Added to_dict method to SparseVector

pgvector · Jun 9, 2024 · 8eaad5c · 8eaad5c
1 parent 33eb182
commit 8eaad5c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py
@@ -49,6 +49,9 @@ def from_dense(cls, value):
     def dim(self):
         return self._dim
 
+    def to_dict(self):
+        return {i: v for i, v in zip(self._indices, self._values)}
+
     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
@@ -26,3 +26,6 @@ def test_repr(self):
 
     def test_dim(self):
         assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).dim() == 6
+
+    def test_to_dict(self):
+        assert SparseVector.from_dense([1, 0, 2, 0, 3, 0]).to_dict() == {0: 1, 2: 2, 4: 3}