Add Zvec Sparse ANN (#1307)

test/python/testann/testdense/testzvec.py

@@ -25,7 +25,7 @@ class TestZvec(DenseTest):
         Test zvec backend with custom settings
         """
 
-        self.runTests("zvec", {"zvec": {"m": 16}})
+        self.runTests("zvec", {"zvec": {"efconstruction": 100, "m": 16, "efsearch": 50}})
 
         # Test invalid file path handled
         with self.assertRaises(FileNotFoundError):

test/python/testann/testsparse/testzvecsparse.py

@@ -0,0 +1,88 @@
+"""
+ZvecSparse ANN module tests
+"""
+
+import os
+import tempfile
+
+from scipy.sparse import csr_matrix
+
+from txtai.ann import SparseANNFactory
+
+from .base import SparseTest
+
+
+class TestZvecSparse(SparseTest):
+    """
+    ZvecSparse ANN tests.
+    """
+
+    def testZvecSparse(self):
+        """
+        Test Sparse zvec backend
+        """
+
+        # Generate test record
+        data = self.generate(1, 30522)
+
+        # Create ANN
+        path = os.path.join(tempfile.gettempdir(), "zvecsparse")
+        ann = SparseANNFactory.create({"backend": "zvecsparse", "dimensions": 30522})
+
+        # Test indexing
+        ann.index(data)
+        ann.append(data)
+
+        # Validate search results
+        self.assertEqual([[(uid, round(score, 5)) for uid, score in x] for x in ann.search(data, 1)], [[(0, 1.0)]])
+
+        # Validate save/load
+        ann.save(path)
+        ann.load(path)
+
+        # Validate count
+        self.assertEqual(ann.count(), 2)
+
+        # Test delete
+        ann.delete([0])
+        self.assertEqual(ann.count(), 1)
+        self.assertEqual(ann.search(data, 1)[0][0][0], 1)
+
+        # Close ANN
+        ann.close()
+
+    def testZvecSparseEmpty(self):
+        """
+        Test ZvecSparse with an empty query
+        """
+
+        # Create ANN with an empty row
+        data = csr_matrix([[1.0, 0.0], [0.0, 0.0]])
+        ann = SparseANNFactory.create({"backend": "zvecsparse", "dimensions": 2})
+        ann.index(data)
+
+        # Validate empty rows are indexed and searching with an empty query doesn't raise an error
+        self.assertEqual(ann.count(), 2)
+        self.assertEqual(len(ann.search(data, 1)), 2)
+
+        # Close ANN
+        ann.close()
+
+    def testZvecSparseScores(self):
+        """
+        Test ZvecSparse scores are exact inner products
+        """
+
+        # Create ANN
+        data = self.generate(50, 30522)
+        ann = SparseANNFactory.create({"backend": "zvecsparse", "dimensions": 30522})
+        ann.index(data)
+
+        # Validate scores against exact inner products
+        expected = (data[:5] @ data.T).toarray()
+        for x, results in enumerate(ann.search(data[:5], 10)):
+            for uid, score in results:
+                self.assertAlmostEqual(score, expected[x, uid], places=5)
+
+        # Close ANN
+        ann.close()

test/python/testoptional.py

@@ -162,6 +162,9 @@ class TestOptional(unittest.TestCase):
         with self.assertRaises(ImportError):
             SparseANNFactory.create({"backend": "pgsparse"})
 
+        with self.assertRaises(ImportError):
+            SparseANNFactory.create({"backend": "zvecsparse"})
+
     def testApi(self):
         """
         Test missing api dependencies

test/python/testscoring/testsparse.py

@@ -205,3 +205,46 @@ class TestSparse(unittest.TestCase):
         results = scoring.search("lottery ticket", 1)
         self.assertGreater(len(results), 0)
         scoring.close()
+
+    def testZvecSparse(self):
+        """
+        Test sparse vectors with a zvec backend
+        """
+
+        # Sparse vectors stored in a zvec index
+        config = {"method": "sparse", "path": "sparse-encoder-testing/splade-bert-tiny-nq", "backend": "zvecsparse"}
+        scoring = ScoringFactory.create(config)
+        scoring.index((uid, {"text": text}, tags) for uid, text, tags in self.data)
+
+        # Run search and validate correct result returned
+        index, _ = scoring.search("lottery ticket", 1)[0]
+        self.assertEqual(index, 4)
+
+        # Validate count
+        self.assertEqual(scoring.count(), len(self.data))
+
+        # Generate temp file path
+        index = os.path.join(tempfile.gettempdir(), "scoring")
+        os.makedirs(index, exist_ok=True)
+
+        # Save scoring instance
+        scoring.save(f"{index}/scoring.zvecsparse.index")
+
+        # Reload scoring instance
+        scoring = ScoringFactory.create(config)
+        scoring.load(f"{index}/scoring.zvecsparse.index")
+
+        # Run search and validate correct result returned
+        index, _ = scoring.search("lottery ticket", 1)[0]
+        self.assertEqual(index, 4)
+
+        # Test delete
+        scoring.delete([4])
+        self.assertEqual(scoring.count(), len(self.data) - 1)
+
+        # Run search after delete
+        index, _ = scoring.search("lottery ticket", 1)[0]
+        self.assertEqual(index, 5)
+
+        # Close scoring
+        scoring.close()