Added from_scipy method to SparseVector

pgvector/utils/sparsevec.py

@@ -20,6 +20,18 @@ def from_dict(cls, d, dim):
         values = [float(v[1]) for v in elements]
         return cls(dim, indices, values)
 
+    @classmethod
+    def from_scipy(cls, value):
+        value = value.tocoo()
+
+        if value.ndim != 1:
+            raise ValueError('expected ndim to be 1')
+
+        dim = value.shape[0]
+        indices = value.coords[0].tolist()
+        values = value.data.tolist()
+        return cls(dim, indices, values)
+
     @classmethod
     def from_dense(cls, value):
         dim = len(value)

requirements.txt

@@ -6,5 +6,6 @@ psycopg[binary]
 psycopg2-binary
 pytest
 pytest-asyncio
+scipy
 SQLAlchemy[asyncio]>=2
 sqlmodel>=0.0.12

tests/test_sparse_vector.py

@@ -1,6 +1,7 @@
 import numpy as np
 from pgvector.utils import SparseVector
 import pytest
+from scipy.sparse import coo_array
 
 
 class TestSparseVector:
@@ -12,6 +13,11 @@ def test_from_dense(self):
     def test_from_coordinates(self):
         assert SparseVector.from_dict({0: 1, 2: 2, 4: 3}, 6).to_list() == [1, 0, 2, 0, 3, 0]
 
+    def test_from_scipy(self):
+        arr = coo_array(np.array([1, 0, 2, 0, 3, 0]))
+        assert SparseVector.from_scipy(arr).to_list() == [1, 0, 2, 0, 3, 0]
+        assert SparseVector.from_scipy(arr.todok()).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])'