Fixed dict constructor for SparseVector [skip ci]

ankane · ankane · commit 0266a2ee38e3 · 2024-06-20T18:39:25.000-07:00
diff --git a/pgvector/utils/sparsevec.py b/pgvector/utils/sparsevec.py
@@ -59,7 +59,7 @@ def _from_dict(self, d, dim):
         if dim is None:
             raise ValueError('dimensions required')
 
-        elements = [(i, v) for i, v in d.items()]
+        elements = [(i, v) for i, v in d.items() if v != 0]
         elements.sort()
         self._dim = int(dim)
         self._indices = [int(v[0]) for v in elements]
diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py
@@ -16,8 +16,9 @@ def test_from_dense_dimensions(self):
         assert str(error.value) == 'dimensions not allowed'
 
     def test_from_dict(self):
-        assert SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6).to_list() == [1, 0, 2, 0, 3, 0]
-        assert SparseVector([1, 0, 2, 0, 3, 0]).indices() == [0, 2, 4]
+        vec = SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6)
+        assert vec.to_list() == [1, 0, 2, 0, 3, 0]
+        assert vec.indices() == [0, 2, 4]
 
     def test_from_dict_no_dimensions(self):
         with pytest.raises(ValueError) as error:
