From 20c9b8e0bee2c29d7b5211272b7421cb09a0af87 Mon Sep 17 00:00:00 2001 From: Ni0ri Date: Thu, 1 Oct 2026 15:53:38 +0200 Subject: [PATCH] Add Zvec Sparse ANN (#1307) --- docs/embeddings/configuration/ann.md | 2 + docs/embeddings/configuration/scoring.md | 12 ++- src/python/txtai/ann/dense/zvec.py | 43 +++++++-- src/python/txtai/ann/sparse/__init__.py | 1 + src/python/txtai/ann/sparse/factory.py | 3 + src/python/txtai/ann/sparse/zvecsparse.py | 35 ++++++++ test/python/testann/testdense/testzvec.py | 2 +- .../testann/testsparse/testzvecsparse.py | 88 +++++++++++++++++++ test/python/testoptional.py | 3 + test/python/testscoring/testsparse.py | 43 +++++++++ 10 files changed, 225 insertions(+), 7 deletions(-) create mode 100644 src/python/txtai/ann/sparse/zvecsparse.py create mode 100644 test/python/testann/testsparse/testzvecsparse.py diff --git a/docs/embeddings/configuration/ann.md b/docs/embeddings/configuration/ann.md index 8f3f8b6b..899fc351 100644 --- a/docs/embeddings/configuration/ann.md +++ b/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. diff --git a/docs/embeddings/configuration/scoring.md b/docs/embeddings/configuration/scoring.md index 876a0843..2767e344 100644 --- a/docs/embeddings/configuration/scoring.md +++ b/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 diff --git a/src/python/txtai/ann/dense/zvec.py b/src/python/txtai/ann/dense/zvec.py index 7d793b8e..46a8041f 100644 --- a/src/python/txtai/ann/dense/zvec.py +++ b/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() diff --git a/src/python/txtai/ann/sparse/__init__.py b/src/python/txtai/ann/sparse/__init__.py index 616cdd4e..a52eddf5 100644 --- a/src/python/txtai/ann/sparse/__init__.py +++ b/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 diff --git a/src/python/txtai/ann/sparse/factory.py b/src/python/txtai/ann/sparse/factory.py index ca12f6dc..addcd0fe 100644 --- a/src/python/txtai/ann/sparse/factory.py +++ b/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) diff --git a/src/python/txtai/ann/sparse/zvecsparse.py b/src/python/txtai/ann/sparse/zvecsparse.py new file mode 100644 index 00000000..b08bd79a --- /dev/null +++ b/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)} diff --git a/test/python/testann/testdense/testzvec.py b/test/python/testann/testdense/testzvec.py index 43ad3c21..c27e687c 100644 --- a/test/python/testann/testdense/testzvec.py +++ b/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): diff --git a/test/python/testann/testsparse/testzvecsparse.py b/test/python/testann/testsparse/testzvecsparse.py new file mode 100644 index 00000000..7084f875 --- /dev/null +++ b/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() diff --git a/test/python/testoptional.py b/test/python/testoptional.py index 116fda7f..c7956d94 100644 --- a/test/python/testoptional.py +++ b/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 diff --git a/test/python/testscoring/testsparse.py b/test/python/testscoring/testsparse.py index 1acb9a3d..cc44151e 100644 --- a/test/python/testscoring/testsparse.py +++ b/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()