From e1579386f2f5dc253b58d1136c187bec5efae056 Mon Sep 17 00:00:00 2001 From: "jiliang.ljl" Date: Thu, 25 Jun 2026 17:45:34 +0800 Subject: [PATCH] feat(ivf-rabitq): add IVF RaBitQ index support Add the core IVF RaBitQ builder, streamer, entity, reformer, on-disk format, and index provider flow. Wire IVF_RABITQ through DB schema validation, proto conversion, collection optimization, C API, and Python bindings. Cover streamer behavior, collection optimize/query paths, parameter serialization, schema validation, proto conversion, and C API usage with tests. --- python/tests/test_collection_ivf_rabitq.py | 212 ++++ python/tests/test_params.py | 103 ++ python/zvec/__init__.py | 4 + python/zvec/__init__.pyi | 4 + python/zvec/model/collection.py | 6 +- python/zvec/model/param/__init__.py | 4 + python/zvec/model/param/__init__.pyi | 61 ++ python/zvec/model/param/query.py | 18 +- python/zvec/model/schema/field_schema.py | 23 +- python/zvec/typing/__init__.pyi | 20 +- src/binding/c/c_api.cc | 215 +++++ src/binding/python/CMakeLists.txt | 2 + .../python/model/param/python_param.cc | 152 +++ src/binding/python/typing/python_type.cc | 3 +- src/core/CMakeLists.txt | 19 + src/core/algorithm/CMakeLists.txt | 16 + .../algorithm/hnsw_rabitq/rabitq_params.h | 4 + src/core/algorithm/ivf_rabitq/CMakeLists.txt | 27 + .../ivf_rabitq/ivf_rabitq_builder.cc | 537 +++++++++++ .../algorithm/ivf_rabitq/ivf_rabitq_builder.h | 91 ++ .../algorithm/ivf_rabitq/ivf_rabitq_context.h | 229 +++++ .../algorithm/ivf_rabitq/ivf_rabitq_entity.cc | 508 ++++++++++ .../algorithm/ivf_rabitq/ivf_rabitq_entity.h | 145 +++ .../ivf_rabitq/ivf_rabitq_index_format.h | 58 ++ .../ivf_rabitq/ivf_rabitq_index_provider.h | 134 +++ .../algorithm/ivf_rabitq/ivf_rabitq_params.h | 46 + .../ivf_rabitq/ivf_rabitq_reformer.cc | 540 +++++++++++ .../ivf_rabitq/ivf_rabitq_reformer.h | 101 ++ .../ivf_rabitq/ivf_rabitq_register.cc | 25 + .../ivf_rabitq/ivf_rabitq_streamer.cc | 397 ++++++++ .../ivf_rabitq/ivf_rabitq_streamer.h | 134 +++ .../algorithm/ivf_rabitq/ivf_rabitq_util.cc | 77 ++ .../algorithm/ivf_rabitq/ivf_rabitq_util.h | 28 + src/core/interface/index_factory.cc | 39 + src/core/interface/index_param.cc | 31 +- .../interface/indexes/ivf_rabitq_index.cc | 346 +++++++ src/core/utility/mmap_file_read_storage.cc | 6 +- .../column/vector_column/engine_helper.hpp | 44 +- src/db/index/common/proto_converter.cc | 35 +- src/db/index/common/proto_converter.h | 5 + src/db/index/common/schema.cc | 14 +- src/db/index/common/type_helper.h | 6 + src/db/index/segment/segment.cc | 2 +- src/db/index/segment/segment_helper.cc | 48 +- src/db/index/segment/segment_helper.h | 2 +- src/db/proto/zvec.proto | 10 + src/include/zvec/c_api.h | 159 +++ src/include/zvec/core/interface/constants.h | 5 +- src/include/zvec/core/interface/index.h | 32 + src/include/zvec/core/interface/index_param.h | 37 + .../core/interface/index_param_builders.h | 34 +- src/include/zvec/db/index_params.h | 88 +- src/include/zvec/db/query_params.h | 24 + src/include/zvec/db/type.h | 1 + tests/c/c_api_test.c | 76 ++ tests/core/algorithm/CMakeLists.txt | 1 + .../core/algorithm/ivf_rabitq/CMakeLists.txt | 34 + .../ivf_rabitq/ivf_rabitq_streamer_test.cc | 912 ++++++++++++++++++ tests/core/interface/CMakeLists.txt | 2 +- tests/db/CMakeLists.txt | 1 + tests/db/collection_test.cc | 90 +- .../index/common/db_proto_converter_test.cc | 63 +- tests/db/index/common/index_params_test.cc | 38 +- tests/db/index/common/schema_test.cc | 94 +- tests/db/index/utils/utils.cc | 8 +- tools/core/CMakeLists.txt | 12 +- tools/core/local_builder.cc | 3 + 67 files changed, 6174 insertions(+), 71 deletions(-) create mode 100644 python/tests/test_collection_ivf_rabitq.py create mode 100644 src/core/algorithm/ivf_rabitq/CMakeLists.txt create mode 100644 src/core/algorithm/ivf_rabitq/ivf_rabitq_builder.cc create mode 100644 src/core/algorithm/ivf_rabitq/ivf_rabitq_builder.h create mode 100644 src/core/algorithm/ivf_rabitq/ivf_rabitq_context.h create mode 100644 src/core/algorithm/ivf_rabitq/ivf_rabitq_entity.cc create mode 100644 src/core/algorithm/ivf_rabitq/ivf_rabitq_entity.h create mode 100644 src/core/algorithm/ivf_rabitq/ivf_rabitq_index_format.h create mode 100644 src/core/algorithm/ivf_rabitq/ivf_rabitq_index_provider.h create mode 100644 src/core/algorithm/ivf_rabitq/ivf_rabitq_params.h create mode 100644 src/core/algorithm/ivf_rabitq/ivf_rabitq_reformer.cc create mode 100644 src/core/algorithm/ivf_rabitq/ivf_rabitq_reformer.h create mode 100644 src/core/algorithm/ivf_rabitq/ivf_rabitq_register.cc create mode 100644 src/core/algorithm/ivf_rabitq/ivf_rabitq_streamer.cc create mode 100644 src/core/algorithm/ivf_rabitq/ivf_rabitq_streamer.h create mode 100644 src/core/algorithm/ivf_rabitq/ivf_rabitq_util.cc create mode 100644 src/core/algorithm/ivf_rabitq/ivf_rabitq_util.h create mode 100644 src/core/interface/indexes/ivf_rabitq_index.cc create mode 100644 tests/core/algorithm/ivf_rabitq/CMakeLists.txt create mode 100644 tests/core/algorithm/ivf_rabitq/ivf_rabitq_streamer_test.cc diff --git a/python/tests/test_collection_ivf_rabitq.py b/python/tests/test_collection_ivf_rabitq.py new file mode 100644 index 000000000..60886bd0b --- /dev/null +++ b/python/tests/test_collection_ivf_rabitq.py @@ -0,0 +1,212 @@ +# 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, + 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}"}, + 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, + ), + FieldSchema("name", DataType.STRING, 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() diff --git a/python/tests/test_params.py b/python/tests/test_params.py index 8f613939f..01772134e 100644 --- a/python/tests/test_params.py +++ b/python/tests/test_params.py @@ -25,11 +25,13 @@ CollectionOption, FlatIndexParam, HnswIndexParam, + IvfRabitqIndexParam, IndexOption, InvertIndexParam, IVFIndexParam, OptimizeOption, HnswQueryParam, + IvfRabitqQueryParam, IVFQueryParam, Query, VectorQuery, @@ -179,6 +181,68 @@ def test_readonly_attributes(self, attr): setattr(param, attr, getattr(param, attr)) +# ---------------------------- +# Ivf Rabitq Index Param Test Case +# ---------------------------- +class TestIvfRabitqIndexParam: + def test_default(self): + param = IvfRabitqIndexParam() + assert param.metric_type == MetricType.IP + assert param.nlist == 1024 + assert param.total_bits == 7 + assert param.sample_count == 0 + assert param.quantize_type == QuantizeType.RABITQ + assert param.type == IndexType.IVF_RABITQ + + def test_custom(self): + param = IvfRabitqIndexParam( + metric_type=MetricType.L2, nlist=128, total_bits=6, sample_count=1000 + ) + assert param.metric_type == MetricType.L2 + assert param.nlist == 128 + assert param.total_bits == 6 + assert param.sample_count == 1000 + assert param.quantize_type == QuantizeType.RABITQ + assert param.type == IndexType.IVF_RABITQ + + def test_to_dict(self): + param = IvfRabitqIndexParam( + metric_type=MetricType.L2, nlist=128, total_bits=6, sample_count=1000 + ) + data = param.to_dict() + assert data["type"] == "IVF_RABITQ" + assert data["metric_type"] == "L2" + assert data["quantize_type"] == "RABITQ" + assert data["nlist"] == 128 + assert data["total_bits"] == 6 + assert data["sample_count"] == 1000 + + def test_vector_schema_accepts_param(self): + param = IvfRabitqIndexParam(metric_type=MetricType.L2, nlist=128) + schema = VectorSchema( + name="embedding", + data_type=DataType.VECTOR_FP32, + dimension=128, + index_param=param, + ) + assert schema.index_param.type == IndexType.IVF_RABITQ + assert schema.index_param.nlist == 128 + + @pytest.mark.parametrize( + "attr", ["metric_type", "nlist", "total_bits", "sample_count", "quantize_type"] + ) + def test_readonly_attributes(self, attr): + param = IvfRabitqIndexParam() + import sys + + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + setattr(param, attr, getattr(param, attr)) + + # ---------------------------- # CollectionOption Test Case # ---------------------------- @@ -360,6 +424,45 @@ def test_readonly_attributes(self): param.is_linear = True +# ---------------------------- +# IvfRabitqQueryParam Test Case +# ---------------------------- +class TestIvfRabitqQueryParam: + def test_default(self): + param = IvfRabitqQueryParam() + assert param is not None + assert param.nprobe == 10 + assert param.is_using_refiner == False + assert param.radius == 0 + assert param.is_linear == False + assert param.type == IndexType.IVF_RABITQ + + def test_custom(self): + param = IvfRabitqQueryParam( + nprobe=20, is_using_refiner=True, radius=30, is_linear=True + ) + assert param.nprobe == 20 + assert param.is_using_refiner == True + assert param.radius == 30 + assert param.is_linear == True + assert param.type == IndexType.IVF_RABITQ + + def test_query_accepts_param(self): + param = IvfRabitqQueryParam(nprobe=20) + query = Query(field_name="embedding", vector=[0.1, 0.2], param=param) + assert query.param == param + assert query.param.type == IndexType.IVF_RABITQ + + def test_readonly_attributes(self): + param = IvfRabitqQueryParam() + if sys.version_info >= (3, 11): + match_pattern = r"(can't set attribute|has no setter|readonly attribute)" + else: + match_pattern = r"can't set attribute" + with pytest.raises(AttributeError, match=match_pattern): + param.nprobe = 10 + + # # ---------------------------- # # IVFQueryParam Test Case # # ---------------------------- diff --git a/python/zvec/__init__.py b/python/zvec/__init__.py index cfda0a4d5..81bf747ba 100644 --- a/python/zvec/__init__.py +++ b/python/zvec/__init__.py @@ -107,6 +107,8 @@ InvertIndexParam, IVFIndexParam, IVFQueryParam, + IvfRabitqIndexParam, + IvfRabitqQueryParam, OptimizeOption, QuantizerParam, VamanaIndexParam, @@ -160,6 +162,7 @@ "InvertIndexParam", "HnswIndexParam", "HnswRabitqIndexParam", + "IvfRabitqIndexParam", "FlatIndexParam", "IVFIndexParam", "DiskAnnIndexParam", @@ -171,6 +174,7 @@ "AlterColumnOption", "HnswQueryParam", "HnswRabitqQueryParam", + "IvfRabitqQueryParam", "IVFQueryParam", "QuantizerParam", "VamanaIndexParam", diff --git a/python/zvec/__init__.pyi b/python/zvec/__init__.pyi index 3e75f931b..5e344f841 100644 --- a/python/zvec/__init__.pyi +++ b/python/zvec/__init__.pyi @@ -29,6 +29,8 @@ from .model.param import ( InvertIndexParam, IVFIndexParam, IVFQueryParam, + IvfRabitqIndexParam, + IvfRabitqQueryParam, OptimizeOption, QuantizerParam, VamanaIndexParam, @@ -75,6 +77,8 @@ __all__: list = [ "IndexOption", "IndexType", "InvertIndexParam", + "IvfRabitqIndexParam", + "IvfRabitqQueryParam", "LogLevel", "LogType", "MetricType", diff --git a/python/zvec/model/collection.py b/python/zvec/model/collection.py index b90caae94..5f0a581b0 100644 --- a/python/zvec/model/collection.py +++ b/python/zvec/model/collection.py @@ -34,6 +34,7 @@ IndexOption, InvertIndexParam, IVFIndexParam, + IvfRabitqIndexParam, OptimizeOption, ) from .param.query import Query @@ -110,6 +111,7 @@ def create_index( index_param: Union[ HnswIndexParam, HnswRabitqIndexParam, + IvfRabitqIndexParam, IVFIndexParam, FlatIndexParam, InvertIndexParam, @@ -119,13 +121,13 @@ def create_index( ) -> None: """Create an index on a field. - Vector index types (HNSW, IVF, FLAT) can only be applied to vector fields. + Vector index types (HNSW, HNSW_RABITQ, IVF, IVF_RABITQ, FLAT) can only be applied to vector fields. Inverted index (`InvertIndexParam`) is for scalar fields. FTS index (`FtsIndexParam`) is for full-text search on STRING fields. Args: field_name (str): Name of the field to index. - index_param (Union[HnswIndexParam, HnswRabitqIndexParam, IVFIndexParam, FlatIndexParam, InvertIndexParam, FtsIndexParam]): + index_param (Union[HnswIndexParam, HnswRabitqIndexParam, IvfRabitqIndexParam, IVFIndexParam, FlatIndexParam, InvertIndexParam, FtsIndexParam]): Index configuration. option (Optional[IndexOption], optional): Index creation options. Defaults to ``IndexOption()``. diff --git a/python/zvec/model/param/__init__.py b/python/zvec/model/param/__init__.py index bcf0f1db5..3539d8bbf 100644 --- a/python/zvec/model/param/__init__.py +++ b/python/zvec/model/param/__init__.py @@ -30,6 +30,8 @@ InvertIndexParam, IVFIndexParam, IVFQueryParam, + IvfRabitqIndexParam, + IvfRabitqQueryParam, OptimizeOption, QuantizerParam, VamanaIndexParam, @@ -53,6 +55,8 @@ "IVFQueryParam", "IndexOption", "InvertIndexParam", + "IvfRabitqIndexParam", + "IvfRabitqQueryParam", "OptimizeOption", "QuantizerParam", "VamanaIndexParam", diff --git a/python/zvec/model/param/__init__.pyi b/python/zvec/model/param/__init__.pyi index 851044a3e..e6c2a3d1e 100644 --- a/python/zvec/model/param/__init__.pyi +++ b/python/zvec/model/param/__init__.pyi @@ -25,6 +25,8 @@ __all__: list[str] = [ "IndexOption", "IndexParam", "InvertIndexParam", + "IvfRabitqIndexParam", + "IvfRabitqQueryParam", "OptimizeOption", "QuantizerParam", "QueryParam", @@ -467,6 +469,65 @@ class HnswRabitqQueryParam(QueryParam): int: Size of the dynamic candidate list during HNSW search. """ +class IvfRabitqIndexParam(VectorIndexParam): + """ + Parameters for configuring an IVF index with RaBitQ quantization. + """ + + def __getstate__(self) -> tuple: ... + def __init__( + self, + metric_type: zvec._zvec.typing.MetricType = ..., + nlist: typing.SupportsInt = 1024, + total_bits: typing.SupportsInt = 7, + sample_count: typing.SupportsInt = 0, + ) -> None: ... + def __repr__(self) -> str: ... + def __setstate__(self, arg0: tuple) -> None: ... + def to_dict(self) -> dict: + """ + Convert to dictionary with all fields + """ + + @property + def nlist(self) -> int: + """ + int: Number of IVF cluster centers. + """ + + @property + def total_bits(self) -> int: + """ + int: Total bits for RaBitQ quantization. + """ + + @property + def sample_count(self) -> int: + """ + int: Sample count for RaBitQ training. + """ + +class IvfRabitqQueryParam(QueryParam): + """ + Query parameters for IVF RaBitQ index. + """ + + def __getstate__(self) -> tuple: ... + def __init__( + self, + nprobe: typing.SupportsInt = 10, + radius: typing.SupportsFloat = 0.0, + is_linear: bool = False, + is_using_refiner: bool = False, + ) -> None: ... + def __repr__(self) -> str: ... + def __setstate__(self, arg0: tuple) -> None: ... + @property + def nprobe(self) -> int: + """ + int: Number of inverted lists to search during IVF RaBitQ query. + """ + class IVFIndexParam(VectorIndexParam): """ diff --git a/python/zvec/model/param/query.py b/python/zvec/model/param/query.py index d0916b3e3..fb752b044 100644 --- a/python/zvec/model/param/query.py +++ b/python/zvec/model/param/query.py @@ -18,7 +18,13 @@ from typing import Optional, Union from ...common import VectorType -from . import FtsQueryParam, HnswQueryParam, HnswRabitqQueryParam, IVFQueryParam +from . import ( + FtsQueryParam, + HnswQueryParam, + HnswRabitqQueryParam, + IVFQueryParam, + IvfRabitqQueryParam, +) __all__ = ["Fts", "Query", "VectorQuery"] @@ -53,7 +59,7 @@ class Query: field_name (str): Name of the field to query. id (Optional[str], optional): Document ID to fetch vector from. Default is None. vector (VectorType, optional): Explicit query vector. Default is None. - param (Optional[Union[HnswQueryParam, HnswRabitqQueryParam, IVFQueryParam, FtsQueryParam]], optional): + param (Optional[Union[HnswQueryParam, HnswRabitqQueryParam, IVFQueryParam, IvfRabitqQueryParam, FtsQueryParam]], optional): Index-specific query parameters. Default is None. fts (Optional[Fts], optional): Full-text search parameters. Default is None. @@ -84,7 +90,13 @@ class Query: id: Optional[str] = None vector: VectorType = None param: Optional[ - Union[HnswQueryParam, HnswRabitqQueryParam, IVFQueryParam, FtsQueryParam] + Union[ + HnswQueryParam, + HnswRabitqQueryParam, + IVFQueryParam, + IvfRabitqQueryParam, + FtsQueryParam, + ] ] = None fts: Optional[Fts] = None diff --git a/python/zvec/model/schema/field_schema.py b/python/zvec/model/schema/field_schema.py index 9b757d6da..0005459ba 100644 --- a/python/zvec/model/schema/field_schema.py +++ b/python/zvec/model/schema/field_schema.py @@ -24,6 +24,7 @@ HnswRabitqIndexParam, InvertIndexParam, IVFIndexParam, + IvfRabitqIndexParam, ) from zvec.typing import DataType @@ -196,9 +197,9 @@ class VectorSchema: data_type (DataType): Vector data type (e.g., VECTOR_FP32, VECTOR_INT8). dimension (int, optional): Dimensionality of the vector. Must be > 0 for dense vectors; may be `None` for sparse vectors. - index_param (Union[HnswIndexParam, IVFIndexParam, FlatIndexParam], optional): + index_param (Union[HnswIndexParam, HnswRabitqIndexParam, IvfRabitqIndexParam, IVFIndexParam, FlatIndexParam], optional): Index configuration for this vector field. Defaults to - ``HnswIndexParam()``. + ``FlatIndexParam()``. Examples: >>> from zvec.typing import DataType @@ -217,7 +218,13 @@ def __init__( data_type: DataType, dimension: Optional[int] = 0, index_param: Optional[ - Union[HnswIndexParam, HnswRabitqIndexParam, FlatIndexParam, IVFIndexParam] + Union[ + HnswIndexParam, + HnswRabitqIndexParam, + IvfRabitqIndexParam, + FlatIndexParam, + IVFIndexParam, + ] ] = None, ): if name is None or not isinstance(name, str): @@ -273,8 +280,14 @@ def dimension(self) -> int: @property def index_param( self, - ) -> Union[HnswIndexParam, HnswRabitqIndexParam, IVFIndexParam, FlatIndexParam]: - """Union[HnswIndexParam, HnswRabitqIndexParam, IVFIndexParam, FlatIndexParam]: Index configuration for the vector.""" + ) -> Union[ + HnswIndexParam, + HnswRabitqIndexParam, + IvfRabitqIndexParam, + IVFIndexParam, + FlatIndexParam, + ]: + """Union[HnswIndexParam, HnswRabitqIndexParam, IvfRabitqIndexParam, IVFIndexParam, FlatIndexParam]: Index configuration for the vector.""" return self._cpp_obj.index_param def __dict__(self) -> dict[str, Any]: diff --git a/python/zvec/typing/__init__.pyi b/python/zvec/typing/__init__.pyi index db3a20104..1a82d2fe3 100644 --- a/python/zvec/typing/__init__.pyi +++ b/python/zvec/typing/__init__.pyi @@ -139,21 +139,30 @@ class IndexType: HNSW + HNSW_RABITQ + IVF + IVF_RABITQ + FLAT + VAMANA + INVERT """ - FLAT: typing.ClassVar[IndexType] # value = + FLAT: typing.ClassVar[IndexType] # value = HNSW: typing.ClassVar[IndexType] # value = + HNSW_RABITQ: typing.ClassVar[IndexType] # value = INVERT: typing.ClassVar[IndexType] # value = - IVF: typing.ClassVar[IndexType] # value = + IVF: typing.ClassVar[IndexType] # value = + IVF_RABITQ: typing.ClassVar[IndexType] # value = UNDEFINED: typing.ClassVar[IndexType] # value = + VAMANA: typing.ClassVar[IndexType] # value = __members__: typing.ClassVar[ dict[str, IndexType] - ] # value = {'UNDEFINED': , 'HNSW': , 'IVF': , 'FLAT': , 'INVERT': } + ] # value = {'UNDEFINED': , 'HNSW': , 'HNSW_RABITQ': , 'IVF': , 'IVF_RABITQ': , 'FLAT': , 'VAMANA': , 'INVERT': } def __eq__(self, other: typing.Any) -> bool: ... def __getstate__(self) -> int: ... @@ -236,15 +245,18 @@ class QuantizeType: INT8 INT4 + + RABITQ """ FP16: typing.ClassVar[QuantizeType] # value = INT4: typing.ClassVar[QuantizeType] # value = INT8: typing.ClassVar[QuantizeType] # value = + RABITQ: typing.ClassVar[QuantizeType] # value = UNDEFINED: typing.ClassVar[QuantizeType] # value = __members__: typing.ClassVar[ dict[str, QuantizeType] - ] # value = {'UNDEFINED': , 'FP16': , 'INT8': , 'INT4': } + ] # value = {'UNDEFINED': , 'FP16': , 'INT8': , 'INT4': , 'RABITQ': } def __eq__(self, other: typing.Any) -> bool: ... def __getstate__(self) -> int: ... diff --git a/src/binding/c/c_api.cc b/src/binding/c/c_api.cc index d1d34428f..5b073fd87 100644 --- a/src/binding/c/c_api.cc +++ b/src/binding/c/c_api.cc @@ -1343,6 +1343,11 @@ zvec_index_params_t *zvec_index_params_create(zvec_index_type_t index_type) { false, // use_soar (default) zvec::QuantizeType::UNDEFINED); break; + case ZVEC_INDEX_TYPE_IVF_RABITQ: + cpp_params = new zvec::IvfRabitqIndexParams( + zvec::MetricType::L2, zvec::core_interface::kDefaultIvfRabitqNlist, + zvec::core_interface::kDefaultRabitqTotalBits, 0); + break; case ZVEC_INDEX_TYPE_VAMANA: cpp_params = new zvec::VamanaIndexParams( @@ -1720,6 +1725,55 @@ zvec_error_code_t zvec_index_params_get_ivf_params(const zvec_index_params_t *pa return ZVEC_OK; } +zvec_error_code_t zvec_index_params_set_ivf_rabitq_params( + zvec_index_params_t *params, int nlist, int total_bits, int sample_count) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not IVF_RABITQ index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_params = reinterpret_cast(params); + auto *ivf_rabitq_params = + dynamic_cast(cpp_params); + if (!ivf_rabitq_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not IVF_RABITQ index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + ivf_rabitq_params->set_nlist(nlist); + ivf_rabitq_params->set_total_bits(total_bits); + ivf_rabitq_params->set_sample_count(sample_count); + return ZVEC_OK; +} + +zvec_error_code_t zvec_index_params_get_ivf_rabitq_params( + const zvec_index_params_t *params, int *out_nlist, int *out_total_bits, + int *out_sample_count) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not IVF_RABITQ index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *cpp_params = reinterpret_cast(params); + auto *ivf_rabitq_params = + dynamic_cast(cpp_params); + if (!ivf_rabitq_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Invalid params or not IVF_RABITQ index type"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + if (out_nlist) { + *out_nlist = ivf_rabitq_params->nlist(); + } + if (out_total_bits) { + *out_total_bits = ivf_rabitq_params->total_bits(); + } + if (out_sample_count) { + *out_sample_count = ivf_rabitq_params->sample_count(); + } + return ZVEC_OK; +} + /** * @brief Set Invert index parameters * @param params Index parameters (must be INVERT type) @@ -2693,6 +2747,10 @@ const char *zvec_index_type_to_string(zvec_index_type_t index_type) { return "IVF"; case ZVEC_INDEX_TYPE_FLAT: return "FLAT"; + case ZVEC_INDEX_TYPE_VAMANA: + return "VAMANA"; + case ZVEC_INDEX_TYPE_IVF_RABITQ: + return "IVF_RABITQ"; case ZVEC_INDEX_TYPE_INVERT: return "INVERT"; case ZVEC_INDEX_TYPE_FTS: @@ -4749,6 +4807,7 @@ void zvec_collection_stats_destroy(zvec_collection_stats_t *stats) { // Users should create type-specific query params: // - HnswQueryParams via zvec_query_params_hnsw_create() // - IVFQueryParams via zvec_query_params_ivf_create() +// - IvfRabitqQueryParams via zvec_query_params_ivf_rabitq_create() // - FlatQueryParams via zvec_query_params_flat_create() // // Each type-specific instance has its own destroy function. @@ -4959,6 +5018,111 @@ bool zvec_query_params_ivf_get_is_using_refiner( return ptr->is_using_refiner(); } +// ============================================================================= +// IvfRabitqQueryParams implementation - wrapper around zvec::IvfRabitqQueryParams +// ============================================================================= + +zvec_ivf_rabitq_query_params_t *zvec_query_params_ivf_rabitq_create( + int nprobe, float radius, bool is_linear, bool is_using_refiner) { + ZVEC_TRY_RETURN_NULL( + "Failed to create IvfRabitqQueryParams", + auto *params = new zvec::IvfRabitqQueryParams( + nprobe, radius, is_linear, is_using_refiner); + return reinterpret_cast(params);) + return nullptr; +} + +void zvec_query_params_ivf_rabitq_destroy( + zvec_ivf_rabitq_query_params_t *params) { + if (params) { + delete reinterpret_cast(params); + } +} + +zvec_error_code_t zvec_query_params_ivf_rabitq_set_nprobe( + zvec_ivf_rabitq_query_params_t *params, int nprobe) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "IVF_RABITQ query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_nprobe(nprobe); + return ZVEC_OK; +} + +int zvec_query_params_ivf_rabitq_get_nprobe( + const zvec_ivf_rabitq_query_params_t *params) { + if (!params) { + return zvec::core_interface::kDefaultIvfRabitqNprobe; + } + auto *ptr = reinterpret_cast(params); + return ptr->nprobe(); +} + +zvec_error_code_t zvec_query_params_ivf_rabitq_set_radius( + zvec_ivf_rabitq_query_params_t *params, float radius) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "IVF_RABITQ query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_radius(radius); + return ZVEC_OK; +} + +float zvec_query_params_ivf_rabitq_get_radius( + const zvec_ivf_rabitq_query_params_t *params) { + if (!params) { + return 0.0f; + } + auto *ptr = reinterpret_cast(params); + return ptr->radius(); +} + +zvec_error_code_t zvec_query_params_ivf_rabitq_set_is_linear( + zvec_ivf_rabitq_query_params_t *params, bool is_linear) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "IVF_RABITQ query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_is_linear(is_linear); + return ZVEC_OK; +} + +bool zvec_query_params_ivf_rabitq_get_is_linear( + const zvec_ivf_rabitq_query_params_t *params) { + if (!params) { + return false; + } + auto *ptr = reinterpret_cast(params); + return ptr->is_linear(); +} + +zvec_error_code_t zvec_query_params_ivf_rabitq_set_is_using_refiner( + zvec_ivf_rabitq_query_params_t *params, bool is_using_refiner) { + if (!params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "IVF_RABITQ query params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(params); + ptr->set_is_using_refiner(is_using_refiner); + return ZVEC_OK; +} + +bool zvec_query_params_ivf_rabitq_get_is_using_refiner( + const zvec_ivf_rabitq_query_params_t *params) { + if (!params) { + return false; + } + auto *ptr = reinterpret_cast(params); + return ptr->is_using_refiner(); +} + // ============================================================================= // FlatQueryParams implementation - wrapper around zvec::FlatQueryParams // ============================================================================= @@ -5416,6 +5580,24 @@ zvec_error_code_t zvec_vector_query_set_ivf_params(zvec_vector_query_t *query, return ZVEC_OK; } +zvec_error_code_t zvec_vector_query_set_ivf_rabitq_params( + zvec_vector_query_t *query, + zvec_ivf_rabitq_query_params_t *ivf_rabitq_params) { + if (!query || !ivf_rabitq_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query or IVF_RABITQ params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *query_ptr = reinterpret_cast(query); + auto *params_ptr = + reinterpret_cast(ivf_rabitq_params); + + query_ptr->target_.query_params_.reset(params_ptr); + + return ZVEC_OK; +} + zvec_error_code_t zvec_vector_query_set_flat_params( zvec_vector_query_t *query, zvec_flat_query_params_t *flat_params) { if (!query || !flat_params) { @@ -5786,6 +5968,24 @@ zvec_error_code_t zvec_group_by_vector_query_set_ivf_params( return ZVEC_OK; } +zvec_error_code_t zvec_group_by_vector_query_set_ivf_rabitq_params( + zvec_group_by_vector_query_t *query, + zvec_ivf_rabitq_query_params_t *ivf_rabitq_params) { + if (!query || !ivf_rabitq_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Query or IVF_RABITQ params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + + auto *query_ptr = reinterpret_cast(query); + auto *params_ptr = + reinterpret_cast(ivf_rabitq_params); + + query_ptr->target_.query_params_.reset(params_ptr); + + return ZVEC_OK; +} + zvec_error_code_t zvec_group_by_vector_query_set_flat_params( zvec_group_by_vector_query_t *query, zvec_flat_query_params_t *flat_params) { if (!query || !flat_params) { @@ -6149,6 +6349,21 @@ zvec_error_code_t zvec_sub_query_set_ivf_params( return ZVEC_OK; } +zvec_error_code_t zvec_sub_query_set_ivf_rabitq_params( + zvec_sub_query_t *query, + zvec_ivf_rabitq_query_params_t *ivf_rabitq_params) { + if (!query || !ivf_rabitq_params) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, + "Sub-vector query or IVF_RABITQ params pointer is null"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + auto *ptr = reinterpret_cast(query); + auto *params_ptr = + reinterpret_cast(ivf_rabitq_params); + ptr->target_.query_params_.reset(params_ptr); + return ZVEC_OK; +} + zvec_error_code_t zvec_sub_query_set_flat_params( zvec_sub_query_t *query, zvec_flat_query_params_t *flat_params) { if (!query || !flat_params) { diff --git a/src/binding/python/CMakeLists.txt b/src/binding/python/CMakeLists.txt index 0db6d75ff..9742ab9b7 100644 --- a/src/binding/python/CMakeLists.txt +++ b/src/binding/python/CMakeLists.txt @@ -53,6 +53,7 @@ if (CMAKE_SYSTEM_NAME STREQUAL "Linux") $ $ $ + $ $ $ $ @@ -90,6 +91,7 @@ elseif (APPLE) -Wl,-force_load,$ -Wl,-force_load,$ -Wl,-force_load,$ + -Wl,-force_load,$ -Wl,-force_load,$ -Wl,-force_load,$ -Wl,-force_load,$ diff --git a/src/binding/python/model/param/python_param.cc b/src/binding/python/model/param/python_param.cc index 43870611c..38ef924df 100644 --- a/src/binding/python/model/param/python_param.cc +++ b/src/binding/python/model/param/python_param.cc @@ -34,6 +34,8 @@ static std::string index_type_to_string(const IndexType type) { return "HNSW"; case IndexType::HNSW_RABITQ: return "HNSW_RABITQ"; + case IndexType::IVF_RABITQ: + return "IVF_RABITQ"; case IndexType::DISKANN: return "DISKANN"; case IndexType::VAMANA: @@ -666,6 +668,88 @@ quantization method that provides high compression with minimal accuracy loss. t[3].cast(), t[4].cast(), t[5].cast()); })); + // binding ivf rabitq index params + py::class_> + ivf_rabitq_params(m, "IvfRabitqIndexParam", R"pbdoc( +Parameters for configuring an IVF index with RaBitQ quantization. + +IVF partitions the vector space into inverted lists. RaBitQ compresses vectors +inside each list for faster scanning with lower memory usage. + +Attributes: + metric_type (MetricType): Distance metric used for similarity computation. + Default is ``MetricType.IP`` (inner product). + nlist (int): Number of IVF cluster centers. Default is 1024. + total_bits (int): Total bits for RaBitQ quantization. Default is 7. + sample_count (int): Sample count for training. 0 means use all vectors. + +Examples: + >>> from zvec.typing import MetricType + >>> params = IvfRabitqIndexParam( + ... metric_type=MetricType.COSINE, + ... nlist=1024, + ... total_bits=7, + ... sample_count=10000 + ... ) + >>> print(params.nlist) + 1024 +)pbdoc"); + ivf_rabitq_params + .def(py::init(), + py::arg("metric_type") = MetricType::IP, + py::arg("nlist") = core_interface::kDefaultIvfRabitqNlist, + py::arg("total_bits") = core_interface::kDefaultRabitqTotalBits, + py::arg("sample_count") = 0) + .def_property_readonly("nlist", &IvfRabitqIndexParams::nlist, + "int: Number of IVF cluster centers.") + .def_property_readonly("total_bits", &IvfRabitqIndexParams::total_bits, + "int: Total bits for RaBitQ quantization.") + .def_property_readonly("sample_count", + &IvfRabitqIndexParams::sample_count, + "int: Sample count for RaBitQ training.") + .def( + "to_dict", + [](const IvfRabitqIndexParams &self) -> py::dict { + py::dict dict; + dict["type"] = index_type_to_string(self.type()); + dict["metric_type"] = metric_type_to_string(self.metric_type()); + dict["quantize_type"] = + quantize_type_to_string(self.quantize_type()); + dict["nlist"] = self.nlist(); + dict["total_bits"] = self.total_bits(); + dict["sample_count"] = self.sample_count(); + return dict; + }, + "Convert to dictionary with all fields") + .def( + "__repr__", + [](const IvfRabitqIndexParams &self) -> std::string { + return "{" + "\"type\":\"" + + index_type_to_string(self.type()) + + "\", \"metric_type\":\"" + + metric_type_to_string(self.metric_type()) + + "\", \"nlist\":" + std::to_string(self.nlist()) + + ", \"total_bits\":" + std::to_string(self.total_bits()) + + ", \"sample_count\":" + std::to_string(self.sample_count()) + + ", \"quantize_type\":\"" + + quantize_type_to_string(self.quantize_type()) + "\"}"; + }) + .def(py::pickle( + [](const IvfRabitqIndexParams &self) { + return py::make_tuple(self.metric_type(), self.nlist(), + self.total_bits(), self.sample_count()); + }, + [](py::tuple t) { + if (t.size() != 4) + throw std::runtime_error( + "Invalid state for IvfRabitqIndexParams"); + return std::make_shared( + t[0].cast(), t[1].cast(), t[2].cast(), + t[3].cast()); + })); + // binding vamana index params py::class_> @@ -1425,6 +1509,74 @@ Constructs an HnswRabitqQueryParam instance. return obj; })); + // binding ivf rabitq query params + py::class_> + ivf_rabitq_query_params(m, "IvfRabitqQueryParam", R"pbdoc( +Query parameters for IVF RaBitQ index. + +Controls how many IVF clusters (`nprobe`) to visit during search. + +Attributes: + type (IndexType): Always ``IndexType.IVF_RABITQ``. + nprobe (int): Number of closest clusters to search. + Higher values improve recall but increase latency. + Default is 10. + radius (float): Search radius for range queries. Default is 0.0. + is_linear (bool): Force linear search. Default is False. + is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False. + +Examples: + >>> params = IvfRabitqQueryParam(nprobe=20) + >>> print(params.nprobe) + 20 +)pbdoc"); + ivf_rabitq_query_params + .def(py::init(), + py::arg("nprobe") = core_interface::kDefaultIvfRabitqNprobe, + py::arg("radius") = 0.0f, py::arg("is_linear") = false, + py::arg("is_using_refiner") = false, + R"pbdoc( +Constructs an IvfRabitqQueryParam instance. + +Args: + nprobe (int, optional): Number of inverted lists to probe during search. + Higher values improve accuracy. Defaults to 10. + radius (float, optional): Search radius for range queries. Default is 0.0. + is_linear (bool, optional): Force linear search. Default is False. + is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False. +)pbdoc") + .def_property_readonly( + "nprobe", + [](const IvfRabitqQueryParams &self) -> int { return self.nprobe(); }, + "int: Number of inverted lists to search during IVF RaBitQ query.") + .def("__repr__", + [](const IvfRabitqQueryParams &self) -> std::string { + return "{" + "\"type\":\"" + + index_type_to_string(self.type()) + + "\", \"nprobe\":" + std::to_string(self.nprobe()) + + ", \"radius\":" + std::to_string(self.radius()) + + ", \"is_linear\":" + std::to_string(self.is_linear()) + + ", \"is_using_refiner\":" + + std::to_string(self.is_using_refiner()) + "}"; + }) + .def(py::pickle( + [](const IvfRabitqQueryParams &self) { + return py::make_tuple(self.nprobe(), self.radius(), + self.is_linear(), self.is_using_refiner()); + }, + [](py::tuple t) { + if (t.size() != 4) + throw std::runtime_error( + "Invalid state for IvfRabitqQueryParams"); + auto obj = std::make_shared(t[0].cast()); + obj->set_radius(t[1].cast()); + obj->set_is_linear(t[2].cast()); + obj->set_is_using_refiner(t[3].cast()); + return obj; + })); + // binding diskann query params py::class_> diff --git a/src/binding/python/typing/python_type.cc b/src/binding/python/typing/python_type.cc index e6b9bd7cd..b6480b03c 100644 --- a/src/binding/python/typing/python_type.cc +++ b/src/binding/python/typing/python_type.cc @@ -98,6 +98,7 @@ Enumeration of supported index types in Zvec. .value("HNSW", IndexType::HNSW) .value("HNSW_RABITQ", IndexType::HNSW_RABITQ) .value("IVF", IndexType::IVF) + .value("IVF_RABITQ", IndexType::IVF_RABITQ) .value("FLAT", IndexType::FLAT) .value("VAMANA", IndexType::VAMANA) .value("INVERT", IndexType::INVERT); @@ -226,4 +227,4 @@ Construct a status with the given code and optional message. }); } -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/CMakeLists.txt b/src/core/CMakeLists.txt index c52c3b8cc..333c81dd0 100644 --- a/src/core/CMakeLists.txt +++ b/src/core/CMakeLists.txt @@ -19,6 +19,24 @@ if(RABITQ_SUPPORTED AND AUTO_DETECT_ARCH) COMPILE_FLAGS "${RABITQ_ARCH_FLAG}" ) endforeach() + + set(IVF_RABITQ_FILES + ivf_rabitq_builder.cc + ivf_rabitq_entity.cc + ivf_rabitq_reformer.cc + ivf_rabitq_register.cc + ivf_rabitq_streamer.cc + ivf_rabitq_util.cc + ) + set(IVF_RABITQ_FILES_FULL ${IVF_RABITQ_FILES}) + list(TRANSFORM IVF_RABITQ_FILES_FULL PREPEND "algorithm/ivf_rabitq/") + foreach(FILE ${IVF_RABITQ_FILES_FULL}) + set_source_files_properties( + ${FILE} + PROPERTIES + COMPILE_FLAGS "${RABITQ_ARCH_FLAG}" + ) + endforeach() endif() # utility/block_heap.cc uses AVX2 intrinsics guarded by __AVX2__. When the @@ -59,6 +77,7 @@ file(GLOB_RECURSE ALL_CORE_SRCS *.cc *.c *.h) # for HNSWRabitqIndex and guards rabitqlib usage with #if RABITQ_SUPPORTED. if(NOT RABITQ_SUPPORTED) list(FILTER ALL_CORE_SRCS EXCLUDE REGEX ".*/algorithm/hnsw_rabitq/.*") + list(FILTER ALL_CORE_SRCS EXCLUDE REGEX ".*/algorithm/ivf_rabitq/.*") endif() # Always exclude algorithm/diskann implementation files from zvec_core. diff --git a/src/core/algorithm/CMakeLists.txt b/src/core/algorithm/CMakeLists.txt index bed772984..b83ed8cc5 100644 --- a/src/core/algorithm/CMakeLists.txt +++ b/src/core/algorithm/CMakeLists.txt @@ -47,6 +47,7 @@ endif() if(RABITQ_SUPPORTED) message(STATUS "BUILD RABITQ") cc_directory(hnsw_rabitq) + cc_directory(ivf_rabitq) else() message(STATUS "NOT BUILD RABITQ") # Empty stub library for unsupported platforms @@ -64,4 +65,19 @@ else() INCS . ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm VERSION "${PROXIMA_ZVEC_VERSION}" ) + + file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/ivf_rabitq_stub.cc + "// Stub implementation for unsupported platforms\n" + "// IVF RaBitQ only supports Linux x86_64\n" + "namespace zvec { namespace core { /* empty namespace for compatibility */ } }\n" + ) + + cc_library( + NAME core_knn_ivf_rabitq + STATIC SHARED STRICT ALWAYS_LINK + SRCS ${CMAKE_CURRENT_BINARY_DIR}/ivf_rabitq_stub.cc + LIBS core_framework + INCS . ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + VERSION "${PROXIMA_ZVEC_VERSION}" + ) endif() diff --git a/src/core/algorithm/hnsw_rabitq/rabitq_params.h b/src/core/algorithm/hnsw_rabitq/rabitq_params.h index c30669f48..8e8e4ff30 100644 --- a/src/core/algorithm/hnsw_rabitq/rabitq_params.h +++ b/src/core/algorithm/hnsw_rabitq/rabitq_params.h @@ -44,6 +44,10 @@ constexpr size_t kDefaultRabitqTotalBits = 7; constexpr int kMinRabitqDimSize = 64; constexpr int kMaxRabitqDimSize = 4095; +// Original vector dimension before converter (e.g., CosineNormalizeConverter) +// modifies it. Used by IVF-RaBitQ to ignore extra dimensions from converter. +static const std::string PARAM_RABITQ_GENERAL_DIMENSION( + "proxima.rabitq.general.dimension"); } // namespace core } // namespace zvec diff --git a/src/core/algorithm/ivf_rabitq/CMakeLists.txt b/src/core/algorithm/ivf_rabitq/CMakeLists.txt new file mode 100644 index 000000000..991d60142 --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/CMakeLists.txt @@ -0,0 +1,27 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +if(AUTO_DETECT_ARCH) + foreach(FILE ${IVF_RABITQ_FILES}) + set_source_files_properties( + ${FILE} + PROPERTIES + COMPILE_FLAGS "${RABITQ_ARCH_FLAG}" + ) + endforeach() +endif() + +if(NOT APPLE) + set(CORE_KNN_IVF_RABITQ_LDFLAGS + "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") +endif() + +cc_library( + NAME core_knn_ivf_rabitq + STATIC SHARED STRICT ALWAYS_LINK + SRCS *.cc + LIBS core_framework rabitqlib core_knn_cluster core_knn_hnsw_rabitq + INCS . ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm + LDFLAGS "${CORE_KNN_IVF_RABITQ_LDFLAGS}" + VERSION "${PROXIMA_ZVEC_VERSION}" + ) diff --git a/src/core/algorithm/ivf_rabitq/ivf_rabitq_builder.cc b/src/core/algorithm/ivf_rabitq/ivf_rabitq_builder.cc new file mode 100644 index 000000000..91bde7136 --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_builder.cc @@ -0,0 +1,537 @@ +// 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. +#include "ivf_rabitq_builder.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "algorithm/hnsw_rabitq/rabitq_converter.h" +#include "algorithm/hnsw_rabitq/rabitq_params.h" +#include "zvec/core/framework/index_error.h" +#include "zvec/core/framework/index_factory.h" +#include "zvec/core/framework/index_memory.h" +#include "ivf_rabitq_params.h" +#include "ivf_rabitq_reformer.h" +#include "ivf_rabitq_util.h" + +namespace zvec { +namespace core { + +namespace { + +struct IvfRabitqVectorEntry { + uint64_t key{0}; + std::vector owned_data; +}; + +struct IvfRabitqBuildData { + IvfRabitqHeader header; + std::vector batch_data_buf; + std::vector ex_data_buf; + std::vector cluster_metas; + std::vector keys_buf; + std::vector mapping_buf; +}; + +int BuildIvfRabitqData(const std::shared_ptr &reformer, + const IndexHolder::Pointer &holder, + std::vector extra_vectors, + IvfRabitqBuildData *build_data) { + if (!reformer || !reformer->loaded()) { + LOG_ERROR("Reformer not loaded"); + return IndexError_NoReady; + } + if (!build_data) { + LOG_ERROR("Null IVF RaBitQ build data"); + return IndexError_InvalidArgument; + } + + *build_data = IvfRabitqBuildData(); + + size_t dimension = reformer->dimension(); + size_t padded_dim = reformer->padded_dim(); + size_t ex_bits = reformer->ex_bits(); + size_t cluster_count = reformer->num_clusters(); + + std::vector all_vectors; + if (holder) { + if (holder->dimension() < dimension) { + LOG_ERROR("Holder dimension=%zu smaller than RaBitQ dimension=%zu", + holder->dimension(), dimension); + return IndexError_Mismatch; + } + auto iter = holder->create_iterator(); + while (iter && iter->is_valid()) { + IvfRabitqVectorEntry entry; + entry.key = iter->key(); + entry.owned_data.assign( + static_cast(iter->data()), + static_cast(iter->data()) + dimension); + all_vectors.push_back(std::move(entry)); + iter->next(); + } + } + + all_vectors.reserve(all_vectors.size() + extra_vectors.size()); + for (auto &entry : extra_vectors) { + all_vectors.push_back(std::move(entry)); + } + + size_t total_count = all_vectors.size(); + if (total_count == 0) { + LOG_WARN("No vectors to build"); + return 0; + } + + LOG_INFO("Building IVF RaBitQ index: %zu vectors, %zu clusters", total_count, + cluster_count); + + std::vector> cluster_assignments(cluster_count); + for (size_t i = 0; i < total_count; ++i) { + uint32_t cid = + reformer->find_nearest_centroid(all_vectors[i].owned_data.data()); + cluster_assignments[cid].push_back(static_cast(i)); + } + + size_t batch_data_per_batch = + rabitqlib::BatchDataMap::data_bytes(padded_dim); + size_t ex_data_per_vector = + rabitqlib::ExDataMap::data_bytes(padded_dim, ex_bits); + + size_t total_batch_data_size = 0; + size_t total_ex_data_size = 0; + std::vector cluster_batch_counts(cluster_count); + + for (size_t cid = 0; cid < cluster_count; ++cid) { + uint32_t vec_count = static_cast(cluster_assignments[cid].size()); + uint32_t num_batches = (vec_count + rabitqlib::fastscan::kBatchSize - 1) / + rabitqlib::fastscan::kBatchSize; + cluster_batch_counts[cid] = num_batches; + total_batch_data_size += num_batches * batch_data_per_batch; + total_ex_data_size += vec_count * ex_data_per_vector; + } + + build_data->batch_data_buf.assign(total_batch_data_size, 0); + build_data->ex_data_buf.assign(total_ex_data_size, 0); + build_data->keys_buf.resize(total_count); + build_data->mapping_buf.resize(total_count); + build_data->cluster_metas.resize(cluster_count); + + size_t batch_data_offset = 0; + size_t ex_data_offset = 0; + uint32_t key_offset = 0; + std::vector rotated_batch(rabitqlib::fastscan::kBatchSize * + padded_dim); + + for (size_t cid = 0; cid < cluster_count; ++cid) { + auto &assignments = cluster_assignments[cid]; + uint32_t vec_count = static_cast(assignments.size()); + + build_data->cluster_metas[cid].batch_data_offset = batch_data_offset; + build_data->cluster_metas[cid].ex_data_offset = ex_data_offset; + build_data->cluster_metas[cid].vector_count = vec_count; + build_data->cluster_metas[cid].batch_count = cluster_batch_counts[cid]; + build_data->cluster_metas[cid].key_offset = key_offset; + + for (uint32_t i = 0; i < vec_count; ++i) { + auto &entry = all_vectors[assignments[i]]; + build_data->keys_buf[key_offset + i] = entry.key; + } + + uint32_t processed = 0; + while (processed < vec_count) { + uint32_t batch_size = + std::min(static_cast(rabitqlib::fastscan::kBatchSize), + vec_count - processed); + + std::memset(rotated_batch.data(), 0, + rabitqlib::fastscan::kBatchSize * padded_dim * sizeof(float)); + for (uint32_t i = 0; i < batch_size; ++i) { + const float *src = + all_vectors[assignments[processed + i]].owned_data.data(); + float *dst = rotated_batch.data() + (i * padded_dim); + reformer->rotate_vector(src, dst); + } + + char *bd_ptr = build_data->batch_data_buf.data() + batch_data_offset; + char *ex_ptr = nullptr; + if (ex_data_per_vector > 0) { + ex_ptr = build_data->ex_data_buf.data() + ex_data_offset; + } + + int ret = reformer->quantize_batch(rotated_batch.data(), + static_cast(cid), batch_size, + bd_ptr, ex_ptr); + if (ret != 0) { + LOG_ERROR("Failed to quantize batch for cluster %zu, ret=%d", cid, ret); + return ret; + } + + batch_data_offset += batch_data_per_batch; + ex_data_offset += batch_size * ex_data_per_vector; + processed += batch_size; + } + + key_offset += vec_count; + } + + std::iota(build_data->mapping_buf.begin(), build_data->mapping_buf.end(), 0U); + std::sort(build_data->mapping_buf.begin(), build_data->mapping_buf.end(), + [&keys_buf = build_data->keys_buf](uint32_t lhs, uint32_t rhs) { + if (keys_buf[lhs] == keys_buf[rhs]) { + return lhs < rhs; + } + return keys_buf[lhs] < keys_buf[rhs]; + }); + + build_data->header.total_vector_count = static_cast(total_count); + build_data->header.cluster_count = static_cast(cluster_count); + build_data->header.dimension = static_cast(dimension); + build_data->header.padded_dim = static_cast(padded_dim); + build_data->header.ex_bits = static_cast(ex_bits); + build_data->header.rotator_type = 0; + build_data->header.metric_type = + (reformer->rabitq_metric_type() == RabitqMetricType::kIP) ? 1 : 0; + build_data->header.batch_data_size = total_batch_data_size; + build_data->header.ex_data_size = total_ex_data_size; + + return 0; +} + +} // namespace + +IvfRabitqBuilder::IvfRabitqBuilder() = default; + +IvfRabitqBuilder::~IvfRabitqBuilder() { + this->cleanup(); +} + +// -------------------------------------------------------------------------- +// init +// -------------------------------------------------------------------------- +int IvfRabitqBuilder::init(const IndexMeta &meta, + const ailego::Params ¶ms) { + if (state_ != INIT) { + LOG_ERROR("IvfRabitqBuilder state wrong. state=%d", state_); + return IndexError_Logic; + } + + meta_ = meta; + params_ = params; + + params.get(PARAM_IVF_RABITQ_NLIST, &nlist_); + if (nlist_ == 0) { + nlist_ = kDefaultIvfRabitqNlist; + } + + params.get(PARAM_RABITQ_TOTAL_BITS, &total_bits_); + if (total_bits_ == 0) { + total_bits_ = kDefaultRabitqTotalBits; + } + if (total_bits_ < 1 || total_bits_ > 9) { + LOG_ERROR("Invalid total_bits=%u, must be in [1, 9]", total_bits_); + return IndexError_InvalidArgument; + } + + params.get(PARAM_RABITQ_SAMPLE_COUNT, &sample_count_); + + int ret = PrepareAndCheckIvfRabitqInternalMeta(meta_, params, &rabitq_meta_, + &metric_name_); + if (ret != 0) { + return ret; + } + + uint32_t dim = rabitq_meta_.dimension(); + LOG_INFO( + "IvfRabitqBuilder initialized: dim=%u, nlist=%u, total_bits=%u, " + "metric=%s, sample_count=%u", + dim, nlist_, total_bits_, metric_name_.c_str(), sample_count_); + + state_ = INITED; + return 0; +} + +// -------------------------------------------------------------------------- +// cleanup +// -------------------------------------------------------------------------- +int IvfRabitqBuilder::cleanup() { + state_ = INIT; + stats_.clear(); + converter_.reset(); + reformer_.reset(); + header_ = IvfRabitqHeader(); + batch_data_buf_.clear(); + ex_data_buf_.clear(); + cluster_metas_.clear(); + keys_buf_.clear(); + mapping_buf_.clear(); + return 0; +} + +// -------------------------------------------------------------------------- +// train — KMeans clustering + rotator training +// -------------------------------------------------------------------------- +int IvfRabitqBuilder::train(IndexThreads::Pointer /*threads*/, + IndexHolder::Pointer holder) { + if (state_ != INITED) { + LOG_ERROR("IvfRabitqBuilder train failed, wrong state=%d", state_); + return IndexError_Runtime; + } + + ailego::ElapsedTime timer; + + // Create and configure RabitqConverter + converter_ = std::make_shared(); + ailego::Params conv_params; + conv_params.set(PARAM_RABITQ_NUM_CLUSTERS, nlist_); + conv_params.set(PARAM_RABITQ_TOTAL_BITS, total_bits_); + if (sample_count_ > 0) { + conv_params.set(PARAM_RABITQ_SAMPLE_COUNT, sample_count_); + } + + int ret = converter_->init(rabitq_meta_, conv_params); + if (ret != 0) { + LOG_ERROR("Failed to init RabitqConverter, ret=%d", ret); + return ret; + } + + // Train KMeans centroids + ret = converter_->train(holder); + if (ret != 0) { + LOG_ERROR("Failed to train RabitqConverter, ret=%d", ret); + return ret; + } + + // Create IvfRabitqReformer via dump+load cycle + auto memory_dumper = IndexFactory::CreateDumper("MemoryDumper"); + memory_dumper->init(ailego::Params()); + std::string file_id = ailego::StringHelper::Concat( + "ivf_rabitq_builder_train_", ailego::Monotime::MilliSeconds(), rand()); + ret = memory_dumper->create(file_id); + if (ret != 0) { + LOG_ERROR("Failed to create memory dumper, ret=%d", ret); + return ret; + } + + ret = converter_->dump(memory_dumper); + if (ret != 0) { + LOG_ERROR("Failed to dump converter, ret=%d", ret); + IndexMemory::Instance()->remove(file_id); + return ret; + } + ret = memory_dumper->close(); + if (ret != 0) { + LOG_ERROR("Failed to close memory dumper, ret=%d", ret); + IndexMemory::Instance()->remove(file_id); + return ret; + } + + reformer_ = std::make_shared(); + ret = reformer_->init(metric_name_); + if (ret != 0) { + LOG_ERROR("Failed to init IvfRabitqReformer, ret=%d", ret); + IndexMemory::Instance()->remove(file_id); + return ret; + } + + auto memory_storage = IndexFactory::CreateStorage("MemoryReadStorage"); + ret = memory_storage->open(file_id, false); + if (ret != 0) { + LOG_ERROR("Failed to open memory storage, ret=%d", ret); + IndexMemory::Instance()->remove(file_id); + return ret; + } + + ret = reformer_->load(memory_storage); + if (ret != 0) { + LOG_ERROR("Failed to load IvfRabitqReformer, ret=%d", ret); + IndexMemory::Instance()->remove(file_id); + return ret; + } + + IndexMemory::Instance()->remove(file_id); + + stats_.set_trained_count(converter_->stats().trained_count()); + stats_.set_trained_costtime(timer.milli_seconds()); + + LOG_INFO("IvfRabitqBuilder training completed: %u clusters, cost %zu ms", + static_cast(reformer_->num_clusters()), + static_cast(timer.milli_seconds())); + + state_ = TRAINED; + return 0; +} + +// -------------------------------------------------------------------------- +// build — assign vectors to centroids, quantize, store results in memory +// -------------------------------------------------------------------------- +int IvfRabitqBuilder::build(IndexThreads::Pointer /*threads*/, + IndexHolder::Pointer holder) { + if (state_ != TRAINED) { + LOG_ERROR("Train the index first before build"); + return IndexError_Runtime; + } + + if (!reformer_ || !reformer_->loaded()) { + LOG_ERROR("Reformer not loaded"); + return IndexError_NoReady; + } + + ailego::ElapsedTime timer; + + IvfRabitqBuildData build_data; + int ret = BuildIvfRabitqData(reformer_, holder, {}, &build_data); + if (ret != 0) { + return ret; + } + + header_ = build_data.header; + batch_data_buf_ = std::move(build_data.batch_data_buf); + ex_data_buf_ = std::move(build_data.ex_data_buf); + cluster_metas_ = std::move(build_data.cluster_metas); + keys_buf_ = std::move(build_data.keys_buf); + mapping_buf_ = std::move(build_data.mapping_buf); + + stats_.set_built_count(header_.total_vector_count); + stats_.set_built_costtime(timer.milli_seconds()); + + LOG_INFO( + "IvfRabitqBuilder build completed: %zu vectors, %zu clusters, " + "batch_data=%zu bytes, ex_data=%zu bytes, cost %zu ms", + static_cast(header_.total_vector_count), + static_cast(header_.cluster_count), + static_cast(header_.batch_data_size), + static_cast(header_.ex_data_size), + static_cast(timer.milli_seconds())); + + state_ = BUILT; + return 0; +} + +// -------------------------------------------------------------------------- +// dump — write meta + reformer + segments to dumper +// -------------------------------------------------------------------------- +int IvfRabitqBuilder::dump(const IndexDumper::Pointer &dumper) { + if (state_ != BUILT) { + LOG_ERROR("Build the index before dump"); + return IndexError_Runtime; + } + + if (!dumper) { + LOG_ERROR("Null dumper"); + return IndexError_InvalidArgument; + } + + ailego::ElapsedTime timer; + + // Write meta + int ret = IndexHelper::SerializeToDumper(meta_, dumper.get()); + if (ret != 0) { + LOG_ERROR("Failed to serialize meta into dumper"); + return ret; + } + + // Write reformer (centroids, rotator) + if (reformer_ && reformer_->loaded()) { + ret = reformer_->dump(dumper); + if (ret != 0) { + LOG_ERROR("Failed to dump reformer"); + return ret; + } + } + + // Write segments + auto dump_segment = [&](const std::string &seg_id, const void *data, + size_t size) -> int { + if (size == 0) { + return 0; + } + uint32_t crc = ailego::Crc32c::Hash(data, size, 0); + size_t align_size = (size + 0x1F) & (~0x1F); + size_t padding_size = align_size - size; + + size_t written = dumper->write(data, size); + if (written != size) { + LOG_ERROR("Failed to write segment %s to dumper", seg_id.c_str()); + return IndexError_WriteData; + } + if (padding_size > 0) { + std::string padding(padding_size, '\0'); + dumper->write(padding.data(), padding_size); + } + return dumper->append(seg_id, size, padding_size, crc); + }; + + ret = dump_segment(IVF_RABITQ_HEADER_SEG_ID, &header_, sizeof(header_)); + if (ret != 0) { + return ret; + } + + if (!batch_data_buf_.empty()) { + ret = dump_segment(IVF_RABITQ_BATCH_DATA_SEG_ID, batch_data_buf_.data(), + batch_data_buf_.size()); + if (ret != 0) { + return ret; + } + } + + if (!ex_data_buf_.empty()) { + ret = dump_segment(IVF_RABITQ_EX_DATA_SEG_ID, ex_data_buf_.data(), + ex_data_buf_.size()); + if (ret != 0) { + return ret; + } + } + + if (!cluster_metas_.empty()) { + ret = dump_segment(IVF_RABITQ_CLUSTER_META_SEG_ID, cluster_metas_.data(), + cluster_metas_.size() * sizeof(IvfRabitqClusterMeta)); + if (ret != 0) { + return ret; + } + } + + if (!keys_buf_.empty()) { + ret = dump_segment(IVF_RABITQ_KEYS_SEG_ID, keys_buf_.data(), + keys_buf_.size() * sizeof(uint64_t)); + if (ret != 0) { + return ret; + } + } + + if (!mapping_buf_.empty()) { + ret = dump_segment(IVF_RABITQ_MAPPING_SEG_ID, mapping_buf_.data(), + mapping_buf_.size() * sizeof(uint32_t)); + if (ret != 0) { + return ret; + } + } + + stats_.set_dumped_count(header_.total_vector_count); + stats_.set_dumped_costtime(timer.milli_seconds()); + + LOG_INFO("IvfRabitqBuilder dump completed, cost %zu ms", + static_cast(timer.milli_seconds())); + + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf_rabitq/ivf_rabitq_builder.h b/src/core/algorithm/ivf_rabitq/ivf_rabitq_builder.h new file mode 100644 index 000000000..2b6040123 --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_builder.h @@ -0,0 +1,91 @@ +// 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. +#pragma once + +#include +#include +#include +#include +#include +#include "ivf_rabitq_index_format.h" + +namespace zvec { +namespace core { + +class IvfRabitqReformer; +class RabitqConverter; + +/*! IVF RaBitQ Builder + * Implements the standard IndexBuilder interface (init → train → build → dump) + * for IVF RaBitQ indexes. + */ +class IvfRabitqBuilder : public IndexBuilder { + public: + IvfRabitqBuilder(); + ~IvfRabitqBuilder() override; + + IvfRabitqBuilder(const IvfRabitqBuilder &) = delete; + IvfRabitqBuilder &operator=(const IvfRabitqBuilder &) = delete; + + //! Initialize the builder + int init(const IndexMeta &meta, const ailego::Params ¶ms) override; + + //! Cleanup the builder + int cleanup() override; + + //! Train the data (KMeans clustering + rotator) + int train(IndexThreads::Pointer threads, + IndexHolder::Pointer holder) override; + + //! Build the index (assign vectors to centroids, quantize) + int build(IndexThreads::Pointer threads, + IndexHolder::Pointer holder) override; + + //! Dump index into file system + int dump(const IndexDumper::Pointer &dumper) override; + + //! Retrieve statistics + const Stats &stats() const override { + return stats_; + } + + private: + enum State { INIT = 0, INITED = 1, TRAINED = 2, BUILT = 3 }; + State state_{INIT}; + Stats stats_; + IndexMeta meta_; + IndexMeta rabitq_meta_; + ailego::Params params_; + + // IVF RaBitQ parameters + uint32_t nlist_{1024}; + uint32_t total_bits_{7}; + uint32_t sample_count_{0}; + std::string metric_name_; + + // Training results + std::shared_ptr converter_; + std::shared_ptr reformer_; + + // Build results (stored in memory until dump) + IvfRabitqHeader header_{}; + std::vector batch_data_buf_; + std::vector ex_data_buf_; + std::vector cluster_metas_; + std::vector keys_buf_; + std::vector mapping_buf_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf_rabitq/ivf_rabitq_context.h b/src/core/algorithm/ivf_rabitq/ivf_rabitq_context.h new file mode 100644 index 000000000..b67212b83 --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_context.h @@ -0,0 +1,229 @@ +// 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. +#pragma once + +#include +#include +#include +#include +#include "zvec/ailego/container/params.h" +#include "zvec/core/framework/index_context.h" +#include "zvec/core/framework/index_document.h" +#include "zvec/core/framework/index_error.h" +#include "zvec/core/framework/index_logger.h" +#include "ivf_rabitq_params.h" + +namespace zvec { +namespace core { + +// Opaque query state for IVF RaBitQ search. +// Holds the rotated query, lazily computed centroid norms, and +// the rabitqlib SplitBatchQuery object (via type-erased shared_ptr). +struct IvfRabitqQueryState { + std::vector rotated_query; + // sqrt(||q_rot - c_rot||^2) for probed centroids, used by g_add. + std::vector centroid_norms; + // Opaque pointer to rabitqlib::SplitBatchQuery, managed by reformer + std::shared_ptr batch_query; +}; + +/*! IVF RaBitQ Context + * Follows the same pattern as IVFSearcherContext: + * - result_heap_: max-heap of size topk_ used during scan + * - mutable_result_heap(): accessor for entity::search_cluster + * - reset_results() / topk_to_result(): lifecycle helpers + */ +class IvfRabitqContext : public IndexContext { + public: + typedef std::shared_ptr Pointer; + + IvfRabitqContext() = default; + ~IvfRabitqContext() override = default; + + // ----------------------------------------------------------------------- + // IndexContext interface + // ----------------------------------------------------------------------- + + //! Update context from ailego params + int update(const ailego::Params ¶ms) override { + params.get(PARAM_IVF_RABITQ_BRUTE_FORCE_THRESHOLD, &bruteforce_threshold_); + params.get(PARAM_IVF_RABITQ_SCAN_RATIO, &scan_ratio_); + if (scan_ratio_ <= 0.0f || scan_ratio_ > 1.0f) { + LOG_ERROR("Invalid params %s=%f", PARAM_IVF_RABITQ_SCAN_RATIO.c_str(), + scan_ratio_); + return IndexError_InvalidArgument; + } + + uint32_t val = 0; + if (params.get(PARAM_IVF_RABITQ_NPROBE, &val)) { + nprobe_ = val; + } + + return 0; + } + + //! Set topk — also sizes the heap + void set_topk(uint32_t k) override { + topk_ = k; + result_heap_.limit(topk_); + result_heap_.set_threshold(this->threshold()); + } + + uint32_t topk() const override { + return topk_; + } + + void set_fetch_vector(bool enable) override { + fetch_vector_ = enable; + } + + bool fetch_vector() const override { + return fetch_vector_; + } + + void reset(void) override { + reset_filter(); + reset_threshold(); + reset_group_by(); + set_fetch_vector(false); + result_heap_.clear(); + result_heap_.set_threshold(this->threshold()); + query_state.rotated_query.clear(); + query_state.centroid_norms.clear(); + query_state.batch_query.reset(); + } + + //! Retrieve search result (first query) + const IndexDocumentList &result(void) const override { + return results_[0]; + } + + //! Retrieve search result with index + const IndexDocumentList &result(size_t idx) const override { + return results_[idx]; + } + + //! Retrieve mutable result with index + IndexDocumentList *mutable_result(size_t idx) override { + if (idx >= results_.size()) { + results_.resize(idx + 1); + } + return &results_[idx]; + } + + // ----------------------------------------------------------------------- + // Heap helpers (same pattern as IVFSearcherContext) + // ----------------------------------------------------------------------- + + //! Accessor for the result heap (used by entity::search_cluster) + IndexDocumentHeap &mutable_result_heap() { + return result_heap_; + } + + //! Reset for a new batch of queries + void reset_results(size_t qnum) { + results_.resize(qnum); + for (size_t i = 0; i < qnum; ++i) { + results_[i].clear(); + } + result_heap_.clear(); + result_heap_.limit(topk_); + result_heap_.set_threshold(this->threshold()); + } + + //! Drain heap → results_[idx], sorted by score ascending (same as IVF) + void topk_to_result(uint32_t idx) { + if (result_heap_.empty()) { + return; + } + if (idx >= results_.size()) { + results_.resize(idx + 1); + } + int sz = std::min(topk_, static_cast(result_heap_.size())); + result_heap_.sort(); + results_[idx].clear(); + for (int i = 0; i < sz; ++i) { + float score = result_heap_[i].score(); + if (score > this->threshold()) { + break; + } + results_[idx].emplace_back(result_heap_[i].key(), score); + } + } + + // ----------------------------------------------------------------------- + // Search parameters + // ----------------------------------------------------------------------- + uint32_t nprobe() const { + return nprobe_; + } + uint32_t max_scan_count() const { + return max_scan_count_; + } + float scan_ratio() const { + return scan_ratio_; + } + uint32_t bruteforce_threshold() const { + return bruteforce_threshold_; + } + + int update_search_limits(uint32_t vector_count, uint32_t cluster_count, + uint32_t *effective_nprobe) { + if (cluster_count == 0) { + LOG_ERROR("Invalid cluster count"); + return IndexError_InvalidFormat; + } + if (!effective_nprobe) { + LOG_ERROR("Invalid effective nprobe output"); + return IndexError_InvalidArgument; + } + + if (nprobe_ > 0) { + *effective_nprobe = std::min(nprobe_, cluster_count); + // Explicit nprobe means fully scanning the selected clusters. + max_scan_count_ = vector_count; + } else { + *effective_nprobe = std::max( + static_cast(std::round(cluster_count * scan_ratio_)), 1u); + *effective_nprobe = std::min(*effective_nprobe, cluster_count); + max_scan_count_ = + static_cast(std::ceil(vector_count * scan_ratio_)); + } + max_scan_count_ = std::max(bruteforce_threshold_, max_scan_count_); + return 0; + } + + void set_search_limits(uint32_t max_scan_count) { + max_scan_count_ = max_scan_count; + } + + // ----------------------------------------------------------------------- + // Per-query state (managed by search loop) + // ----------------------------------------------------------------------- + IvfRabitqQueryState query_state; + + private: + IndexDocumentHeap result_heap_; + std::vector results_; + + uint32_t topk_{10}; + uint32_t nprobe_{10}; + uint32_t max_scan_count_{0}; + float scan_ratio_{kDefaultIvfRabitqScanRatio}; + uint32_t bruteforce_threshold_{kDefaultIvfRabitqBruteForceThreshold}; + bool fetch_vector_{false}; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf_rabitq/ivf_rabitq_entity.cc b/src/core/algorithm/ivf_rabitq/ivf_rabitq_entity.cc new file mode 100644 index 000000000..f8e4d53de --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_entity.cc @@ -0,0 +1,508 @@ +// 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. + +#include "ivf_rabitq_entity.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "zvec/core/framework/index_error.h" +#include "ivf_rabitq_params.h" + +namespace zvec { +namespace core { + +namespace { + +void ExtractCompactCode(const uint8_t *batch_code, size_t padded_dim, + uint32_t lane, uint8_t *compact_code) { + uint32_t perm_pos = 0; + uint32_t perm_lane = lane % 16U; + for (uint32_t i = 0; i < rabitqlib::fastscan::kPerm0.size(); ++i) { + if (static_cast(rabitqlib::fastscan::kPerm0[i]) == perm_lane) { + perm_pos = i; + break; + } + } + + bool high_half = lane >= 16U; + size_t code_bytes = padded_dim / 8; + for (size_t col = 0; col < code_bytes; ++col) { + const uint8_t *block = batch_code + (col * rabitqlib::fastscan::kBatchSize); + uint8_t high_bits = + high_half ? ((block[perm_pos] >> 4) & 0x0F) : (block[perm_pos] & 0x0F); + uint8_t low_bits = high_half ? ((block[perm_pos + 16] >> 4) & 0x0F) + : (block[perm_pos + 16] & 0x0F); + compact_code[col] = static_cast((high_bits << 4) | low_bits); + } +} + +} // namespace + +int IvfRabitqEntity::load(IndexStorage::Pointer storage) { + if (!storage) { + LOG_ERROR("Invalid storage for IvfRabitqEntity::load"); + return IndexError_InvalidArgument; + } + + // Read header + auto header_seg = storage->get(IVF_RABITQ_HEADER_SEG_ID); + if (!header_seg) { + LOG_ERROR("Failed to get segment %s", IVF_RABITQ_HEADER_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + IndexStorage::MemoryBlock hdr_block; + size_t read_size = header_seg->read(0, hdr_block, sizeof(IvfRabitqHeader)); + if (read_size != sizeof(IvfRabitqHeader)) { + LOG_ERROR("Failed to read IvfRabitqHeader"); + return IndexError_InvalidFormat; + } + memcpy(&header_, hdr_block.data(), sizeof(IvfRabitqHeader)); + + // Read cluster metas + auto cluster_meta_seg = storage->get(IVF_RABITQ_CLUSTER_META_SEG_ID); + if (!cluster_meta_seg) { + LOG_ERROR("Failed to get segment %s", + IVF_RABITQ_CLUSTER_META_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + size_t meta_size = header_.cluster_count * sizeof(IvfRabitqClusterMeta); + IndexStorage::MemoryBlock meta_block; + read_size = cluster_meta_seg->read(0, meta_block, meta_size); + if (read_size != meta_size) { + LOG_ERROR("Failed to read cluster metas"); + return IndexError_InvalidFormat; + } + cluster_metas_.resize(header_.cluster_count); + memcpy(cluster_metas_.data(), meta_block.data(), meta_size); + + // Read batch data + auto batch_seg = storage->get(IVF_RABITQ_BATCH_DATA_SEG_ID); + if (!batch_seg) { + LOG_ERROR("Failed to get segment %s", IVF_RABITQ_BATCH_DATA_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + read_size = batch_seg->read(0, batch_data_block_, header_.batch_data_size); + if (read_size != header_.batch_data_size) { + LOG_ERROR("Failed to read batch data: got %zu, expected %zu", read_size, + (size_t)header_.batch_data_size); + return IndexError_InvalidFormat; + } + batch_data_ = reinterpret_cast(batch_data_block_.data()); + + // Read ex data + ex_data_ = nullptr; + if (header_.ex_data_size > 0) { + auto ex_seg = storage->get(IVF_RABITQ_EX_DATA_SEG_ID); + if (!ex_seg) { + LOG_ERROR("Failed to get segment %s", IVF_RABITQ_EX_DATA_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + read_size = ex_seg->read(0, ex_data_block_, header_.ex_data_size); + if (read_size != header_.ex_data_size) { + LOG_ERROR("Failed to read ex data"); + return IndexError_InvalidFormat; + } + ex_data_ = reinterpret_cast(ex_data_block_.data()); + } + + // Read keys + auto keys_seg = storage->get(IVF_RABITQ_KEYS_SEG_ID); + if (!keys_seg) { + LOG_ERROR("Failed to get segment %s", IVF_RABITQ_KEYS_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + size_t keys_size = header_.total_vector_count * sizeof(uint64_t); + read_size = keys_seg->read(0, keys_block_, keys_size); + if (read_size != keys_size) { + LOG_ERROR("Failed to read keys"); + return IndexError_InvalidFormat; + } + keys_ = reinterpret_cast(keys_block_.data()); + + int ret = load_key_order_mapping(storage); + if (ret != 0) { + return ret; + } + ret = init_quantized_vector_layout(); + if (ret != 0) { + return ret; + } + + LOG_INFO( + "IvfRabitqEntity loaded: total_vectors=%u, clusters=%u, " + "padded_dim=%u, ex_bits=%u, batch_data_size=%zu, ex_data_size=%zu", + header_.total_vector_count, header_.cluster_count, header_.padded_dim, + header_.ex_bits, (size_t)header_.batch_data_size, + (size_t)header_.ex_data_size); + + // Resolve function pointer once at load time + ip_func_ = rabitqlib::select_excode_ipfunc(header_.ex_bits); + + return 0; +} + +int IvfRabitqEntity::search_cluster(uint32_t cluster_id, + const IvfRabitqQueryState &query_state, + size_t padded_dim, size_t ex_bits, + IndexDocumentHeap *heap) const { + return search_cluster_impl(cluster_id, query_state, padded_dim, + ex_bits, nullptr, heap); +} + +int IvfRabitqEntity::search_cluster(uint32_t cluster_id, + const IvfRabitqQueryState &query_state, + size_t padded_dim, size_t ex_bits, + const IndexFilter &filter, + IndexDocumentHeap *heap) const { + return search_cluster_impl(cluster_id, query_state, padded_dim, ex_bits, + &filter, heap); +} + +uint64_t IvfRabitqEntity::get_key(size_t id) const { + if (id >= header_.total_vector_count || !keys_) { + return std::numeric_limits::max(); + } + return keys_[id]; +} + +const void *IvfRabitqEntity::get_vector(size_t id) const { + int ret = materialize_quantized_vector(id, &quantized_vector_buffer_); + if (ret != 0) { + return nullptr; + } + return quantized_vector_buffer_.data(); +} + +int IvfRabitqEntity::get_vector(size_t id, + IndexStorage::MemoryBlock &block) const { + return materialize_quantized_vector(id, block); +} + +const void *IvfRabitqEntity::get_vector_by_key(uint64_t key) const { + uint32_t id = key_to_id(key); + if (id == std::numeric_limits::max()) { + return nullptr; + } + return get_vector(id); +} + +int IvfRabitqEntity::get_vector_by_key(uint64_t key, + IndexStorage::MemoryBlock &block) const { + uint32_t id = key_to_id(key); + if (id == std::numeric_limits::max()) { + return IndexError_NoExist; + } + return get_vector(id, block); +} + +uint32_t IvfRabitqEntity::key_to_id(uint64_t key) const { + if (!keys_ || !mapping_) { + return std::numeric_limits::max(); + } + + uint32_t start = 0U; + uint32_t end = header_.total_vector_count; + const void *data = nullptr; + while (start < end) { + uint32_t idx = start + (end - start) / 2U; + if (mapping_->read(idx * sizeof(uint32_t), &data, sizeof(uint32_t)) != + sizeof(uint32_t)) { + LOG_ERROR("Failed to read mapping segment, idx=%u", idx); + return std::numeric_limits::max(); + } + uint32_t local_id = *reinterpret_cast(data); + if (local_id >= header_.total_vector_count) { + LOG_ERROR("Invalid mapping local_id=%u, vector_count=%u", local_id, + header_.total_vector_count); + return std::numeric_limits::max(); + } + uint64_t local_key = keys_[local_id]; + if (local_key < key) { + start = idx + 1U; + } else if (local_key > key) { + end = idx; + } else { + return local_id; + } + } + return std::numeric_limits::max(); +} + +int IvfRabitqEntity::load_key_order_mapping(IndexStorage::Pointer storage) { + mapping_.reset(); + + if (header_.total_vector_count == 0) { + return 0; + } + if (!keys_) { + LOG_ERROR("Missing keys for mapping"); + return IndexError_InvalidFormat; + } + + size_t mapping_size = + static_cast(header_.total_vector_count) * sizeof(uint32_t); + mapping_ = storage->get(IVF_RABITQ_MAPPING_SEG_ID); + if (!mapping_) { + LOG_ERROR("Failed to get segment %s", IVF_RABITQ_MAPPING_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + if (mapping_->data_size() != mapping_size) { + LOG_ERROR("Invalid mapping segment size=%zu, expected=%zu", + mapping_->data_size(), mapping_size); + return IndexError_InvalidFormat; + } + return 0; +} + +int IvfRabitqEntity::init_quantized_vector_layout() { + quantized_vector_buffer_.clear(); + quantized_vector_element_size_ = 0; + compact_code_size_ = header_.padded_dim / 8; + bin_data_size_ = rabitqlib::BinDataMap::data_bytes(header_.padded_dim); + batch_data_size_ = + rabitqlib::BatchDataMap::data_bytes(header_.padded_dim); + ex_data_size_ = rabitqlib::ExDataMap::data_bytes(header_.padded_dim, + header_.ex_bits); + quantized_vector_element_size_ = bin_data_size_ + ex_data_size_; + + if (header_.total_vector_count == 0) { + return 0; + } + if (!batch_data_) { + LOG_ERROR("Missing batch data for quantized vectors"); + return IndexError_InvalidFormat; + } + if (header_.ex_bits > 0 && !ex_data_) { + LOG_ERROR("Missing extra-bit data for quantized vectors"); + return IndexError_InvalidFormat; + } + + return 0; +} + +const IvfRabitqClusterMeta *IvfRabitqEntity::find_cluster_meta( + size_t id) const { + auto it = + std::upper_bound(cluster_metas_.begin(), cluster_metas_.end(), id, + [](size_t value, const IvfRabitqClusterMeta &meta) { + return value < meta.key_offset; + }); + if (it == cluster_metas_.begin()) { + return nullptr; + } + --it; + if (id >= static_cast(it->key_offset) + it->vector_count) { + return nullptr; + } + return &(*it); +} + +int IvfRabitqEntity::materialize_quantized_vector(size_t id, + std::string *buffer) const { + if (!buffer) { + return IndexError_InvalidArgument; + } + if (id >= header_.total_vector_count || quantized_vector_element_size_ == 0) { + return IndexError_NoExist; + } + + buffer->resize(quantized_vector_element_size_); + char *dst = &(*buffer)[0]; + return materialize_quantized_vector(id, dst); +} + +int IvfRabitqEntity::materialize_quantized_vector( + size_t id, IndexStorage::MemoryBlock &block) const { + if (id >= header_.total_vector_count || quantized_vector_element_size_ == 0) { + return IndexError_NoExist; + } + + void *data = ailego_malloc(quantized_vector_element_size_); + if (!data) { + return IndexError_NoMemory; + } + + int ret = materialize_quantized_vector(id, static_cast(data)); + if (ret != 0) { + ailego_free(data); + return ret; + } + block = IndexStorage::MemoryBlock::MakeOwned(data, + quantized_vector_element_size_); + return 0; +} + +int IvfRabitqEntity::materialize_quantized_vector(size_t id, char *dst) const { + if (!dst) { + return IndexError_InvalidArgument; + } + if (id >= header_.total_vector_count || quantized_vector_element_size_ == 0) { + return IndexError_NoExist; + } + if (!batch_data_) { + LOG_ERROR("Missing batch data for quantized vectors"); + return IndexError_InvalidFormat; + } + if (header_.ex_bits > 0 && !ex_data_) { + LOG_ERROR("Missing extra-bit data for quantized vectors"); + return IndexError_InvalidFormat; + } + + const auto *meta = find_cluster_meta(id); + if (!meta) { + LOG_ERROR("Failed to locate cluster for vector id=%zu", id); + return IndexError_InvalidFormat; + } + + size_t local_id = id - meta->key_offset; + uint32_t batch_id = static_cast(local_id) / + static_cast(rabitqlib::fastscan::kBatchSize); + uint32_t lane = static_cast(local_id) % + static_cast(rabitqlib::fastscan::kBatchSize); + + const char *batch_ptr = batch_data_ + meta->batch_data_offset + + (static_cast(batch_id) * batch_data_size_); + rabitqlib::ConstBatchDataMap batch_map(batch_ptr, header_.padded_dim); + + ExtractCompactCode(batch_map.bin_code(), header_.padded_dim, lane, + reinterpret_cast(dst)); + std::memcpy(dst + compact_code_size_, batch_map.f_add() + lane, + sizeof(float)); + std::memcpy(dst + compact_code_size_ + sizeof(float), + batch_map.f_rescale() + lane, sizeof(float)); + std::memcpy(dst + compact_code_size_ + (sizeof(float) * 2), + batch_map.f_error() + lane, sizeof(float)); + + if (ex_data_size_ > 0) { + const char *cluster_ex = ex_data_ + meta->ex_data_offset; + std::memcpy(dst + bin_data_size_, cluster_ex + (local_id * ex_data_size_), + ex_data_size_); + } + + return 0; +} + +// Cluster scan with lower-bound pruning — mirrors rabitqlib IVF::scan_one_batch +template +int IvfRabitqEntity::search_cluster_impl(uint32_t cluster_id, + const IvfRabitqQueryState &query_state, + size_t padded_dim, size_t ex_bits, + const IndexFilter *filter, + IndexDocumentHeap *heap) const { + if (cluster_id >= cluster_metas_.size()) { + LOG_ERROR("Invalid cluster_id=%u", cluster_id); + return IndexError_OutOfRange; + } + + const auto &meta = cluster_metas_[cluster_id]; + if (meta.vector_count == 0) { + return 0; + } + + const auto *bq = static_cast *>( + query_state.batch_query.get()); + if (!bq) { + LOG_ERROR("Null batch_query in query_state"); + return IndexError_InvalidArgument; + } + + const size_t batch_stride = + rabitqlib::BatchDataMap::data_bytes(padded_dim); + const size_t ex_stride = + rabitqlib::ExDataMap::data_bytes(padded_dim, ex_bits); + + const char *cur_batch = batch_data_ + meta.batch_data_offset; + const char *cur_ex = nullptr; + if (ex_bits > 0) { + if (!ex_data_) { + LOG_ERROR("Missing extra-bit data for ex_bits=%zu", ex_bits); + return IndexError_InvalidFormat; + } + cur_ex = ex_data_ + meta.ex_data_offset; + } + const uint64_t *cur_keys = keys_ + meta.key_offset; + + uint32_t remaining = meta.vector_count; + uint32_t offset = 0; + + while (remaining > 0) { + const uint32_t batch_size = + std::min(remaining, (uint32_t)rabitqlib::fastscan::kBatchSize); + + // Step 1: 1-bit estimation for all 32 vectors at once (cheap, SIMD) + std::array est_dist; + std::array low_dist; + std::array ip_x0_qr; + + rabitqlib::split_batch_estdist(cur_batch, *bq, padded_dim, est_dist.data(), + low_dist.data(), ip_x0_qr.data(), + true /* use_hacc */); + + if (ex_bits == 0) { + // 1-bit only: insert all directly + for (uint32_t i = 0; i < batch_size; ++i) { + uint64_t key = cur_keys[offset + i]; + if constexpr (HasFilter) { + if ((*filter)(key)) { + continue; + } + } + heap->emplace(key, est_dist[i]); + } + } else { + // Step 2: lower-bound pruning + extra-bit boosting (same as rabitqlib) + // distk: current worst distance in heap, updated after each insertion + float distk = heap->full() ? heap->begin()->score() + : std::numeric_limits::max(); + + for (uint32_t i = 0; i < batch_size; ++i) { + if (low_dist[i] < distk) { + uint64_t key = cur_keys[offset + i]; + if constexpr (HasFilter) { + if ((*filter)(key)) { + continue; + } + } + // Extra-bit boosting (expensive, but only when needed) + const float ex_dist = rabitqlib::split_distance_boosting( + cur_ex + i * ex_stride, ip_func_, *bq, padded_dim, ex_bits, + ip_x0_qr[i]); + heap->emplace(key, ex_dist); + if (heap->full()) { + distk = heap->begin()->score(); + } + } + } + } + + cur_batch += batch_stride; + if (ex_bits > 0) { + cur_ex += batch_size * ex_stride; + } + offset += batch_size; + remaining -= batch_size; + } + + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf_rabitq/ivf_rabitq_entity.h b/src/core/algorithm/ivf_rabitq/ivf_rabitq_entity.h new file mode 100644 index 000000000..7b032bc9d --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_entity.h @@ -0,0 +1,145 @@ +// 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. +#pragma once + +#include +#include +#include +#include +#include +#include "zvec/core/framework/index_document.h" +#include "zvec/core/framework/index_filter.h" +#include "zvec/core/framework/index_storage.h" +#include "ivf_rabitq_context.h" +#include "ivf_rabitq_index_format.h" + +namespace zvec { +namespace core { + +/*! IVF RaBitQ Entity + * Stores quantized data (batch-packed binary codes + extra-bit codes + keys) + * for all clusters. Provides cluster scanning with lower-bound pruning. + */ +class IvfRabitqEntity { + public: + typedef std::shared_ptr Pointer; + + IvfRabitqEntity() = default; + ~IvfRabitqEntity() = default; + + //! Load entity data from storage + int load(IndexStorage::Pointer storage); + + //! Scan one cluster and insert qualifying candidates into heap. + //! Uses 1-bit lower bound to prune before expensive extra-bit boosting. + //! Same interface as IVFEntity::search for consistency. + int search_cluster(uint32_t cluster_id, + const IvfRabitqQueryState &query_state, size_t padded_dim, + size_t ex_bits, IndexDocumentHeap *heap) const; + + int search_cluster(uint32_t cluster_id, + const IvfRabitqQueryState &query_state, size_t padded_dim, + size_t ex_bits, const IndexFilter &filter, + IndexDocumentHeap *heap) const; + + uint32_t cluster_count() const { + return header_.cluster_count; + } + uint32_t total_vector_count() const { + return header_.total_vector_count; + } + uint32_t padded_dim() const { + return header_.padded_dim; + } + uint8_t ex_bits() const { + return header_.ex_bits; + } + + const IvfRabitqClusterMeta &cluster_meta(uint32_t cluster_id) const { + return cluster_metas_[cluster_id]; + } + + size_t quantized_vector_element_size() const { + return quantized_vector_element_size_; + } + + uint64_t get_key(size_t id) const; + + const void *get_vector(size_t id) const; + + int get_vector(size_t id, IndexStorage::MemoryBlock &block) const; + + const void *get_vector_by_key(uint64_t key) const; + + int get_vector_by_key(uint64_t key, IndexStorage::MemoryBlock &block) const; + + int materialize_quantized_vector(size_t id, std::string *buffer) const; + + int materialize_quantized_vector(size_t id, + IndexStorage::MemoryBlock &block) const; + + uint32_t key_to_id(uint64_t key) const; + + const uint32_t *get_key_order_mapping() const { + if (!mapping_) { + return nullptr; + } + const void *data = nullptr; + size_t size = total_vector_count() * sizeof(uint32_t); + if (mapping_->read(0, &data, size) != size) { + return nullptr; + } + return static_cast(data); + } + + private: + IvfRabitqHeader header_; + std::vector cluster_metas_; + + IndexStorage::MemoryBlock batch_data_block_; + IndexStorage::MemoryBlock ex_data_block_; + IndexStorage::MemoryBlock keys_block_; + IndexStorage::Segment::Pointer mapping_; + + const char *batch_data_{nullptr}; + const char *ex_data_{nullptr}; + const uint64_t *keys_{nullptr}; + mutable std::string quantized_vector_buffer_; + size_t quantized_vector_element_size_{0}; + size_t compact_code_size_{0}; + size_t bin_data_size_{0}; + size_t batch_data_size_{0}; + size_t ex_data_size_{0}; + + // Resolved once at load time to avoid per-cluster function-pointer lookup + rabitqlib::ex_ipfunc ip_func_{nullptr}; + + template + int search_cluster_impl(uint32_t cluster_id, + const IvfRabitqQueryState &query_state, + size_t padded_dim, size_t ex_bits, + const IndexFilter *filter, + IndexDocumentHeap *heap) const; + + const IvfRabitqClusterMeta *find_cluster_meta(size_t id) const; + + int materialize_quantized_vector(size_t id, char *dst) const; + + int init_quantized_vector_layout(); + + int load_key_order_mapping(IndexStorage::Pointer storage); +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf_rabitq/ivf_rabitq_index_format.h b/src/core/algorithm/ivf_rabitq/ivf_rabitq_index_format.h new file mode 100644 index 000000000..838d2a92e --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_index_format.h @@ -0,0 +1,58 @@ +// 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. +#pragma once + +#include +#include + +namespace zvec { +namespace core { + +struct IvfRabitqHeader { + uint32_t total_vector_count; + uint32_t cluster_count; + uint32_t dimension; + uint32_t padded_dim; + uint8_t ex_bits; + uint8_t rotator_type; + uint8_t metric_type; // 0=L2, 1=IP + uint8_t padding1; + uint64_t batch_data_size; + uint64_t ex_data_size; + uint32_t reserve[6]; + + IvfRabitqHeader() { + memset(static_cast(this), 0, sizeof(IvfRabitqHeader)); + } +}; +static_assert(sizeof(IvfRabitqHeader) % 32 == 0, + "IvfRabitqHeader must be 32-byte aligned"); + +struct IvfRabitqClusterMeta { + uint64_t batch_data_offset; + uint64_t ex_data_offset; + uint32_t vector_count; + uint32_t batch_count; + uint32_t key_offset; + uint32_t reserve; + + IvfRabitqClusterMeta() { + memset(static_cast(this), 0, sizeof(IvfRabitqClusterMeta)); + } +}; +static_assert(sizeof(IvfRabitqClusterMeta) == 32, + "IvfRabitqClusterMeta must be 32 bytes"); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf_rabitq/ivf_rabitq_index_provider.h b/src/core/algorithm/ivf_rabitq/ivf_rabitq_index_provider.h new file mode 100644 index 000000000..a9c26ca35 --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_index_provider.h @@ -0,0 +1,134 @@ +// 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. +#pragma once + +#include +#include +#include +#include +#include +#include "zvec/core/framework/index_provider.h" +#include "ivf_rabitq_entity.h" + +namespace zvec { +namespace core { + +/*! IVF RaBitQ IndexProvider + */ +class IvfRabitqIndexProvider : public IndexProvider { + public: + IvfRabitqIndexProvider(const IndexMeta &, IvfRabitqEntity::Pointer entity, + const std::string &owner) + : entity_(std::move(entity)), owner_class_(owner) {} + + IvfRabitqIndexProvider(const IvfRabitqIndexProvider &) = delete; + IvfRabitqIndexProvider &operator=(const IvfRabitqIndexProvider &) = delete; + + //! Create a new iterator + Iterator::Pointer create_iterator(void) override { + return Iterator::Pointer(new (std::nothrow) Iterator(entity_)); + } + + //! Retrieve count of vectors + size_t count(void) const override { + return entity_->total_vector_count(); + } + + //! Retrieve dimension of vector + size_t dimension(void) const override { + return entity_->padded_dim(); + } + + //! Retrieve type of vector + IndexMeta::DataType data_type(void) const override { + return IndexMeta::DataType::DT_UNDEFINED; + } + + //! Retrieve vector size in bytes + size_t element_size(void) const override { + return entity_->quantized_vector_element_size(); + } + + //! Retrieve a vector using a primary key + const void *get_vector(uint64_t key) const override { + uint32_t id = entity_->key_to_id(key); + if (id == std::numeric_limits::max()) { + return nullptr; + } + int ret = entity_->materialize_quantized_vector(id, &vector_buffer_); + if (ret != 0) { + return nullptr; + } + return vector_buffer_.data(); + } + + //! Retrieve a vector using a primary key + int get_vector(const uint64_t key, + IndexStorage::MemoryBlock &block) const override { + uint32_t id = entity_->key_to_id(key); + if (id == std::numeric_limits::max()) { + return IndexError_NoExist; + } + return entity_->materialize_quantized_vector(id, block); + } + + //! Retrieve the owner class + const std::string &owner_class(void) const override { + return owner_class_; + } + + private: + class Iterator : public IndexProvider::Iterator { + public: + explicit Iterator(const IvfRabitqEntity::Pointer &entity) + : entity_(entity) {} + + //! Retrieve pointer of data + const void *data(void) const override { + int ret = entity_->materialize_quantized_vector(id_, &vector_buffer_); + if (ret != 0) { + return nullptr; + } + return vector_buffer_.data(); + } + + //! Test if the iterator is valid + bool is_valid(void) const override { + return id_ < entity_->total_vector_count(); + } + + //! Retrieve primary key + uint64_t key(void) const override { + return entity_->get_key(id_); + } + + //! Next iterator + void next(void) override { + ++id_; + } + + private: + IvfRabitqEntity::Pointer entity_; + mutable std::string vector_buffer_; + size_t id_{0}; + }; + + private: + IvfRabitqEntity::Pointer entity_; + mutable std::string vector_buffer_; + std::string owner_class_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf_rabitq/ivf_rabitq_params.h b/src/core/algorithm/ivf_rabitq/ivf_rabitq_params.h new file mode 100644 index 000000000..120b22783 --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_params.h @@ -0,0 +1,46 @@ +// 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. +#pragma once + +#include + +namespace zvec { +namespace core { + +// IVF RabitQ specific params +static const std::string PARAM_IVF_RABITQ_NLIST("proxima.ivf_rabitq.nlist"); +static const std::string PARAM_IVF_RABITQ_NPROBE("proxima.ivf_rabitq.nprobe"); +static const std::string PARAM_IVF_RABITQ_SCAN_RATIO( + "proxima.ivf_rabitq.scan_ratio"); +static const std::string PARAM_IVF_RABITQ_BRUTE_FORCE_THRESHOLD( + "proxima.ivf_rabitq.brute_force_threshold"); + +// Segment IDs +static const std::string IVF_RABITQ_HEADER_SEG_ID{"ivf_rabitq.header"}; +static const std::string IVF_RABITQ_BATCH_DATA_SEG_ID{"ivf_rabitq.batch_data"}; +static const std::string IVF_RABITQ_EX_DATA_SEG_ID{"ivf_rabitq.ex_data"}; +static const std::string IVF_RABITQ_CLUSTER_META_SEG_ID{ + "ivf_rabitq.cluster_meta"}; +static const std::string IVF_RABITQ_KEYS_SEG_ID{"ivf_rabitq.keys"}; +static const std::string IVF_RABITQ_MAPPING_SEG_ID{"ivf_rabitq.mapping"}; +static const std::string IVF_RABITQ_CENTROID_SEG_ID{"ivf_rabitq.centroid"}; + +// Defaults +constexpr uint32_t kDefaultIvfRabitqNlist = 1024; +constexpr uint32_t kDefaultIvfRabitqNprobe = 10; +constexpr float kDefaultIvfRabitqScanRatio = 0.1f; +constexpr uint32_t kDefaultIvfRabitqBruteForceThreshold = 1000; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf_rabitq/ivf_rabitq_reformer.cc b/src/core/algorithm/ivf_rabitq/ivf_rabitq_reformer.cc new file mode 100644 index 000000000..1e8b92edc --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_reformer.cc @@ -0,0 +1,540 @@ +// 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. + +#include "ivf_rabitq_reformer.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "algorithm/hnsw_rabitq/rabitq_utils.h" +#include "core/algorithm/cluster/linear_seeker.h" +#include "zvec/core/framework/index_error.h" +#include "zvec/core/framework/index_factory.h" +#include "zvec/core/framework/index_features.h" +#include "zvec/core/framework/index_metric.h" +#include "zvec/core/framework/index_storage.h" + +namespace zvec { +namespace core { + +// All rabitqlib types are confined to this translation unit via pimpl. +struct IvfRabitqReformer::Impl { + // RaBitQ parameters + size_t num_clusters{0}; + size_t ex_bits{0}; + size_t dimension{0}; + size_t padded_dim{0}; + bool is_loaded{false}; + + // Original centroids: num_clusters * dimension (for nearest centroid search) + std::vector centroids; + // Rotated centroids: num_clusters * padded_dim (for quantization) + std::vector rotated_centroids; + + rabitqlib::RotatorType rotator_type{rabitqlib::RotatorType::FhtKacRotator}; + std::unique_ptr> rotator; + + // Precomputed configs for quantization + rabitqlib::quant::RabitqConfig query_config; + rabitqlib::quant::RabitqConfig build_config; + rabitqlib::MetricType metric_type{rabitqlib::METRIC_L2}; + + // LinearSeeker for SIMD-optimized nearest centroid search + LinearSeeker::Pointer centroid_seeker; + CoherentIndexFeatures::Pointer centroid_features; + IndexMetric::Pointer centroid_metric; + IndexMetric::MatrixDistance centroid_distance_func; + + // Translate string metric to rabitqlib enum + static rabitqlib::MetricType metric_from_string(const std::string &name) { + if (name == "InnerProduct" || name == "Cosine") { + return rabitqlib::METRIC_IP; + } + return rabitqlib::METRIC_L2; + } + + // Translate rabitqlib enum to local enum + static RabitqMetricType to_local(rabitqlib::MetricType m) { + return m == rabitqlib::METRIC_IP ? RabitqMetricType::kIP + : RabitqMetricType::kL2; + } +}; + +IvfRabitqReformer::IvfRabitqReformer() : impl_(std::make_unique()) {} + +IvfRabitqReformer::~IvfRabitqReformer() = default; + +size_t IvfRabitqReformer::num_clusters() const { + return impl_->num_clusters; +} + +size_t IvfRabitqReformer::dimension() const { + return impl_->dimension; +} + +size_t IvfRabitqReformer::padded_dim() const { + return impl_->padded_dim; +} + +size_t IvfRabitqReformer::ex_bits() const { + return impl_->ex_bits; +} + +RabitqMetricType IvfRabitqReformer::rabitq_metric_type() const { + return Impl::to_local(impl_->metric_type); +} + +bool IvfRabitqReformer::loaded() const { + return impl_->is_loaded; +} + +int IvfRabitqReformer::init(const std::string &metric_name) { + impl_->metric_type = Impl::metric_from_string(metric_name); + LOG_DEBUG("IvfRabitqReformer init done. metric_name=%s metric_type=%d", + metric_name.c_str(), static_cast(impl_->metric_type)); + return 0; +} + +int IvfRabitqReformer::load(IndexStorage::Pointer storage) { + if (!storage) { + LOG_ERROR("Invalid storage for load"); + return IndexError_InvalidArgument; + } + + auto segment = storage->get(RABITQ_CONVERTER_SEG_ID); + if (!segment) { + LOG_ERROR("Failed to get segment %s", RABITQ_CONVERTER_SEG_ID.c_str()); + return IndexError_InvalidFormat; + } + + size_t offset = 0; + RabitqConverterHeader header; + IndexStorage::MemoryBlock block; + + size_t size = segment->read(offset, block, sizeof(header)); + if (size != sizeof(header)) { + LOG_ERROR("Failed to read header"); + return IndexError_InvalidFormat; + } + memcpy(static_cast(&header), block.data(), sizeof(header)); + impl_->dimension = header.dim; + impl_->padded_dim = header.padded_dim; + impl_->ex_bits = header.ex_bits; + impl_->num_clusters = header.num_clusters; + impl_->rotator_type = + static_cast(header.rotator_type); + offset += sizeof(header); + + // Read rotated centroids + size_t rotated_centroids_size = + sizeof(float) * header.num_clusters * header.padded_dim; + size = segment->read(offset, block, rotated_centroids_size); + if (size != rotated_centroids_size) { + LOG_ERROR("Failed to read rotated centroids"); + return IndexError_InvalidFormat; + } + impl_->rotated_centroids.resize(header.num_clusters * header.padded_dim); + memcpy(impl_->rotated_centroids.data(), block.data(), rotated_centroids_size); + offset += size; + + // Read original centroids + size_t centroids_size = sizeof(float) * header.num_clusters * header.dim; + size = segment->read(offset, block, centroids_size); + if (size != centroids_size) { + LOG_ERROR("Failed to read centroids"); + return IndexError_InvalidFormat; + } + impl_->centroids.resize(header.num_clusters * header.dim); + memcpy(impl_->centroids.data(), block.data(), centroids_size); + offset += size; + + // Read rotator + size_t rotator_size = header.rotator_size; + size = segment->read(offset, block, rotator_size); + if (size != rotator_size) { + LOG_ERROR("Failed to read rotator"); + return IndexError_InvalidFormat; + } + impl_->rotator.reset(rabitqlib::choose_rotator( + impl_->dimension, impl_->rotator_type, impl_->padded_dim)); + impl_->rotator->load(reinterpret_cast(block.data())); + offset += size; + + // Precompute configs + impl_->query_config = rabitqlib::quant::faster_config( + impl_->padded_dim, rabitqlib::SplitSingleQuery::kNumBits); + impl_->build_config = + rabitqlib::quant::faster_config(impl_->padded_dim, impl_->ex_bits + 1); + + // Initialize LinearSeeker for SIMD-optimized centroid search + IndexMeta centroid_meta; + centroid_meta.set_data_type(IndexMeta::DataType::DT_FP32); + centroid_meta.set_dimension(static_cast(impl_->dimension)); + // Cosine maps to IP internally; use InnerProduct for centroid search + // since input vectors are already normalized by CosineNormalizeConverter + centroid_meta.set_metric(impl_->metric_type == rabitqlib::METRIC_L2 + ? "SquaredEuclidean" + : "InnerProduct", + 0, ailego::Params()); + + impl_->centroid_features = std::make_shared(); + impl_->centroid_features->mount(centroid_meta, impl_->centroids.data(), + impl_->centroids.size() * sizeof(float)); + + impl_->centroid_seeker = std::make_shared(); + int ret = impl_->centroid_seeker->init(centroid_meta); + if (ret != 0) { + LOG_ERROR("Failed to init centroid seeker, ret=%d", ret); + return ret; + } + ret = impl_->centroid_seeker->mount(impl_->centroid_features); + if (ret != 0) { + LOG_ERROR("Failed to mount centroid features, ret=%d", ret); + return ret; + } + impl_->centroid_metric = + IndexFactory::CreateMetric(centroid_meta.metric_name()); + if (!impl_->centroid_metric) { + LOG_ERROR("Failed to create centroid metric %s", + centroid_meta.metric_name().c_str()); + return IndexError_Unsupported; + } + ret = impl_->centroid_metric->init(centroid_meta, + centroid_meta.metric_params()); + if (ret != 0) { + LOG_ERROR("Failed to init centroid metric, ret=%d", ret); + return ret; + } + impl_->centroid_distance_func = impl_->centroid_metric->distance_matrix(1, 1); + if (!impl_->centroid_distance_func) { + LOG_ERROR("Failed to get centroid distance function"); + return IndexError_Unsupported; + } + + impl_->is_loaded = true; + + LOG_INFO( + "IvfRabitqReformer loaded: dimension=%zu, padded_dim=%zu, " + "ex_bits=%zu, num_clusters=%zu, rotator_type=%d", + impl_->dimension, impl_->padded_dim, impl_->ex_bits, impl_->num_clusters, + (int)impl_->rotator_type); + return 0; +} + +int IvfRabitqReformer::dump(const IndexStorage::Pointer &storage) { + if (!storage) { + LOG_ERROR("Null storage"); + return IndexError_InvalidArgument; + } + if (!impl_->is_loaded || impl_->rotated_centroids.empty() || + impl_->centroids.empty()) { + LOG_ERROR("No centroids to dump"); + return IndexError_NoReady; + } + + auto align_size = [](size_t s) -> size_t { return (s + 0x1F) & (~0x1F); }; + + size_t header_size = sizeof(RabitqConverterHeader); + size_t rotated_centroids_size = + impl_->rotated_centroids.size() * sizeof(float); + size_t centroids_size = impl_->centroids.size() * sizeof(float); + size_t rotator_size = impl_->rotator->dump_bytes(); + size_t data_size = + header_size + rotated_centroids_size + centroids_size + rotator_size; + size_t total_size = align_size(data_size); + + int ret = storage->append(RABITQ_CONVERTER_SEG_ID, total_size); + if (ret != 0) { + LOG_ERROR("Failed to append segment %s, ret=%d", + RABITQ_CONVERTER_SEG_ID.c_str(), ret); + return ret; + } + + auto segment = storage->get(RABITQ_CONVERTER_SEG_ID); + if (!segment) { + LOG_ERROR("Failed to get segment %s", RABITQ_CONVERTER_SEG_ID.c_str()); + return IndexError_ReadData; + } + + size_t offset = 0; + + RabitqConverterHeader header; + header.dim = static_cast(impl_->dimension); + header.padded_dim = static_cast(impl_->padded_dim); + header.num_clusters = static_cast(impl_->num_clusters); + header.ex_bits = static_cast(impl_->ex_bits); + header.rotator_type = static_cast(impl_->rotator_type); + header.rotator_size = static_cast(rotator_size); + + size_t written = segment->write(offset, &header, header_size); + if (written != header_size) { + LOG_ERROR("Failed to write header"); + return IndexError_WriteData; + } + offset += header_size; + + written = segment->write(offset, impl_->rotated_centroids.data(), + rotated_centroids_size); + if (written != rotated_centroids_size) { + LOG_ERROR("Failed to write rotated centroids"); + return IndexError_WriteData; + } + offset += rotated_centroids_size; + + written = segment->write(offset, impl_->centroids.data(), centroids_size); + if (written != centroids_size) { + LOG_ERROR("Failed to write centroids"); + return IndexError_WriteData; + } + offset += centroids_size; + + std::vector buffer(rotator_size); + impl_->rotator->save(buffer.data()); + written = segment->write(offset, buffer.data(), rotator_size); + if (written != rotator_size) { + LOG_ERROR("Failed to write rotator data"); + return IndexError_WriteData; + } + + LOG_INFO("IvfRabitqReformer dumped to storage: %zu bytes", data_size); + return 0; +} + +int IvfRabitqReformer::dump(const IndexDumper::Pointer &dumper) { + if (!dumper) { + LOG_ERROR("Null dumper"); + return IndexError_InvalidArgument; + } + if (!impl_->is_loaded || impl_->rotated_centroids.empty() || + impl_->centroids.empty()) { + LOG_ERROR("No centroids to dump"); + return IndexError_NoReady; + } + + size_t dumped_size = 0; + int ret = dump_rabitq_centroids( + dumper, impl_->dimension, impl_->padded_dim, impl_->ex_bits, + impl_->num_clusters, impl_->rotator_type, impl_->rotated_centroids, + impl_->centroids, impl_->rotator, &dumped_size); + if (ret != 0) { + return ret; + } + + LOG_INFO("IvfRabitqReformer dump completed: %zu bytes", dumped_size); + return 0; +} + +int IvfRabitqReformer::create_query_state(const float *query, + IvfRabitqQueryState *state) const { + if (!impl_->is_loaded) { + LOG_ERROR("Reformer not loaded yet"); + return IndexError_NoReady; + } + if (!query || !state) { + LOG_ERROR("Invalid arguments for create_query_state"); + return IndexError_InvalidArgument; + } + + // Rotate query + state->rotated_query.resize(impl_->padded_dim); + impl_->rotator->rotate(query, state->rotated_query.data()); + + // Create SplitBatchQuery object for FastScan + auto batch_query = std::make_shared>( + state->rotated_query.data(), impl_->padded_dim, impl_->ex_bits, + impl_->metric_type, true /* use_hacc */); + state->batch_query = batch_query; + + state->centroid_norms.assign(impl_->num_clusters, + std::numeric_limits::quiet_NaN()); + + return 0; +} + +int IvfRabitqReformer::prepare_for_cluster(uint32_t centroid_id, + IvfRabitqQueryState *state) const { + if (!impl_->is_loaded || !state || !state->batch_query) { + return IndexError_NoReady; + } + if (centroid_id >= impl_->num_clusters) { + LOG_ERROR("Invalid centroid_id=%u, num_clusters=%zu", centroid_id, + impl_->num_clusters); + return IndexError_OutOfRange; + } + + auto *bq = static_cast *>( + state->batch_query.get()); + + float centroid_norm = state->centroid_norms[centroid_id]; + if (!std::isfinite(centroid_norm)) { + centroid_norm = std::sqrt(rabitqlib::euclidean_sqr( + state->rotated_query.data(), + impl_->rotated_centroids.data() + (centroid_id * impl_->padded_dim), + impl_->padded_dim)); + state->centroid_norms[centroid_id] = centroid_norm; + } + + if (impl_->metric_type == rabitqlib::METRIC_L2) { + bq->set_g_add(centroid_norm); + } else { + // Compute dot_product on demand for the selected centroid only + float ip = rabitqlib::dot_product( + state->rotated_query.data(), + impl_->rotated_centroids.data() + (centroid_id * impl_->padded_dim), + impl_->padded_dim); + bq->set_g_add(centroid_norm, ip); + } + + return 0; +} + +int IvfRabitqReformer::rotate_vector(const float *input, float *output) const { + if (!impl_->is_loaded || !impl_->rotator) { + return IndexError_NoReady; + } + impl_->rotator->rotate(input, output); + return 0; +} + +int IvfRabitqReformer::quantize_batch(const float *rotated_data, + uint32_t centroid_id, size_t num_points, + char *batch_data, char *ex_data) const { + if (!impl_->is_loaded) { + return IndexError_NoReady; + } + if (centroid_id >= impl_->num_clusters) { + LOG_ERROR("Invalid centroid_id=%u, num_clusters=%zu", centroid_id, + impl_->num_clusters); + return IndexError_OutOfRange; + } + + const float *rotated_centroid = + impl_->rotated_centroids.data() + (centroid_id * impl_->padded_dim); + + rabitqlib::quant::quantize_split_batch( + rotated_data, rotated_centroid, num_points, impl_->padded_dim, + impl_->ex_bits, batch_data, ex_data, impl_->metric_type, + impl_->build_config); + + return 0; +} + +uint32_t IvfRabitqReformer::find_nearest_centroid(const float *vector) const { + if (!impl_->is_loaded || impl_->num_clusters == 0) { + return 0; + } + + Seeker::Document doc; + int ret = impl_->centroid_seeker->seek( + vector, impl_->dimension * sizeof(float), &doc); + if (ret != 0) { + LOG_ERROR("Failed to seek nearest centroid, ret=%d", ret); + return 0; + } + return doc.index; +} + +int IvfRabitqReformer::select_probe_centroids( + const float *query, size_t nprobe, std::vector *centroids) const { + return select_probe_centroids(query, nprobe, nullptr, centroids); +} + +int IvfRabitqReformer::select_probe_centroids( + const float *query, size_t nprobe, IvfRabitqQueryState *state, + std::vector *centroids) const { + if (!impl_->is_loaded || impl_->num_clusters == 0) { + return IndexError_NoReady; + } + if (!query || !centroids) { + return IndexError_InvalidArgument; + } + if (!impl_->centroid_distance_func) { + LOG_ERROR("Centroid distance function is not initialized"); + return IndexError_NoReady; + } + + size_t probe_count = std::min(nprobe, impl_->num_clusters); + centroids->clear(); + if (probe_count == 0) { + return 0; + } + if (probe_count == 1) { + Seeker::Document doc; + int ret = impl_->centroid_seeker->seek( + query, impl_->dimension * sizeof(float), &doc); + if (ret != 0) { + LOG_ERROR("Failed to seek probe centroid, ret=%d", ret); + return ret; + } + centroids->push_back(doc.index); + if (state && impl_->metric_type == rabitqlib::METRIC_L2 && + state->centroid_norms.size() == impl_->num_clusters) { + state->centroid_norms[doc.index] = std::sqrt(std::max(doc.score, 0.0f)); + } + return 0; + } + + struct CentroidDist { + uint32_t id; + float dist; + }; + + std::vector centroid_dists(impl_->num_clusters); + for (size_t i = 0; i < impl_->num_clusters; ++i) { + float score = 0.0f; + impl_->centroid_distance_func(impl_->centroid_features->element(i), query, + impl_->dimension, &score); + centroid_dists[i].id = static_cast(i); + centroid_dists[i].dist = score; + } + + auto compare_dist = [](const CentroidDist &a, const CentroidDist &b) { + if (a.dist == b.dist) { + return a.id < b.id; + } + return a.dist < b.dist; + }; + auto top_end = centroid_dists.begin() + probe_count; + if (top_end != centroid_dists.end()) { + std::nth_element(centroid_dists.begin(), top_end, centroid_dists.end(), + compare_dist); + } + std::sort(centroid_dists.begin(), top_end, compare_dist); + + centroids->reserve(probe_count); + for (size_t i = 0; i < probe_count; ++i) { + uint32_t centroid_id = centroid_dists[i].id; + centroids->push_back(centroid_id); + if (state && impl_->metric_type == rabitqlib::METRIC_L2 && + state->centroid_norms.size() == impl_->num_clusters) { + float dist = std::max(centroid_dists[i].dist, 0.0f); + state->centroid_norms[centroid_id] = std::sqrt(dist); + } + } + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf_rabitq/ivf_rabitq_reformer.h b/src/core/algorithm/ivf_rabitq/ivf_rabitq_reformer.h new file mode 100644 index 000000000..bc2488fd5 --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_reformer.h @@ -0,0 +1,101 @@ +// 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. +#pragma once + +#include +#include +#include "algorithm/hnsw_rabitq/rabitq_params.h" +#include "zvec/core/framework/index_dumper.h" +#include "zvec/core/framework/index_storage.h" +#include "ivf_rabitq_context.h" + +namespace zvec { +namespace core { + +/*! IVF RaBitQ Reformer + * Manages centroids, rotator, and quantization for IVF RaBitQ. + * Provides methods for query transformation, centroid assignment, + * and batch quantization for cluster building. + * + * All rabitqlib types are hidden behind a pimpl to avoid leaking rabitqlib + * headers to consumers of this class. + */ +class IvfRabitqReformer { + public: + typedef std::shared_ptr Pointer; + + IvfRabitqReformer(); + ~IvfRabitqReformer(); + + // Non-copyable + IvfRabitqReformer(const IvfRabitqReformer &) = delete; + IvfRabitqReformer &operator=(const IvfRabitqReformer &) = delete; + + //! Initialize with metric parameters + int init(const std::string &metric_name); + + //! Load from storage (reads rabitq.converter segment) + int load(IndexStorage::Pointer storage); + + //! Dump to storage (writes rabitq.converter segment) + int dump(const IndexStorage::Pointer &storage); + + //! Dump to dumper (writes rabitq.converter segment via IndexDumper) + int dump(const IndexDumper::Pointer &dumper); + + //! Create a query state for searching + int create_query_state(const float *query, IvfRabitqQueryState *state) const; + + //! Prepare query state for a specific centroid cluster + //! This sets g_add / g_error in the batch_query for the given centroid + int prepare_for_cluster(uint32_t centroid_id, + IvfRabitqQueryState *state) const; + + //! Rotate a single vector: input[dimension] -> output[padded_dim] + int rotate_vector(const float *input, float *output) const; + + //! Quantize a batch of rotated vectors for a specific centroid. + //! rotated_data: num_points vectors of padded_dim dimension each + //! centroid_id: which centroid to use as residual reference + //! num_points: up to fastscan::kBatchSize (32) + //! batch_data: output batch-packed data + //! ex_data: output extra-bit data for each vector + int quantize_batch(const float *rotated_data, uint32_t centroid_id, + size_t num_points, char *batch_data, char *ex_data) const; + + //! Find nearest centroid for a vector (using original centroids, brute-force) + uint32_t find_nearest_centroid(const float *vector) const; + + //! Select top-n probe centroids using the same metric as build assignment. + int select_probe_centroids(const float *query, size_t nprobe, + std::vector *centroids) const; + int select_probe_centroids(const float *query, size_t nprobe, + IvfRabitqQueryState *state, + std::vector *centroids) const; + + //! Accessors + size_t num_clusters() const; + size_t dimension() const; + size_t padded_dim() const; + size_t ex_bits() const; + RabitqMetricType rabitq_metric_type() const; + bool loaded() const; + + private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf_rabitq/ivf_rabitq_register.cc b/src/core/algorithm/ivf_rabitq/ivf_rabitq_register.cc new file mode 100644 index 000000000..0259014d0 --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_register.cc @@ -0,0 +1,25 @@ +// 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. +#include "zvec/core/framework/index_factory.h" +#include "ivf_rabitq_builder.h" +#include "ivf_rabitq_streamer.h" + +namespace zvec { +namespace core { + +INDEX_FACTORY_REGISTER_BUILDER(IvfRabitqBuilder); +INDEX_FACTORY_REGISTER_STREAMER(IvfRabitqStreamer); + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf_rabitq/ivf_rabitq_streamer.cc b/src/core/algorithm/ivf_rabitq/ivf_rabitq_streamer.cc new file mode 100644 index 000000000..5fe4580c7 --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_streamer.cc @@ -0,0 +1,397 @@ +// 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. + +#include "ivf_rabitq_streamer.h" +#include +#include +#include +#include +#include +#include "zvec/core/framework/index_error.h" +#include "zvec/core/framework/index_helper.h" +#include "zvec/core/framework/index_meta.h" +#include "ivf_rabitq_context.h" +#include "ivf_rabitq_entity.h" +#include "ivf_rabitq_index_provider.h" +#include "ivf_rabitq_params.h" +#include "ivf_rabitq_reformer.h" +#include "ivf_rabitq_util.h" + +namespace zvec { +namespace core { + +// -------------------------------------------------------------------------- +// init +// -------------------------------------------------------------------------- +int IvfRabitqStreamer::init(const IndexMeta &meta, + const ailego::Params ¶ms) { + meta_ = meta; + params_ = params; + + // Parse IVF RaBitQ specific params + params.get(PARAM_IVF_RABITQ_NPROBE, &nprobe_); + + params.get(PARAM_IVF_RABITQ_SCAN_RATIO, &scan_ratio_); + if (scan_ratio_ <= 0.0f || scan_ratio_ > 1.0f) { + LOG_ERROR("Invalid params %s=%f", PARAM_IVF_RABITQ_SCAN_RATIO.c_str(), + scan_ratio_); + return IndexError_InvalidArgument; + } + + params.get(PARAM_IVF_RABITQ_BRUTE_FORCE_THRESHOLD, &brute_force_threshold_); + + int ret = PrepareAndCheckIvfRabitqInternalMeta(meta_, params, &rabitq_meta_, + &metric_name_); + if (ret != 0) { + return ret; + } + + uint32_t dim = rabitq_meta_.dimension(); + state_ = STATE_INITED; + + LOG_INFO( + "IvfRabitqStreamer initialized: dim=%u, nprobe=%u, metric=%s, " + "scan_ratio=%.3f, bf_threshold=%u", + dim, nprobe_, metric_name_.c_str(), scan_ratio_, brute_force_threshold_); + + return 0; +} + +// -------------------------------------------------------------------------- +// open +// -------------------------------------------------------------------------- +int IvfRabitqStreamer::open(IndexStorage::Pointer storage) { + if (!storage) { + LOG_ERROR("Invalid storage"); + return IndexError_InvalidArgument; + } + if (state_ != STATE_INITED) { + LOG_ERROR("Streamer not initialized"); + return IndexError_NoReady; + } + + int ret = IndexHelper::DeserializeFromStorage(storage.get(), &meta_); + if (ret != 0) { + LOG_ERROR("Failed to deserialize meta from storage"); + return ret; + } + ret = PrepareAndCheckIvfRabitqInternalMeta(meta_, params_, &rabitq_meta_, + &metric_name_); + if (ret != 0) { + return ret; + } + + storage_ = std::move(storage); + ret = load_index(storage_); + if (ret != 0) { + LOG_ERROR("Failed to load index, ret=%d", ret); + return ret; + } + + state_ = STATE_LOADED; + return 0; +} + +// -------------------------------------------------------------------------- +// load_index +// -------------------------------------------------------------------------- +int IvfRabitqStreamer::load_index(IndexStorage::Pointer storage) { + ailego::ElapsedTime timer; + + // Load reformer + reformer_ = std::make_shared(); + int ret = reformer_->init(metric_name_); + if (ret != 0) { + LOG_ERROR("Failed to init IvfRabitqReformer, ret=%d", ret); + return ret; + } + ret = reformer_->load(storage); + if (ret != 0) { + LOG_ERROR("Failed to load IvfRabitqReformer, ret=%d", ret); + return ret; + } + if (reformer_->dimension() != rabitq_meta_.dimension()) { + LOG_ERROR("RaBitQ dimension mismatch: reformer=%zu, meta=%zu", + reformer_->dimension(), + static_cast(rabitq_meta_.dimension())); + return IndexError_Mismatch; + } + + // Load entity data + entity_ = std::make_shared(); + ret = entity_->load(storage); + if (ret != 0) { + LOG_ERROR("Failed to load IvfRabitqEntity, ret=%d", ret); + return ret; + } + + stats_.set_loaded_count(entity_->total_vector_count()); + stats_.set_loaded_costtime(timer.milli_seconds()); + + LOG_INFO("IvfRabitqStreamer loaded: %u vectors, %u clusters, cost %zu ms", + entity_->total_vector_count(), entity_->cluster_count(), + static_cast(timer.milli_seconds())); + + return 0; +} + +// -------------------------------------------------------------------------- +// create_context +// -------------------------------------------------------------------------- +IndexStreamer::Context::Pointer IvfRabitqStreamer::create_context() const { + if (state_ != STATE_LOADED) { + LOG_ERROR("Load the index first before create context"); + return Context::Pointer(); + } + + auto *ctx = new (std::nothrow) IvfRabitqContext(); + if (!ctx) { + LOG_ERROR("Failed to allocate IvfRabitqContext"); + return Context::Pointer(); + } + ailego::Params defaults; + defaults.set(PARAM_IVF_RABITQ_NPROBE, nprobe_); + defaults.set(PARAM_IVF_RABITQ_SCAN_RATIO, scan_ratio_); + defaults.set(PARAM_IVF_RABITQ_BRUTE_FORCE_THRESHOLD, brute_force_threshold_); + ctx->update(defaults); + return Context::Pointer(ctx); +} + +IndexProvider::Pointer IvfRabitqStreamer::create_provider(void) const { + if (state_ != STATE_LOADED) { + LOG_ERROR("Load the index first before create provider"); + return Provider::Pointer(); + } + if (!entity_ || entity_->quantized_vector_element_size() == 0) { + LOG_ERROR("Quantized vectors are not available for IVF RaBitQ provider"); + return Provider::Pointer(); + } + + auto *provider = new (std::nothrow) + IvfRabitqIndexProvider(meta_, entity_, "IvfRabitqStreamer"); + if (!provider) { + LOG_ERROR("Failed to alloc IvfRabitqIndexProvider"); + return Provider::Pointer(); + } + return Provider::Pointer(provider); +} + +const void *IvfRabitqStreamer::get_vector(uint64_t key) const { + if (!entity_) { + return nullptr; + } + return entity_->get_vector_by_key(key); +} + +int IvfRabitqStreamer::get_vector(const uint64_t key, + IndexStorage::MemoryBlock &block) const { + if (!entity_) { + return IndexError_NoReady; + } + return entity_->get_vector_by_key(key, block); +} + +// -------------------------------------------------------------------------- +// search_bf_impl (brute force - search all clusters, single query) +// -------------------------------------------------------------------------- +int IvfRabitqStreamer::search_bf_impl(const void *query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) const { + return search_impl_internal(query, qmeta, 1, context, true); +} + +// -------------------------------------------------------------------------- +// search_bf_impl (brute force - search all clusters, multi query) +// -------------------------------------------------------------------------- +int IvfRabitqStreamer::search_bf_impl(const void *query, + const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context) const { + return search_impl_internal(query, qmeta, count, context, true); +} + +// -------------------------------------------------------------------------- +// search_impl (single query) +// -------------------------------------------------------------------------- +int IvfRabitqStreamer::search_impl(const void *query, + const IndexQueryMeta &qmeta, + Context::Pointer &context) const { + return search_impl(query, qmeta, 1, context); +} + +// -------------------------------------------------------------------------- +// search_impl +// -------------------------------------------------------------------------- +int IvfRabitqStreamer::search_impl(const void *query, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + return search_impl_internal(query, qmeta, count, context, false); +} + +int IvfRabitqStreamer::search_impl_internal(const void *query, + const IndexQueryMeta &qmeta, + uint32_t count, + Context::Pointer &context, + bool force_brute_force) const { + if (!query) { + LOG_ERROR("Null query"); + return IndexError_InvalidArgument; + } + if (!reformer_ || !reformer_->loaded()) { + LOG_ERROR("Reformer not loaded for search"); + return IndexError_NoReady; + } + if (!entity_) { + LOG_ERROR("Entity not loaded for search"); + return IndexError_NoReady; + } + + IvfRabitqContext *ctx = dynamic_cast(context.get()); + if (!ctx) { + LOG_ERROR("Invalid context type"); + return IndexError_Cast; + } + + bool brute_force = force_brute_force || entity_->total_vector_count() <= + ctx->bruteforce_threshold(); + + uint32_t topk = ctx->topk(); + if (topk == 0) { + topk = 10; + } + uint32_t nprobe = 0; + uint32_t max_scan = 0; + if (brute_force) { + nprobe = entity_->cluster_count(); + max_scan = entity_->total_vector_count(); + ctx->set_search_limits(max_scan); + } else { + int ret = ctx->update_search_limits(entity_->total_vector_count(), + entity_->cluster_count(), &nprobe); + if (ret != 0) { + return ret; + } + max_scan = ctx->max_scan_count(); + } + + size_t padded_dim = reformer_->padded_dim(); + size_t ex_bits = reformer_->ex_bits(); + size_t dimension = reformer_->dimension(); + if (qmeta.dimension() != meta_.dimension() || + qmeta.data_type() != meta_.data_type() || + qmeta.element_size() != meta_.element_size()) { + LOG_ERROR("Unsupported query meta"); + return IndexError_Mismatch; + } + if (qmeta.dimension() < dimension) { + LOG_ERROR("Query dimension=%zu smaller than RaBitQ dimension=%zu", + static_cast(qmeta.dimension()), dimension); + return IndexError_Mismatch; + } + + // Reset results and heap for all queries + ctx->reset_results(count); + + for (uint32_t q = 0; q < count; ++q) { + const float *q_vec = reinterpret_cast( + static_cast(query) + + (static_cast(q) * qmeta.element_size())); + + // Create query state (rotate query and prepare per-query scan state) + IvfRabitqQueryState query_state; + int ret = reformer_->create_query_state(q_vec, &query_state); + if (ret != 0) { + LOG_ERROR("Failed to create query state, ret=%d", ret); + return ret; + } + + // Select probe centroids with the same metric used for build assignment. + std::vector probe_centroids; + ret = reformer_->select_probe_centroids(q_vec, nprobe, &query_state, + &probe_centroids); + if (ret != 0) { + LOG_ERROR("Failed to select probe centroids, ret=%d", ret); + return ret; + } + + // Use context heap for online distk-gated pruning + IndexDocumentHeap &heap = ctx->mutable_result_heap(); + heap.clear(); + heap.limit(topk); + const auto &filter = ctx->filter(); + + uint32_t scanned = 0; + + for (uint32_t p = 0; p < probe_centroids.size() && scanned < max_scan; + ++p) { + uint32_t cid = probe_centroids[p]; + + ret = reformer_->prepare_for_cluster(cid, &query_state); + if (ret != 0) { + LOG_ERROR("Failed to prepare for cluster %u, ret=%d", cid, ret); + continue; + } + + // Scan cluster with 1-bit lower-bound pruning before extra-bit boosting + if (!filter.is_valid()) { + ret = entity_->search_cluster(cid, query_state, padded_dim, ex_bits, + &heap); + } else { + ret = entity_->search_cluster(cid, query_state, padded_dim, ex_bits, + filter, &heap); + } + if (ret != 0) { + LOG_ERROR("Failed to search cluster %u, ret=%d", cid, ret); + continue; + } + + scanned += entity_->cluster_meta(cid).vector_count; + } + + // Drain heap into sorted result list for query q + ctx->topk_to_result(q); + } + + return 0; +} + +int IvfRabitqStreamer::unload() { + reformer_.reset(); + entity_.reset(); + storage_.reset(); + stats_.set_loaded_count(0UL); + stats_.set_loaded_costtime(0UL); + stats_.clear_attributes(); + state_ = STATE_INITED; + + return 0; +} + +// -------------------------------------------------------------------------- +// cleanup +// -------------------------------------------------------------------------- +int IvfRabitqStreamer::cleanup() { + LOG_INFO("IvfRabitqStreamer cleanup"); + + this->unload(); + params_.clear(); + nprobe_ = kDefaultIvfRabitqNprobe; + scan_ratio_ = kDefaultIvfRabitqScanRatio; + brute_force_threshold_ = kDefaultIvfRabitqBruteForceThreshold; + state_ = STATE_INIT; + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf_rabitq/ivf_rabitq_streamer.h b/src/core/algorithm/ivf_rabitq/ivf_rabitq_streamer.h new file mode 100644 index 000000000..f0642260f --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_streamer.h @@ -0,0 +1,134 @@ +// 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. +#pragma once + +#include +#include +#include "zvec/core/framework/index_streamer.h" +#include "ivf_rabitq_params.h" + +namespace zvec { +namespace core { + +class IvfRabitqReformer; +class IvfRabitqEntity; + +/*! IVF RaBitQ Streamer + * Combines IVF partitioning with RaBitQ quantization for fast approximate + * nearest neighbor search over a built index. + */ +class IvfRabitqStreamer : public IndexStreamer { + public: + IvfRabitqStreamer() = default; + ~IvfRabitqStreamer() override = default; + + //! Initialize Streamer + int init(const IndexMeta &meta, const ailego::Params ¶ms) override; + + //! Open index from storage + int open(IndexStorage::Pointer storage) override; + + //! Flush index + int flush(uint64_t /*check_point*/) override { + return 0; + } + + //! Close index + int close() override { + return this->unload(); + } + + //! Cleanup Streamer + int cleanup() override; + + //! Unload index + int unload() override; + + //! Retrieve statistics + const Stats &stats(void) const override { + return stats_; + } + + //! Retrieve meta + const IndexMeta &meta(void) const override { + return meta_; + } + + //! Create a search context + Context::Pointer create_context(void) const override; + + //! Create a new iterator + IndexProvider::Pointer create_provider(void) const override; + + //! Similarity search + int search_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const override; + + //! Similarity search (single query) + int search_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + + //! Brute force search (for ground truth generation) + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + Context::Pointer &context) const override; + int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context) const override; + + //! Fetch vector by key + const void *get_vector(uint64_t key) const override; + + int get_vector(const uint64_t key, + IndexStorage::MemoryBlock &block) const override; + + //! Fetch vector by id + const void *get_vector_by_id(uint32_t id) const override { + return get_vector(id); + } + + int get_vector_by_id(const uint32_t id, + IndexStorage::MemoryBlock &block) const override { + return get_vector(id, block); + } + + private: + //! Load index from storage + int load_index(IndexStorage::Pointer storage); + + int search_impl_internal(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, Context::Pointer &context, + bool force_brute_force) const; + + //! Internal state + enum State { STATE_INIT = 0, STATE_INITED = 1, STATE_LOADED = 2 }; + + IndexMeta meta_; + IndexMeta rabitq_meta_; + ailego::Params params_; + IndexStorage::Pointer storage_; + Stats stats_; + State state_{STATE_INIT}; + + // Core components + std::shared_ptr reformer_; + std::shared_ptr entity_; + + // Parameters + uint32_t nprobe_{kDefaultIvfRabitqNprobe}; + float scan_ratio_{kDefaultIvfRabitqScanRatio}; + uint32_t brute_force_threshold_{kDefaultIvfRabitqBruteForceThreshold}; + std::string metric_name_; +}; + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf_rabitq/ivf_rabitq_util.cc b/src/core/algorithm/ivf_rabitq/ivf_rabitq_util.cc new file mode 100644 index 000000000..03a5b7063 --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_util.cc @@ -0,0 +1,77 @@ +// 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. +#include "ivf_rabitq_util.h" +#include +#include "algorithm/hnsw_rabitq/rabitq_params.h" +#include "zvec/core/framework/index_error.h" + +namespace zvec { +namespace core { + +int PrepareAndCheckIvfRabitqInternalMeta(const IndexMeta &meta, + const ailego::Params ¶ms, + IndexMeta *rabitq_meta, + std::string *metric_name) { + if (!rabitq_meta || !metric_name) { + LOG_ERROR("Invalid IVF RaBitQ meta arguments"); + return IndexError_InvalidArgument; + } + + *rabitq_meta = meta; + *metric_name = meta.metric_name(); + if (metric_name->empty()) { + LOG_ERROR("Meta metric is empty"); + return IndexError_InvalidArgument; + } + + // Keep public meta unchanged. The internal RaBitQ meta describes only the + // dimensions consumed by centroids, rotator and quantized data. + uint32_t general_dimension = 0; + params.get(PARAM_RABITQ_GENERAL_DIMENSION, &general_dimension); + if (*metric_name == "Cosine" && general_dimension == 0) { + LOG_ERROR("%s not set for Cosine IVF RaBitQ", + PARAM_RABITQ_GENERAL_DIMENSION.c_str()); + return IndexError_InvalidArgument; + } + if (general_dimension > meta.dimension()) { + LOG_ERROR("Invalid general dimension=%zu, meta dimension=%zu", + static_cast(general_dimension), + static_cast(meta.dimension())); + return IndexError_InvalidArgument; + } + if (general_dimension > 0) { + rabitq_meta->set_dimension(general_dimension); + } + + if (*metric_name == "Cosine") { + rabitq_meta->set_metric("InnerProduct", 0, ailego::Params()); + *metric_name = "InnerProduct"; + } + + if (rabitq_meta->data_type() != IndexMeta::DataType::DT_FP32) { + LOG_ERROR("IVF RaBitQ only supports FP32 data type"); + return IndexError_Unsupported; + } + + uint32_t dim = rabitq_meta->dimension(); + if (dim < kMinRabitqDimSize || dim > kMaxRabitqDimSize) { + LOG_ERROR("Invalid dimension=%u, must be in [%d, %d]", dim, + kMinRabitqDimSize, kMaxRabitqDimSize); + return IndexError_InvalidArgument; + } + return 0; +} + +} // namespace core +} // namespace zvec diff --git a/src/core/algorithm/ivf_rabitq/ivf_rabitq_util.h b/src/core/algorithm/ivf_rabitq/ivf_rabitq_util.h new file mode 100644 index 000000000..3441d41cf --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_util.h @@ -0,0 +1,28 @@ +// 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. +#pragma once + +#include +#include "zvec/core/framework/index_meta.h" + +namespace zvec { +namespace core { + +int PrepareAndCheckIvfRabitqInternalMeta(const IndexMeta &meta, + const ailego::Params ¶ms, + IndexMeta *rabitq_meta, + std::string *metric_name); + +} // namespace core +} // namespace zvec diff --git a/src/core/interface/index_factory.cc b/src/core/interface/index_factory.cc index 5d9dea13e..e5f1a553e 100644 --- a/src/core/interface/index_factory.cc +++ b/src/core/interface/index_factory.cc @@ -47,6 +47,8 @@ Index::Pointer IndexFactory::CreateAndInitIndex(const BaseIndexParam ¶m) { ptr = std::make_shared(); } else if (param.index_type == IndexType::kHNSWRabitq) { ptr = std::make_shared(); + } else if (param.index_type == IndexType::kIVFRabitq) { + ptr = std::make_shared(); #if DISKANN_SUPPORTED } else if (param.index_type == IndexType::kDiskAnn) { ptr = std::make_shared(); @@ -121,6 +123,15 @@ BaseIndexParam::Pointer IndexFactory::DeserializeIndexParamFromJson( } return param; } + case IndexType::kIVFRabitq: { + IVFRabitqIndexParam::Pointer param = + std::make_shared(); + if (!param->DeserializeFromJson(json_str)) { + LOG_ERROR("Failed to deserialize ivf_rabitq index param"); + return nullptr; + } + return param; + } case IndexType::kVamana: { VamanaIndexParam::Pointer param = std::make_shared(); if (!param->DeserializeFromJson(json_str)) { @@ -187,6 +198,11 @@ std::string IndexFactory::QueryParamSerializeToJson(const QueryParamType ¶m, json_obj.set("ef_search", ailego::JsonValue(param.ef_search)); } index_type = IndexType::kHNSWRabitq; + } else if constexpr (std::is_same_v) { + if (!omit_empty_value || param.nprobe != 0) { + json_obj.set("nprobe", ailego::JsonValue(param.nprobe)); + } + index_type = IndexType::kIVFRabitq; } else if constexpr (std::is_same_v) { if (!omit_empty_value || param.ef_search != 0) { json_obj.set("ef_search", ailego::JsonValue(param.ef_search)); @@ -313,6 +329,17 @@ typename QueryParamType::Pointer IndexFactory::QueryParamDeserializeFromJson( return nullptr; } return param; + } else if (index_type == IndexType::kIVFRabitq) { + auto param = std::make_shared(); + if (!parse_common_fields(param)) { + return nullptr; + } + if (!extract_value_from_json(json_obj, "nprobe", param->nprobe, + tmp_json_value)) { + LOG_ERROR("Failed to deserialize nprobe"); + return nullptr; + } + return param; } else if (index_type == IndexType::kVamana) { auto param = std::make_shared(); if (!parse_common_fields(param)) { @@ -373,6 +400,12 @@ typename QueryParamType::Pointer IndexFactory::QueryParamDeserializeFromJson( LOG_ERROR("Failed to deserialize ef_search"); return nullptr; } + } else if constexpr (std::is_same_v) { + if (!extract_value_from_json(json_obj, "nprobe", param->nprobe, + tmp_json_value)) { + LOG_ERROR("Failed to deserialize nprobe"); + return nullptr; + } } else if constexpr (std::is_same_v) { if (!extract_value_from_json(json_obj, "ef_search", param->ef_search, tmp_json_value)) { @@ -411,5 +444,11 @@ template std::string IndexFactory::QueryParamSerializeToJson( const VamanaQueryParam ¶m, bool omit_empty_value); template VamanaQueryParam::Pointer IndexFactory::QueryParamDeserializeFromJson< VamanaQueryParam>(const std::string &json_str); +template std::string +IndexFactory::QueryParamSerializeToJson( + const IVFRabitqQueryParam ¶m, bool omit_empty_value); +template IVFRabitqQueryParam::Pointer +IndexFactory::QueryParamDeserializeFromJson( + const std::string &json_str); } // namespace zvec::core_interface diff --git a/src/core/interface/index_param.cc b/src/core/interface/index_param.cc index 5d75276fd..017fb4439 100644 --- a/src/core/interface/index_param.cc +++ b/src/core/interface/index_param.cc @@ -183,6 +183,35 @@ ailego::JsonObject HNSWRabitqIndexParam::SerializeToJsonObject( return json_obj; } +bool IVFRabitqIndexParam::DeserializeFromJsonObject( + const ailego::JsonObject &json_obj) { + if (!BaseIndexParam::DeserializeFromJsonObject(json_obj)) { + return false; + } + + if (index_type != IndexType::kIVFRabitq) { + LOG_ERROR("index_type is not kIVFRabitq"); + return false; + } + + DESERIALIZE_VALUE_FIELD(json_obj, nlist); + DESERIALIZE_VALUE_FIELD(json_obj, total_bits); + DESERIALIZE_VALUE_FIELD(json_obj, sample_count); + + return true; +} + +ailego::JsonObject IVFRabitqIndexParam::SerializeToJsonObject( + bool omit_empty_value) const { + auto json_obj = BaseIndexParam::SerializeToJsonObject(omit_empty_value); + json_obj.set("nlist", ailego::JsonValue(nlist)); + json_obj.set("total_bits", ailego::JsonValue(total_bits)); + if (!omit_empty_value || sample_count != 0) { + json_obj.set("sample_count", ailego::JsonValue(sample_count)); + } + return json_obj; +} + ailego::JsonObject VamanaIndexParam::SerializeToJsonObject( bool omit_empty_value) const { auto json_obj = BaseIndexParam::SerializeToJsonObject(omit_empty_value); @@ -340,4 +369,4 @@ bool QuantizerParam::DeserializeFromJsonObject( } // namespace core_interface -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/interface/indexes/ivf_rabitq_index.cc b/src/core/interface/indexes/ivf_rabitq_index.cc new file mode 100644 index 000000000..8c0f8c84a --- /dev/null +++ b/src/core/interface/indexes/ivf_rabitq_index.cc @@ -0,0 +1,346 @@ +// 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. + +#include +#include +#include +#include +#include "zvec/ailego/io/file.h" +#include "zvec/core/framework/index_error.h" +#include "holder_builder.h" + +#if RABITQ_SUPPORTED +#include "algorithm/hnsw_rabitq/hnsw_rabitq_params.h" +#include "algorithm/hnsw_rabitq/rabitq_params.h" +#include "algorithm/ivf_rabitq/ivf_rabitq_params.h" +#include "algorithm/ivf_rabitq/ivf_rabitq_streamer.h" +#endif + +namespace zvec::core_interface { + +int IVFRabitqIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) { +#if !RABITQ_SUPPORTED + (void)param; + LOG_ERROR("RaBitQ is not supported on this platform (Linux x86_64 only)"); + return core::IndexError_Unsupported; +#else + param_ = dynamic_cast(param); + + if (is_sparse_) { + LOG_ERROR("Sparse index is not supported for IVF RaBitQ"); + return core::IndexError_Runtime; + } + + if (param.dimension < core::kMinRabitqDimSize || + param.dimension > core::kMaxRabitqDimSize) { + LOG_ERROR("Unsupported dimension: %d, must be in [%d, %d]", param.dimension, + core::kMinRabitqDimSize, core::kMaxRabitqDimSize); + return core::IndexError_Unsupported; + } + + proxima_index_params_.set(core::PARAM_IVF_RABITQ_NLIST, param_.nlist); + proxima_index_params_.set(core::PARAM_RABITQ_TOTAL_BITS, param_.total_bits); + proxima_index_params_.set(core::PARAM_RABITQ_SAMPLE_COUNT, + param_.sample_count); + // Pass original dimension so builder can ignore extra dim from converter + proxima_index_params_.set(core::PARAM_RABITQ_GENERAL_DIMENSION, + input_vector_meta_.dimension()); + + // Create builder (for train/build/dump) + builder_ = core::IndexFactory::CreateBuilder("IvfRabitqBuilder"); + if (ailego_unlikely(!builder_)) { + LOG_ERROR("Failed to create IvfRabitqBuilder"); + return core::IndexError_Runtime; + } + if (ailego_unlikely( + builder_->init(proxima_index_meta_, proxima_index_params_) != 0)) { + LOG_ERROR("Failed to init IvfRabitqBuilder"); + return core::IndexError_Runtime; + } + + // Create streamer (for search) + auto streamer = std::make_shared(); + streamer_ = streamer; + if (ailego_unlikely(!streamer_)) { + LOG_ERROR("Failed to create IvfRabitqStreamer"); + return core::IndexError_Runtime; + } + if (ailego_unlikely( + streamer_->init(proxima_index_meta_, proxima_index_params_) != 0)) { + LOG_ERROR("Failed to init IvfRabitqStreamer"); + return core::IndexError_Runtime; + } + return 0; +#endif // RABITQ_SUPPORTED +} + +int IVFRabitqIndex::Open(const std::string &file_path, + StorageOptions storage_options) { +#if !RABITQ_SUPPORTED + (void)file_path; + (void)storage_options; + LOG_ERROR("RaBitQ is not supported on this platform (Linux x86_64 only)"); + return core::IndexError_Unsupported; +#else + ailego::Params storage_params; + file_path_ = file_path; + is_read_only_ = storage_options.read_only; + + // Use MMapFileReadStorage (consistent with IVFIndex, read-only after builder + // dump) + storage_ = core::IndexFactory::CreateStorage("MMapFileReadStorage"); + if (storage_ == nullptr) { + LOG_ERROR("Failed to create MMapFileReadStorage"); + return core::IndexError_Runtime; + } + int ret = storage_->init(storage_params); + if (ret != 0) { + LOG_ERROR("Failed to init MMapFileReadStorage, path: %s, err: %s", + file_path_.c_str(), core::IndexError::What(ret)); + return ret; + } + + if (is_read_only_ || !storage_options.create_new) { + ret = storage_->open(file_path_, false); + if (ret != 0) { + LOG_ERROR("Failed to open storage, path: %s, err: %s", file_path_.c_str(), + core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + if (streamer_ == nullptr || streamer_->open(storage_) != 0) { + LOG_ERROR("Failed to open streamer, path: %s", file_path_.c_str()); + return core::IndexError_Runtime; + } + is_trained_ = true; + } + is_open_ = true; + return 0; +#endif // RABITQ_SUPPORTED +} + +int IVFRabitqIndex::GenerateHolder() { +#if !RABITQ_SUPPORTED + LOG_ERROR("RaBitQ is not supported on this platform (Linux x86_64 only)"); + return core::IndexError_Unsupported; +#else + return BuildMultiPassHolder(param_.data_type, param_.dimension, doc_cache_, + converter_, &holder_); +#endif // RABITQ_SUPPORTED +} + +int IVFRabitqIndex::Add(const VectorData &vector, uint32_t doc_id) { +#if !RABITQ_SUPPORTED + (void)vector; + (void)doc_id; + LOG_ERROR("RaBitQ is not supported on this platform (Linux x86_64 only)"); + return core::IndexError_Unsupported; +#else + if (is_trained_) { + LOG_ERROR("this IVF RaBitQ index is trained"); + return core::IndexError_Runtime; + } + if (!std::holds_alternative(vector.vector)) { + LOG_ERROR("Invalid vector data"); + return core::IndexError_Runtime; + } + const DenseVector &dense_vector = std::get(vector.vector); + std::string out_vector_buffer = std::string( + static_cast(dense_vector.data), + input_vector_meta_.dimension() * input_vector_meta_.unit_size()); + + std::lock_guard lock(mutex_); + while (doc_cache_.size() <= doc_id) { + std::string fake_data( + input_vector_meta_.dimension() * input_vector_meta_.unit_size(), 0); + doc_cache_.push_back(std::make_pair(kInvalidKey, fake_data)); + } + doc_cache_[doc_id] = std::make_pair(doc_id, out_vector_buffer); + return 0; +#endif // RABITQ_SUPPORTED +} + +int IVFRabitqIndex::Train() { +#if !RABITQ_SUPPORTED + LOG_ERROR("RaBitQ is not supported on this platform (Linux x86_64 only)"); + return core::IndexError_Unsupported; +#else + int ret = GenerateHolder(); + if (ret != 0) { + LOG_ERROR("Failed to generate holder"); + return core::IndexError_Runtime; + } + + // Train centroids + rotator + ret = builder_->train(holder_); + if (ret != 0) { + LOG_ERROR("Failed to train IVF RaBitQ index, ret=%d", ret); + return core::IndexError_Runtime; + } + + // Build index (assign vectors to centroids, quantize) + ret = builder_->build(holder_); + if (ret != 0) { + LOG_ERROR("Failed to build IVF RaBitQ index, ret=%d", ret); + return core::IndexError_Runtime; + } + + // Dump to file + auto dumper = core::IndexFactory::CreateDumper("FileDumper"); + if (!dumper) { + LOG_ERROR("Failed to create FileDumper"); + return core::IndexError_Runtime; + } + ret = dumper->create(file_path_); + if (ret != 0) { + LOG_ERROR("Failed to create dumper at path: %s", file_path_.c_str()); + return core::IndexError_Runtime; + } + ret = builder_->dump(dumper); + if (ret != 0) { + LOG_ERROR("Failed to dump IVF RaBitQ index, ret=%d", ret); + return core::IndexError_Runtime; + } + dumper->close(); + + // Reopen storage + streamer + ret = storage_->open(file_path_, false); + if (ret != 0) { + LOG_ERROR("Failed to open storage, path: %s, err: %s", file_path_.c_str(), + core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + if (streamer_ == nullptr || streamer_->open(storage_) != 0) { + LOG_ERROR("Failed to open streamer, path: %s", file_path_.c_str()); + return core::IndexError_Runtime; + } + + is_trained_ = true; + return 0; +#endif // RABITQ_SUPPORTED +} + +int IVFRabitqIndex::Merge(const std::vector &indexes, + const IndexFilter &filter, + const MergeOptions &options) { +#if !RABITQ_SUPPORTED + (void)indexes; + (void)filter; + (void)options; + LOG_ERROR("RaBitQ is not supported on this platform (Linux x86_64 only)"); + return core::IndexError_Unsupported; +#else + int ret = Index::Merge(indexes, filter, options); + if (ret != 0) { + return ret; + } + + auto dumper = core::IndexFactory::CreateDumper("FileDumper"); + if (!dumper) { + LOG_ERROR("Failed to create FileDumper"); + return core::IndexError_Runtime; + } + ret = dumper->create(file_path_); + if (ret != 0) { + LOG_ERROR("Failed to create dumper at path: %s", file_path_.c_str()); + return core::IndexError_Runtime; + } + ret = builder_->dump(dumper); + if (ret != 0) { + LOG_ERROR("Failed to dump IVF RaBitQ index, ret=%d", ret); + return core::IndexError_Runtime; + } + ret = dumper->close(); + if (ret != 0) { + LOG_ERROR("Failed to close dumper at path: %s", file_path_.c_str()); + return core::IndexError_Runtime; + } + + ret = storage_->open(file_path_, false); + if (ret != 0) { + LOG_ERROR("Failed to open storage, path: %s, err: %s", file_path_.c_str(), + core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + if (streamer_ == nullptr || streamer_->open(storage_) != 0) { + LOG_ERROR("Failed to open streamer, path: %s", file_path_.c_str()); + return core::IndexError_Runtime; + } + + is_trained_ = true; + return 0; +#endif // RABITQ_SUPPORTED +} + +int IVFRabitqIndex::_dense_fetch(const uint32_t doc_id, + VectorDataBuffer *vector_data_buffer) { +#if !RABITQ_SUPPORTED + (void)doc_id; + (void)vector_data_buffer; + LOG_ERROR("RaBitQ is not supported on this platform (Linux x86_64 only)"); + return core::IndexError_Unsupported; +#else + if (is_trained_) { + return Index::_dense_fetch(doc_id, vector_data_buffer); + } else { + DenseVectorBuffer dense_vector_buffer; + std::string &out_vector_buffer = dense_vector_buffer.data; + out_vector_buffer = doc_cache_[doc_id].second; + vector_data_buffer->vector_buffer = std::move(dense_vector_buffer); + return 0; + } +#endif // RABITQ_SUPPORTED +} + +int IVFRabitqIndex::_prepare_for_search( + const VectorData & /*vector_data*/, + const BaseIndexQueryParam::Pointer &search_param, + core::IndexContext::Pointer &context) { +#if !RABITQ_SUPPORTED + (void)search_param; + (void)context; + LOG_ERROR("RaBitQ is not supported on this platform (Linux x86_64 only)"); + return core::IndexError_Unsupported; +#else + const auto &ivf_rabitq_param = + std::dynamic_pointer_cast(search_param); + + context->set_topk(search_param->topk); + context->set_fetch_vector(search_param->fetch_vector); + if (search_param->filter) { + context->set_filter(std::move(*search_param->filter)); + } + if (search_param->radius > 0.0f) { + context->set_threshold(search_param->radius); + } + if (ivf_rabitq_param && ivf_rabitq_param->nprobe > 0) { + ailego::Params params; + params.set(core::PARAM_IVF_RABITQ_NPROBE, ivf_rabitq_param->nprobe); + context->update(params); + } + return 0; +#endif // RABITQ_SUPPORTED +} + +int IVFRabitqIndex::_get_coarse_search_topk( + const BaseIndexQueryParam::Pointer &search_param) { +#if !RABITQ_SUPPORTED + (void)search_param; + LOG_ERROR("RaBitQ is not supported on this platform (Linux x86_64 only)"); + return -1; +#else + return search_param->topk; +#endif // RABITQ_SUPPORTED +} + +} // namespace zvec::core_interface diff --git a/src/core/utility/mmap_file_read_storage.cc b/src/core/utility/mmap_file_read_storage.cc index 04605bf55..e90968769 100644 --- a/src/core/utility/mmap_file_read_storage.cc +++ b/src/core/utility/mmap_file_read_storage.cc @@ -237,7 +237,9 @@ class MMapFileReadStorage : public IndexStorage { } int close(void) override { - file_ptr_->close(); + if (file_ptr_) { + file_ptr_->close(); + } file_ptr_ = nullptr; segments_.clear(); return 0; @@ -294,4 +296,4 @@ class MMapFileReadStorage : public IndexStorage { INDEX_FACTORY_REGISTER_STORAGE(MMapFileReadStorage); } // namespace core -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/db/index/column/vector_column/engine_helper.hpp b/src/db/index/column/vector_column/engine_helper.hpp index 7fa0c16d3..0a4563de7 100644 --- a/src/db/index/column/vector_column/engine_helper.hpp +++ b/src/db/index/column/vector_column/engine_helper.hpp @@ -207,6 +207,25 @@ class ProximaEngineHelper { return std::move(ivf_query_param); } + case IndexType::IVF_RABITQ: { + auto ivf_rabitq_query_param_result = + _build_common_query_param( + query_params); + if (!ivf_rabitq_query_param_result.has_value()) { + return tl::make_unexpected(Status::InvalidArgument( + "failed to build query param: " + + ivf_rabitq_query_param_result.error().message())); + } + auto &ivf_rabitq_query_param = ivf_rabitq_query_param_result.value(); + if (query_params.query_params) { + auto db_ivf_rabitq_query_params = + dynamic_cast( + query_params.query_params.get()); + ivf_rabitq_query_param->nprobe = db_ivf_rabitq_query_params->nprobe(); + } + return std::move(ivf_rabitq_query_param); + } + case IndexType::DISKANN: { auto diskann_query_param_result = _build_common_query_param( @@ -446,6 +465,29 @@ class ProximaEngineHelper { return index_param_builder->Build(); } + case IndexType::IVF_RABITQ: { + auto index_param_builder_result = _build_common_index_param< + IvfRabitqIndexParams, core_interface::IVFRabitqIndexParamBuilder>( + field_schema); + if (!index_param_builder_result.has_value()) { + return tl::make_unexpected(Status::InvalidArgument( + "failed to build index param: " + + index_param_builder_result.error().message())); + } + auto index_param_builder = index_param_builder_result.value(); + + auto db_index_params = dynamic_cast( + field_schema.index_params().get()); + index_param_builder->WithNlist(db_index_params->nlist()); + index_param_builder->WithTotalBits(db_index_params->total_bits()); + index_param_builder->WithSampleCount(db_index_params->sample_count()); + index_param_builder->WithProvider( + db_index_params->raw_vector_provider()); + index_param_builder->WithReformer(db_index_params->rabitq_reformer()); + + return index_param_builder->Build(); + } + case IndexType::DISKANN: { auto index_param_builder_result = _build_common_index_param( + MetricTypeCodeBook::Get(params_pb.base().metric_type()), + params_pb.nlist(), params_pb.total_bits(), params_pb.sample_count()); + return params; +} + +proto::IvfRabitqIndexParams ProtoConverter::ToPb( + const IvfRabitqIndexParams *params) { + proto::IvfRabitqIndexParams params_pb; + params_pb.mutable_base()->set_metric_type( + MetricTypeCodeBook::Get(params->metric_type())); + params_pb.mutable_base()->set_quantize_type( + QuantizeTypeCodeBook::Get(params->quantize_type())); + params_pb.set_nlist(params->nlist()); + params_pb.set_total_bits(params->total_bits()); + params_pb.set_sample_count(params->sample_count()); + return params_pb; +} + // FlatIndexParams FlatIndexParams::OPtr ProtoConverter::FromPb( const proto::FlatIndexParams ¶ms_pb) { @@ -276,6 +298,8 @@ IndexParams::Ptr ProtoConverter::FromPb(const proto::IndexParams ¶ms_pb) { return ProtoConverter::FromPb(params_pb.flat()); } else if (params_pb.has_hnsw_rabitq()) { return ProtoConverter::FromPb(params_pb.hnsw_rabitq()); + } else if (params_pb.has_ivf_rabitq()) { + return ProtoConverter::FromPb(params_pb.ivf_rabitq()); } else if (params_pb.has_diskann()) { return ProtoConverter::FromPb(params_pb.diskann()); } else if (params_pb.has_vamana()) { @@ -345,6 +369,15 @@ proto::IndexParams ProtoConverter::ToPb(const IndexParams *params) { } break; } + case IndexType::IVF_RABITQ: { + auto ivf_rabitq_params = + dynamic_cast(params); + if (ivf_rabitq_params) { + params_pb.mutable_ivf_rabitq()->CopyFrom( + ProtoConverter::ToPb(ivf_rabitq_params)); + } + break; + } case IndexType::DISKANN: { auto diskann_params = dynamic_cast(params); if (diskann_params) { @@ -436,4 +469,4 @@ proto::SegmentMeta ProtoConverter::ToPb(const SegmentMeta &meta) { return meta_pb; } -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/db/index/common/proto_converter.h b/src/db/index/common/proto_converter.h index e0f015cb1..1bd5be628 100644 --- a/src/db/index/common/proto_converter.h +++ b/src/db/index/common/proto_converter.h @@ -30,6 +30,11 @@ struct ProtoConverter { const proto::HnswRabitqIndexParams ¶ms_pb); static proto::HnswRabitqIndexParams ToPb(const HnswRabitqIndexParams *params); + // IvfRabitqIndexParams + static IvfRabitqIndexParams::OPtr FromPb( + const proto::IvfRabitqIndexParams ¶ms_pb); + static proto::IvfRabitqIndexParams ToPb(const IvfRabitqIndexParams *params); + // FlatIndexParams static FlatIndexParams::OPtr FromPb(const proto::FlatIndexParams ¶ms_pb); static proto::FlatIndexParams ToPb(const FlatIndexParams *params); diff --git a/src/db/index/common/schema.cc b/src/db/index/common/schema.cc index 9471eee90..8b4d60f0b 100644 --- a/src/db/index/common/schema.cc +++ b/src/db/index/common/schema.cc @@ -55,8 +55,9 @@ std::unordered_set support_sparse_vector_type = { }; std::unordered_set support_dense_vector_index = { - IndexType::FLAT, IndexType::HNSW, IndexType::HNSW_RABITQ, - IndexType::IVF, IndexType::DISKANN, IndexType::VAMANA}; + IndexType::FLAT, IndexType::HNSW, IndexType::HNSW_RABITQ, + IndexType::IVF, IndexType::IVF_RABITQ, IndexType::DISKANN, + IndexType::VAMANA}; std::unordered_set support_sparse_vector_index = {IndexType::FLAT, IndexType::HNSW}; @@ -136,23 +137,24 @@ Status FieldSchema::validate() const { } } - if (index_params_->type() == IndexType::HNSW_RABITQ) { + if (index_params_->type() == IndexType::HNSW_RABITQ || + index_params_->type() == IndexType::IVF_RABITQ) { if (dimension_ < kMinRabitqDimSize || dimension_ > kMaxRabitqDimSize) { return Status::InvalidArgument( - "schema validate failed: HNSW_RABITQ index only support " + "schema validate failed: RabitQ index only support " "dimension in [", kMinRabitqDimSize, ", ", kMaxRabitqDimSize, "]"); } if (data_type_ != DataType::VECTOR_FP32) { return Status::InvalidArgument( - "schema validate failed: HNSW_RABITQ index only support FP32 " + "schema validate failed: RabitQ index only support FP32 " "data types"); } auto metric_type = vector_index_params->metric_type(); if (metric_type != MetricType::L2 && metric_type != MetricType::IP && metric_type != MetricType::COSINE) { return Status::InvalidArgument( - "schema validate failed: HNSW_RABITQ index only support " + "schema validate failed: RabitQ index only support " "L2/IP/COSINE metric"); } #if !RABITQ_SUPPORTED diff --git a/src/db/index/common/type_helper.h b/src/db/index/common/type_helper.h index d24f1ee52..dd93dcd19 100644 --- a/src/db/index/common/type_helper.h +++ b/src/db/index/common/type_helper.h @@ -46,6 +46,8 @@ struct IndexTypeCodeBook { return IndexType::HNSW; case proto::IT_HNSW_RABITQ: return IndexType::HNSW_RABITQ; + case proto::IT_IVF_RABITQ: + return IndexType::IVF_RABITQ; case proto::IT_FLAT: return IndexType::FLAT; case proto::IT_IVF: @@ -71,6 +73,8 @@ struct IndexTypeCodeBook { return proto::IT_HNSW; case IndexType::HNSW_RABITQ: return proto::IT_HNSW_RABITQ; + case IndexType::IVF_RABITQ: + return proto::IT_IVF_RABITQ; case IndexType::FLAT: return proto::IT_FLAT; case IndexType::IVF: @@ -96,6 +100,8 @@ struct IndexTypeCodeBook { return "HNSW"; case IndexType::HNSW_RABITQ: return "HNSW_RABITQ"; + case IndexType::IVF_RABITQ: + return "IVF_RABITQ"; case IndexType::FLAT: return "FLAT"; case IndexType::IVF: diff --git a/src/db/index/segment/segment.cc b/src/db/index/segment/segment.cc index e21ececa2..4d83dd572 100644 --- a/src/db/index/segment/segment.cc +++ b/src/db/index/segment/segment.cc @@ -1781,7 +1781,7 @@ Status SegmentImpl::create_vector_index( field_with_new_index_params->set_index_params(index_params); // For RABITQ, PrepareQuantizeField trains a reformer with - // raw_vector_provider and attaches it to a cloned HnswRabitqIndexParams. + // raw_vector_provider and attaches it to cloned RabitQ index params. // For other quantize types, field_with_new_index_params is reused as-is. std::shared_ptr field_for_quantize; { diff --git a/src/db/index/segment/segment_helper.cc b/src/db/index/segment/segment_helper.cc index 1238a4865..b24411a8d 100644 --- a/src/db/index/segment/segment_helper.cc +++ b/src/db/index/segment/segment_helper.cc @@ -759,8 +759,8 @@ namespace { // Only the first indexer's file is reused as the merge base; the remaining // indexers are merged in via Merge(). Reuse is restricted to streaming -// indexes (HNSW, FLAT). Builder-rebuild indexes (IVF, VAMANA) and HNSW_RABITQ -// fall back to the full-rebuild merge. +// indexes (HNSW, FLAT). Builder-rebuild indexes (IVF, VAMANA) and HNSW_RABITQ, +// IVF_RABITQ fall back to the full-rebuild merge. bool CanReuseFirstIndexer(const std::vector &indexers, const FieldSchema &output_field, const IndexFilter::Ptr &filter) { @@ -876,10 +876,13 @@ Status SegmentHelper::PrepareQuantizeField( "raw_vector_provider is required for RABITQ training"); } - auto rabitq_params = std::dynamic_pointer_cast( - vector_index_params->clone()); - if (!rabitq_params) { - return Status::InternalError("Expect HnswRabitqIndexParams"); + auto cloned_index_params = vector_index_params->clone(); + auto hnsw_rabitq_params = + std::dynamic_pointer_cast(cloned_index_params); + auto ivf_rabitq_params = + std::dynamic_pointer_cast(cloned_index_params); + if (!hnsw_rabitq_params && !ivf_rabitq_params) { + return Status::InternalError("Expect RabitQ index params"); } auto converter = core::IndexFactory::CreateConverter("RabitqConverter"); @@ -897,13 +900,23 @@ Status SegmentHelper::PrepareQuantizeField( .value(), false), 0, ailego::Params{}); + int total_bits = 0; + int num_clusters = 0; + int sample_count = 0; + if (hnsw_rabitq_params) { + total_bits = hnsw_rabitq_params->total_bits(); + num_clusters = hnsw_rabitq_params->num_clusters(); + sample_count = hnsw_rabitq_params->sample_count(); + } else { + total_bits = ivf_rabitq_params->total_bits(); + num_clusters = ivf_rabitq_params->nlist(); + sample_count = ivf_rabitq_params->sample_count(); + } + ailego::Params converter_params; - converter_params.set(core::PARAM_RABITQ_TOTAL_BITS, - rabitq_params->total_bits()); - converter_params.set(core::PARAM_RABITQ_NUM_CLUSTERS, - rabitq_params->num_clusters()); - converter_params.set(core::PARAM_RABITQ_SAMPLE_COUNT, - rabitq_params->sample_count()); + converter_params.set(core::PARAM_RABITQ_TOTAL_BITS, total_bits); + converter_params.set(core::PARAM_RABITQ_NUM_CLUSTERS, num_clusters); + converter_params.set(core::PARAM_RABITQ_SAMPLE_COUNT, sample_count); if (int ret = converter->init(index_meta, converter_params); ret != 0) { return Status::InternalError("Failed to init rabitq converter:", ret); } @@ -914,10 +927,15 @@ Status SegmentHelper::PrepareQuantizeField( if (int ret = converter->to_reformer(&reformer); ret != 0) { return Status::InternalError("Failed to to get rabitq reformer:", ret); } - rabitq_params->set_rabitq_reformer(reformer); - rabitq_params->set_raw_vector_provider(raw_vector_provider); + if (hnsw_rabitq_params) { + hnsw_rabitq_params->set_rabitq_reformer(reformer); + hnsw_rabitq_params->set_raw_vector_provider(raw_vector_provider); + } else { + ivf_rabitq_params->set_rabitq_reformer(reformer); + ivf_rabitq_params->set_raw_vector_provider(raw_vector_provider); + } - field_clone->set_index_params(rabitq_params); + field_clone->set_index_params(cloned_index_params); *out_field = field_clone; return Status::OK(); #endif diff --git a/src/db/index/segment/segment_helper.h b/src/db/index/segment/segment_helper.h index ab8c3a9df..905b1926b 100644 --- a/src/db/index/segment/segment_helper.h +++ b/src/db/index/segment/segment_helper.h @@ -272,7 +272,7 @@ class SegmentHelper { // Returns a FieldSchema clone whose index_params is ready for building the // quantize indexer. - // - RABITQ: clones HnswRabitqIndexParams, trains a RabitqConverter against + // - RABITQ: clones RabitQ index params, trains a RabitqConverter against // `raw_vector_provider`, and attaches the resulting reformer and raw // provider to the cloned params. // - Other quantize types: clones the field with its current index_params diff --git a/src/db/proto/zvec.proto b/src/db/proto/zvec.proto index f2c18f5ad..f8c885c1b 100644 --- a/src/db/proto/zvec.proto +++ b/src/db/proto/zvec.proto @@ -62,6 +62,8 @@ enum IndexType { IT_VAMANA = 5; // Proxima DiskAnn Index IT_DISKANN = 6; + // Proxima IVF RABITQ Index + IT_IVF_RABITQ = 7; // Invert Index IT_INVERT = 10; // Full-Text Search Index @@ -121,6 +123,13 @@ message HnswRabitqIndexParams { int32 sample_count = 6; } +message IvfRabitqIndexParams { + BaseIndexParams base = 1; + int32 nlist = 2; + int32 total_bits = 3; + int32 sample_count = 4; +} + message FlatIndexParams { BaseIndexParams base = 1; } @@ -168,6 +177,7 @@ message IndexParams { VamanaIndexParams vamana = 6; FtsIndexParams fts = 7; DiskAnnIndexParams diskann = 8; + IvfRabitqIndexParams ivf_rabitq = 9; }; }; diff --git a/src/include/zvec/c_api.h b/src/include/zvec/c_api.h index 98828d3b9..152dc5fc1 100644 --- a/src/include/zvec/c_api.h +++ b/src/include/zvec/c_api.h @@ -827,6 +827,7 @@ typedef uint32_t zvec_index_type_t; #define ZVEC_INDEX_TYPE_IVF 2 #define ZVEC_INDEX_TYPE_FLAT 3 #define ZVEC_INDEX_TYPE_VAMANA 6 +#define ZVEC_INDEX_TYPE_IVF_RABITQ 7 #define ZVEC_INDEX_TYPE_INVERT 10 #define ZVEC_INDEX_TYPE_FTS 11 @@ -856,6 +857,7 @@ typedef uint32_t zvec_quantize_type_t; #define ZVEC_QUANTIZE_TYPE_FP16 1 #define ZVEC_QUANTIZE_TYPE_INT8 2 #define ZVEC_QUANTIZE_TYPE_INT4 3 +#define ZVEC_QUANTIZE_TYPE_RABITQ 4 // ============================================================================= // Collection Structures (Opaque Pointer Pattern) @@ -1065,6 +1067,29 @@ ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_index_params_get_ivf_params( const zvec_index_params_t *params, int *out_n_list, int *out_n_iters, bool *out_use_soar); +/** + * @brief Set IVF RaBitQ specific parameters + * @param params Index parameters (must be IVF_RABITQ type) + * @param nlist Number of cluster centers + * @param total_bits Total bits for RaBitQ quantization + * @param sample_count Sample count for training, 0 means use all vectors + * @return ZVEC_OK on success, error code on failure + */ +ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_index_params_set_ivf_rabitq_params( + zvec_index_params_t *params, int nlist, int total_bits, int sample_count); + +/** + * @brief Get IVF RaBitQ parameters (all at once) + * @param params Index parameters (must be IVF_RABITQ type) + * @param out_nlist Output parameter for nlist + * @param out_total_bits Output parameter for total_bits + * @param out_sample_count Output parameter for sample_count + * @return ZVEC_OK on success, error code on failure + */ +ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_index_params_get_ivf_rabitq_params( + const zvec_index_params_t *params, int *out_nlist, int *out_total_bits, + int *out_sample_count); + /** * @brief Get invert index parameters (all at once) * @param params Index parameters (must not be NULL) @@ -1140,6 +1165,16 @@ typedef struct zvec_hnsw_query_params_t zvec_hnsw_query_params_t; */ typedef struct zvec_ivf_query_params_t zvec_ivf_query_params_t; +/** + * @brief IVF RaBitQ query parameters handle (opaque pointer) + * + * Internally maps to zvec::IvfRabitqQueryParams* (raw pointer). + * Created by zvec_query_params_ivf_rabitq_create() and destroyed by + * zvec_query_params_ivf_rabitq_destroy(). Caller owns the pointer and must + * explicitly destroy it. + */ +typedef struct zvec_ivf_rabitq_query_params_t zvec_ivf_rabitq_query_params_t; + /** * @brief Flat query parameters handle (opaque pointer) * @@ -1428,6 +1463,100 @@ zvec_query_params_ivf_set_is_using_refiner(zvec_ivf_query_params_t *params, ZVEC_EXPORT bool ZVEC_CALL zvec_query_params_ivf_get_is_using_refiner( const zvec_ivf_query_params_t *params); +// ----------------------------------------------------------------------------- +// zvec_ivf_rabitq_query_params_t (IVF RaBitQ Query Parameters) +// ----------------------------------------------------------------------------- + +/** + * @brief Create IVF RaBitQ query parameters + * @param nprobe Number of clusters to probe (default: 10) + * @param radius Search radius (default: 0.0) + * @param is_linear Whether linear search (default: false) + * @param is_using_refiner Whether using refiner (default: false) + * @return zvec_ivf_rabitq_query_params_t* Pointer to the newly created IVF + * RaBitQ query parameters + */ +ZVEC_EXPORT zvec_ivf_rabitq_query_params_t *ZVEC_CALL +zvec_query_params_ivf_rabitq_create(int nprobe, float radius, bool is_linear, + bool is_using_refiner); + +/** + * @brief Destroy IVF RaBitQ query parameters + * @param params IVF RaBitQ query parameters pointer + */ +ZVEC_EXPORT void ZVEC_CALL +zvec_query_params_ivf_rabitq_destroy(zvec_ivf_rabitq_query_params_t *params); + +/** + * @brief Set number of probe clusters + * @param params IVF RaBitQ query parameters pointer + * @param nprobe Number of probe clusters + * @return zvec_error_code_t Error code + */ +ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_query_params_ivf_rabitq_set_nprobe( + zvec_ivf_rabitq_query_params_t *params, int nprobe); + +/** + * @brief Get number of probe clusters + * @param params IVF RaBitQ query parameters pointer + * @return int Number of probe clusters + */ +ZVEC_EXPORT int ZVEC_CALL zvec_query_params_ivf_rabitq_get_nprobe( + const zvec_ivf_rabitq_query_params_t *params); + +/** + * @brief Set search radius + * @param params IVF RaBitQ query parameters pointer + * @param radius Search radius + * @return zvec_error_code_t Error code + */ +ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_query_params_ivf_rabitq_set_radius( + zvec_ivf_rabitq_query_params_t *params, float radius); + +/** + * @brief Get search radius + * @param params IVF RaBitQ query parameters pointer + * @return float Search radius + */ +ZVEC_EXPORT float ZVEC_CALL zvec_query_params_ivf_rabitq_get_radius( + const zvec_ivf_rabitq_query_params_t *params); + +/** + * @brief Set linear search mode + * @param params IVF RaBitQ query parameters pointer + * @param is_linear Whether linear search + * @return zvec_error_code_t Error code + */ +ZVEC_EXPORT zvec_error_code_t ZVEC_CALL +zvec_query_params_ivf_rabitq_set_is_linear( + zvec_ivf_rabitq_query_params_t *params, bool is_linear); + +/** + * @brief Get linear search mode + * @param params IVF RaBitQ query parameters pointer + * @return bool Whether linear search + */ +ZVEC_EXPORT bool ZVEC_CALL zvec_query_params_ivf_rabitq_get_is_linear( + const zvec_ivf_rabitq_query_params_t *params); + +/** + * @brief Set whether to use refiner + * @param params IVF RaBitQ query parameters pointer + * @param is_using_refiner Whether to use refiner + * @return zvec_error_code_t Error code + */ +ZVEC_EXPORT zvec_error_code_t ZVEC_CALL +zvec_query_params_ivf_rabitq_set_is_using_refiner( + zvec_ivf_rabitq_query_params_t *params, bool is_using_refiner); + +/** + * @brief Get whether to use refiner + * @param params IVF RaBitQ query parameters pointer + * @return bool Whether to use refiner + */ +ZVEC_EXPORT bool ZVEC_CALL zvec_query_params_ivf_rabitq_get_is_using_refiner( + const zvec_ivf_rabitq_query_params_t *params); + // ----------------------------------------------------------------------------- // zvec_flat_query_params_t (Flat Query Parameters) // ----------------------------------------------------------------------------- @@ -1815,6 +1944,16 @@ ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_vector_query_set_hnsw_params( ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_vector_query_set_ivf_params( zvec_vector_query_t *query, zvec_ivf_query_params_t *ivf_params); +/** + * @brief Set IVF RaBitQ query parameters (takes ownership) + * @param query Vector query pointer + * @param ivf_rabitq_params IVF RaBitQ query parameters pointer + * @return zvec_error_code_t Error code + */ +ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_vector_query_set_ivf_rabitq_params( + zvec_vector_query_t *query, + zvec_ivf_rabitq_query_params_t *ivf_rabitq_params); + /** * @brief Set Flat query parameters (takes ownership) * @param query Vector query pointer @@ -2104,6 +2243,17 @@ ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_group_by_vector_query_set_ivf_params(zvec_group_by_vector_query_t *query, zvec_ivf_query_params_t *ivf_params); +/** + * @brief Set IVF RaBitQ query parameters (takes ownership) + * @param query Group by vector query pointer + * @param ivf_rabitq_params IVF RaBitQ query parameters pointer + * @return zvec_error_code_t Error code + */ +ZVEC_EXPORT zvec_error_code_t ZVEC_CALL +zvec_group_by_vector_query_set_ivf_rabitq_params( + zvec_group_by_vector_query_t *query, + zvec_ivf_rabitq_query_params_t *ivf_rabitq_params); + /** * @brief Set Flat query parameters (takes ownership) * @param query Group by vector query pointer @@ -2367,6 +2517,15 @@ ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_sub_query_set_hnsw_params( ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_sub_query_set_ivf_params( zvec_sub_query_t *query, zvec_ivf_query_params_t *ivf_params); +/** + * @brief Set IVF RaBitQ query parameters (takes ownership) + * @param query Sub-query pointer + * @param ivf_rabitq_params IVF RaBitQ query parameters pointer + * @return zvec_error_code_t Error code + */ +ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_sub_query_set_ivf_rabitq_params( + zvec_sub_query_t *query, zvec_ivf_rabitq_query_params_t *ivf_rabitq_params); + /** * @brief Set Flat query parameters (takes ownership) * @param query Sub-query pointer diff --git a/src/include/zvec/core/interface/constants.h b/src/include/zvec/core/interface/constants.h index 1f39daa0b..5d6f2f855 100644 --- a/src/include/zvec/core/interface/constants.h +++ b/src/include/zvec/core/interface/constants.h @@ -35,8 +35,11 @@ constexpr static bool kDefaultVamanaSaturateGraph = false; constexpr const uint32_t kDefaultRabitqTotalBits = 7; constexpr const uint32_t kDefaultRabitqNumClusters = 16; +constexpr const uint32_t kDefaultIvfRabitqNlist = 1024; +constexpr const uint32_t kDefaultIvfRabitqNprobe = 10; + constexpr const uint32_t kDefaultDiskAnnMaxDegree = 100; constexpr const uint32_t kDefaultDiskAnnListSize = 200; constexpr const uint32_t kDefaultDiskAnnPqChunkNum = 16; -} // namespace zvec::core_interface \ No newline at end of file +} // namespace zvec::core_interface diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index e9ed39051..2350500e0 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -16,6 +16,7 @@ #include #include +#include #include #include #include @@ -356,6 +357,37 @@ class HNSWRabitqIndex : public Index { HNSWRabitqIndexParam param_{}; }; +class IVFRabitqIndex : public Index { + public: + IVFRabitqIndex() = default; + + protected: + int CreateAndInitStreamer(const BaseIndexParam ¶m) override; + + int _prepare_for_search(const VectorData &query, + const BaseIndexQueryParam::Pointer &search_param, + core::IndexContext::Pointer &context) override; + int _get_coarse_search_topk( + const BaseIndexQueryParam::Pointer &search_param) override; + + int Add(const VectorData &vector, uint32_t doc_id) override; + int Train() override; + int Open(const std::string &file_path, + StorageOptions storage_options) override; + int _dense_fetch(const uint32_t doc_id, + VectorDataBuffer *vector_data_buffer) override; + int Merge(const std::vector &indexes, + const IndexFilter &filter, const MergeOptions &options) override; + int GenerateHolder(); + + private: + IVFRabitqIndexParam param_{}; + std::mutex mutex_{}; + std::vector> doc_cache_; + core::IndexHolder::Pointer holder_{}; + std::string file_path_; +}; + class DiskAnnIndex : public Index { public: DiskAnnIndex() = default; diff --git a/src/include/zvec/core/interface/index_param.h b/src/include/zvec/core/interface/index_param.h index 5d4e8a206..d233450ed 100644 --- a/src/include/zvec/core/interface/index_param.h +++ b/src/include/zvec/core/interface/index_param.h @@ -69,6 +69,7 @@ enum class IndexType { kIVF, // it's actual a two-layer index kHNSW, kHNSWRabitq, + kIVFRabitq, kDiskAnn, kVamana, }; @@ -432,6 +433,42 @@ struct HNSWRabitqIndexParam : public BaseIndexParam { bool omit_empty_value = false) const override; }; +struct IVFRabitqIndexParam : public BaseIndexParam { + using Pointer = std::shared_ptr; + + // IVF parameters + int nlist = kDefaultIvfRabitqNlist; + + // Rabitq parameters + int total_bits = kDefaultRabitqTotalBits; + int sample_count = 0; + core::IndexProvider::Pointer provider = nullptr; + core::IndexReformer::Pointer reformer = nullptr; + + IVFRabitqIndexParam() : BaseIndexParam(IndexType::kIVFRabitq) {} + + IVFRabitqIndexParam(int nlist) + : BaseIndexParam(IndexType::kIVFRabitq), nlist(nlist) {} + + IVFRabitqIndexParam(MetricType metric, int dim, int nlist) + : BaseIndexParam(IndexType::kIVFRabitq, metric, dim), nlist(nlist) {} + + protected: + bool DeserializeFromJsonObject(const ailego::JsonObject &json_obj) override; + ailego::JsonObject SerializeToJsonObject( + bool omit_empty_value = false) const override; +}; + +struct IVFRabitqQueryParam : public BaseIndexQueryParam { + using Pointer = std::shared_ptr; + + uint32_t nprobe = kDefaultIvfRabitqNprobe; + + BaseIndexQueryParam::Pointer Clone() const override { + return std::make_shared(*this); + } +}; + struct DiskAnnIndexParam : public BaseIndexParam { using Pointer = std::shared_ptr; diff --git a/src/include/zvec/core/interface/index_param_builders.h b/src/include/zvec/core/interface/index_param_builders.h index d88057a93..7df99134b 100644 --- a/src/include/zvec/core/interface/index_param_builders.h +++ b/src/include/zvec/core/interface/index_param_builders.h @@ -204,6 +204,38 @@ class HNSWRabitqIndexParamBuilder } }; +class IVFRabitqIndexParamBuilder + : public BaseIndexParamBuilder { + public: + IVFRabitqIndexParamBuilder() = default; + IVFRabitqIndexParamBuilder &WithNlist(int nlist) { + param->nlist = nlist; + return *this; + } + IVFRabitqIndexParamBuilder &WithTotalBits(int total_bits) { + param->total_bits = total_bits; + return *this; + } + IVFRabitqIndexParamBuilder &WithSampleCount(int sample_count) { + param->sample_count = sample_count; + return *this; + } + IVFRabitqIndexParamBuilder &WithReformer( + core::IndexReformer::Pointer reformer) { + param->reformer = std::move(reformer); + return *this; + } + IVFRabitqIndexParamBuilder &WithProvider( + core::IndexProvider::Pointer provider) { + param->provider = std::move(provider); + return *this; + } + std::shared_ptr Build() override { + return param; + } +}; + class DiskAnnIndexParamBuilder : public BaseIndexParamBuilder { @@ -511,4 +543,4 @@ class SCANNIndexParamBuilder { }; } // namespace predefined -} // namespace zvec::core_interface \ No newline at end of file +} // namespace zvec::core_interface diff --git a/src/include/zvec/db/index_params.h b/src/include/zvec/db/index_params.h index 1a79e4b5f..98806f2b3 100644 --- a/src/include/zvec/db/index_params.h +++ b/src/include/zvec/db/index_params.h @@ -54,7 +54,8 @@ class IndexParams { bool is_vector_index_type() const { return type_ == IndexType::FLAT || type_ == IndexType::HNSW || type_ == IndexType::HNSW_RABITQ || type_ == IndexType::IVF || - type_ == IndexType::DISKANN || type_ == IndexType::VAMANA; + type_ == IndexType::IVF_RABITQ || type_ == IndexType::DISKANN || + type_ == IndexType::VAMANA; } IndexType type() const { @@ -397,6 +398,91 @@ class HnswRabitqIndexParams : public VectorIndexParams { core::IndexReformer::Pointer rabitq_reformer_; }; +class IvfRabitqIndexParams : public VectorIndexParams { + public: + IvfRabitqIndexParams(MetricType metric_type, + int nlist = core_interface::kDefaultIvfRabitqNlist, + int total_bits = core_interface::kDefaultRabitqTotalBits, + int sample_count = 0) + : VectorIndexParams(IndexType::IVF_RABITQ, metric_type, + QuantizeType::RABITQ), + nlist_(nlist), + total_bits_(total_bits), + sample_count_(sample_count) {} + + using OPtr = std::shared_ptr; + + Ptr clone() const override { + auto obj = std::make_shared( + metric_type_, nlist_, total_bits_, sample_count_); + obj->set_rabitq_reformer(rabitq_reformer_); + obj->set_raw_vector_provider(raw_vector_provider_); + return obj; + } + + std::string to_string() const override { + auto base_str = vector_index_params_to_string("IvfRabitqIndexParams", + metric_type_, quantize_type_); + std::ostringstream oss; + oss << base_str << ",nlist:" << nlist_ << ",total_bits:" << total_bits_ + << ",sample_count:" << sample_count_ << "}"; + return oss.str(); + } + + bool operator==(const IndexParams &other) const override { + if (type() != other.type()) { + return false; + } + auto &rhs = dynamic_cast(other); + return metric_type() == rhs.metric_type() && + quantize_type_ == rhs.quantize_type_ && nlist_ == rhs.nlist_ && + total_bits_ == rhs.total_bits_ && sample_count_ == rhs.sample_count_; + } + + void set_nlist(int nlist) { + nlist_ = nlist; + } + int nlist() const { + return nlist_; + } + + void set_total_bits(int total_bits) { + total_bits_ = total_bits; + } + int total_bits() const { + return total_bits_; + } + + void set_sample_count(int sample_count) { + sample_count_ = sample_count; + } + int sample_count() const { + return sample_count_; + } + + void set_raw_vector_provider( + core::IndexProvider::Pointer raw_vector_provider) { + raw_vector_provider_ = std::move(raw_vector_provider); + } + core::IndexProvider::Pointer raw_vector_provider() const { + return raw_vector_provider_; + } + + void set_rabitq_reformer(core::IndexReformer::Pointer rabitq_reformer) { + rabitq_reformer_ = std::move(rabitq_reformer); + } + core::IndexReformer::Pointer rabitq_reformer() const { + return rabitq_reformer_; + } + + private: + int nlist_; + int total_bits_; + int sample_count_; + core::IndexProvider::Pointer raw_vector_provider_; + core::IndexReformer::Pointer rabitq_reformer_; +}; + class FlatIndexParams : public VectorIndexParams { public: FlatIndexParams(MetricType metric_type, diff --git a/src/include/zvec/db/query_params.h b/src/include/zvec/db/query_params.h index da90b7b93..8ba5256a8 100644 --- a/src/include/zvec/db/query_params.h +++ b/src/include/zvec/db/query_params.h @@ -170,6 +170,30 @@ class HnswRabitqQueryParams : public QueryParams { int ef_; }; +class IvfRabitqQueryParams : public QueryParams { + public: + IvfRabitqQueryParams(int nprobe = core_interface::kDefaultIvfRabitqNprobe, + float radius = 0.0f, bool is_linear = false, + bool is_using_refiner = false) + : QueryParams(IndexType::IVF_RABITQ), nprobe_(nprobe) { + set_radius(radius); + set_is_linear(is_linear); + set_is_using_refiner(is_using_refiner); + } + + ~IvfRabitqQueryParams() override = default; + + int nprobe() const { + return nprobe_; + } + void set_nprobe(int nprobe) { + nprobe_ = nprobe; + } + + private: + int nprobe_; +}; + class FlatQueryParams : public QueryParams { public: FlatQueryParams(bool is_using_refiner = false, float scale_factor = 10) diff --git a/src/include/zvec/db/type.h b/src/include/zvec/db/type.h index b98a0d066..2dde300b2 100644 --- a/src/include/zvec/db/type.h +++ b/src/include/zvec/db/type.h @@ -28,6 +28,7 @@ enum class IndexType : uint32_t { HNSW_RABITQ = 4, DISKANN = 5, VAMANA = 6, + IVF_RABITQ = 7, INVERT = 10, FTS = 11, }; diff --git a/tests/c/c_api_test.c b/tests/c/c_api_test.c index 791636eeb..c5b84890a 100644 --- a/tests/c/c_api_test.c +++ b/tests/c/c_api_test.c @@ -3482,11 +3482,30 @@ void test_index_params_functions(void) { TEST_ASSERT(n_iters == 10); TEST_ASSERT(use_soar == false); // Default is false + // Test IVF RaBitQ index params + zvec_index_params_t *ivf_rabitq_params = + zvec_index_params_create(ZVEC_INDEX_TYPE_IVF_RABITQ); + TEST_ASSERT(ivf_rabitq_params != NULL); + TEST_ASSERT(zvec_index_params_get_type(ivf_rabitq_params) == + ZVEC_INDEX_TYPE_IVF_RABITQ); + TEST_ASSERT(zvec_index_params_get_metric_type(ivf_rabitq_params) == + ZVEC_METRIC_TYPE_L2); + TEST_ASSERT(zvec_index_params_get_quantize_type(ivf_rabitq_params) == + ZVEC_QUANTIZE_TYPE_RABITQ); + + int nlist, total_bits, sample_count; + zvec_index_params_get_ivf_rabitq_params(ivf_rabitq_params, &nlist, + &total_bits, &sample_count); + TEST_ASSERT(nlist == 1024); + TEST_ASSERT(total_bits == 7); + TEST_ASSERT(sample_count == 0); + // Cleanup zvec_index_params_destroy(hnsw_params); zvec_index_params_destroy(invert_params); zvec_index_params_destroy(flat_params); zvec_index_params_destroy(ivf_params); + zvec_index_params_destroy(ivf_rabitq_params); TEST_END(); } @@ -3705,6 +3724,26 @@ void test_index_params_api_functions(void) { TEST_ASSERT(n_iters == 20); TEST_ASSERT(use_soar == true); + // Test zvec_index_params_create for IVF_RABITQ + zvec_index_params_t *ivf_rabitq_params = + zvec_index_params_create(ZVEC_INDEX_TYPE_IVF_RABITQ); + TEST_ASSERT(ivf_rabitq_params != NULL); + TEST_ASSERT(zvec_index_params_get_type(ivf_rabitq_params) == + ZVEC_INDEX_TYPE_IVF_RABITQ); + TEST_ASSERT(zvec_index_params_get_metric_type(ivf_rabitq_params) == + ZVEC_METRIC_TYPE_L2); + TEST_ASSERT(zvec_index_params_get_quantize_type(ivf_rabitq_params) == + ZVEC_QUANTIZE_TYPE_RABITQ); + + // Test zvec_index_params_set_ivf_rabitq_params + zvec_index_params_set_ivf_rabitq_params(ivf_rabitq_params, 256, 6, 4096); + int nlist, total_bits, sample_count; + zvec_index_params_get_ivf_rabitq_params(ivf_rabitq_params, &nlist, + &total_bits, &sample_count); + TEST_ASSERT(nlist == 256); + TEST_ASSERT(total_bits == 6); + TEST_ASSERT(sample_count == 4096); + // Test zvec_index_params_create for INVERT zvec_index_params_t *invert_params = zvec_index_params_create(ZVEC_INDEX_TYPE_INVERT); @@ -3732,6 +3771,7 @@ void test_index_params_api_functions(void) { // Cleanup zvec_index_params_destroy(hnsw_params); zvec_index_params_destroy(ivf_params); + zvec_index_params_destroy(ivf_rabitq_params); zvec_index_params_destroy(invert_params); zvec_index_params_destroy(flat_params); @@ -3796,6 +3836,11 @@ void test_query_params_functions(void) { zvec_query_params_ivf_create(10, true, 1.5f); TEST_ASSERT(ivf_params != NULL); + // Test IVF RaBitQ query parameters + zvec_ivf_rabitq_query_params_t *ivf_rabitq_params = + zvec_query_params_ivf_rabitq_create(10, 0.2f, false, true); + TEST_ASSERT(ivf_rabitq_params != NULL); + // Test Flat query parameters zvec_flat_query_params_t *flat_params = zvec_query_params_flat_create(false, 2.0f); @@ -3848,6 +3893,29 @@ void test_query_params_functions(void) { is_using_refiner = zvec_query_params_ivf_get_is_using_refiner(ivf_params); TEST_ASSERT(is_using_refiner == false); + // Test IVF RaBitQ-specific parameters + err = zvec_query_params_ivf_rabitq_set_nprobe(ivf_rabitq_params, 18); + TEST_ASSERT(err == ZVEC_OK); + TEST_ASSERT(zvec_query_params_ivf_rabitq_get_nprobe(ivf_rabitq_params) == 18); + + // Test IVF RaBitQ common parameters (radius, is_linear, is_using_refiner) + err = zvec_query_params_ivf_rabitq_set_radius(ivf_rabitq_params, 0.6f); + TEST_ASSERT(err == ZVEC_OK); + radius = zvec_query_params_ivf_rabitq_get_radius(ivf_rabitq_params); + TEST_ASSERT(radius == 0.6f); + + err = zvec_query_params_ivf_rabitq_set_is_linear(ivf_rabitq_params, true); + TEST_ASSERT(err == ZVEC_OK); + is_linear = zvec_query_params_ivf_rabitq_get_is_linear(ivf_rabitq_params); + TEST_ASSERT(is_linear == true); + + err = zvec_query_params_ivf_rabitq_set_is_using_refiner(ivf_rabitq_params, + false); + TEST_ASSERT(err == ZVEC_OK); + is_using_refiner = + zvec_query_params_ivf_rabitq_get_is_using_refiner(ivf_rabitq_params); + TEST_ASSERT(is_using_refiner == false); + // Test Flat scale factor setting err = zvec_query_params_flat_set_scale_factor(flat_params, 3.0f); TEST_ASSERT(err == ZVEC_OK); @@ -3900,12 +3968,14 @@ void test_query_params_functions(void) { // Test destruction of valid parameters zvec_query_params_hnsw_destroy(hnsw_params); zvec_query_params_ivf_destroy(ivf_params); + zvec_query_params_ivf_rabitq_destroy(ivf_rabitq_params); zvec_query_params_flat_destroy(flat_params); zvec_query_params_vamana_destroy(vamana_params); // Test boundary cases - null pointer handling zvec_query_params_hnsw_destroy(NULL); zvec_query_params_ivf_destroy(NULL); + zvec_query_params_ivf_rabitq_destroy(NULL); zvec_query_params_flat_destroy(NULL); zvec_query_params_vamana_destroy(NULL); @@ -3914,6 +3984,8 @@ void test_query_params_functions(void) { TEST_ASSERT(err == ZVEC_ERROR_INVALID_ARGUMENT); err = zvec_query_params_ivf_set_radius(NULL, 0.5f); TEST_ASSERT(err == ZVEC_ERROR_INVALID_ARGUMENT); + err = zvec_query_params_ivf_rabitq_set_radius(NULL, 0.5f); + TEST_ASSERT(err == ZVEC_ERROR_INVALID_ARGUMENT); err = zvec_query_params_flat_set_radius(NULL, 0.5f); TEST_ASSERT(err == ZVEC_ERROR_INVALID_ARGUMENT); err = zvec_query_params_vamana_set_ef_search(NULL, 100); @@ -3922,13 +3994,17 @@ void test_query_params_functions(void) { // Test default values for getters with NULL TEST_ASSERT(zvec_query_params_hnsw_get_radius(NULL) == 0.0f); TEST_ASSERT(zvec_query_params_ivf_get_radius(NULL) == 0.0f); + TEST_ASSERT(zvec_query_params_ivf_rabitq_get_radius(NULL) == 0.0f); TEST_ASSERT(zvec_query_params_flat_get_radius(NULL) == 0.0f); TEST_ASSERT(zvec_query_params_hnsw_get_is_linear(NULL) == false); TEST_ASSERT(zvec_query_params_ivf_get_is_linear(NULL) == false); + TEST_ASSERT(zvec_query_params_ivf_rabitq_get_is_linear(NULL) == false); TEST_ASSERT(zvec_query_params_flat_get_is_linear(NULL) == false); TEST_ASSERT(zvec_query_params_hnsw_get_is_using_refiner(NULL) == false); TEST_ASSERT(zvec_query_params_ivf_get_is_using_refiner(NULL) == false); + TEST_ASSERT(zvec_query_params_ivf_rabitq_get_is_using_refiner(NULL) == false); TEST_ASSERT(zvec_query_params_flat_get_is_using_refiner(NULL) == false); + TEST_ASSERT(zvec_query_params_ivf_rabitq_get_nprobe(NULL) == 10); TEST_ASSERT(zvec_query_params_vamana_get_ef_search(NULL) == 200); TEST_ASSERT(zvec_query_params_vamana_get_radius(NULL) == 0.0f); TEST_ASSERT(zvec_query_params_vamana_get_is_linear(NULL) == false); diff --git a/tests/core/algorithm/CMakeLists.txt b/tests/core/algorithm/CMakeLists.txt index 1c7e16e91..614d52d6e 100644 --- a/tests/core/algorithm/CMakeLists.txt +++ b/tests/core/algorithm/CMakeLists.txt @@ -18,4 +18,5 @@ endif() if(RABITQ_SUPPORTED) cc_directories(hnsw_rabitq) + cc_directories(ivf_rabitq) endif() diff --git a/tests/core/algorithm/ivf_rabitq/CMakeLists.txt b/tests/core/algorithm/ivf_rabitq/CMakeLists.txt new file mode 100644 index 000000000..89f185f5b --- /dev/null +++ b/tests/core/algorithm/ivf_rabitq/CMakeLists.txt @@ -0,0 +1,34 @@ +include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) +include(${PROJECT_ROOT_DIR}/cmake/option.cmake) + +if(APPLE) + set(APPLE_FRAMEWORK_LIBS + -framework CoreFoundation + -framework CoreGraphics + -framework CoreData + -framework CoreText + -framework Security + -framework Foundation + -Wl,-U,_MallocExtension_ReleaseFreeMemory + -Wl,-U,_ProfilerStart + -Wl,-U,_ProfilerStop + -Wl,-U,_RegisterThriftProtocol + ) +endif() + +file(GLOB_RECURSE ALL_TEST_SRCS *_test.cc) + +foreach(CC_SRCS ${ALL_TEST_SRCS}) + get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE) + cc_gtest( + NAME ${CC_TARGET} + STRICT + LIBS zvec_ailego core_framework core_utility core_metric core_quantizer core_knn_ivf_rabitq core_knn_hnsw_rabitq core_knn_ivf core_knn_flat core_knn_cluster + ${CMAKE_THREAD_LIBS_INIT} + ${CMAKE_DL_LIBS} + SRCS ${CC_SRCS} + INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm/ivf_rabitq ${PROJECT_ROOT_DIR}/src/core/algorithm/hnsw_rabitq + LDFLAGS ${APPLE_FRAMEWORK_LIBS} + ) + cc_test_suite(ivf_rabitq ${CC_TARGET}) +endforeach() diff --git a/tests/core/algorithm/ivf_rabitq/ivf_rabitq_streamer_test.cc b/tests/core/algorithm/ivf_rabitq/ivf_rabitq_streamer_test.cc new file mode 100644 index 000000000..9db913883 --- /dev/null +++ b/tests/core/algorithm/ivf_rabitq/ivf_rabitq_streamer_test.cc @@ -0,0 +1,912 @@ +// 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. + +#include "ivf_rabitq_streamer.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include "zvec/ailego/container/params.h" +#include "zvec/ailego/logger/logger.h" +#include "zvec/ailego/utility/file_helper.h" +#include "zvec/core/framework/index_factory.h" +#include "zvec/core/framework/index_holder.h" +#include "zvec/core/framework/index_memory.h" +#include "zvec/core/framework/index_streamer.h" +#include "ivf_rabitq_builder.h" +#include "ivf_rabitq_context.h" +#include "ivf_rabitq_params.h" +#include "ivf_rabitq_reformer.h" +#include "ivf_rabitq_util.h" +#include "rabitq_converter.h" +#include "rabitq_params.h" + +using namespace std; +using namespace zvec::ailego; + +namespace zvec { +namespace core { + +constexpr size_t static kDim = 128; + +class IvfRabitqStreamerTest : public testing::Test { + protected: + void SetUp(void) override; + void TearDown(void) override; + + static std::string dir_; + static shared_ptr index_meta_ptr_; +}; + +std::string IvfRabitqStreamerTest::dir_("ivfRabitqStreamerTest"); +shared_ptr IvfRabitqStreamerTest::index_meta_ptr_; + +namespace { + +void FillTestVector(size_t seed, NumericalVector *vec) { + ASSERT_NE(nullptr, vec); + for (size_t j = 0; j < vec->dimension(); ++j) { + size_t value = (seed + 1) * 1103515245ULL + (j + 1) * 12345ULL + + (seed + 17) * (j + 31); + value ^= (value >> 11); + (*vec)[j] = static_cast(value % 1000003ULL) / 1000003.0f; + } +} + +IndexHolder::Pointer BuildHolder(size_t dim, size_t doc_cnt) { + auto holder = + make_shared>(dim); + for (size_t i = 0; i < doc_cnt; i++) { + NumericalVector vec(dim); + FillTestVector(i, &vec); + EXPECT_TRUE(holder->emplace(i, vec)); + } + return holder; +} + +void BuildIndexFile(const IndexMeta &meta, const IndexHolder::Pointer &holder, + const ailego::Params ¶ms, const std::string &path) { + auto builder = std::make_shared(); + ASSERT_EQ(0, builder->init(meta, params)); + ASSERT_EQ(0, builder->train(nullptr, holder)); + ASSERT_EQ(0, builder->build(nullptr, holder)); + + auto dumper = IndexFactory::CreateDumper("FileDumper"); + ASSERT_NE(nullptr, dumper); + ASSERT_EQ(0, dumper->create(path)); + ASSERT_EQ(0, builder->dump(dumper)); + ASSERT_EQ(0, dumper->close()); +} + +void BuildAndOpenStreamer(const IndexMeta &meta, + const IndexHolder::Pointer &holder, + const ailego::Params ¶ms, const std::string &path, + IndexStreamer::Pointer *streamer) { + ASSERT_NE(nullptr, streamer); + BuildIndexFile(meta, holder, params, path); + + auto storage = IndexFactory::CreateStorage("MMapFileReadStorage"); + ASSERT_NE(nullptr, storage); + ailego::Params stg_params; + ASSERT_EQ(0, storage->init(stg_params)); + ASSERT_EQ(0, storage->open(path, false)); + + *streamer = std::make_shared(); + ASSERT_EQ(0, (*streamer)->init(meta, params)); + ASSERT_EQ(0, (*streamer)->open(storage)); +} + +void RunBuildSearchAndReopen(const IndexMeta &meta, + const IndexHolder::Pointer &holder, + const void *query, + const IndexQueryMeta &query_meta, + const ailego::Params ¶ms, + const std::string &path) { + BuildIndexFile(meta, holder, params, path); + + { + IndexStreamer::Pointer streamer = std::make_shared(); + ASSERT_EQ(0, streamer->init(meta, params)); + auto storage = IndexFactory::CreateStorage("MMapFileReadStorage"); + ASSERT_NE(nullptr, storage); + ASSERT_EQ(0, storage->init(ailego::Params())); + ASSERT_EQ(0, storage->open(path, false)); + ASSERT_EQ(0, streamer->open(storage)); + + auto context = streamer->create_context(); + context->set_topk(10); + ailego::Params search_params; + search_params.set(PARAM_IVF_RABITQ_NPROBE, 16U); + ASSERT_EQ(0, context->update(search_params)); + ASSERT_EQ(0, streamer->search_impl(query, query_meta, 1, context)); + EXPECT_GT(context->result(0).size(), 0UL); + + ASSERT_EQ(0, streamer->flush(0UL)); + ASSERT_EQ(0, streamer->close()); + } + + { + IndexStreamer::Pointer streamer = std::make_shared(); + ASSERT_EQ(0, streamer->init(meta, params)); + auto storage = IndexFactory::CreateStorage("MMapFileReadStorage"); + ASSERT_NE(nullptr, storage); + ASSERT_EQ(0, storage->init(ailego::Params())); + ASSERT_EQ(0, storage->open(path, false)); + ASSERT_EQ(0, streamer->open(storage)); + + auto context = streamer->create_context(); + context->set_topk(10); + ailego::Params search_params; + search_params.set(PARAM_IVF_RABITQ_NPROBE, 16U); + ASSERT_EQ(0, context->update(search_params)); + ASSERT_EQ(0, streamer->search_impl(query, query_meta, 1, context)); + EXPECT_GT(context->result(0).size(), 0UL); + } +} + +void BuildLoadedReformer(const IndexMeta &rabitq_meta, + const std::string &metric_name, + const IndexHolder::Pointer &holder, + const std::string &file_id, + std::shared_ptr *out_reformer) { + ASSERT_NE(nullptr, out_reformer); + + RabitqConverter converter; + ailego::Params params; + params.set(PARAM_RABITQ_TOTAL_BITS, 7U); + params.set(PARAM_RABITQ_NUM_CLUSTERS, 16U); + ASSERT_EQ(0, converter.init(rabitq_meta, params)); + ASSERT_EQ(0, converter.train(holder)); + + auto memory_dumper = IndexFactory::CreateDumper("MemoryDumper"); + ASSERT_NE(nullptr, memory_dumper); + ASSERT_EQ(0, memory_dumper->init(ailego::Params())); + ASSERT_EQ(0, memory_dumper->create(file_id)); + ASSERT_EQ(0, converter.dump(memory_dumper)); + ASSERT_EQ(0, memory_dumper->close()); + + auto reformer = std::make_shared(); + ASSERT_EQ(0, reformer->init(metric_name)); + auto memory_storage = IndexFactory::CreateStorage("MemoryReadStorage"); + ASSERT_NE(nullptr, memory_storage); + ASSERT_EQ(0, memory_storage->open(file_id, false)); + ASSERT_EQ(0, reformer->load(memory_storage)); + IndexMemory::Instance()->remove(file_id); + + *out_reformer = reformer; +} + +} // namespace + +void IvfRabitqStreamerTest::SetUp(void) { + index_meta_ptr_.reset(new (nothrow) + IndexMeta(IndexMeta::DataType::DT_FP32, kDim)); + index_meta_ptr_->set_metric("SquaredEuclidean", 0, ailego::Params()); +} + +void IvfRabitqStreamerTest::TearDown(void) { + ailego::FileHelper::RemovePath(dir_.c_str()); +} + +TEST_F(IvfRabitqStreamerTest, TestBuildAndSearch) { + auto holder = + make_shared>(kDim); + size_t doc_cnt = 1000UL; + for (size_t i = 0; i < doc_cnt; i++) { + NumericalVector vec(kDim); + for (size_t j = 0; j < kDim; ++j) { + vec[j] = static_cast(i * kDim + j) / 1000.0f; + } + ASSERT_TRUE(holder->emplace(i, vec)); + } + + ailego::Params params; + params.set("proxima.ivf_rabitq.nlist", 32U); + params.set(PARAM_RABITQ_TOTAL_BITS, 7U); + IndexStreamer::Pointer streamer; + BuildAndOpenStreamer(*index_meta_ptr_, holder, params, + dir_ + "/TestBuildAndSearch", &streamer); + ASSERT_NE(nullptr, streamer); + + // Search + NumericalVector query_vec(kDim); + for (size_t j = 0; j < kDim; ++j) { + query_vec[j] = static_cast(j) / 1000.0f; + } + + IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, kDim); + auto context = streamer->create_context(); + context->set_topk(10); + ailego::Params search_params; + search_params.set("proxima.ivf_rabitq.nprobe", 8U); + context->update(search_params); + + ASSERT_EQ(0, streamer->search_impl(query_vec.data(), query_meta, 1, context)); + + const auto &result = context->result(0); + ASSERT_GT(result.size(), 0UL); + ASSERT_LE(result.size(), 10UL); + + // The nearest vector should be key=0 (same direction as query) + LOG_INFO("IVF RabitQ search returned %zu results, top key=%lu, score=%f", + result.size(), result[0].key(), result[0].score()); +} + +TEST_F(IvfRabitqStreamerTest, TestCreateProvider) { + auto holder = BuildHolder(kDim, 300UL); + + ailego::Params params; + params.set(PARAM_IVF_RABITQ_NLIST, 16U); + params.set(PARAM_RABITQ_TOTAL_BITS, 4U); + IndexStreamer::Pointer streamer; + BuildAndOpenStreamer(*index_meta_ptr_, holder, params, + dir_ + "/TestCreateProvider", &streamer); + ASSERT_NE(nullptr, streamer); + + auto provider = streamer->create_provider(); + ASSERT_NE(nullptr, provider); + EXPECT_EQ(holder->count(), provider->count()); + EXPECT_EQ(kDim, provider->dimension()); + EXPECT_EQ(IndexMeta::DataType::DT_UNDEFINED, provider->data_type()); + size_t expected_element_size = + rabitqlib::BinDataMap::data_bytes(kDim) + + rabitqlib::ExDataMap::data_bytes(kDim, 3UL); + EXPECT_EQ(expected_element_size, provider->element_size()); + + size_t seen = 0; + auto iter = provider->create_iterator(); + ASSERT_NE(nullptr, iter); + while (iter->is_valid()) { + ASSERT_NE(nullptr, iter->data()); + + IndexStorage::MemoryBlock block; + EXPECT_EQ(0, provider->get_vector(iter->key(), block)); + ASSERT_NE(nullptr, block.data()); + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_HEAP_SCRATCH, block.type_); + EXPECT_EQ( + 0, std::memcmp(iter->data(), block.data(), provider->element_size())); + + ++seen; + iter->next(); + } + EXPECT_EQ(provider->count(), seen); + + IndexStorage::MemoryBlock block; + ASSERT_EQ(0, streamer->get_vector_by_id(17, block)); + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_HEAP_SCRATCH, block.type_); + IndexStorage::MemoryBlock provider_block; + ASSERT_EQ(0, provider->get_vector(17, provider_block)); + EXPECT_EQ(IndexStorage::MemoryBlock::MBT_HEAP_SCRATCH, provider_block.type_); + EXPECT_EQ(0, std::memcmp(provider_block.data(), block.data(), + provider->element_size())); +} + +TEST_F(IvfRabitqStreamerTest, TestTotalBitsOneBuildSearchAndReopen) { + auto holder = BuildHolder(kDim, 1000UL); + NumericalVector query_vec(kDim); + FillTestVector(17, &query_vec); + + ailego::Params params; + params.set(PARAM_IVF_RABITQ_NLIST, 32U); + params.set(PARAM_IVF_RABITQ_BRUTE_FORCE_THRESHOLD, 0U); + params.set(PARAM_RABITQ_TOTAL_BITS, 1U); + + IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, kDim); + RunBuildSearchAndReopen(*index_meta_ptr_, holder, query_vec.data(), + query_meta, params, + dir_ + "/TestTotalBitsOneBuildSearchAndReopen"); +} + +TEST_F(IvfRabitqStreamerTest, TestInnerProductBuildSearchAndReopen) { + IndexMeta meta(IndexMeta::DataType::DT_FP32, kDim); + meta.set_metric("InnerProduct", 0, ailego::Params()); + + auto holder = BuildHolder(kDim, 1000UL); + NumericalVector query_vec(kDim); + FillTestVector(31, &query_vec); + + ailego::Params params; + params.set(PARAM_IVF_RABITQ_NLIST, 32U); + params.set(PARAM_IVF_RABITQ_BRUTE_FORCE_THRESHOLD, 0U); + params.set(PARAM_RABITQ_TOTAL_BITS, 7U); + + IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, kDim); + RunBuildSearchAndReopen(meta, holder, query_vec.data(), query_meta, params, + dir_ + "/TestInnerProductBuildSearchAndReopen"); +} + +TEST_F(IvfRabitqStreamerTest, TestCosineBuildSearchAndReopen) { + IndexMeta raw_meta(IndexMeta::DataType::DT_FP32, kDim); + raw_meta.set_metric("Cosine", 0, ailego::Params()); + + auto raw_holder = BuildHolder(kDim, 1000UL); + auto converter = IndexFactory::CreateConverter("CosineNormalizeConverter"); + ASSERT_NE(nullptr, converter); + ASSERT_EQ(0, converter->init(raw_meta, ailego::Params())); + IndexMeta transformed_meta = converter->meta(); + ASSERT_EQ(kDim + 1, transformed_meta.dimension()); + ASSERT_EQ(0, converter->transform(raw_holder)); + auto holder = converter->result(); + ASSERT_NE(nullptr, holder); + + NumericalVector raw_query_vec(kDim); + FillTestVector(43, &raw_query_vec); + auto reformer = + IndexFactory::CreateReformer(transformed_meta.reformer_name()); + ASSERT_NE(nullptr, reformer); + ASSERT_EQ(0, reformer->init(transformed_meta.reformer_params())); + + IndexQueryMeta raw_query_meta(IndexMeta::DataType::DT_FP32, kDim); + IndexQueryMeta transformed_query_meta; + std::string transformed_query; + ASSERT_EQ(0, + reformer->transform(raw_query_vec.data(), raw_query_meta, + &transformed_query, &transformed_query_meta)); + ASSERT_EQ(kDim + 1, transformed_query_meta.dimension()); + + ailego::Params params; + params.set(PARAM_IVF_RABITQ_NLIST, 32U); + params.set(PARAM_IVF_RABITQ_BRUTE_FORCE_THRESHOLD, 0U); + params.set(PARAM_RABITQ_TOTAL_BITS, 7U); + params.set(PARAM_RABITQ_GENERAL_DIMENSION, static_cast(kDim)); + + RunBuildSearchAndReopen(transformed_meta, holder, transformed_query.data(), + transformed_query_meta, params, + dir_ + "/TestCosineBuildSearchAndReopen"); +} + +TEST_F(IvfRabitqStreamerTest, TestContextResetClearsSearchState) { + IvfRabitqContext context; + context.set_topk(3); + context.set_fetch_vector(true); + context.set_filter([](uint64_t key) { return key == 7; }); + context.set_threshold(0.5f); + context.query_state.rotated_query.push_back(1.0f); + context.query_state.centroid_norms.push_back(2.0f); + context.query_state.batch_query = std::make_shared(1); + + ASSERT_TRUE(context.fetch_vector()); + ASSERT_TRUE(context.filter().is_valid()); + ASSERT_NE(std::numeric_limits::max(), context.threshold()); + + context.reset(); + + EXPECT_FALSE(context.fetch_vector()); + EXPECT_FALSE(context.filter().is_valid()); + EXPECT_EQ(std::numeric_limits::max(), context.threshold()); + EXPECT_TRUE(context.query_state.rotated_query.empty()); + EXPECT_TRUE(context.query_state.centroid_norms.empty()); + EXPECT_EQ(nullptr, context.query_state.batch_query); +} + +TEST_F(IvfRabitqStreamerTest, TestProbeCentroidsUseAssignmentMetric) { + IndexMeta meta(IndexMeta::DataType::DT_FP32, kDim); + meta.set_metric("InnerProduct", 0, ailego::Params()); + auto holder = BuildHolder(kDim, 1000UL); + + std::shared_ptr reformer; + BuildLoadedReformer(meta, "InnerProduct", holder, + dir_ + "/TestProbeCentroidsUseAssignmentMetric", + &reformer); + ASSERT_NE(nullptr, reformer); + + NumericalVector query_vec(kDim); + FillTestVector(71, &query_vec); + + uint32_t assigned = reformer->find_nearest_centroid(query_vec.data()); + std::vector probe_centroids; + ASSERT_EQ(0, reformer->select_probe_centroids(query_vec.data(), 1, + &probe_centroids)); + ASSERT_EQ(1UL, probe_centroids.size()); + EXPECT_EQ(assigned, probe_centroids[0]); +} + +TEST_F(IvfRabitqStreamerTest, TestCosineProbeCentroidsUseInnerProductMetric) { + IndexMeta raw_meta(IndexMeta::DataType::DT_FP32, kDim); + raw_meta.set_metric("Cosine", 0, ailego::Params()); + auto raw_holder = BuildHolder(kDim, 1000UL); + + auto converter = IndexFactory::CreateConverter("CosineNormalizeConverter"); + ASSERT_NE(nullptr, converter); + ASSERT_EQ(0, converter->init(raw_meta, ailego::Params())); + IndexMeta transformed_meta = converter->meta(); + ASSERT_EQ(0, converter->transform(raw_holder)); + auto holder = converter->result(); + ASSERT_NE(nullptr, holder); + + ailego::Params params; + params.set(PARAM_RABITQ_GENERAL_DIMENSION, static_cast(kDim)); + IndexMeta rabitq_meta; + std::string metric_name; + ASSERT_EQ(0, PrepareAndCheckIvfRabitqInternalMeta( + transformed_meta, params, &rabitq_meta, &metric_name)); + ASSERT_EQ("InnerProduct", metric_name); + ASSERT_EQ(kDim, rabitq_meta.dimension()); + + std::shared_ptr reformer; + BuildLoadedReformer(rabitq_meta, metric_name, holder, + dir_ + "/TestCosineProbeCentroidsUseInnerProductMetric", + &reformer); + ASSERT_NE(nullptr, reformer); + + NumericalVector raw_query_vec(kDim); + FillTestVector(79, &raw_query_vec); + auto query_reformer = + IndexFactory::CreateReformer(transformed_meta.reformer_name()); + ASSERT_NE(nullptr, query_reformer); + ASSERT_EQ(0, query_reformer->init(transformed_meta.reformer_params())); + + IndexQueryMeta raw_query_meta(IndexMeta::DataType::DT_FP32, kDim); + IndexQueryMeta transformed_query_meta; + std::string transformed_query; + ASSERT_EQ(0, query_reformer->transform(raw_query_vec.data(), raw_query_meta, + &transformed_query, + &transformed_query_meta)); + + const float *query = + reinterpret_cast(transformed_query.data()); + uint32_t assigned = reformer->find_nearest_centroid(query); + std::vector probe_centroids; + ASSERT_EQ(0, reformer->select_probe_centroids(query, 1, &probe_centroids)); + ASSERT_EQ(1UL, probe_centroids.size()); + EXPECT_EQ(assigned, probe_centroids[0]); +} + +TEST_F(IvfRabitqStreamerTest, TestIncrementalAddUnsupported) { + NumericalVector query_vec(kDim); + FillTestVector(83, &query_vec); + IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, kDim); + + { + IndexStreamer::Pointer streamer = std::make_shared(); + ailego::Params params; + params.set(PARAM_IVF_RABITQ_NLIST, 32U); + params.set(PARAM_RABITQ_TOTAL_BITS, 7U); + ASSERT_EQ(0, streamer->init(*index_meta_ptr_, params)); + + auto storage = IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_NE(nullptr, storage); + ASSERT_EQ(0, storage->init(ailego::Params())); + ASSERT_EQ(0, + storage->open(dir_ + "/TestUnsupportedFlushBeforeBuild", true)); + ASSERT_NE(0, streamer->open(storage)); + + auto context = streamer->create_context(); + EXPECT_EQ(IndexError_NotImplemented, + streamer->add_impl(1, query_vec.data(), query_meta, context)); + EXPECT_EQ(0, streamer->flush(0UL)); + } + + { + auto holder = BuildHolder(kDim, 1000UL); + ailego::Params params; + params.set(PARAM_IVF_RABITQ_NLIST, 32U); + params.set(PARAM_RABITQ_TOTAL_BITS, 7U); + IndexStreamer::Pointer streamer; + BuildAndOpenStreamer(*index_meta_ptr_, holder, params, + dir_ + "/TestUnsupportedAddAfterBuild", &streamer); + ASSERT_NE(nullptr, streamer); + + auto context = streamer->create_context(); + EXPECT_EQ(IndexError_NotImplemented, + streamer->add_impl(1001, query_vec.data(), query_meta, context)); + EXPECT_EQ(0, streamer->flush(0UL)); + } +} + +TEST_F(IvfRabitqStreamerTest, TestSearchScanLimitParams) { + auto holder = + make_shared>(kDim); + size_t doc_cnt = 1000UL; + for (size_t i = 0; i < doc_cnt; i++) { + NumericalVector vec(kDim); + for (size_t j = 0; j < kDim; ++j) { + vec[j] = static_cast((i + 1) * (j + 3)) / 1000.0f; + } + ASSERT_TRUE(holder->emplace(i, vec)); + } + + ailego::Params params; + params.set(PARAM_IVF_RABITQ_NLIST, 32U); + params.set(PARAM_IVF_RABITQ_NPROBE, 0U); + params.set(PARAM_IVF_RABITQ_SCAN_RATIO, 0.05f); + params.set(PARAM_IVF_RABITQ_BRUTE_FORCE_THRESHOLD, 80U); + params.set(PARAM_RABITQ_TOTAL_BITS, 7U); + IndexStreamer::Pointer streamer; + BuildAndOpenStreamer(*index_meta_ptr_, holder, params, + dir_ + "/TestSearchScanLimitParams", &streamer); + ASSERT_NE(nullptr, streamer); + + NumericalVector query_vec(kDim); + for (size_t j = 0; j < kDim; ++j) { + query_vec[j] = static_cast(j + 3) / 1000.0f; + } + IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, kDim); + + auto base_context = streamer->create_context(); + auto *context = dynamic_cast(base_context.get()); + ASSERT_NE(nullptr, context); + context->set_topk(10); + + ASSERT_EQ( + 0, streamer->search_impl(query_vec.data(), query_meta, 1, base_context)); + EXPECT_EQ(80U, context->max_scan_count()); + EXPECT_GT(context->result(0).size(), 0UL); + + ailego::Params nprobe_params; + nprobe_params.set(PARAM_IVF_RABITQ_NPROBE, 8U); + nprobe_params.set(PARAM_IVF_RABITQ_BRUTE_FORCE_THRESHOLD, 80U); + ASSERT_EQ(0, context->update(nprobe_params)); + ASSERT_EQ( + 0, streamer->search_impl(query_vec.data(), query_meta, 1, base_context)); + EXPECT_EQ(static_cast(doc_cnt), context->max_scan_count()); + + ailego::Params bf_params; + bf_params.set(PARAM_IVF_RABITQ_NPROBE, 1U); + bf_params.set(PARAM_IVF_RABITQ_BRUTE_FORCE_THRESHOLD, + static_cast(doc_cnt)); + ASSERT_EQ(0, context->update(bf_params)); + ASSERT_EQ( + 0, streamer->search_impl(query_vec.data(), query_meta, 1, base_context)); + EXPECT_EQ(static_cast(doc_cnt), context->max_scan_count()); +} + +TEST_F(IvfRabitqStreamerTest, TestSearchWithFilter) { + auto holder = + make_shared>(kDim); + size_t doc_cnt = 1000UL; + for (size_t i = 0; i < doc_cnt; i++) { + NumericalVector vec(kDim); + for (size_t j = 0; j < kDim; ++j) { + vec[j] = static_cast(i * kDim + j) / 1000.0f; + } + ASSERT_TRUE(holder->emplace(i, vec)); + } + + ailego::Params params; + params.set("proxima.ivf_rabitq.nlist", 32U); + params.set(PARAM_RABITQ_TOTAL_BITS, 7U); + IndexStreamer::Pointer streamer; + BuildAndOpenStreamer(*index_meta_ptr_, holder, params, + dir_ + "/TestSearchWithFilter", &streamer); + ASSERT_NE(nullptr, streamer); + + NumericalVector query_vec(kDim); + for (size_t j = 0; j < kDim; ++j) { + query_vec[j] = static_cast(j) / 1000.0f; + } + + IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, kDim); + auto context = streamer->create_context(); + context->set_topk(10); + ailego::Params search_params; + search_params.set("proxima.ivf_rabitq.nprobe", 16U); + context->update(search_params); + + ASSERT_EQ(0, streamer->search_impl(query_vec.data(), query_meta, 1, context)); + const auto &unfiltered_result = context->result(0); + ASSERT_GT(unfiltered_result.size(), 0UL); + uint64_t filtered_key = unfiltered_result[0].key(); + + context->set_filter( + [filtered_key](uint64_t key) { return key == filtered_key; }); + ASSERT_EQ(0, streamer->search_impl(query_vec.data(), query_meta, 1, context)); + const auto &filtered_result = context->result(0); + auto iter = std::find_if(filtered_result.begin(), filtered_result.end(), + [filtered_key](const IndexDocument &doc) { + return doc.key() == filtered_key; + }); + EXPECT_EQ(filtered_result.end(), iter); + + context->reset(); + context->set_topk(10); + ASSERT_EQ(0, context->update(search_params)); + ASSERT_EQ(0, streamer->search_impl(query_vec.data(), query_meta, 1, context)); + const auto &reset_result = context->result(0); + auto reset_iter = std::find_if(reset_result.begin(), reset_result.end(), + [filtered_key](const IndexDocument &doc) { + return doc.key() == filtered_key; + }); + EXPECT_NE(reset_result.end(), reset_iter); +} + +TEST_F(IvfRabitqStreamerTest, TestRecallQuality) { + auto holder = + make_shared>(kDim); + size_t doc_cnt = 2000UL; + srand(42); + for (size_t i = 0; i < doc_cnt; i++) { + NumericalVector vec(kDim); + for (size_t j = 0; j < kDim; ++j) { + vec[j] = static_cast(rand()) / static_cast(RAND_MAX); + } + ASSERT_TRUE(holder->emplace(i, vec)); + } + + ailego::Params params; + params.set("proxima.ivf_rabitq.nlist", 64U); + params.set(PARAM_RABITQ_TOTAL_BITS, 7U); + IndexStreamer::Pointer streamer; + BuildAndOpenStreamer(*index_meta_ptr_, holder, params, dir_ + "/TestRecall", + &streamer); + ASSERT_NE(nullptr, streamer); + + // Use high nprobe (all clusters) as ground truth, low nprobe as test + size_t topk = 10; + int total_hits = 0; + int total_cnts = 0; + size_t query_cnt = 100; + IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, kDim); + + auto gt_ctx = streamer->create_context(); + auto test_ctx = streamer->create_context(); + gt_ctx->set_topk(topk); + test_ctx->set_topk(topk); + + // Ground truth: search all clusters + ailego::Params gt_params; + gt_params.set("proxima.ivf_rabitq.nprobe", 64U); + gt_ctx->update(gt_params); + + // Test: search with limited nprobe + ailego::Params test_search_params; + test_search_params.set("proxima.ivf_rabitq.nprobe", 16U); + test_ctx->update(test_search_params); + + srand(123); + for (size_t q = 0; q < query_cnt; q++) { + NumericalVector query_vec(kDim); + for (size_t j = 0; j < kDim; ++j) { + query_vec[j] = static_cast(rand()) / static_cast(RAND_MAX); + } + + ASSERT_EQ(0, + streamer->search_impl(query_vec.data(), query_meta, 1, gt_ctx)); + ASSERT_EQ(0, + streamer->search_impl(query_vec.data(), query_meta, 1, test_ctx)); + + auto >_result = gt_ctx->result(0); + auto &test_result = test_ctx->result(0); + + ASSERT_GT(gt_result.size(), 0UL); + ASSERT_GT(test_result.size(), 0UL); + + for (size_t k = 0; k < test_result.size(); ++k) { + total_cnts++; + for (size_t j = 0; j < gt_result.size(); ++j) { + if (gt_result[j].key() == test_result[k].key()) { + total_hits++; + break; + } + } + } + } + + float recall = total_hits * 1.0f / total_cnts; + LOG_INFO("IVF RabitQ recall@%zu (nprobe=16 vs nprobe=64): %.2f%% (%d/%d)", + topk, recall * 100.0f, total_hits, total_cnts); + EXPECT_GT(recall, 0.50f) << "Recall too low"; +} + +TEST_F(IvfRabitqStreamerTest, TestCloseAndReopen) { + auto holder = + make_shared>(kDim); + size_t doc_cnt = 500UL; + for (size_t i = 0; i < doc_cnt; i++) { + NumericalVector vec(kDim); + for (size_t j = 0; j < kDim; ++j) { + vec[j] = static_cast(i); + } + ASSERT_TRUE(holder->emplace(i, vec)); + } + + ailego::Params params; + params.set("proxima.ivf_rabitq.nlist", 16U); + params.set(PARAM_RABITQ_TOTAL_BITS, 4U); + + std::string path = dir_ + "/TestCloseReopen"; + BuildIndexFile(*index_meta_ptr_, holder, params, path); + + // Open and close the built index once before reopening it. + { + IndexStreamer::Pointer streamer = std::make_shared(); + ASSERT_EQ(0, streamer->init(*index_meta_ptr_, params)); + auto storage = IndexFactory::CreateStorage("MMapFileReadStorage"); + ASSERT_NE(nullptr, storage); + ASSERT_EQ(0, storage->init(ailego::Params())); + ASSERT_EQ(0, storage->open(path, false)); + ASSERT_EQ(0, streamer->open(storage)); + ASSERT_EQ(0, streamer->flush(0UL)); + ASSERT_EQ(0, streamer->close()); + } + + // Reopen and search + { + IndexStreamer::Pointer streamer2 = std::make_shared(); + ASSERT_EQ(0, streamer2->init(*index_meta_ptr_, params)); + auto storage = IndexFactory::CreateStorage("MMapFileReadStorage"); + ASSERT_NE(nullptr, storage); + ASSERT_EQ(0, storage->init(ailego::Params())); + ASSERT_EQ(0, storage->open(path, false)); + ASSERT_EQ(0, streamer2->open(storage)); + + NumericalVector query_vec(kDim); + for (size_t j = 0; j < kDim; ++j) { + query_vec[j] = 10.0f; + } + + IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, kDim); + auto context = streamer2->create_context(); + context->set_topk(5); + ailego::Params search_params; + search_params.set("proxima.ivf_rabitq.nprobe", 8U); + context->update(search_params); + + ASSERT_EQ(0, + streamer2->search_impl(query_vec.data(), query_meta, 1, context)); + const auto &result = context->result(0); + ASSERT_GT(result.size(), 0UL); + // Key 10 should be the nearest + ASSERT_EQ(10UL, result[0].key()); + } +} + +TEST_F(IvfRabitqStreamerTest, TestCosineCloseAndReopen) { + IndexMeta raw_meta(IndexMeta::DataType::DT_FP32, kDim); + raw_meta.set_metric("Cosine", 0, ailego::Params()); + + auto raw_holder = + make_shared>(kDim); + size_t doc_cnt = 1000UL; + std::vector> raw_vectors; + raw_vectors.reserve(doc_cnt); + for (size_t i = 0; i < doc_cnt; i++) { + NumericalVector vec(kDim); + for (size_t j = 0; j < kDim; ++j) { + vec[j] = static_cast(((i + 3) * (j + 5)) % 97 + 1) / 97.0f; + } + raw_vectors.push_back(vec); + ASSERT_TRUE(raw_holder->emplace(i, raw_vectors.back())); + } + + auto converter = IndexFactory::CreateConverter("CosineNormalizeConverter"); + ASSERT_NE(nullptr, converter); + ASSERT_EQ(0, converter->init(raw_meta, ailego::Params())); + IndexMeta transformed_meta = converter->meta(); + ASSERT_EQ(kDim + 1, transformed_meta.dimension()); + ASSERT_EQ(0, converter->transform(raw_holder)); + auto holder = converter->result(); + ASSERT_NE(nullptr, holder); + + ailego::Params params; + params.set("proxima.ivf_rabitq.nlist", 32U); + params.set(PARAM_RABITQ_TOTAL_BITS, 7U); + params.set(PARAM_RABITQ_GENERAL_DIMENSION, static_cast(kDim)); + + std::string path = dir_ + "/TestCosineCloseAndReopen"; + BuildIndexFile(transformed_meta, holder, params, path); + + { + IndexStreamer::Pointer streamer = std::make_shared(); + ASSERT_EQ(0, streamer->init(transformed_meta, params)); + ASSERT_EQ(kDim + 1, streamer->meta().dimension()); + ASSERT_EQ("Cosine", streamer->meta().metric_name()); + auto storage = IndexFactory::CreateStorage("MMapFileReadStorage"); + ASSERT_NE(nullptr, storage); + ASSERT_EQ(0, storage->init(ailego::Params())); + ASSERT_EQ(0, storage->open(path, false)); + ASSERT_EQ(0, streamer->open(storage)); + ASSERT_EQ(0, streamer->flush(0UL)); + ASSERT_EQ(0, streamer->close()); + } + + IndexStreamer::Pointer streamer2 = std::make_shared(); + ASSERT_EQ(0, streamer2->init(transformed_meta, params)); + ASSERT_EQ(kDim + 1, streamer2->meta().dimension()); + ASSERT_EQ("Cosine", streamer2->meta().metric_name()); + auto storage = IndexFactory::CreateStorage("MMapFileReadStorage"); + ASSERT_NE(nullptr, storage); + ASSERT_EQ(0, storage->init(ailego::Params())); + ASSERT_EQ(0, storage->open(path, false)); + ASSERT_EQ(0, streamer2->open(storage)); + + auto reformer = + IndexFactory::CreateReformer(transformed_meta.reformer_name()); + ASSERT_NE(nullptr, reformer); + ASSERT_EQ(0, reformer->init(transformed_meta.reformer_params())); + + IndexQueryMeta raw_query_meta(IndexMeta::DataType::DT_FP32, kDim); + IndexQueryMeta transformed_query_meta; + std::string transformed_query0; + std::string transformed_query1; + ASSERT_EQ(0, + reformer->transform(raw_vectors[17].data(), raw_query_meta, + &transformed_query0, &transformed_query_meta)); + IndexQueryMeta transformed_query_meta1; + ASSERT_EQ(0, + reformer->transform(raw_vectors[101].data(), raw_query_meta, + &transformed_query1, &transformed_query_meta1)); + ASSERT_EQ(transformed_query_meta.dimension(), + transformed_query_meta1.dimension()); + ASSERT_EQ(kDim + 1, transformed_query_meta.dimension()); + + std::vector query_buffer(2 * transformed_query_meta.dimension()); + std::memcpy(query_buffer.data(), transformed_query0.data(), + transformed_query0.size()); + std::memcpy(query_buffer.data() + transformed_query_meta.dimension(), + transformed_query1.data(), transformed_query1.size()); + + auto context = streamer2->create_context(); + context->set_topk(10); + ailego::Params search_params; + search_params.set("proxima.ivf_rabitq.nprobe", 16U); + context->update(search_params); + + ASSERT_EQ(0, streamer2->search_impl(query_buffer.data(), + transformed_query_meta, 2, context)); + EXPECT_GT(context->result(0).size(), 0UL); + EXPECT_GT(context->result(1).size(), 0UL); +} + +TEST_F(IvfRabitqStreamerTest, TestSmallDataset) { + // Test with a small but valid dataset. nlist must be significantly smaller + // than doc_count because rabitqlib's f_rescale = 1/||residual|| becomes NaN + // when a cluster has only 1 vector (residual ≈ 0). Use 200 vectors / 32 + // clusters ≈ 6 vectors per cluster to stay well within the valid range. + auto holder = + make_shared>(kDim); + size_t doc_cnt = 200UL; + srand(7); + for (size_t i = 0; i < doc_cnt; i++) { + NumericalVector vec(kDim); + for (size_t j = 0; j < kDim; ++j) { + vec[j] = static_cast(rand()) / static_cast(RAND_MAX); + } + ASSERT_TRUE(holder->emplace(i, vec)); + } + + ailego::Params params; + params.set("proxima.ivf_rabitq.nlist", 32U); + params.set(PARAM_RABITQ_TOTAL_BITS, 3U); + IndexStreamer::Pointer streamer; + BuildAndOpenStreamer(*index_meta_ptr_, holder, params, + dir_ + "/TestSmallDataset", &streamer); + ASSERT_NE(nullptr, streamer); + + NumericalVector query_vec(kDim); + srand(99); + for (size_t j = 0; j < kDim; ++j) { + query_vec[j] = static_cast(rand()) / static_cast(RAND_MAX); + } + + IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, kDim); + auto context = streamer->create_context(); + context->set_topk(5); + ailego::Params search_params; + search_params.set("proxima.ivf_rabitq.nprobe", 32U); + context->update(search_params); + + ASSERT_EQ(0, streamer->search_impl(query_vec.data(), query_meta, 1, context)); + const auto &result = context->result(0); + ASSERT_GT(result.size(), 0UL); + LOG_INFO("TestSmallDataset: %zu results, top key=%lu, score=%f", + result.size(), result[0].key(), result[0].score()); +} + +} // namespace core +} // namespace zvec diff --git a/tests/core/interface/CMakeLists.txt b/tests/core/interface/CMakeLists.txt index 83bf8d95e..db719df92 100644 --- a/tests/core/interface/CMakeLists.txt +++ b/tests/core/interface/CMakeLists.txt @@ -8,7 +8,7 @@ foreach(CC_SRCS ${ALL_TEST_SRCS}) NAME ${CC_TARGET} STRICT LIBS zvec_ailego core_framework core_metric core_interface core_knn_flat core_utility core_quantizer sparsehash core_knn_hnsw core_mix_reducer - core_knn_flat_sparse core_knn_hnsw_sparse core_knn_ivf core_knn_hnsw_rabitq core_knn_vamana + core_knn_flat_sparse core_knn_hnsw_sparse core_knn_ivf core_knn_hnsw_rabitq core_knn_ivf_rabitq core_knn_vamana SRCS ${CC_SRCS} INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm ) diff --git a/tests/db/CMakeLists.txt b/tests/db/CMakeLists.txt index bfad2bc57..b179eed48 100644 --- a/tests/db/CMakeLists.txt +++ b/tests/db/CMakeLists.txt @@ -32,6 +32,7 @@ foreach(CC_SRCS ${ALL_TEST_SRCS}) core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_rabitq + core_knn_ivf_rabitq core_knn_hnsw_sparse core_knn_ivf core_knn_diskann diff --git a/tests/db/collection_test.cc b/tests/db/collection_test.cc index 74407af6e..42e1e60f2 100644 --- a/tests/db/collection_test.cc +++ b/tests/db/collection_test.cc @@ -211,10 +211,10 @@ TEST_F(CollectionTest, Feature_OpenReadOnly_WithReadOnlyLockFile) { // Use std::filesystem to set read-only permissions (cross-platform) std::error_code ec; - fs::permissions(lock_path, - fs::perms::owner_read | fs::perms::group_read | - fs::perms::others_read, - fs::perm_options::replace, ec); + fs::permissions( + lock_path, + fs::perms::owner_read | fs::perms::group_read | fs::perms::others_read, + fs::perm_options::replace, ec); ASSERT_FALSE(ec) << "Failed to set read-only permissions: " << ec.message(); // Open with read_only=true should succeed even with read-only LOCK file @@ -2858,6 +2858,9 @@ TEST_F(CollectionTest, Feature_Optimize_Repeated) { run_repeated_optimize_test( enable_mmap, std::make_shared(MetricType::IP, 7, 256, 16, 200, 0)); + run_repeated_optimize_test( + enable_mmap, + std::make_shared(MetricType::IP, 32, 7, 0)); #endif } } @@ -5353,6 +5356,85 @@ TEST_F(CollectionTest, Feature_Optimize_HNSW_RABITQ) { } #endif +#if RABITQ_SUPPORTED +TEST_F(CollectionTest, Feature_Optimize_IVF_RABITQ) { + auto func = [](MetricType metric_type, int concurrency) { + FileHelper::RemoveDirectory(col_path); + + int doc_count = 1000; + + // create simple schema with only FP32 dense vector for IVF_RABITQ + auto schema = std::make_shared("demo"); + schema->set_max_doc_count_per_segment(MAX_DOC_COUNT_PER_SEGMENT); + + auto ivf_rabitq_params = + std::make_shared(metric_type, 32, 7, 0); + schema->add_field(std::make_shared( + "dense_fp32", DataType::VECTOR_FP32, 128, false, ivf_rabitq_params)); + + auto options = CollectionOptions{false, true, 64 * 1024 * 1024}; + auto collection = TestHelper::CreateCollectionWithDoc( + col_path, *schema, options, 0, doc_count, false); + + auto check_doc = [&]() { + for (int i = 0; i < doc_count; i++) { + auto expect_doc = TestHelper::CreateDoc(i, *schema); + auto result = collection->Fetch({expect_doc.pk()}); + ASSERT_TRUE(result.has_value()); + ASSERT_EQ(result.value().size(), 1); + ASSERT_EQ(result.value().count(expect_doc.pk()), 1); + auto doc = result.value()[expect_doc.pk()]; + ASSERT_NE(doc, nullptr); + if (metric_type != MetricType::COSINE) { + if (*doc != expect_doc) { + std::cout << " doc:" << doc->to_detail_string() << std::endl; + std::cout << "expect_doc:" << expect_doc.to_detail_string() + << std::endl; + } + ASSERT_EQ(*doc, expect_doc); + } + } + }; + + check_doc(); + std::cout << "check success 1" << std::endl; + + ASSERT_TRUE(collection->Flush().ok()); + auto stats = collection->Stats().value(); + ASSERT_EQ(stats.doc_count, doc_count); + ASSERT_EQ(stats.index_completeness["dense_fp32"], 0); + + auto s = collection->Optimize(OptimizeOptions{concurrency}); + if (!s.ok()) { + std::cout << s.message() << std::endl; + } + ASSERT_TRUE(s.ok()); + + stats = collection->Stats().value(); + ASSERT_EQ(stats.doc_count, doc_count); + ASSERT_EQ(stats.index_completeness["dense_fp32"], 1); + + check_doc(); + std::cout << "check success 2" << std::endl; + + collection.reset(); + auto result = Collection::Open(col_path, options); + ASSERT_TRUE(result.has_value()); + collection = std::move(result.value()); + + check_doc(); + std::cout << "check success 3" << std::endl; + }; + + func(MetricType::L2, 0); + func(MetricType::L2, 4); + func(MetricType::IP, 0); + func(MetricType::IP, 4); + func(MetricType::COSINE, 0); + func(MetricType::COSINE, 4); +} +#endif + #if DISKANN_SUPPORTED TEST_F(CollectionTest, Feature_Optimize_DiskAnn) { auto func = [](MetricType metric_type, int concurrency) { diff --git a/tests/db/index/common/db_proto_converter_test.cc b/tests/db/index/common/db_proto_converter_test.cc index 9c71c3c89..833804341 100644 --- a/tests/db/index/common/db_proto_converter_test.cc +++ b/tests/db/index/common/db_proto_converter_test.cc @@ -114,6 +114,35 @@ TEST(ConverterTest, IVFIndexParamsConversion) { EXPECT_EQ(pb_result.base().quantize_type(), proto::QT_INT4); } +#if RABITQ_SUPPORTED +TEST(ConverterTest, IvfRabitqIndexParamsConversion) { + proto::IvfRabitqIndexParams ivf_rabitq_pb; + auto *base_params = ivf_rabitq_pb.mutable_base(); + base_params->set_metric_type(proto::MT_COSINE); + base_params->set_quantize_type(proto::QT_RABITQ); + ivf_rabitq_pb.set_nlist(32); + ivf_rabitq_pb.set_total_bits(1); + ivf_rabitq_pb.set_sample_count(5); + + auto ivf_rabitq_params = ProtoConverter::FromPb(ivf_rabitq_pb); + ASSERT_NE(ivf_rabitq_params, nullptr); + EXPECT_EQ(ivf_rabitq_params->metric_type(), MetricType::COSINE); + EXPECT_EQ(ivf_rabitq_params->quantize_type(), QuantizeType::RABITQ); + EXPECT_EQ(ivf_rabitq_params->type(), IndexType::IVF_RABITQ); + EXPECT_EQ(ivf_rabitq_params->nlist(), 32); + EXPECT_EQ(ivf_rabitq_params->total_bits(), 1); + EXPECT_EQ(ivf_rabitq_params->sample_count(), 5); + + IvfRabitqIndexParams original_params(MetricType::IP, 64, 7, 3); + auto pb_result = ProtoConverter::ToPb(&original_params); + EXPECT_EQ(pb_result.base().metric_type(), proto::MT_IP); + EXPECT_EQ(pb_result.base().quantize_type(), proto::QT_RABITQ); + EXPECT_EQ(pb_result.nlist(), 64); + EXPECT_EQ(pb_result.total_bits(), 7); + EXPECT_EQ(pb_result.sample_count(), 3); +} +#endif + TEST(ConverterTest, IndexParamsConversion) { // Test conversion from protobuf to C++ IndexParams for HNSW proto::IndexParams index_pb; @@ -184,6 +213,38 @@ TEST(ConverterTest, IndexParamsConversion) { EXPECT_EQ(pb_result3.base().metric_type(), proto::MT_COSINE); EXPECT_EQ(pb_result3.n_list(), 256); +#if RABITQ_SUPPORTED + // Test conversion from protobuf to C++ IndexParams for IVF_RABITQ + proto::IndexParams index_pb5; + auto *ivf_rabitq_pb = index_pb5.mutable_ivf_rabitq(); + auto *base_params5 = ivf_rabitq_pb->mutable_base(); + base_params5->set_metric_type(proto::MT_IP); + base_params5->set_quantize_type(proto::QT_RABITQ); + ivf_rabitq_pb->set_nlist(32); + ivf_rabitq_pb->set_total_bits(1); + ivf_rabitq_pb->set_sample_count(5); + + auto index_params5 = ProtoConverter::FromPb(index_pb5); + ASSERT_NE(index_params5, nullptr); + EXPECT_EQ(index_params5->type(), IndexType::IVF_RABITQ); + auto ivf_rabitq_cast = + std::dynamic_pointer_cast(index_params5); + ASSERT_NE(ivf_rabitq_cast, nullptr); + EXPECT_EQ(ivf_rabitq_cast->metric_type(), MetricType::IP); + EXPECT_EQ(ivf_rabitq_cast->quantize_type(), QuantizeType::RABITQ); + EXPECT_EQ(ivf_rabitq_cast->nlist(), 32); + EXPECT_EQ(ivf_rabitq_cast->total_bits(), 1); + EXPECT_EQ(ivf_rabitq_cast->sample_count(), 5); + + IvfRabitqIndexParams ivf_rabitq_original(MetricType::COSINE, 64, 7, 3); + auto pb_result5 = ProtoConverter::ToPb(&ivf_rabitq_original); + EXPECT_EQ(pb_result5.base().metric_type(), proto::MT_COSINE); + EXPECT_EQ(pb_result5.base().quantize_type(), proto::QT_RABITQ); + EXPECT_EQ(pb_result5.nlist(), 64); + EXPECT_EQ(pb_result5.total_bits(), 7); + EXPECT_EQ(pb_result5.sample_count(), 3); +#endif + // Test conversion from protobuf to C++ IndexParams for INVERT proto::IndexParams index_pb4; auto *invert_pb = index_pb4.mutable_invert(); @@ -547,4 +608,4 @@ TEST(ConverterTest, IVFIndexParamsWithEnableRotate) { EXPECT_FALSE(pb2.base().quantizer_param().enable_rotate()); auto restored2 = ProtoConverter::FromPb(pb2); EXPECT_FALSE(restored2->quantizer_param().enable_rotate()); -} \ No newline at end of file +} diff --git a/tests/db/index/common/index_params_test.cc b/tests/db/index/common/index_params_test.cc index d5a85aeb9..9f6645aa4 100644 --- a/tests/db/index/common/index_params_test.cc +++ b/tests/db/index/common/index_params_test.cc @@ -163,6 +163,42 @@ TEST(IndexParamsTest, IVFIndexParams) { EXPECT_EQ(params.n_list(), 64); } +#if RABITQ_SUPPORTED +TEST(IndexParamsTest, IvfRabitqIndexParams) { + IvfRabitqIndexParams params(MetricType::COSINE, 32, 1, 5); + EXPECT_EQ(params.type(), IndexType::IVF_RABITQ); + EXPECT_EQ(params.metric_type(), MetricType::COSINE); + EXPECT_EQ(params.quantize_type(), QuantizeType::RABITQ); + EXPECT_EQ(params.nlist(), 32); + EXPECT_EQ(params.total_bits(), 1); + EXPECT_EQ(params.sample_count(), 5); + + auto cloned = params.clone(); + EXPECT_NE(cloned.get(), ¶ms); + EXPECT_EQ(cloned->type(), IndexType::IVF_RABITQ); + EXPECT_TRUE(*cloned == params); + + IvfRabitqIndexParams same_params(MetricType::COSINE, 32, 1, 5); + IvfRabitqIndexParams diff_metric(MetricType::IP, 32, 1, 5); + IvfRabitqIndexParams diff_nlist(MetricType::COSINE, 64, 1, 5); + IvfRabitqIndexParams diff_total_bits(MetricType::COSINE, 32, 7, 5); + IvfRabitqIndexParams diff_sample_count(MetricType::COSINE, 32, 1, 0); + + EXPECT_TRUE(params == same_params); + EXPECT_FALSE(params == diff_metric); + EXPECT_FALSE(params == diff_nlist); + EXPECT_FALSE(params == diff_total_bits); + EXPECT_FALSE(params == diff_sample_count); + + params.set_nlist(48); + params.set_total_bits(7); + params.set_sample_count(0); + EXPECT_EQ(params.nlist(), 48); + EXPECT_EQ(params.total_bits(), 7); + EXPECT_EQ(params.sample_count(), 0); +} +#endif + TEST(IndexParamsTest, DefaultVectorIndexParams) { // Test default vector index params EXPECT_EQ(DefaultVectorIndexParams.type(), IndexType::FLAT); @@ -278,4 +314,4 @@ TEST(IndexParamsTest, QuantizerParamWithVectorIndex) { IVFIndexParams params2(MetricType::IP, 128, 10, false, QuantizeType::INT8); EXPECT_FALSE(params == params2); } -} \ No newline at end of file +} diff --git a/tests/db/index/common/schema_test.cc b/tests/db/index/common/schema_test.cc index e45f22150..a1e73d54a 100644 --- a/tests/db/index/common/schema_test.cc +++ b/tests/db/index/common/schema_test.cc @@ -878,9 +878,8 @@ TEST(FieldSchemaTest, HnswRabitqIndexValidation_Dimension) { auto status = field.validate(); EXPECT_FALSE(status.ok()) << "Dimension 63 should not be supported with HNSW_RABITQ"; - EXPECT_NE( - status.message().find("HNSW_RABITQ index only support dimension in"), - std::string::npos) + EXPECT_NE(status.message().find("RabitQ index only support dimension in"), + std::string::npos) << "Error message should mention dimension range, got: " << status.message(); } @@ -905,9 +904,8 @@ TEST(FieldSchemaTest, HnswRabitqIndexValidation_Dimension) { auto status = field.validate(); EXPECT_FALSE(status.ok()) << "Dimension 4096 should not be supported with HNSW_RABITQ"; - EXPECT_NE( - status.message().find("HNSW_RABITQ index only support dimension in"), - std::string::npos) + EXPECT_NE(status.message().find("RabitQ index only support dimension in"), + std::string::npos) << "Error message should mention dimension range, got: " << status.message(); } @@ -936,6 +934,86 @@ TEST(FieldSchemaTest, HnswRabitqIndexValidation_Dimension) { << status.message(); } } + +TEST(FieldSchemaTest, IvfRabitqIndexValidationMetricTypes) { + // Test supported combinations: FP32 + (L2/IP/COSINE) + for (auto metric_type : + {MetricType::L2, MetricType::IP, MetricType::COSINE}) { + auto index_params = + std::make_shared(metric_type, 32, 1, 0); + FieldSchema field("vector_field", DataType::VECTOR_FP32, 128, false, + index_params); + auto status = field.validate(); + EXPECT_TRUE(status.ok()) + << "FP32 + IVF_RABITQ metric should be supported, but got error: " + << status.message(); + } + + { + auto index_params = + std::make_shared(MetricType::MIPSL2, 32, 1, 0); + FieldSchema field("vector_field", DataType::VECTOR_FP32, 128, false, + index_params); + auto status = field.validate(); + EXPECT_FALSE(status.ok()) + << "FP32 + MIPSL2 should not be supported with IVF_RABITQ"; + } +} + +TEST(FieldSchemaTest, IvfRabitqIndexValidationDimensionAndDataTypes) { + { + auto index_params = + std::make_shared(MetricType::L2, 32, 1, 0); + FieldSchema field("vector_field", DataType::VECTOR_FP32, 63, false, + index_params); + auto status = field.validate(); + EXPECT_FALSE(status.ok()) + << "Dimension 63 should not be supported with IVF_RABITQ"; + EXPECT_NE(status.message().find("RabitQ index only support dimension in"), + std::string::npos) + << "Error message should mention dimension range, got: " + << status.message(); + } + + { + auto index_params = + std::make_shared(MetricType::L2, 32, 1, 0); + FieldSchema field("vector_field", DataType::VECTOR_FP32, 64, false, + index_params); + auto status = field.validate(); + EXPECT_TRUE(status.ok()) + << "Dimension 64 should be supported, but got error: " + << status.message(); + } + + { + auto index_params = + std::make_shared(MetricType::L2, 32, 1, 0); + FieldSchema field("vector_field", DataType::VECTOR_FP32, 4096, false, + index_params); + auto status = field.validate(); + EXPECT_FALSE(status.ok()) + << "Dimension 4096 should not be supported with IVF_RABITQ"; + EXPECT_NE(status.message().find("RabitQ index only support dimension in"), + std::string::npos) + << "Error message should mention dimension range, got: " + << status.message(); + } + + { + auto index_params = + std::make_shared(MetricType::L2, 32, 1, 0); + FieldSchema field("vector_field", DataType::VECTOR_FP16, 128, false, + index_params); + auto status = field.validate(); + EXPECT_FALSE(status.ok()) << "FP16 should not be supported with IVF_RABITQ"; + EXPECT_NE( + status.message().find("RabitQ index only support FP32 data types"), + std::string::npos) + << "Error message should mention FP32 support only, got: " + << status.message(); + } +} #endif TEST(FieldSchemaTest, HnswRabitqIndexValidation_UnsupportedDataTypes) { @@ -951,7 +1029,7 @@ TEST(FieldSchemaTest, HnswRabitqIndexValidation_UnsupportedDataTypes) { EXPECT_FALSE(status.ok()) << "FP16 should not be supported with HNSW_RABITQ"; EXPECT_NE( - status.message().find("HNSW_RABITQ index only support FP32 data type"), + status.message().find("RabitQ index only support FP32 data types"), std::string::npos) << "Error message should mention FP32 support only, got: " << status.message(); @@ -967,7 +1045,7 @@ TEST(FieldSchemaTest, HnswRabitqIndexValidation_UnsupportedDataTypes) { EXPECT_FALSE(status.ok()) << "INT8 should not be supported with HNSW_RABITQ"; EXPECT_NE( - status.message().find("HNSW_RABITQ index only support FP32 data type"), + status.message().find("RabitQ index only support FP32 data types"), std::string::npos) << "Error message should mention FP32 support only, got: " << status.message(); diff --git a/tests/db/index/utils/utils.cc b/tests/db/index/utils/utils.cc index 827660b25..0ca338172 100644 --- a/tests/db/index/utils/utils.cc +++ b/tests/db/index/utils/utils.cc @@ -112,12 +112,12 @@ CollectionSchema::Ptr TestHelper::CreateNormalSchema( "dense_int8", DataType::VECTOR_INT8, 128, false, std::make_shared(MetricType::IP))); - // IVF, HNSW_RABITQ and DISKANN do not support sparse vectors, always use - // Flat for sparse fields in those cases. + // IVF, RabitQ and DISKANN do not support sparse vectors, always use Flat for + // sparse fields in those cases. auto supports_sparse = [](const IndexParams::Ptr ¶ms) { auto type = params->type(); return type != IndexType::IVF && type != IndexType::HNSW_RABITQ && - type != IndexType::DISKANN; + type != IndexType::IVF_RABITQ && type != IndexType::DISKANN; }; IndexParams::Ptr sparse_index_params; @@ -801,4 +801,4 @@ arrow::Status TestHelper::WriteTestFile(const std::string &filepath, ARROW_RETURN_NOT_OK(out->Close()); return arrow::Status::OK(); -} \ No newline at end of file +} diff --git a/tools/core/CMakeLists.txt b/tools/core/CMakeLists.txt index 86a3c4eac..1de295a00 100644 --- a/tools/core/CMakeLists.txt +++ b/tools/core/CMakeLists.txt @@ -14,7 +14,7 @@ cc_binary( STRICT PACKED SRCS local_builder.cc INCS ${PROJECT_ROOT_DIR}/src/core/ - LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_vamana core_knn_cluster core_knn_ivf core_interface core_knn_diskann + LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_ivf_rabitq core_knn_vamana core_knn_cluster core_knn_ivf core_interface core_knn_diskann ) cc_binary( @@ -22,7 +22,7 @@ cc_binary( STRICT PACKED SRCS recall.cc INCS ${PROJECT_ROOT_DIR}/src/core/ - LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_vamana core_knn_cluster core_knn_ivf roaring core_interface core_knn_diskann + LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_ivf_rabitq core_knn_vamana core_knn_cluster core_knn_ivf roaring core_interface core_knn_diskann ) cc_binary( @@ -30,7 +30,7 @@ cc_binary( STRICT PACKED SRCS bench.cc INCS ${PROJECT_ROOT_DIR}/src/core/ - LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_vamana core_knn_cluster core_knn_ivf roaring core_interface core_knn_diskann + LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_ivf_rabitq core_knn_vamana core_knn_cluster core_knn_ivf roaring core_interface core_knn_diskann ) @@ -39,7 +39,7 @@ cc_binary( STRICT PACKED SRCS recall_original.cc INCS ${PROJECT_ROOT_DIR}/src/core/ - LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_vamana core_knn_cluster core_knn_ivf roaring core_interface core_knn_diskann + LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_ivf_rabitq core_knn_vamana core_knn_cluster core_knn_ivf roaring core_interface core_knn_diskann ) cc_binary( @@ -47,7 +47,7 @@ cc_binary( STRICT PACKED SRCS bench_original.cc INCS ${PROJECT_ROOT_DIR}/src/core/ - LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_vamana core_knn_cluster core_knn_ivf roaring core_interface core_knn_diskann + LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_ivf_rabitq core_knn_vamana core_knn_cluster core_knn_ivf roaring core_interface core_knn_diskann ) cc_binary( @@ -55,5 +55,5 @@ cc_binary( STRICT PACKED SRCS local_builder_original.cc INCS ${PROJECT_ROOT_DIR}/src/core/ - LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_vamana core_knn_cluster core_knn_ivf core_interface core_knn_diskann + LIBS gflags yaml-cpp magic_enum core_framework core_metric core_quantizer core_utility core_knn_flat core_knn_flat_sparse core_knn_hnsw core_knn_hnsw_sparse core_knn_hnsw_rabitq core_knn_ivf_rabitq core_knn_vamana core_knn_cluster core_knn_ivf core_interface core_knn_diskann ) diff --git a/tools/core/local_builder.cc b/tools/core/local_builder.cc index f8ea9ca51..58a79f945 100644 --- a/tools/core/local_builder.cc +++ b/tools/core/local_builder.cc @@ -21,6 +21,7 @@ #include #include "algorithm/flat/flat_utility.h" #include "algorithm/hnsw_rabitq/hnsw_rabitq_params.h" +#include "algorithm/hnsw_rabitq/rabitq_params.h" #if RABITQ_SUPPORTED #include "algorithm/hnsw_rabitq/hnsw_rabitq_streamer.h" #include "algorithm/hnsw_rabitq/rabitq_converter.h" @@ -1049,6 +1050,8 @@ int do_build(YAML::Node &config_root, YAML::Node &config_common) { for (auto ¶m : id_map_param_list) { params.set(param, !g_disable_id_map); } + // Pass original dimension for Cosine support (before converter modifies it) + params.set(PARAM_RABITQ_GENERAL_DIMENSION, input_meta.dimension()); // INIT int ret =