feat: support ivf rabitq

.github/workflows/03-macos-linux-build.yml

@@ -177,7 +177,7 @@ jobs:
         shell: bash
 
       # ------------------------------------------------------------------ #
-      # DiskAnn libaio round: install libaio and re-run tests 
+      # DiskAnn libaio round: install libaio and re-run tests
       # ------------------------------------------------------------------ #
       - name: Install libaio runtime
         if: matrix.platform == 'linux-x64'

python/tests/test_collection_ivf_rabitq.py

@@ -0,0 +1,261 @@
+# Copyright 2025-present the zvec project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import platform
+import sys
+
+import pytest
+import zvec
+from zvec import (
+    Collection,
+    CollectionOption,
+    CollectionSchema,
+    DataType,
+    Doc,
+    FieldSchema,
+    IndexType,
+    InvertIndexParam,
+    IvfRabitqIndexParam,
+    IvfRabitqQueryParam,
+    MetricType,
+    OptimizeOption,
+    Query,
+    QuantizeType,
+    VectorSchema,
+)
+
+pytestmark = pytest.mark.skipif(
+    not (sys.platform == "linux" and platform.machine() in ("x86_64", "AMD64")),
+    reason="IVF RaBitQ only supported on Linux x86_64",
+)
+
+DIMENSION = 128
+NLIST = 16
+NPROBE = 16
+TOPK = 10
+SMALL_DOC_COUNT = 64
+INDEXED_DOC_COUNT = 1200
+
+
+def _vector_for(doc_id: int) -> list[float]:
+    return [float(doc_id)] * DIMENSION
+
+
+def _make_docs(count: int) -> list[Doc]:
+    return [
+        Doc(
+            id=str(doc_id),
+            fields={
+                "id": doc_id,
+                "name": f"test_doc_{doc_id}",
+                "group_id": doc_id % 4,
+            },
+            vectors={"embedding": _vector_for(doc_id)},
+        )
+        for doc_id in range(count)
+    ]
+
+
+def _build_schema(name: str) -> CollectionSchema:
+    return CollectionSchema(
+        name=name,
+        fields=[
+            FieldSchema(
+                "id",
+                DataType.INT64,
+                nullable=False,
+                index_param=InvertIndexParam(enable_range_optimization=True),
+            ),
+            FieldSchema("name", DataType.STRING, nullable=False),
+            FieldSchema("group_id", DataType.INT64, nullable=False),
+        ],
+        vectors=[
+            VectorSchema(
+                "embedding",
+                DataType.VECTOR_FP32,
+                dimension=DIMENSION,
+                index_param=IvfRabitqIndexParam(
+                    metric_type=MetricType.L2,
+                    nlist=NLIST,
+                    total_bits=7,
+                    sample_count=0,
+                ),
+            ),
+        ],
+    )
+
+
+def _create_collection(tmp_path_factory, schema_name: str) -> Collection:
+    path = tmp_path_factory.mktemp("zvec_ivf_rabitq") / schema_name
+    return zvec.create_and_open(
+        path=str(path),
+        schema=_build_schema(schema_name),
+        option=CollectionOption(read_only=False, enable_mmap=True),
+    )
+
+
+def _insert_docs(coll: Collection, docs: list[Doc]) -> None:
+    batch_size = 512
+    for offset in range(0, len(docs), batch_size):
+        batch = docs[offset : offset + batch_size]
+        result = coll.insert(batch)
+        assert len(result) == len(batch)
+        for item in result:
+            assert item.ok()
+    assert coll.stats.doc_count == len(docs)
+
+
+def _query(
+    coll: Collection,
+    query_vector: list[float],
+    *,
+    filter: str | None = None,
+    include_vector: bool = False,
+) -> list[Doc]:
+    query = Query(
+        field_name="embedding",
+        vector=query_vector,
+        param=IvfRabitqQueryParam(nprobe=NPROBE),
+    )
+    result = coll.query(
+        queries=query,
+        topk=TOPK,
+        filter=filter,
+        include_vector=include_vector,
+    )
+    assert result is not None
+    assert 0 < len(result) <= TOPK
+    return result
+
+
+class TestIvfRabitqCollection:
+    def test_schema_round_trip(self, tmp_path_factory):
+        coll = _create_collection(tmp_path_factory, "ivf_rabitq_schema_rt")
+        try:
+            vector_schema = coll.schema.vector("embedding")
+            assert vector_schema is not None
+            assert vector_schema.data_type == DataType.VECTOR_FP32
+            assert vector_schema.dimension == DIMENSION
+
+            index_param = vector_schema.index_param
+            assert index_param.type == IndexType.IVF_RABITQ
+            assert index_param.metric_type == MetricType.L2
+            assert index_param.quantize_type == QuantizeType.RABITQ
+            assert index_param.nlist == NLIST
+            assert index_param.total_bits == 7
+            assert index_param.sample_count == 0
+            assert coll.stats.index_completeness["embedding"] == 1
+        finally:
+            coll.destroy()
+
+    def test_insert_fetch_and_query_writer_segment(self, tmp_path_factory):
+        coll = _create_collection(tmp_path_factory, "ivf_rabitq_writer_query")
+        try:
+            docs = _make_docs(SMALL_DOC_COUNT)
+            _insert_docs(coll, docs)
+
+            fetched = coll.fetch(ids=["7", "42"])
+            assert set(fetched.keys()) == {"7", "42"}
+            assert fetched["42"].field("name") == "test_doc_42"
+            assert fetched["42"].vector("embedding") == _vector_for(42)
+
+            result = _query(coll, docs[42].vector("embedding"))
+            assert result[0].id == "42"
+        finally:
+            coll.destroy()
+
+    def test_optimize_reopen_and_query(self, tmp_path_factory):
+        schema_name = "ivf_rabitq_optimize_reopen"
+        collection_path = tmp_path_factory.mktemp("zvec_ivf_rabitq") / schema_name
+        option = CollectionOption(read_only=False, enable_mmap=True)
+        coll = zvec.create_and_open(
+            path=str(collection_path),
+            schema=_build_schema(schema_name),
+            option=option,
+        )
+        reopened = None
+        try:
+            docs = _make_docs(INDEXED_DOC_COUNT)
+            _insert_docs(coll, docs)
+
+            coll.optimize(OptimizeOption())
+            assert coll.stats.doc_count == INDEXED_DOC_COUNT
+            assert coll.stats.index_completeness["embedding"] == 1
+
+            query_vector = docs[42].vector("embedding")
+            result = _query(coll, query_vector)
+            assert result[0].id == "42"
+
+            filtered = _query(coll, query_vector, filter="id < 50", include_vector=True)
+            assert filtered[0].id == "42"
+            for doc in filtered:
+                assert doc.field("id") < 50
+            assert len(filtered[0].vector("embedding")) == DIMENSION
+
+            path = str(collection_path)
+            del coll
+            coll = None
+
+            reopened = zvec.open(path=path, option=option)
+            assert reopened.stats.doc_count == INDEXED_DOC_COUNT
+            reopened_result = _query(reopened, query_vector)
+            assert reopened_result[0].id == "42"
+        finally:
+            if reopened is not None:
+                reopened.destroy()
+            elif coll is not None:
+                coll.destroy()
+
+    def test_group_by_after_optimize(self, tmp_path_factory):
+        coll = _create_collection(tmp_path_factory, "ivf_rabitq_group_by")
+        try:
+            docs = _make_docs(INDEXED_DOC_COUNT)
+            _insert_docs(coll, docs)
+            coll.optimize(OptimizeOption())
+
+            query = Query(
+                field_name="embedding",
+                vector=docs[42].vector("embedding"),
+                param=IvfRabitqQueryParam(nprobe=NPROBE),
+            )
+            results = coll.group_by_query(
+                query,
+                group_by_field_name="group_id",
+                group_count=4,
+                topk_per_group=2,
+            )
+            assert {int(group.group_by_value) for group in results} == set(range(4))
+            for group in results:
+                assert 1 <= len(group.docs) <= 2
+                assert all(
+                    doc.field("group_id") == int(group.group_by_value)
+                    for doc in group.docs
+                )
+
+            filtered = coll.group_by_query(
+                query,
+                group_by_field_name="group_id",
+                group_count=4,
+                topk_per_group=2,
+                filter="id < 8",
+                include_vector=True,
+            )
+            assert {int(group.group_by_value) for group in filtered} == set(range(4))
+            for group in filtered:
+                for doc in group.docs:
+                    assert doc.field("id") < 8
+                    assert len(doc.vector("embedding")) == DIMENSION
+        finally:
+            coll.destroy()