Added from_coordinates method to SparseVector

ankane · ankane · commit df3427dace3c · 2024-06-04T15:36:51.000-07:00
diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py
@@ -12,6 +12,16 @@ def __init__(self, dim, indices, values):
     def __repr__(self):
         return f'SparseVector({self._dim}, {self._indices}, {self._values})'
 
+    @classmethod
+    def from_coordinates(cls, coordinates, dim):
+        if isinstance(coordinates, dict):
+            coordinates = coordinates.items()
+        elements = [(i, v) for i, v in coordinates]
+        elements.sort()
+        indices = [int(v[0]) for v in elements]
+        values = [float(v[1]) for v in elements]
+        return cls(dim, indices, values)
+
     @classmethod
     def from_dense(cls, value):
         dim = len(value)
diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py
@@ -8,6 +8,10 @@ def test_from_dense(self):
         assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).to_list() == [1, 0, 2, 0, 3, 0]
         assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).to_numpy().tolist() == [1, 0, 2, 0, 3, 0]
 
+    def test_from_coordinates(self):
+        assert SparseVector.from_coordinates({0: 1, 2: 2, 4: 3}, 6).to_list() == [1, 0, 2, 0, 3, 0]
+        assert SparseVector.from_coordinates([(0, 1), (2, 2), (4, 3)], 6).to_list() == [1, 0, 2, 0, 3, 0]
+
     def test_repr(self):
         assert repr(SparseVector.from_dense([1, 0, 2, 0, 3, 0])) == 'SparseVector(6, [0, 2, 4], [1.0, 2.0, 3.0])'
         assert str(SparseVector.from_dense([1, 0, 2, 0, 3, 0])) == 'SparseVector(6, [0, 2, 4], [1.0, 2.0, 3.0])'
