Improved SparseVector code [skip ci]

ankane · ankane · commit b1d2cdb17c8a · 2024-06-01T20:02:37.000-07:00
diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py
@@ -13,9 +13,10 @@ def to_db_value(value):
 
 class SparseVector:
     def __init__(self, dim, indices, values):
-        self._dim = dim
-        self._indices = indices
-        self._values = values
+        # TODO improve
+        self._dim = int(dim)
+        self._indices = [int(i) for i in indices]
+        self._values = [float(v) for v in values]
 
     def __repr__(self):
         return f'SparseVector({self._dim}, {self._indices}, {self._values})'
@@ -25,7 +26,7 @@ def from_dense(value):
             value = value.tolist()
         dim = len(value)
         indices = [i for i, v in enumerate(value) if v != 0]
-        values = [value[i] for i in indices]
+        values = [float(value[i]) for i in indices]
         return SparseVector(dim, indices, values)
 
     def dim(self):
diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py
@@ -9,5 +9,5 @@ def test_from_dense(self):
         assert SparseVector.from_dense([1, 2, 3]).to_numpy().tolist() == [1, 2, 3]
 
     def test_repr(self):
-        assert repr(SparseVector.from_dense([1, 2, 3])) == 'SparseVector(3, [0, 1, 2], [1, 2, 3])'
-        assert str(SparseVector.from_dense([1, 2, 3])) == 'SparseVector(3, [0, 1, 2], [1, 2, 3])'
+        assert repr(SparseVector.from_dense([1, 2, 3])) == 'SparseVector(3, [0, 1, 2], [1.0, 2.0, 3.0])'
+        assert str(SparseVector.from_dense([1, 2, 3])) == 'SparseVector(3, [0, 1, 2], [1.0, 2.0, 3.0])'
