Add Zvec Sparse ANN (#1307)

docs/embeddings/configuration/ann.md

@@ -139,7 +139,9 @@ The [turbovec](https://github.com/RyanCodrai/turbovec) backend is a k-nearest ne
 
 ```yaml
 zvec:
+    efconstruction: ef_construction param for HnswIndexParam (int) - defaults to 200
     m: number of HNSW links per element (int) - defaults to 50
+    efsearch: ef search param for HnswQueryParam (int) - defaults to 300
 ```
 
 The [zvec](https://github.com/alibaba/zvec) backend is an embedded, path-based vector index.

docs/embeddings/configuration/scoring.md

@@ -36,7 +36,7 @@ batch: Sets the transform batch size
 encodebatch: Sets the encode batch size
 vectors: additional model init args
 encodeargs: additional encode() args
-backend: ivfsparse|pgsparse
+backend: ivfsparse|pgsparse|zvecsparse
 ```
 
 Sparse vector scoring options. The sparse scoring instance combines a sparse vector model with a sparse approximate nearest neighbor index (ANN). This method supports both vector normalization and score normalization.
@@ -66,6 +66,16 @@ Inverted file (IVF) index with flat vector file storage and sparse array support
 
 Sparse ANN backed by Postgres. Supports same options as the [pgvector](../ann/#pgvector) ANN.
 
+#### zvecsparse
+```yaml
+zvecsparse:
+  efconstruction: ef_construction param for HnswIndexParam (int) - defaults to 200
+  m: number of HNSW links per element (int) - defaults to 50
+  efsearch: ef search param for HnswQueryParam (int) - defaults to 300
+```
+
+Sparse ANN backed by [zvec](https://github.com/alibaba/zvec). Stores sparse vectors in an embedded, path-based vector index. Requires the [ann](../../../install/#ann) extras package.
+
 ## terms
 ```yaml
 terms: boolean|dict

src/python/txtai/ann/dense/zvec.py

@@ -49,6 +49,7 @@ class Zvec(ANN):
         self.close()
 
         # Lookup index settings
+        efconstruction = self.setting("efconstruction", 200)
         m = self.setting("m", 50)
 
         # Create collection
@@ -58,9 +59,9 @@ class Zvec(ANN):
             name="txtai",
             vectors=zvec.VectorSchema(
                 "embedding",
-                zvec.DataType.VECTOR_FP32,
+                self.datatype(),
                 self.config["dimensions"],
-                index_param=zvec.HnswIndexParam(metric_type=zvec.MetricType.IP, m=m),
+                index_param=zvec.HnswIndexParam(metric_type=zvec.MetricType.IP, m=m, ef_construction=efconstruction),
             ),
         )
         self.backend = zvec.create_and_open(path=self.path, schema=schema)
@@ -68,13 +69,19 @@ class Zvec(ANN):
         # Add items - position in embeddings is used as the id
         self.insert(embeddings, 0)
 
+        # Build HNSW index
+        self.backend.optimize()
+
         # Add id offset and index build metadata
         self.config["offset"] = embeddings.shape[0]
-        self.metadata({"m": m, "zvec": zvec.__version__})
+        self.metadata({"efconstruction": efconstruction, "m": m, "zvec": zvec.__version__})
 
     def append(self, embeddings):
         self.insert(embeddings, self.config["offset"])
 
+        # Add new items to HNSW index
+        self.backend.optimize()
+
         # Update id offset and index metadata
         self.config["offset"] += embeddings.shape[0]
         self.metadata()
@@ -85,10 +92,13 @@ class Zvec(ANN):
             self.backend.delete([str(uid) for uid in ids])
 
     def search(self, queries, limit):
+        # Lookup search settings
+        param = zvec.HnswQueryParam(ef=self.setting("efsearch", 300))
+
         results = []
         for query in queries:
             matches = self.backend.query(
-                zvec.Query(field_name="embedding", vector=query.tolist()),
+                zvec.Query(field_name="embedding", vector=self.prepare(query), param=param),
                 topk=limit,
             )
             results.append([(int(match.id), float(match.score)) for match in matches])
@@ -129,5 +139,28 @@ class Zvec(ANN):
             for start in range(0, embeddings.shape[0], 1024):
                 batch = embeddings[start : start + 1024]
                 self.backend.insert(
-                    [zvec.Doc(id=str(offset + start + uid), vectors={"embedding": embedding.tolist()}) for uid, embedding in enumerate(batch)]
+                    [zvec.Doc(id=str(offset + start + uid), vectors={"embedding": self.prepare(embedding)}) for uid, embedding in enumerate(batch)]
                 )
+
+    def datatype(self):
+        """
+        Gets the vector data type.
+
+        Returns:
+            zvec vector data type
+        """
+
+        return zvec.DataType.VECTOR_FP32
+
+    def prepare(self, data):
+        """
+        Prepares data for the vector field.
+
+        Args:
+            data: input data
+
+        Returns:
+            data formatted for zvec
+        """
+
+        return data.tolist()

src/python/txtai/ann/sparse/__init__.py

@@ -5,3 +5,4 @@ Sparse ANN imports
 from .factory import SparseANNFactory
 from .ivfsparse import IVFSparse
 from .pgsparse import PGSparse
+from .zvecsparse import ZvecSparse

src/python/txtai/ann/sparse/factory.py

@@ -7,6 +7,7 @@ from ...util import Resolver
 from ..base import ANN
 from .ivfsparse import IVFSparse
 from .pgsparse import PGSparse
+from .zvecsparse import ZvecSparse
 
 
 class SparseANNFactory:
@@ -35,6 +36,8 @@ class SparseANNFactory:
             ann = IVFSparse(config)
         elif backend == "pgsparse":
             ann = PGSparse(config)
+        elif backend == "zvecsparse":
+            ann = ZvecSparse(config)
         else:
             ann = SparseANNFactory.resolve(backend, config)
 

src/python/txtai/ann/sparse/zvecsparse.py

@@ -0,0 +1,35 @@
+"""
+ZvecSparse module
+"""
+
+# Conditional import
+try:
+    import zvec
+except ImportError:
+    pass
+
+from ..dense import Zvec
+
+
+class ZvecSparse(Zvec):
+    """
+    Builds a Sparse ANN index using the zvec library.
+    """
+
+    def search(self, queries, limit):
+        # Lookup search settings
+        param = zvec.HnswQueryParam(ef=self.setting("efsearch", 300))
+
+        results = []
+        for query in queries:
+            # zvec requires at least one non-zero value per query
+            matches = self.backend.query(zvec.Query(field_name="embedding", vector=self.prepare(query), param=param), topk=limit) if query.nnz else []
+            results.append([(int(match.id), float(match.score)) for match in matches])
+
+        return results
+
+    def datatype(self):
+        return zvec.DataType.SPARSE_VECTOR_FP32
+
+    def prepare(self, data):
+        return {int(index): float(value) for index, value in zip(data.indices, data.data)}

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()