Added to_coo method to SparseVector

ankane · ankane · commit 8be15b29e900 · 2024-06-09T14:54:21.000-07:00
diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py
@@ -52,6 +52,12 @@ def dim(self):
     def to_dict(self):
         return dict(zip(self._indices, self._values))
 
+    def to_coo(self):
+        from scipy.sparse import coo_array
+
+        coords = ([0] * len(self._indices), self._indices)
+        return coo_array((self._values, coords), shape=(1, self._dim))
+
     def to_list(self):
         vec = [0.0] * self._dim
         for i, v in zip(self._indices, self._values):
diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py
@@ -27,3 +27,6 @@ def test_dim(self):
 
     def test_to_dict(self):
         assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).to_dict() == {0: 1, 2: 2, 4: 3}
+
+    def test_to_coo(self):
+        assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).to_coo().toarray().tolist() == [[1, 0, 2, 0, 3, 0]]
