diff --git a/.github/workflows/03-macos-linux-build.yml b/.github/workflows/03-macos-linux-build.yml index d74257679..b68f47a94 100644 --- a/.github/workflows/03-macos-linux-build.yml +++ b/.github/workflows/03-macos-linux-build.yml @@ -177,7 +177,7 @@ jobs: shell: bash # ------------------------------------------------------------------ # - # DiskAnn libaio round: install libaio and re-run tests + # DiskAnn libaio round: install libaio and re-run tests # ------------------------------------------------------------------ # - name: Install libaio runtime if: matrix.platform == 'linux-x64' diff --git a/python/tests/test_collection_ivf_rabitq.py b/python/tests/test_collection_ivf_rabitq.py new file mode 100644 index 000000000..5c115a2cd --- /dev/null +++ b/python/tests/test_collection_ivf_rabitq.py @@ -0,0 +1,261 @@ +# Copyright 2025-present the zvec project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import platform +import sys + +import pytest +import zvec +from zvec import ( + Collection, + CollectionOption, + CollectionSchema, + DataType, + Doc, + FieldSchema, + IndexType, + InvertIndexParam, + IvfRabitqIndexParam, + IvfRabitqQueryParam, + MetricType, + OptimizeOption, + Query, + QuantizeType, + VectorSchema, +) + +pytestmark = pytest.mark.skipif( + not (sys.platform == "linux" and platform.machine() in ("x86_64", "AMD64")), + reason="IVF RaBitQ only supported on Linux x86_64", +) + +DIMENSION = 128 +NLIST = 16 +NPROBE = 16 +TOPK = 10 +SMALL_DOC_COUNT = 64 +INDEXED_DOC_COUNT = 1200 + + +def _vector_for(doc_id: int) -> list[float]: + return [float(doc_id)] * DIMENSION + + +def _make_docs(count: int) -> list[Doc]: + return [ + Doc( + id=str(doc_id), + fields={ + "id": doc_id, + "name": f"test_doc_{doc_id}", + "group_id": doc_id % 4, + }, + vectors={"embedding": _vector_for(doc_id)}, + ) + for doc_id in range(count) + ] + + +def _build_schema(name: str) -> CollectionSchema: + return CollectionSchema( + name=name, + fields=[ + FieldSchema( + "id", + DataType.INT64, + nullable=False, + index_param=InvertIndexParam(enable_range_optimization=True), + ), + FieldSchema("name", DataType.STRING, nullable=False), + FieldSchema("group_id", DataType.INT64, nullable=False), + ], + vectors=[ + VectorSchema( + "embedding", + DataType.VECTOR_FP32, + dimension=DIMENSION, + index_param=IvfRabitqIndexParam( + metric_type=MetricType.L2, + nlist=NLIST, + total_bits=7, + sample_count=0, + ), + ), + ], + ) + + +def _create_collection(tmp_path_factory, schema_name: str) -> Collection: + path = tmp_path_factory.mktemp("zvec_ivf_rabitq") / schema_name + return zvec.create_and_open( + path=str(path), + schema=_build_schema(schema_name), + option=CollectionOption(read_only=False, enable_mmap=True), + ) + + +def _insert_docs(coll: Collection, docs: list[Doc]) -> None: + batch_size = 512 + for offset in range(0, len(docs), batch_size): + batch = docs[offset : offset + batch_size] + result = coll.insert(batch) + assert len(result) == len(batch) + for item in result: + assert item.ok() + assert coll.stats.doc_count == len(docs) + + +def _query( + coll: Collection, + query_vector: list[float], + *, + filter: str | None = None, + include_vector: bool = False, +) -> list[Doc]: + query = Query( + field_name="embedding", + vector=query_vector, + param=IvfRabitqQueryParam(nprobe=NPROBE), + ) + result = coll.query( + queries=query, + topk=TOPK, + filter=filter, + include_vector=include_vector, + ) + assert result is not None + assert 0 < len(result) <= TOPK + return result + + +class TestIvfRabitqCollection: + def test_schema_round_trip(self, tmp_path_factory): + coll = _create_collection(tmp_path_factory, "ivf_rabitq_schema_rt") + try: + vector_schema = coll.schema.vector("embedding") + assert vector_schema is not None + assert vector_schema.data_type == DataType.VECTOR_FP32 + assert vector_schema.dimension == DIMENSION + + index_param = vector_schema.index_param + assert index_param.type == IndexType.IVF_RABITQ + assert index_param.metric_type == MetricType.L2 + assert index_param.quantize_type == QuantizeType.RABITQ + assert index_param.nlist == NLIST + assert index_param.total_bits == 7 + assert index_param.sample_count == 0 + assert coll.stats.index_completeness["embedding"] == 1 + finally: + coll.destroy() + + def test_insert_fetch_and_query_writer_segment(self, tmp_path_factory): + coll = _create_collection(tmp_path_factory, "ivf_rabitq_writer_query") + try: + docs = _make_docs(SMALL_DOC_COUNT) + _insert_docs(coll, docs) + + fetched = coll.fetch(ids=["7", "42"]) + assert set(fetched.keys()) == {"7", "42"} + assert fetched["42"].field("name") == "test_doc_42" + assert fetched["42"].vector("embedding") == _vector_for(42) + + result = _query(coll, docs[42].vector("embedding")) + assert result[0].id == "42" + finally: + coll.destroy() + + def test_optimize_reopen_and_query(self, tmp_path_factory): + schema_name = "ivf_rabitq_optimize_reopen" + collection_path = tmp_path_factory.mktemp("zvec_ivf_rabitq") / schema_name + option = CollectionOption(read_only=False, enable_mmap=True) + coll = zvec.create_and_open( + path=str(collection_path), + schema=_build_schema(schema_name), + option=option, + ) + reopened = None + try: + docs = _make_docs(INDEXED_DOC_COUNT) + _insert_docs(coll, docs) + + coll.optimize(OptimizeOption()) + assert coll.stats.doc_count == INDEXED_DOC_COUNT + assert coll.stats.index_completeness["embedding"] == 1 + + query_vector = docs[42].vector("embedding") + result = _query(coll, query_vector) + assert result[0].id == "42" + + filtered = _query(coll, query_vector, filter="id < 50", include_vector=True) + assert filtered[0].id == "42" + for doc in filtered: + assert doc.field("id") < 50 + assert len(filtered[0].vector("embedding")) == DIMENSION + + path = str(collection_path) + del coll + coll = None + + reopened = zvec.open(path=path, option=option) + assert reopened.stats.doc_count == INDEXED_DOC_COUNT + reopened_result = _query(reopened, query_vector) + assert reopened_result[0].id == "42" + finally: + if reopened is not None: + reopened.destroy() + elif coll is not None: + coll.destroy() + + def test_group_by_after_optimize(self, tmp_path_factory): + coll = _create_collection(tmp_path_factory, "ivf_rabitq_group_by") + try: + docs = _make_docs(INDEXED_DOC_COUNT) + _insert_docs(coll, docs) + coll.optimize(OptimizeOption()) + + query = Query( + field_name="embedding", + vector=docs[42].vector("embedding"), + param=IvfRabitqQueryParam(nprobe=NPROBE), + ) + results = coll.group_by_query( + query, + group_by_field_name="group_id", + group_count=4, + topk_per_group=2, + ) + assert {int(group.group_by_value) for group in results} == set(range(4)) + for group in results: + assert 1 <= len(group.docs) <= 2 + assert all( + doc.field("group_id") == int(group.group_by_value) + for doc in group.docs + ) + + filtered = coll.group_by_query( + query, + group_by_field_name="group_id", + group_count=4, + topk_per_group=2, + filter="id < 8", + include_vector=True, + ) + assert {int(group.group_by_value) for group in filtered} == set(range(4)) + for group in filtered: + for doc in group.docs: + assert doc.field("id") < 8 + assert len(doc.vector("embedding")) == DIMENSION + finally: + coll.destroy() diff --git a/python/tests/test_params.py b/python/tests/test_params.py index f83c7cdb5..ca9df0fdc 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,60 @@ 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.scale_factor == 10.0 + assert param.type == IndexType.IVF_RABITQ + + def test_custom(self): + param = IvfRabitqQueryParam( + nprobe=20, + is_using_refiner=True, + radius=30, + is_linear=True, + scale_factor=3.5, + ) + assert param.nprobe == 20 + assert param.is_using_refiner == True + assert param.radius == 30 + assert param.is_linear == True + assert param.scale_factor == 3.5 + 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_pickle_roundtrip(self): + import pickle + + param = IvfRabitqQueryParam(nprobe=20, is_using_refiner=True, scale_factor=3.5) + restored = pickle.loads(pickle.dumps(param)) + assert restored.nprobe == 20 + assert restored.is_using_refiner == True + assert restored.scale_factor == 3.5 + + 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/tests/test_typing.py b/python/tests/test_typing.py index d566d5efc..f930d07fe 100644 --- a/python/tests/test_typing.py +++ b/python/tests/test_typing.py @@ -104,6 +104,7 @@ def test_data_type_has_member(member): "HNSW_RABITQ", "DISKANN", "VAMANA", + "IVF_RABITQ", "INVERT", "FTS", ], @@ -117,7 +118,7 @@ def test_io_backend_type_has_member(member): assert member in IOBackendType.__members__ -@pytest.mark.parametrize("member", ["FP16", "INT8", "INT4", "UNDEFINED"]) +@pytest.mark.parametrize("member", ["FP16", "INT8", "INT4", "RABITQ", "UNDEFINED"]) def test_quantize_type_has_member(member): assert member in QuantizeType.__members__ diff --git a/python/zvec/__init__.py b/python/zvec/__init__.py index 2dbcc172f..af26b3254 100644 --- a/python/zvec/__init__.py +++ b/python/zvec/__init__.py @@ -103,6 +103,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 09e828dae..74368f092 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, @@ -94,6 +96,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 a002d0de5..058552dd4 100644 --- a/python/zvec/model/collection.py +++ b/python/zvec/model/collection.py @@ -35,6 +35,7 @@ IndexOption, InvertIndexParam, IVFIndexParam, + IvfRabitqIndexParam, OptimizeOption, ) from .param.query import Query @@ -116,6 +117,7 @@ def create_index( index_param: Union[ HnswIndexParam, HnswRabitqIndexParam, + IvfRabitqIndexParam, IVFIndexParam, FlatIndexParam, InvertIndexParam, @@ -125,13 +127,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 a8c3b6560..d81c00524 100644 --- a/python/zvec/model/param/__init__.pyi +++ b/python/zvec/model/param/__init__.pyi @@ -27,6 +27,8 @@ __all__: list[str] = [ "IndexOption", "IndexParam", "InvertIndexParam", + "IvfRabitqIndexParam", + "IvfRabitqQueryParam", "OptimizeOption", "QuantizerParam", "QueryParam", @@ -587,6 +589,72 @@ 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, + scale_factor: typing.SupportsFloat = 10.0, + ) -> 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. + """ + + @property + def scale_factor(self) -> float: + """ + float: Candidate expansion factor used by the refiner. + """ + class IVFIndexParam(VectorIndexParam): """ diff --git a/python/zvec/model/param/query.py b/python/zvec/model/param/query.py index fd46d2caa..6c75b096e 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 fab03f4b9..5a63e84c0 100644 --- a/python/zvec/typing/__init__.pyi +++ b/python/zvec/typing/__init__.pyi @@ -184,6 +184,8 @@ class IndexType: IVF + IVF_RABITQ + FLAT HNSW_RABITQ @@ -204,11 +206,12 @@ class IndexType: HNSW_RABITQ: typing.ClassVar[IndexType] # value = INVERT: 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': , 'HNSW_RABITQ': , 'DISKANN': , 'VAMANA': , 'INVERT': , 'FTS': } + ] # value = {'UNDEFINED': , 'HNSW': , 'IVF': , 'IVF_RABITQ': , 'FLAT': , 'HNSW_RABITQ': , 'DISKANN': , 'VAMANA': , 'INVERT': , 'FTS': } 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 1dfce56f2..7402a4c2d 100644 --- a/src/binding/c/c_api.cc +++ b/src/binding/c/c_api.cc @@ -1364,6 +1364,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( @@ -1848,6 +1853,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) @@ -2827,6 +2881,8 @@ const char *zvec_index_type_to_string(zvec_index_type_t index_type) { return "DISKANN"; 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: @@ -4883,6 +4939,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. @@ -5187,6 +5244,132 @@ 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_scale_factor( + zvec_ivf_rabitq_query_params_t *params, float scale_factor) { + 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_scale_factor(scale_factor); + return ZVEC_OK; +} + +float zvec_query_params_ivf_rabitq_get_scale_factor( + const zvec_ivf_rabitq_query_params_t *params) { + if (!params) { + return 10.0f; + } + auto *ptr = reinterpret_cast(params); + return ptr->scale_factor(); +} + +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 // ============================================================================= @@ -5646,6 +5829,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) { @@ -6034,6 +6235,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) { @@ -6416,6 +6635,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 7524502be..2a7b2aff9 100644 --- a/src/binding/python/CMakeLists.txt +++ b/src/binding/python/CMakeLists.txt @@ -38,6 +38,7 @@ if (CMAKE_SYSTEM_NAME STREQUAL "Linux") $ $ $ + $ $ $ $ @@ -65,6 +66,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 3c3ff623b..feae5eaa1 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: @@ -704,6 +706,89 @@ 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_> @@ -1463,6 +1548,91 @@ 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. + scale_factor (float): Candidate expansion factor used by the refiner. + Default is 10.0. + +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, py::arg("scale_factor") = 10.0f, + 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. + scale_factor (float, optional): Candidate expansion factor used by the + refiner. Default is 10.0. +)pbdoc") + .def_property_readonly( + "nprobe", + [](const IvfRabitqQueryParams &self) -> int { return self.nprobe(); }, + "int: Number of inverted lists to search during IVF RaBitQ query.") + .def_property_readonly( + "scale_factor", + [](const IvfRabitqQueryParams &self) -> float { + return self.scale_factor(); + }, + "float: Candidate expansion factor used by the refiner.") + .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()) + + ", \"scale_factor\":" + + std::to_string(self.scale_factor()) + "}"; + }) + .def(py::pickle( + [](const IvfRabitqQueryParams &self) { + return py::make_tuple(self.nprobe(), self.radius(), + self.is_linear(), self.is_using_refiner(), + self.scale_factor()); + }, + [](py::tuple t) { + if (t.size() != 4 && t.size() != 5) { + 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()); + if (t.size() == 5) { + obj->set_scale_factor(t[4].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 d20cac755..7d512085a 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("UNDEFINED", IndexType::UNDEFINED) .value("HNSW", IndexType::HNSW) .value("IVF", IndexType::IVF) + .value("IVF_RABITQ", IndexType::IVF_RABITQ) .value("FLAT", IndexType::FLAT) .value("HNSW_RABITQ", IndexType::HNSW_RABITQ) .value("DISKANN", IndexType::DISKANN) @@ -245,4 +246,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 d875176f2..c9491a930 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 @@ -58,6 +76,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() # Exclude algorithm/diskann implementation files from zvec_core when not diff --git a/src/core/algorithm/CMakeLists.txt b/src/core/algorithm/CMakeLists.txt index f874eba62..d87fe55d9 100644 --- a/src/core/algorithm/CMakeLists.txt +++ b/src/core/algorithm/CMakeLists.txt @@ -53,6 +53,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 @@ -70,4 +71,33 @@ 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" + ) + + if(MSVC) + # MSVC cannot create an import library for an empty DLL. Keep the + # cross-platform static target name expected by the Python binding. + cc_library( + NAME core_knn_ivf_rabitq + STATIC 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}" + ) + add_library(core_knn_ivf_rabitq_static ALIAS core_knn_ivf_rabitq) + else() + 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() 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..b02546bf3 --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_builder.cc @@ -0,0 +1,550 @@ +// 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", + static_cast(state_)); + return IndexError_Logic; + } + + meta_ = meta; + params_ = params; + + int configured_nlist = static_cast(kDefaultIvfRabitqNlist); + if (params.get(PARAM_IVF_RABITQ_NLIST, &configured_nlist) && + configured_nlist <= 0) { + LOG_ERROR("Invalid nlist=%d, must be greater than 0", configured_nlist); + return IndexError_InvalidArgument; + } + nlist_ = static_cast(configured_nlist); + + 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=%zu, must be in [1, 9]", (size_t)total_bits_); + return IndexError_InvalidArgument; + } + + int configured_sample_count = 0; + if (params.get(PARAM_RABITQ_SAMPLE_COUNT, &configured_sample_count) && + configured_sample_count < 0) { + LOG_ERROR("Invalid sample_count=%d, must be greater than or equal to 0", + configured_sample_count); + return IndexError_InvalidArgument; + } + sample_count_ = static_cast(configured_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=%zu, nlist=%zu, total_bits=%zu, " + "metric=%s, sample_count=%zu", + (size_t)dim, (size_t)nlist_, (size_t)total_bits_, metric_name_.c_str(), + (size_t)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", + static_cast(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: %zu clusters, cost %zu ms", + (size_t)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..181dd2876 --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_context.h @@ -0,0 +1,361 @@ +// 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 +#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 + * - group_state_: lazily allocated only for group-by search + */ +class IvfRabitqContext : public IndexContext { + public: + typedef std::shared_ptr Pointer; + using GroupTopkHeaps = std::map; + + 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; + } + + int64_t val = 0; + if (params.get(PARAM_IVF_RABITQ_NPROBE, &val)) { + if (val < 0) { + LOG_ERROR("Invalid nprobe, must be greater than or equal to 0"); + return IndexError_InvalidArgument; + } + nprobe_ = static_cast(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(); + group_state_.reset(); + 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]; + } + + const IndexGroupDocumentList &group_result(void) const override { + return group_result(0); + } + + const IndexGroupDocumentList &group_result(size_t idx) const override { + static const IndexGroupDocumentList kEmpty; + if (!group_state_ || idx >= group_state_->results.size()) { + return kEmpty; + } + return group_state_->results[idx]; + } + + IndexGroupDocumentList *mutable_group_result(void) override { + return mutable_group_result(0); + } + + IndexGroupDocumentList *mutable_group_result(size_t idx) override { + if (!group_state_) { + return nullptr; + } + if (idx >= group_state_->results.size()) { + group_state_->results.resize(idx + 1); + } + return &group_state_->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()); + } + + void reset_group_results(size_t qnum) { + if (!group_state_) { + return; + } + group_state_->results.resize(qnum); + for (auto &result : group_state_->results) { + result.clear(); + } + group_state_->heaps.clear(); + } + + //! 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); + } + } + + void topk_to_group_result(uint32_t idx) { + if (!group_state_ || idx >= group_state_->results.size()) { + return; + } + + auto &result = group_state_->results[idx]; + result.clear(); + + std::vector> ranked_groups; + ranked_groups.reserve(group_state_->heaps.size()); + for (auto &entry : group_state_->heaps) { + auto &heap = entry.second; + heap.sort(); + if (!heap.empty()) { + ranked_groups.emplace_back(&entry.first, heap[0].score()); + } + } + std::sort(ranked_groups.begin(), ranked_groups.end(), + [](const auto &lhs, const auto &rhs) { + if (lhs.second != rhs.second) { + return lhs.second < rhs.second; + } + return *lhs.first < *rhs.first; + }); + + const size_t group_count = std::min( + static_cast(group_state_->group_num), ranked_groups.size()); + result.reserve(group_count); + for (size_t i = 0; i < group_count; ++i) { + const std::string &group_id = *ranked_groups[i].first; + auto heap_it = group_state_->heaps.find(group_id); + if (heap_it == group_state_->heaps.end()) { + continue; + } + auto &heap = heap_it->second; + result.emplace_back(); + auto &group = result.back(); + group.set_group_id(group_id); + auto *docs = group.mutable_docs(); + const size_t doc_count = + std::min(static_cast(group_state_->group_topk), heap.size()); + docs->reserve(doc_count); + for (size_t j = 0; j < doc_count; ++j) { + if (heap[j].score() > this->threshold()) { + break; + } + docs->emplace_back(heap[j].key(), heap[j].score()); + } + } + } + + bool group_by_search() const { + return group_state_ != nullptr; + } + + void set_group_params(uint32_t group_num, uint32_t group_topk) override { + if (group_num == 0 || group_topk == 0) { + group_state_.reset(); + return; + } + if (!group_state_) { + group_state_ = std::make_unique(); + } + group_state_->group_num = group_num; + group_state_->group_topk = group_topk; + group_state_->heaps.clear(); + group_state_->results.clear(); + } + + uint32_t group_topk() const { + return group_state_ ? group_state_->group_topk : 0; + } + + GroupTopkHeaps &group_topk_heaps() { + return group_state_->heaps; + } + + // ----------------------------------------------------------------------- + // 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: + struct GroupSearchState { + uint32_t group_num{0}; + uint32_t group_topk{0}; + GroupTopkHeaps heaps; + std::vector results; + }; + + IndexDocumentHeap result_heap_; + std::vector results_; + std::unique_ptr group_state_; + + 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..0492cb85d --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_entity.cc @@ -0,0 +1,799 @@ +// 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 "algorithm/hnsw_rabitq/rabitq_params.h" +#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; + } + if (header_seg->data_size() != sizeof(IvfRabitqHeader)) { + LOG_ERROR("Invalid IVF RaBitQ header segment size=%zu, expected=%zu", + header_seg->data_size(), sizeof(IvfRabitqHeader)); + 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)); + int ret = validate_header(); + if (ret != 0) { + return ret; + } + + // 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 = + static_cast(header_.cluster_count) * sizeof(IvfRabitqClusterMeta); + if (cluster_meta_seg->data_size() != meta_size) { + LOG_ERROR("Invalid cluster meta segment size=%zu, expected=%zu", + cluster_meta_seg->data_size(), meta_size); + return IndexError_InvalidFormat; + } + 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; + } + size_t batch_data_size = static_cast(header_.batch_data_size); + if (static_cast(batch_data_size) != header_.batch_data_size || + batch_seg->data_size() != batch_data_size) { + LOG_ERROR("Invalid batch data segment size=%zu, expected=%zu", + batch_seg->data_size(), batch_data_size); + return IndexError_InvalidFormat; + } + read_size = batch_seg->read(0, batch_data_block_, batch_data_size); + if (read_size != batch_data_size) { + LOG_ERROR("Failed to read batch data: got %zu, expected %zu", read_size, + 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; + } + size_t ex_data_size = static_cast(header_.ex_data_size); + if (static_cast(ex_data_size) != header_.ex_data_size || + ex_seg->data_size() != ex_data_size) { + LOG_ERROR("Invalid extra-bit data segment size=%zu, expected=%zu", + ex_seg->data_size(), ex_data_size); + return IndexError_InvalidFormat; + } + read_size = ex_seg->read(0, ex_data_block_, ex_data_size); + if (read_size != 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 = + static_cast(header_.total_vector_count) * sizeof(uint64_t); + if (keys_seg->data_size() != keys_size) { + LOG_ERROR("Invalid keys segment size=%zu, expected=%zu", + keys_seg->data_size(), keys_size); + return IndexError_InvalidFormat; + } + 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()); + + ret = load_key_order_mapping(storage); + if (ret != 0) { + return ret; + } + ret = init_quantized_vector_layout(); + if (ret != 0) { + return ret; + } + ret = validate_cluster_metas(); + if (ret != 0) { + return ret; + } + + LOG_INFO( + "IvfRabitqEntity loaded: total_vectors=%zu, clusters=%zu, " + "padded_dim=%zu, ex_bits=%zu, batch_data_size=%zu, ex_data_size=%zu", + (size_t)header_.total_vector_count, (size_t)header_.cluster_count, + (size_t)header_.padded_dim, (size_t)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); + if (header_.ex_bits > 0 && !ip_func_) { + LOG_ERROR("Unsupported IVF RaBitQ ex_bits=%zu", (size_t)header_.ex_bits); + return IndexError_InvalidFormat; + } + + 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); +} + +int IvfRabitqEntity::search_cluster_group_by( + uint32_t cluster_id, const IvfRabitqQueryState &query_state, + size_t padded_dim, size_t ex_bits, const IndexGroupBy &group_by, + uint32_t group_topk, float threshold, + IvfRabitqContext::GroupTopkHeaps *heaps) const { + return search_cluster_group_by_impl( + cluster_id, query_state, padded_dim, ex_bits, nullptr, group_by, + group_topk, threshold, heaps); +} + +int IvfRabitqEntity::search_cluster_group_by( + uint32_t cluster_id, const IvfRabitqQueryState &query_state, + size_t padded_dim, size_t ex_bits, const IndexFilter &filter, + const IndexGroupBy &group_by, uint32_t group_topk, float threshold, + IvfRabitqContext::GroupTopkHeaps *heaps) const { + return search_cluster_group_by_impl(cluster_id, query_state, padded_dim, + ex_bits, &filter, group_by, + group_topk, threshold, heaps); +} + +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=%zu", (size_t)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=%zu, vector_count=%zu", + (size_t)local_id, (size_t)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(); +} + +uint32_t IvfRabitqEntity::get_cluster_id(uint32_t id) const { + const auto *meta = find_cluster_meta(id); + if (!meta) { + return std::numeric_limits::max(); + } + return static_cast(meta - cluster_metas_.data()); +} + +int IvfRabitqEntity::validate_header() const { + if (header_.magic != kIvfRabitqIndexMagic) { + LOG_ERROR("Invalid IVF RaBitQ index magic: got=%zu, expected=%zu", + (size_t)header_.magic, (size_t)kIvfRabitqIndexMagic); + return IndexError_InvalidFormat; + } + if (header_.version != kIvfRabitqIndexVersion) { + LOG_ERROR("Unsupported IVF RaBitQ index version: got=%zu, expected=%zu", + (size_t)header_.version, (size_t)kIvfRabitqIndexVersion); + return IndexError_InvalidFormat; + } + if (header_.dimension < kMinRabitqDimSize || + header_.dimension > kMaxRabitqDimSize) { + LOG_ERROR("Invalid IVF RaBitQ dimension=%zu", (size_t)header_.dimension); + return IndexError_InvalidFormat; + } + uint32_t expected_padded_dim = ((header_.dimension + 63U) / 64U) * 64U; + if (header_.padded_dim != expected_padded_dim) { + LOG_ERROR("Invalid IVF RaBitQ padded_dim=%zu, expected=%zu", + (size_t)header_.padded_dim, (size_t)expected_padded_dim); + return IndexError_InvalidFormat; + } + if (header_.total_vector_count == 0 || header_.cluster_count == 0 || + header_.cluster_count > header_.total_vector_count) { + LOG_ERROR("Invalid IVF RaBitQ vector_count=%zu, cluster_count=%zu", + (size_t)header_.total_vector_count, + (size_t)header_.cluster_count); + return IndexError_InvalidFormat; + } + if (header_.ex_bits > 8) { + LOG_ERROR("Invalid IVF RaBitQ ex_bits=%zu", (size_t)header_.ex_bits); + return IndexError_InvalidFormat; + } + return 0; +} + +int IvfRabitqEntity::validate_cluster_metas() const { + size_t expected_key_offset = 0; + for (size_t cluster_id = 0; cluster_id < cluster_metas_.size(); + ++cluster_id) { + const auto &meta = cluster_metas_[cluster_id]; + size_t expected_batch_count = (static_cast(meta.vector_count) + + rabitqlib::fastscan::kBatchSize - 1) / + rabitqlib::fastscan::kBatchSize; + size_t cluster_batch_size = + static_cast(meta.batch_count) * batch_data_size_; + size_t cluster_ex_size = + static_cast(meta.vector_count) * ex_data_size_; + if (meta.batch_count != expected_batch_count || + meta.key_offset != expected_key_offset || + meta.key_offset > header_.total_vector_count || + meta.vector_count > header_.total_vector_count - meta.key_offset || + meta.batch_data_offset > header_.batch_data_size || + cluster_batch_size > header_.batch_data_size - meta.batch_data_offset || + meta.ex_data_offset > header_.ex_data_size || + cluster_ex_size > header_.ex_data_size - meta.ex_data_offset) { + LOG_ERROR("Invalid IVF RaBitQ cluster range, cluster_id=%zu", cluster_id); + return IndexError_InvalidFormat; + } + expected_key_offset += meta.vector_count; + } + if (expected_key_offset != header_.total_vector_count) { + LOG_ERROR("Invalid IVF RaBitQ cluster key coverage=%zu, expected=%zu", + expected_key_offset, (size_t)header_.total_vector_count); + return IndexError_InvalidFormat; + } + return 0; +} + +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; +} + +int IvfRabitqEntity::compute_distances(uint32_t cluster_id, + const std::vector &ids, + const IvfRabitqQueryState &query_state, + size_t padded_dim, size_t ex_bits, + IndexDocumentList *documents) const { + if (!documents) { + return IndexError_InvalidArgument; + } + if (cluster_id >= cluster_metas_.size()) { + return IndexError_OutOfRange; + } + const auto *batch_query = + static_cast *>( + query_state.batch_query.get()); + if (!batch_query) { + return IndexError_InvalidArgument; + } + + const auto &meta = cluster_metas_[cluster_id]; + std::vector sorted_ids(ids); + std::sort(sorted_ids.begin(), sorted_ids.end()); + + const size_t batch_stride = + rabitqlib::BatchDataMap::data_bytes(padded_dim); + const size_t ex_stride = + rabitqlib::ExDataMap::data_bytes(padded_dim, ex_bits); + std::array est_dist; + std::array low_dist; + std::array ip_x0_qr; + uint32_t loaded_batch = std::numeric_limits::max(); + + documents->clear(); + documents->reserve(sorted_ids.size()); + for (uint32_t id : sorted_ids) { + if (static_cast(id) < meta.key_offset || + static_cast(id) >= meta.key_offset + meta.vector_count) { + return IndexError_InvalidArgument; + } + const uint32_t local_id = id - static_cast(meta.key_offset); + const uint32_t batch_id = + local_id / static_cast(rabitqlib::fastscan::kBatchSize); + const uint32_t lane = + local_id % static_cast(rabitqlib::fastscan::kBatchSize); + if (batch_id != loaded_batch) { + const char *batch_data = batch_data_ + meta.batch_data_offset + + (static_cast(batch_id) * batch_stride); + rabitqlib::split_batch_estdist(batch_data, *batch_query, padded_dim, + est_dist.data(), low_dist.data(), + ip_x0_qr.data(), true /* use_hacc */); + loaded_batch = batch_id; + } + + float distance = est_dist[lane]; + if (ex_bits > 0) { + const char *ex_data = ex_data_ + meta.ex_data_offset + + (static_cast(local_id) * ex_stride); + distance = rabitqlib::split_distance_boosting( + ex_data, ip_func_, *batch_query, padded_dim, ex_bits, ip_x0_qr[lane]); + } + documents->emplace_back(keys_[id], distance); + } + 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=%zu", (size_t)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; +} + +template +int IvfRabitqEntity::search_cluster_group_by_impl( + uint32_t cluster_id, const IvfRabitqQueryState &query_state, + size_t padded_dim, size_t ex_bits, const IndexFilter *filter, + const IndexGroupBy &group_by, uint32_t group_topk, float threshold, + IvfRabitqContext::GroupTopkHeaps *heaps) const { + if (cluster_id >= cluster_metas_.size()) { + return IndexError_OutOfRange; + } + if (!heaps || !group_by.is_valid() || group_topk == 0) { + return IndexError_InvalidArgument; + } + + const auto &meta = cluster_metas_[cluster_id]; + if (meta.vector_count == 0) { + return 0; + } + + const auto *batch_query = + static_cast *>( + query_state.batch_query.get()); + if (!batch_query) { + 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 *batch_data = batch_data_ + meta.batch_data_offset; + const char *ex_data = ex_bits > 0 ? ex_data_ + meta.ex_data_offset : nullptr; + const uint64_t *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); + std::array est_dist; + std::array low_dist; + std::array ip_x0_qr; + rabitqlib::split_batch_estdist(batch_data, *batch_query, padded_dim, + est_dist.data(), low_dist.data(), + ip_x0_qr.data(), true /* use_hacc */); + + for (uint32_t i = 0; i < batch_size; ++i) { + const uint64_t key = keys[offset + i]; + if constexpr (HasFilter) { + if ((*filter)(key)) { + continue; + } + } + + const std::string group_id = group_by(key); + auto &heap = (*heaps)[group_id]; + if (heap.empty()) { + heap.limit(group_topk); + heap.set_threshold(threshold); + } + + if (ex_bits == 0) { + heap.emplace(key, est_dist[i]); + continue; + } + + const float group_threshold = + heap.full() ? heap.begin()->score() : threshold; + if (low_dist[i] >= group_threshold) { + continue; + } + const float distance = rabitqlib::split_distance_boosting( + ex_data + (static_cast(i) * ex_stride), ip_func_, + *batch_query, padded_dim, ex_bits, ip_x0_qr[i]); + heap.emplace(key, distance); + } + + batch_data += batch_stride; + if (ex_bits > 0) { + ex_data += static_cast(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..9b4b73bfd --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_entity.h @@ -0,0 +1,181 @@ +// 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; + + int search_cluster_group_by(uint32_t cluster_id, + const IvfRabitqQueryState &query_state, + size_t padded_dim, size_t ex_bits, + const IndexGroupBy &group_by, uint32_t group_topk, + float threshold, + IvfRabitqContext::GroupTopkHeaps *heaps) const; + + int search_cluster_group_by(uint32_t cluster_id, + const IvfRabitqQueryState &query_state, + size_t padded_dim, size_t ex_bits, + const IndexFilter &filter, + const IndexGroupBy &group_by, uint32_t group_topk, + float threshold, + IvfRabitqContext::GroupTopkHeaps *heaps) const; + + int compute_distances(uint32_t cluster_id, const std::vector &ids, + const IvfRabitqQueryState &query_state, + size_t padded_dim, size_t ex_bits, + IndexDocumentList *documents) 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; + } + uint32_t dimension() const { + return header_.dimension; + } + 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; + + uint32_t get_cluster_id(uint32_t id) 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; + + template + int search_cluster_group_by_impl( + uint32_t cluster_id, const IvfRabitqQueryState &query_state, + size_t padded_dim, size_t ex_bits, const IndexFilter *filter, + const IndexGroupBy &group_by, uint32_t group_topk, float threshold, + IvfRabitqContext::GroupTopkHeaps *heaps) 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 validate_header() const; + + int validate_cluster_metas() const; + + 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..66e486ee6 --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_index_format.h @@ -0,0 +1,65 @@ +// 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 { + +inline constexpr uint32_t kIvfRabitqIndexMagic = 0x49565251U; // "IVRQ" +inline constexpr uint32_t kIvfRabitqIndexVersion = 1U; + +struct IvfRabitqHeader { + uint32_t magic; + uint32_t version; + 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[4]; + + IvfRabitqHeader() { + memset(static_cast(this), 0, sizeof(IvfRabitqHeader)); + magic = kIvfRabitqIndexMagic; + version = kIvfRabitqIndexVersion; + } +}; +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..31f445815 --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_reformer.cc @@ -0,0 +1,606 @@ +// 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 { + +namespace { + +struct ReformerLayout { + size_t rotated_centroids_size{0}; + size_t centroids_size{0}; + size_t rotator_size{0}; +}; + +int ValidateReformerHeader(const RabitqConverterHeader &header, + size_t segment_size, ReformerLayout *layout) { + if (!layout) { + return IndexError_InvalidArgument; + } + if (header.dim < kMinRabitqDimSize || header.dim > kMaxRabitqDimSize) { + LOG_ERROR("Invalid IVF RaBitQ reformer dimension=%zu", (size_t)header.dim); + return IndexError_InvalidFormat; + } + uint32_t expected_padded_dim = ((header.dim + 63U) / 64U) * 64U; + if (header.padded_dim != expected_padded_dim || header.num_clusters == 0 || + header.ex_bits > 8 || + header.rotator_type > + static_cast(rabitqlib::RotatorType::FhtKacRotator)) { + LOG_ERROR( + "Invalid IVF RaBitQ reformer header: padded_dim=%zu, clusters=%zu, " + "ex_bits=%zu, rotator_type=%zu", + (size_t)header.padded_dim, (size_t)header.num_clusters, + (size_t)header.ex_bits, (size_t)header.rotator_type); + return IndexError_InvalidFormat; + } + + layout->rotated_centroids_size = static_cast(header.num_clusters) * + header.padded_dim * sizeof(float); + layout->centroids_size = + static_cast(header.num_clusters) * header.dim * sizeof(float); + if (header.rotator_type == + static_cast(rabitqlib::RotatorType::MatrixRotator)) { + layout->rotator_size = + static_cast(header.dim) * header.padded_dim * sizeof(float); + } else { + layout->rotator_size = header.padded_dim / 2U; + } + if (header.rotator_size != layout->rotator_size) { + LOG_ERROR("Invalid IVF RaBitQ rotator size=%zu, expected=%zu", + (size_t)header.rotator_size, layout->rotator_size); + return IndexError_InvalidFormat; + } + + size_t expected_size = sizeof(RabitqConverterHeader) + + layout->rotated_centroids_size + + layout->centroids_size + layout->rotator_size; + if (expected_size != segment_size) { + LOG_ERROR("Invalid IVF RaBitQ reformer segment size=%zu, expected=%zu", + segment_size, expected_size); + return IndexError_InvalidFormat; + } + return 0; +} + +} // namespace + +// 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)); + ReformerLayout layout; + int ret = ValidateReformerHeader(header, segment->data_size(), &layout); + if (ret != 0) { + return ret; + } + 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 = segment->read(offset, block, layout.rotated_centroids_size); + if (size != layout.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(), + layout.rotated_centroids_size); + offset += size; + + // Read original centroids + size = segment->read(offset, block, layout.centroids_size); + if (size != layout.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(), layout.centroids_size); + offset += size; + + // Read rotator + size = segment->read(offset, block, layout.rotator_size); + if (size != layout.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)); + if (!impl_->rotator) { + LOG_ERROR("Failed to create IVF RaBitQ rotator"); + return IndexError_InvalidFormat; + } + 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(); + 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, + static_cast(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=%zu, num_clusters=%zu", (size_t)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=%zu, num_clusters=%zu", (size_t)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..8ec4bdfdd --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_streamer.cc @@ -0,0 +1,638 @@ +// 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 "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 + int64_t configured_nprobe = static_cast(kDefaultIvfRabitqNprobe); + if (params.get(PARAM_IVF_RABITQ_NPROBE, &configured_nprobe) && + configured_nprobe < 0) { + LOG_ERROR("Invalid nprobe, must be greater than or equal to 0"); + return IndexError_InvalidArgument; + } + nprobe_ = static_cast(configured_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=%zu, nprobe=%zu, metric=%s, " + "scan_ratio=%.3f, bf_threshold=%zu", + (size_t)dim, (size_t)nprobe_, metric_name_.c_str(), scan_ratio_, + (size_t)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; + + // Validate the main header and all entity segment ranges before loading the + // reformer, whose persisted dimensions control allocations and rotator setup. + entity_ = std::make_shared(); + int ret = entity_->load(storage); + if (ret != 0) { + LOG_ERROR("Failed to load IvfRabitqEntity, ret=%d", ret); + return ret; + } + if (entity_->dimension() != rabitq_meta_.dimension()) { + LOG_ERROR("RaBitQ dimension mismatch: entity=%zu, meta=%zu", + (size_t)entity_->dimension(), + static_cast(rabitq_meta_.dimension())); + return IndexError_InvalidFormat; + } + + // Load reformer + reformer_ = std::make_shared(); + 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() != entity_->dimension() || + reformer_->padded_dim() != entity_->padded_dim() || + reformer_->ex_bits() != entity_->ex_bits() || + reformer_->num_clusters() != entity_->cluster_count()) { + LOG_ERROR( + "IVF RaBitQ entity and reformer metadata mismatch: " + "dimension=%zu/%zu, padded_dim=%zu/%zu, ex_bits=%zu/%zu, " + "clusters=%zu/%zu", + (size_t)entity_->dimension(), reformer_->dimension(), + (size_t)entity_->padded_dim(), reformer_->padded_dim(), + (size_t)entity_->ex_bits(), reformer_->ex_bits(), + (size_t)entity_->cluster_count(), reformer_->num_clusters()); + return IndexError_InvalidFormat; + } + + stats_.set_loaded_count(entity_->total_vector_count()); + stats_.set_loaded_costtime(timer.milli_seconds()); + + LOG_INFO("IvfRabitqStreamer loaded: %zu vectors, %zu clusters, cost %zu ms", + (size_t)entity_->total_vector_count(), + (size_t)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 { + if (context && context->group_by().is_valid()) { + return search_group_by_impl_internal(query, qmeta, 1, context, true); + } + 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 { + if (context && context->group_by().is_valid()) { + return search_group_by_impl_internal(query, qmeta, count, context, true); + } + 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 { + if (context && context->group_by().is_valid()) { + return search_group_by_impl_internal(query, qmeta, count, context, false); + } + 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 %zu, ret=%d", (size_t)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 %zu, ret=%d", (size_t)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::search_group_by_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() || !entity_) { + return IndexError_NoReady; + } + + auto *ctx = dynamic_cast(context.get()); + if (!ctx) { + return IndexError_Cast; + } + if (!ctx->group_by_search() || !ctx->group_by().is_valid()) { + LOG_ERROR("Invalid group-by state"); + return IndexError_InvalidArgument; + } + + const bool brute_force = force_brute_force || entity_->total_vector_count() <= + ctx->bruteforce_threshold(); + 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(); + } + + const size_t padded_dim = reformer_->padded_dim(); + const size_t ex_bits = reformer_->ex_bits(); + const size_t dimension = reformer_->dimension(); + if (qmeta.dimension() != meta_.dimension() || + qmeta.data_type() != meta_.data_type() || + qmeta.element_size() != meta_.element_size() || + qmeta.dimension() < dimension) { + return IndexError_Mismatch; + } + + ctx->reset_group_results(count); + for (uint32_t q = 0; q < count; ++q) { + const float *query_vector = reinterpret_cast( + static_cast(query) + + (static_cast(q) * qmeta.element_size())); + IvfRabitqQueryState query_state; + int ret = reformer_->create_query_state(query_vector, &query_state); + if (ret != 0) { + return ret; + } + + std::vector probe_centroids; + ret = reformer_->select_probe_centroids(query_vector, nprobe, &query_state, + &probe_centroids); + if (ret != 0) { + return ret; + } + + auto &heaps = ctx->group_topk_heaps(); + heaps.clear(); + const auto &filter = ctx->filter(); + uint32_t scanned = 0; + for (uint32_t cluster_id : probe_centroids) { + if (scanned >= max_scan) { + break; + } + ret = reformer_->prepare_for_cluster(cluster_id, &query_state); + if (ret != 0) { + LOG_ERROR("Failed to prepare for cluster %zu, ret=%d", + (size_t)cluster_id, ret); + continue; + } + + if (filter.is_valid()) { + ret = entity_->search_cluster_group_by( + cluster_id, query_state, padded_dim, ex_bits, filter, + ctx->group_by(), ctx->group_topk(), ctx->threshold(), &heaps); + } else { + ret = entity_->search_cluster_group_by( + cluster_id, query_state, padded_dim, ex_bits, ctx->group_by(), + ctx->group_topk(), ctx->threshold(), &heaps); + } + if (ret != 0) { + LOG_ERROR("Failed to search cluster %zu, ret=%d", (size_t)cluster_id, + ret); + continue; + } + scanned += entity_->cluster_meta(cluster_id).vector_count; + } + ctx->topk_to_group_result(q); + } + return 0; +} + +int IvfRabitqStreamer::search_bf_by_p_keys_impl( + const void *query, const std::vector> &p_keys, + const IndexQueryMeta &qmeta, uint32_t count, + Context::Pointer &context) const { + if (!query || p_keys.size() != count) { + return IndexError_InvalidArgument; + } + if (!reformer_ || !reformer_->loaded() || !entity_) { + return IndexError_NoReady; + } + if (qmeta.dimension() != meta_.dimension() || + qmeta.data_type() != meta_.data_type() || + qmeta.element_size() != meta_.element_size()) { + return IndexError_Mismatch; + } + + auto *ctx = dynamic_cast(context.get()); + if (!ctx) { + return IndexError_Cast; + } + const bool has_group_by = ctx->group_by_search(); + if (has_group_by && !ctx->group_by().is_valid()) { + return IndexError_InvalidArgument; + } + if (has_group_by) { + ctx->reset_group_results(count); + } else { + ctx->reset_results(count); + } + + const size_t padded_dim = reformer_->padded_dim(); + const size_t ex_bits = reformer_->ex_bits(); + const uint32_t topk = ctx->topk() == 0 ? 10 : ctx->topk(); + const auto &filter = ctx->filter(); + for (uint32_t q = 0; q < count; ++q) { + const float *query_vector = reinterpret_cast( + static_cast(query) + + (static_cast(q) * qmeta.element_size())); + IvfRabitqQueryState query_state; + int ret = reformer_->create_query_state(query_vector, &query_state); + if (ret != 0) { + return ret; + } + + std::map> cluster_ids; + for (uint64_t key : p_keys[q]) { + if (filter.is_valid() && filter(key)) { + continue; + } + const uint32_t id = entity_->key_to_id(key); + if (id == std::numeric_limits::max()) { + continue; + } + const uint32_t cluster_id = entity_->get_cluster_id(id); + if (cluster_id == std::numeric_limits::max()) { + continue; + } + cluster_ids[cluster_id].push_back(id); + } + + auto &result_heap = ctx->mutable_result_heap(); + if (!has_group_by) { + result_heap.clear(); + result_heap.limit(topk); + result_heap.set_threshold(ctx->threshold()); + } else { + ctx->group_topk_heaps().clear(); + } + + for (const auto &entry : cluster_ids) { + const uint32_t cluster_id = entry.first; + ret = reformer_->prepare_for_cluster(cluster_id, &query_state); + if (ret != 0) { + return ret; + } + IndexDocumentList documents; + ret = entity_->compute_distances(cluster_id, entry.second, query_state, + padded_dim, ex_bits, &documents); + if (ret != 0) { + return ret; + } + for (const auto &document : documents) { + if (!has_group_by) { + result_heap.emplace(document.key(), document.score()); + continue; + } + const std::string group_id = ctx->group_by()(document.key()); + auto &heap = ctx->group_topk_heaps()[group_id]; + if (heap.empty()) { + heap.limit(ctx->group_topk()); + heap.set_threshold(ctx->threshold()); + } + heap.emplace(document.key(), document.score()); + } + } + + if (has_group_by) { + ctx->topk_to_group_result(q); + } else { + 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..f0f083e25 --- /dev/null +++ b/src/core/algorithm/ivf_rabitq/ivf_rabitq_streamer.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 "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; + + int search_bf_by_p_keys_impl(const void *query, + const std::vector> &p_keys, + 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; + + int search_group_by_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..7d39632f5 --- /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=%zu, must be in [%d, %d]", (size_t)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.cc b/src/core/interface/index.cc index 56bf9f7f9..6c7db8a09 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -557,10 +557,11 @@ int Index::Search(const VectorData &vector_data, return core::IndexError_Runtime; } - if (_prepare_for_search(vector_data, search_param, context) != 0) { + int prepare_ret = _prepare_for_search(vector_data, search_param, context); + if (prepare_ret != 0) { LOG_ERROR("Failed to prepare for search"); context->reset(); - return core::IndexError_Runtime; + return prepare_ret; } if (is_sparse_) { 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 29ce8f4cc..992120c72 100644 --- a/src/core/interface/index_param.cc +++ b/src/core/interface/index_param.cc @@ -66,6 +66,20 @@ BaseIndexQueryParam::Pointer HNSWRabitqQueryParam::Clone() const { return std::make_shared(*this); } +IVFRabitqQueryParam::IVFRabitqQueryParam() = default; +IVFRabitqQueryParam::IVFRabitqQueryParam(const IVFRabitqQueryParam &) = default; +IVFRabitqQueryParam::IVFRabitqQueryParam(IVFRabitqQueryParam &&) noexcept = + default; +IVFRabitqQueryParam &IVFRabitqQueryParam::operator=( + const IVFRabitqQueryParam &) = default; +IVFRabitqQueryParam &IVFRabitqQueryParam::operator=( + IVFRabitqQueryParam &&) noexcept = default; +IVFRabitqQueryParam::~IVFRabitqQueryParam() = default; + +BaseIndexQueryParam::Pointer IVFRabitqQueryParam::Clone() const { + return std::make_shared(*this); +} + IVFQueryParam::IVFQueryParam() = default; IVFQueryParam::IVFQueryParam(const IVFQueryParam &) = default; IVFQueryParam::IVFQueryParam(IVFQueryParam &&) noexcept = default; @@ -147,6 +161,20 @@ HNSWRabitqIndexParam &HNSWRabitqIndexParam::operator=(HNSWRabitqIndexParam &&) = default; HNSWRabitqIndexParam::~HNSWRabitqIndexParam() = default; +IVFRabitqIndexParam::IVFRabitqIndexParam() + : BaseIndexParam(IndexType::kIVFRabitq) {} +IVFRabitqIndexParam::IVFRabitqIndexParam(int nlist) + : BaseIndexParam(IndexType::kIVFRabitq), nlist(nlist) {} +IVFRabitqIndexParam::IVFRabitqIndexParam(MetricType metric, int dim, int nlist) + : BaseIndexParam(IndexType::kIVFRabitq, metric, dim), nlist(nlist) {} +IVFRabitqIndexParam::IVFRabitqIndexParam(const IVFRabitqIndexParam &) = default; +IVFRabitqIndexParam::IVFRabitqIndexParam(IVFRabitqIndexParam &&) = default; +IVFRabitqIndexParam &IVFRabitqIndexParam::operator=( + const IVFRabitqIndexParam &) = default; +IVFRabitqIndexParam &IVFRabitqIndexParam::operator=(IVFRabitqIndexParam &&) = + default; +IVFRabitqIndexParam::~IVFRabitqIndexParam() = default; + ailego::JsonObject BaseIndexParam::SerializeToJsonObject( bool omit_empty_value) const { ailego::JsonObject json_obj; @@ -312,6 +340,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); 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..6d6092727 --- /dev/null +++ b/src/core/interface/indexes/ivf_rabitq_index.cc @@ -0,0 +1,351 @@ +// 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_.nlist <= 0) { + LOG_ERROR("nlist must be greater than 0, got %d", param_.nlist); + return core::IndexError_InvalidArgument; + } + if (param_.sample_count < 0) { + LOG_ERROR("sample_count must be greater than or equal to 0, got %d", + param_.sample_count); + return core::IndexError_InvalidArgument; + } + + 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 + (void)doc_id; + (void)vector_data_buffer; + LOG_ERROR("Fetch is not supported for IVF RaBitQ index"); + return core::IndexError_Unsupported; +#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); + + if (!ivf_rabitq_param) { + LOG_ERROR("Invalid search param type, expected IVFRabitqQueryParam"); + return core::IndexError_InvalidArgument; + } + if (ivf_rabitq_param->nprobe == 0) { + LOG_ERROR("nprobe must be greater than 0"); + return core::IndexError_InvalidArgument; + } + if (search_param->fetch_vector) { + LOG_ERROR("fetch_vector is not supported for IVF RaBitQ index"); + return core::IndexError_Unsupported; + } + + _set_group_by_on_context(search_param, context); + context->set_topk(search_param->topk); + context->set_fetch_vector(false); + if (search_param->filter) { + context->set_filter(std::move(*search_param->filter)); + } + if (search_param->radius > 0.0f) { + context->set_threshold(search_param->radius); + } + ailego::Params params; + params.set(core::PARAM_IVF_RABITQ_NPROBE, ivf_rabitq_param->nprobe); + return context->update(params); +#endif // RABITQ_SUPPORTED +} + +} // namespace zvec::core_interface diff --git a/src/db/index/column/vector_column/combined_vector_column_indexer.cc b/src/db/index/column/vector_column/combined_vector_column_indexer.cc index 5e3fe8066..ef6bbeeb8 100644 --- a/src/db/index/column/vector_column/combined_vector_column_indexer.cc +++ b/src/db/index/column/vector_column/combined_vector_column_indexer.cc @@ -339,6 +339,9 @@ Result CombinedVectorColumnIndexer::Search( } else if (q_params->type() == IndexType::IVF) { scale_factor = std::dynamic_pointer_cast(q_params)->scale_factor(); + } else if (q_params->type() == IndexType::IVF_RABITQ) { + scale_factor = std::dynamic_pointer_cast(q_params) + ->scale_factor(); } need_refine = true; } diff --git a/src/db/index/column/vector_column/engine_helper.hpp b/src/db/index/column/vector_column/engine_helper.hpp index dec1177cc..5d1524901 100644 --- a/src/db/index/column/vector_column/engine_helper.hpp +++ b/src/db/index/column/vector_column/engine_helper.hpp @@ -217,6 +217,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( @@ -456,6 +475,26 @@ 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()); + + 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/query.cc b/src/db/index/common/query.cc index ca94a2db9..3f20aa943 100644 --- a/src/db/index/common/query.cc +++ b/src/db/index/common/query.cc @@ -178,6 +178,18 @@ Status QueryTarget::validate(const FieldSchema *schema, IndexTypeCodeBook::AsString(schema->index_type()), " but got ", IndexTypeCodeBook::AsString(query_params->type())); } + if (query_params && query_params->type() == IndexType::IVF_RABITQ) { + auto ivf_rabitq_params = + std::dynamic_pointer_cast(query_params); + if (!ivf_rabitq_params) { + return Status::InvalidArgument( + "Invalid query: IVF_RABITQ index requires IvfRabitqQueryParams"); + } + if (ivf_rabitq_params->nprobe() <= 0) { + return Status::InvalidArgument( + "Invalid query: IVF_RABITQ nprobe must be greater than 0"); + } + } return Status::OK(); } diff --git a/src/db/index/common/schema.cc b/src/db/index/common/schema.cc index 532696803..fe5a1398f 100644 --- a/src/db/index/common/schema.cc +++ b/src/db/index/common/schema.cc @@ -56,8 +56,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}; @@ -153,30 +154,32 @@ Status FieldSchema::validate() const { support_dense_vector_index.end()) { return Status::InvalidArgument( "schema validate failed: dense_vector's index_params only " - "support FLAT|HNSW|HNSW_RABITQ|IVF|DISKANN|VAMANA index, but " + "support FLAT|HNSW|HNSW_RABITQ|IVF|IVF_RABITQ|DISKANN|VAMANA " + "index, but " "field[", name_, "]'s index_type is ", IndexTypeCodeBook::AsString(index_params_->type())); } } - 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 @@ -197,6 +200,26 @@ Status FieldSchema::validate() const { } } + if (index_params_->type() == IndexType::IVF_RABITQ) { + auto ivf_rabitq_params = + std::dynamic_pointer_cast(index_params_); + if (!ivf_rabitq_params) { + return Status::InvalidArgument( + "schema validate failed: IVF_RABITQ index requires " + "IvfRabitqIndexParams"); + } + if (ivf_rabitq_params->nlist() <= 0) { + return Status::InvalidArgument( + "schema validate failed: IVF_RABITQ nlist must be greater than " + "0"); + } + if (ivf_rabitq_params->sample_count() < 0) { + return Status::InvalidArgument( + "schema validate failed: IVF_RABITQ sample_count must be " + "greater than or equal to 0"); + } + } + if (index_params_->type() == IndexType::DISKANN) { // DiskAnn requires Linux x86_64/i686/i386. The CMake variable // DISKANN_SUPPORTED (defined in the top-level CMakeLists.txt) is the 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 a4f2e5a94..bfae09128 100644 --- a/src/db/index/segment/segment.cc +++ b/src/db/index/segment/segment.cc @@ -1769,10 +1769,10 @@ Status SegmentImpl::create_vector_index( block.set_max_doc_id(meta()->max_doc_id()); block.set_doc_count(meta()->doc_count()); new_segment_meta->add_persisted_block(block); - if (vector_index_params->quantize_type() == QuantizeType::RABITQ) { + if (vector_index_params->type() == IndexType::HNSW_RABITQ) { raw_vector_provider = vector_indexer.value()->create_index_provider(); } - } else { + } else if (vector_index_params->type() == IndexType::HNSW_RABITQ) { raw_vector_provider = vector_indexers_[column][0]->create_index_provider(); } @@ -1780,9 +1780,9 @@ Status SegmentImpl::create_vector_index( auto field_with_new_index_params = std::make_shared(*field); 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. - // For other quantize types, field_with_new_index_params is reused as-is. + // For HNSW_RABITQ, PrepareQuantizeField trains a reformer with + // raw_vector_provider. IVF_RABITQ trains inside its builder. For other + // quantize types, field_with_new_index_params is reused as-is. std::shared_ptr field_for_quantize; { auto s = SegmentHelper::PrepareQuantizeField(*field_with_new_index_params, diff --git a/src/db/index/segment/segment_helper.cc b/src/db/index/segment/segment_helper.cc index 1238a4865..e2fc6d8e4 100644 --- a/src/db/index/segment/segment_helper.cc +++ b/src/db/index/segment/segment_helper.cc @@ -698,11 +698,11 @@ Status SegmentHelper::ReduceVectorIndex( concurrency, &vector_indexer); CHECK_RETURN_STATUS(s); - // The training step (for RABITQ) and the subsequent quantize merge both - // rely on the raw provider held by the flat indexer, so its Close() is - // deferred until after the quantize indexer is written. + // HNSW_RABITQ training relies on the raw provider held by the flat + // indexer, so its Close() is deferred until after the quantize indexer + // is written. core::IndexProvider::Pointer raw_vector_provider; - if (vector_index_params->quantize_type() == QuantizeType::RABITQ) { + if (vector_index_params->type() == IndexType::HNSW_RABITQ) { raw_vector_provider = vector_indexer->create_index_provider(); } @@ -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) { @@ -871,6 +871,11 @@ Status SegmentHelper::PrepareQuantizeField( return Status::NotSupported( "RabitQ is not supported on this platform (Linux x86_64 only)"); #else + if (vector_index_params->type() == IndexType::IVF_RABITQ) { + *out_field = field_clone; + return Status::OK(); + } + if (raw_vector_provider == nullptr) { return Status::InvalidArgument( "raw_vector_provider is required for RABITQ training"); diff --git a/src/db/index/segment/segment_helper.h b/src/db/index/segment/segment_helper.h index ab8c3a9df..e2c5ccb29 100644 --- a/src/db/index/segment/segment_helper.h +++ b/src/db/index/segment/segment_helper.h @@ -272,13 +272,15 @@ class SegmentHelper { // Returns a FieldSchema clone whose index_params is ready for building the // quantize indexer. - // - RABITQ: clones HnswRabitqIndexParams, trains a RabitqConverter against - // `raw_vector_provider`, and attaches the resulting reformer and raw + // - HNSW_RABITQ: clones index params, trains a RabitqConverter against + // `raw_vector_provider`, and attaches the resulting reformer and // provider to the cloned params. + // - IVF_RABITQ: uses the cloned field unchanged because its builder owns + // KMeans and RaBitQ training. // - Other quantize types: clones the field with its current index_params // unchanged. - // `raw_vector_provider` must remain alive until the quantize indexer has - // been flushed; it may be null for non-RABITQ cases. + // `raw_vector_provider` must remain alive until an HNSW_RABITQ quantize + // indexer has been flushed; it may be null for other cases. static Status PrepareQuantizeField( const FieldSchema &field, const core::IndexProvider::Pointer &raw_vector_provider, 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 2d048689b..7be0bb4ef 100644 --- a/src/include/zvec/c_api.h +++ b/src/include/zvec/c_api.h @@ -874,6 +874,7 @@ typedef uint32_t zvec_index_type_t; #define ZVEC_INDEX_TYPE_HNSW_RABITQ 4 #define ZVEC_INDEX_TYPE_DISKANN 5 #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 @@ -903,6 +904,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) @@ -1148,6 +1150,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) @@ -1242,6 +1267,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) * @@ -1540,6 +1575,118 @@ 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 candidate expansion factor used by the refiner + * @param params IVF RaBitQ query parameters pointer + * @param scale_factor Candidate expansion factor + * @return zvec_error_code_t Error code + */ +ZVEC_EXPORT zvec_error_code_t ZVEC_CALL +zvec_query_params_ivf_rabitq_set_scale_factor( + zvec_ivf_rabitq_query_params_t *params, float scale_factor); + +/** + * @brief Get candidate expansion factor used by the refiner + * @param params IVF RaBitQ query parameters pointer + * @return float Candidate expansion factor + */ +ZVEC_EXPORT float ZVEC_CALL zvec_query_params_ivf_rabitq_get_scale_factor( + 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) // ----------------------------------------------------------------------------- @@ -2016,6 +2163,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 @@ -2314,6 +2471,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 @@ -2588,6 +2756,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 3006718e6..b8c4ec1ee 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 @@ -350,6 +351,35 @@ class ZVEC_CORE_API HNSWRabitqIndex : public Index { HNSWRabitqIndexParam param_{}; }; +class ZVEC_CORE_API 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 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 ZVEC_CORE_API 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 877e6e3a0..c8dd53bf9 100644 --- a/src/include/zvec/core/interface/index_param.h +++ b/src/include/zvec/core/interface/index_param.h @@ -73,6 +73,7 @@ enum class IndexType { kHNSWRabitq, kDiskAnn, kVamana, + kIVFRabitq, }; enum class IVFSearchMethod { kBF, kHNSW }; @@ -457,6 +458,46 @@ struct ZVEC_CORE_API HNSWRabitqIndexParam : public BaseIndexParam { bool omit_empty_value = false) const override; }; +struct ZVEC_CORE_API IVFRabitqIndexParam : public BaseIndexParam { + using Pointer = std::shared_ptr; + + // IVF parameters + int nlist = kDefaultIvfRabitqNlist; + + // Rabitq parameters + int total_bits = kDefaultRabitqTotalBits; + int sample_count = 0; + + IVFRabitqIndexParam(); + explicit IVFRabitqIndexParam(int nlist); + IVFRabitqIndexParam(MetricType metric, int dim, int nlist); + IVFRabitqIndexParam(const IVFRabitqIndexParam &); + IVFRabitqIndexParam(IVFRabitqIndexParam &&); + IVFRabitqIndexParam &operator=(const IVFRabitqIndexParam &); + IVFRabitqIndexParam &operator=(IVFRabitqIndexParam &&); + ~IVFRabitqIndexParam() override; + + protected: + bool DeserializeFromJsonObject(const ailego::JsonObject &json_obj) override; + ailego::JsonObject SerializeToJsonObject( + bool omit_empty_value = false) const override; +}; + +struct ZVEC_CORE_API IVFRabitqQueryParam : public BaseIndexQueryParam { + using Pointer = std::shared_ptr; + + IVFRabitqQueryParam(); + IVFRabitqQueryParam(const IVFRabitqQueryParam &); + IVFRabitqQueryParam(IVFRabitqQueryParam &&) noexcept; + IVFRabitqQueryParam &operator=(const IVFRabitqQueryParam &); + IVFRabitqQueryParam &operator=(IVFRabitqQueryParam &&) noexcept; + ~IVFRabitqQueryParam() override; + + uint32_t nprobe = kDefaultIvfRabitqNprobe; + + BaseIndexQueryParam::Pointer Clone() const override; +}; + struct ZVEC_CORE_API 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..e7f365b9e 100644 --- a/src/include/zvec/core/interface/index_param_builders.h +++ b/src/include/zvec/core/interface/index_param_builders.h @@ -204,6 +204,28 @@ 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; + } + std::shared_ptr Build() override { + return param; + } +}; + class DiskAnnIndexParamBuilder : public BaseIndexParamBuilder { @@ -511,4 +533,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 5509e1ead..e71082b33 100644 --- a/src/include/zvec/db/index_params.h +++ b/src/include/zvec/db/index_params.h @@ -55,7 +55,8 @@ class ZVEC_API 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 { @@ -398,6 +399,71 @@ class ZVEC_API HnswRabitqIndexParams : public VectorIndexParams { core::IndexReformer::Pointer rabitq_reformer_; }; +class ZVEC_API 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 { + return std::make_shared(metric_type_, nlist_, + total_bits_, sample_count_); + } + + 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_; + } + + private: + int nlist_; + int total_bits_; + int sample_count_; +}; + class ZVEC_API 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 9ef417b16..60b7a8094 100644 --- a/src/include/zvec/db/query_params.h +++ b/src/include/zvec/db/query_params.h @@ -171,6 +171,40 @@ class ZVEC_API HnswRabitqQueryParams : public QueryParams { int ef_; }; +class ZVEC_API IvfRabitqQueryParams : public QueryParams { + public: + IvfRabitqQueryParams(int nprobe = core_interface::kDefaultIvfRabitqNprobe, + float radius = 0.0f, bool is_linear = false, + bool is_using_refiner = false, + float scale_factor = 10.0f) + : QueryParams(IndexType::IVF_RABITQ), nprobe_(nprobe) { + set_radius(radius); + set_is_linear(is_linear); + set_is_using_refiner(is_using_refiner); + set_scale_factor(scale_factor); + } + + ~IvfRabitqQueryParams() override = default; + + int nprobe() const { + return nprobe_; + } + void set_nprobe(int nprobe) { + nprobe_ = nprobe; + } + + float scale_factor() const { + return scale_factor_; + } + void set_scale_factor(float scale_factor) { + scale_factor_ = scale_factor; + } + + private: + int nprobe_; + float scale_factor_{10.0f}; +}; + class ZVEC_API 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 7365dcf1a..9a26bd09e 100644 --- a/tests/c/c_api_test.c +++ b/tests/c/c_api_test.c @@ -3492,12 +3492,31 @@ void test_index_params_functions(void) { TEST_ASSERT(zvec_index_params_get_diskann_list_size(diskann_params) == 50); TEST_ASSERT(zvec_index_params_get_diskann_pq_chunk_num(diskann_params) == 0); + // 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(diskann_params); + zvec_index_params_destroy(ivf_rabitq_params); TEST_END(); } @@ -3716,6 +3735,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); @@ -3758,6 +3797,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); zvec_index_params_destroy(diskann_params); @@ -3792,6 +3832,14 @@ void test_utility_functions(void) { index_type_str = zvec_index_type_to_string(ZVEC_INDEX_TYPE_INVERT); TEST_ASSERT(index_type_str != NULL); + index_type_str = zvec_index_type_to_string(ZVEC_INDEX_TYPE_HNSW_RABITQ); + TEST_ASSERT(index_type_str != NULL && + strcmp(index_type_str, "HNSW_RABITQ") == 0); + + index_type_str = zvec_index_type_to_string(ZVEC_INDEX_TYPE_IVF_RABITQ); + TEST_ASSERT(index_type_str != NULL && + strcmp(index_type_str, "IVF_RABITQ") == 0); + TEST_END(); } @@ -3823,6 +3871,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); @@ -3880,6 +3933,35 @@ 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 scale factor setting + err = zvec_query_params_ivf_rabitq_set_scale_factor(ivf_rabitq_params, 3.5f); + TEST_ASSERT(err == ZVEC_OK); + TEST_ASSERT( + zvec_query_params_ivf_rabitq_get_scale_factor(ivf_rabitq_params) == 3.5f); + + // 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); @@ -3955,6 +4037,7 @@ 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); zvec_query_params_diskann_destroy(diskann_params); @@ -3963,6 +4046,7 @@ void test_query_params_functions(void) { // 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); zvec_query_params_diskann_destroy(NULL); @@ -3972,6 +4056,10 @@ 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_ivf_rabitq_set_scale_factor(NULL, 2.0f); + 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); @@ -3984,15 +4072,20 @@ 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_ivf_rabitq_get_scale_factor(NULL) == 10.0f); TEST_ASSERT(zvec_query_params_flat_get_radius(NULL) == 0.0f); TEST_ASSERT(zvec_query_params_diskann_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_diskann_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..60a2f9840 --- /dev/null +++ b/tests/core/algorithm/ivf_rabitq/ivf_rabitq_streamer_test.cc @@ -0,0 +1,1091 @@ +// 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_entity.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, TestParamValidation) { + auto init_builder = [&](const ailego::Params ¶ms) { + auto builder = std::make_shared(); + return builder->init(*index_meta_ptr_, params); + }; + + ailego::Params params; + params.set(PARAM_IVF_RABITQ_NLIST, 0); + EXPECT_EQ(IndexError_InvalidArgument, init_builder(params)); + + params.set(PARAM_IVF_RABITQ_NLIST, 1025); + EXPECT_EQ(0, init_builder(params)); + + params.set(PARAM_IVF_RABITQ_NLIST, 32); + params.set(PARAM_RABITQ_SAMPLE_COUNT, -1); + EXPECT_EQ(IndexError_InvalidArgument, init_builder(params)); + + IvfRabitqContext context; + ailego::Params search_params; + search_params.set(PARAM_IVF_RABITQ_NPROBE, 1025); + EXPECT_EQ(0, context.update(search_params)); + + auto init_streamer = [&](const ailego::Params &streamer_params) { + auto streamer = std::make_shared(); + return streamer->init(*index_meta_ptr_, streamer_params); + }; + EXPECT_EQ(0, init_streamer(search_params)); + + search_params.set(PARAM_IVF_RABITQ_NPROBE, -1); + EXPECT_EQ(IndexError_InvalidArgument, context.update(search_params)); + EXPECT_EQ(IndexError_InvalidArgument, init_streamer(search_params)); + + search_params.set(PARAM_IVF_RABITQ_NPROBE, 0); + EXPECT_EQ(0, init_streamer(search_params)); +} + +TEST_F(IvfRabitqStreamerTest, TestRejectInvalidFormatMagicAndVersion) { + auto holder = BuildHolder(kDim, 64); + ailego::Params params; + params.set(PARAM_IVF_RABITQ_NLIST, 16U); + params.set(PARAM_RABITQ_TOTAL_BITS, 4U); + + auto builder = std::make_shared(); + ASSERT_EQ(0, builder->init(*index_meta_ptr_, params)); + ASSERT_EQ(0, builder->train(nullptr, holder)); + ASSERT_EQ(0, builder->build(nullptr, holder)); + + const std::string file_id{"ivf_rabitq_invalid_format"}; + IndexMemory::Instance()->remove(file_id); + auto dumper = IndexFactory::CreateDumper("MemoryDumper"); + ASSERT_NE(nullptr, dumper); + ASSERT_EQ(0, dumper->init(ailego::Params())); + ASSERT_EQ(0, dumper->create(file_id)); + ASSERT_EQ(0, builder->dump(dumper)); + ASSERT_EQ(0, dumper->close()); + + auto storage = IndexFactory::CreateStorage("MemoryReadStorage"); + ASSERT_NE(nullptr, storage); + ASSERT_EQ(0, storage->init(ailego::Params())); + ASSERT_EQ(0, storage->open(file_id, false)); + + auto header_segment = storage->get(IVF_RABITQ_HEADER_SEG_ID); + ASSERT_NE(nullptr, header_segment); + const void *header_data = nullptr; + ASSERT_EQ(sizeof(IvfRabitqHeader), + header_segment->read(0, &header_data, sizeof(IvfRabitqHeader))); + auto *header = const_cast( + static_cast(header_data)); + ASSERT_EQ(kIvfRabitqIndexMagic, header->magic); + ASSERT_EQ(kIvfRabitqIndexVersion, header->version); + + header->magic = 0; + IvfRabitqEntity invalid_magic_entity; + EXPECT_EQ(IndexError_InvalidFormat, invalid_magic_entity.load(storage)); + + header->magic = kIvfRabitqIndexMagic; + header->version = kIvfRabitqIndexVersion + 1; + IvfRabitqEntity invalid_version_entity; + EXPECT_EQ(IndexError_InvalidFormat, invalid_version_entity.load(storage)); + + ASSERT_EQ(0, storage->close()); + IndexMemory::Instance()->remove(file_id); +} + +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=%zu, score=%f", + result.size(), (size_t)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.set_group_by([](uint64_t key) { return std::to_string(key % 3U); }); + context.set_group_params(/*group_num=*/3, /*group_topk=*/2); + 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_TRUE(context.group_by_search()); + ASSERT_NE(std::numeric_limits::max(), context.threshold()); + + context.reset(); + + EXPECT_FALSE(context.fetch_vector()); + EXPECT_FALSE(context.filter().is_valid()); + EXPECT_FALSE(context.group_by_search()); + EXPECT_FALSE(context.group_by().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, TestSearchByPrimaryKeys) { + constexpr size_t kDocCount = 256; + auto holder = BuildHolder(kDim, kDocCount); + + ailego::Params params; + params.set(PARAM_IVF_RABITQ_NLIST, 16U); + params.set(PARAM_RABITQ_TOTAL_BITS, 1U); + IndexStreamer::Pointer streamer; + BuildAndOpenStreamer(*index_meta_ptr_, holder, params, + dir_ + "/TestSearchByPrimaryKeys", &streamer); + ASSERT_NE(nullptr, streamer); + + NumericalVector query_vec0(kDim); + NumericalVector query_vec1(kDim); + FillTestVector(37, &query_vec0); + FillTestVector(42, &query_vec1); + std::vector query_vectors(kDim * 2); + memcpy(query_vectors.data(), query_vec0.data(), kDim * sizeof(float)); + memcpy(query_vectors.data() + kDim, query_vec1.data(), kDim * sizeof(float)); + IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, kDim); + auto context = streamer->create_context(); + context->set_topk(3); + + const std::vector> keys{{5, 37, 91, 9999}, + {7, 42, 100, 9999}}; + ASSERT_EQ(0, streamer->search_bf_by_p_keys_impl(query_vectors.data(), keys, + query_meta, 2, context)); + ASSERT_EQ(3UL, context->result(0).size()); + EXPECT_EQ(37UL, context->result(0)[0].key()); + for (const auto &doc : context->result(0)) { + EXPECT_NE(keys[0].end(), + std::find(keys[0].begin(), keys[0].end(), doc.key())); + } + ASSERT_EQ(3UL, context->result(1).size()); + EXPECT_EQ(42UL, context->result(1)[0].key()); + for (const auto &doc : context->result(1)) { + EXPECT_NE(keys[1].end(), + std::find(keys[1].begin(), keys[1].end(), doc.key())); + } + + context->set_filter([](uint64_t key) { return key == 37 || key == 42; }); + ASSERT_EQ(0, streamer->search_bf_by_p_keys_impl(query_vectors.data(), keys, + query_meta, 2, context)); + ASSERT_EQ(2UL, context->result(0).size()); + ASSERT_EQ(2UL, context->result(1).size()); + for (uint32_t q = 0; q < 2; ++q) { + for (const auto &doc : context->result(q)) { + EXPECT_NE(37UL, doc.key()); + EXPECT_NE(42UL, doc.key()); + } + } +} + +TEST_F(IvfRabitqStreamerTest, TestGroupBySearchByPrimaryKeys) { + constexpr size_t kDocCount = 256; + auto holder = BuildHolder(kDim, kDocCount); + + ailego::Params params; + params.set(PARAM_IVF_RABITQ_NLIST, 16U); + params.set(PARAM_RABITQ_TOTAL_BITS, 7U); + IndexStreamer::Pointer streamer; + BuildAndOpenStreamer(*index_meta_ptr_, holder, params, + dir_ + "/TestGroupBySearchByPrimaryKeys", &streamer); + ASSERT_NE(nullptr, streamer); + + NumericalVector query_vec(kDim); + FillTestVector(37, &query_vec); + IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, kDim); + auto context = streamer->create_context(); + context->set_group_by([](uint64_t key) { return std::to_string(key % 3U); }); + context->set_group_params(/*group_num=*/3, /*group_topk=*/2); + context->set_topk(10); + + const std::vector> keys{ + {3, 4, 5, 36, 37, 38, 90, 91, 92}}; + ASSERT_EQ(0, streamer->search_bf_by_p_keys_impl(query_vec.data(), keys, + query_meta, 1, context)); + + const auto &groups = context->group_result(0); + ASSERT_EQ(3UL, groups.size()); + for (const auto &group : groups) { + ASSERT_EQ(2UL, group.docs().size()); + const uint64_t expected_group = std::stoull(group.group_id()); + for (const auto &doc : group.docs()) { + EXPECT_EQ(expected_group, doc.key() % 3U); + } + } +} + +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=%zu, score=%f", + result.size(), (size_t)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 65c80ffbb..7aa446f11 100644 --- a/tests/core/interface/CMakeLists.txt +++ b/tests/core/interface/CMakeLists.txt @@ -8,8 +8,8 @@ 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_diskann + core_knn_flat_sparse core_knn_hnsw_sparse core_knn_ivf core_knn_hnsw_rabitq core_knn_ivf_rabitq core_knn_vamana core_knn_diskann SRCS ${CC_SRCS} INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm ) -endforeach() \ No newline at end of file +endforeach() diff --git a/tests/core/interface/index_group_by_test.cc b/tests/core/interface/index_group_by_test.cc index bb41938fb..508674dfa 100644 --- a/tests/core/interface/index_group_by_test.cc +++ b/tests/core/interface/index_group_by_test.cc @@ -44,6 +44,7 @@ struct GroupByCase { bool is_sparse = false; uint32_t dimension = kDimension; bool with_refiner = false; + int expected_error = 0; }; std::shared_ptr> AllPks() { @@ -181,6 +182,31 @@ BaseIndexQueryParam::Pointer HnswRabitqQuery(bool fetch_vector = false, } return builder.build(); } + +BaseIndexParam::Pointer DenseIvfRabitqParam(uint32_t dimension) { + return IVFRabitqIndexParamBuilder() + .WithMetricType(MetricType::kInnerProduct) + .WithDataType(DataType::DT_FP32) + .WithDimension(dimension) + .WithIsSparse(false) + .WithNlist(4) + .WithTotalBits(7) + .Build(); +} + +BaseIndexQueryParam::Pointer IvfRabitqQuery(bool is_linear = false, + bool with_bf_pks = false, + bool fetch_vector = false) { + auto query = std::make_shared(); + query->topk = kSearchTopk; + query->nprobe = 4; + query->is_linear = is_linear; + query->fetch_vector = fetch_vector; + if (with_bf_pks) { + query->bf_pks = AllPks(); + } + return query; +} #endif #if DISKANN_SUPPORTED @@ -264,7 +290,11 @@ class GroupByInterfaceTest : public ::testing::Test { SearchResult result; const int ret = index->Search(query.data, query_param, &result); if (expect_error) { - ASSERT_NE(0, ret) << test_case.name; + if (test_case.expected_error != 0) { + ASSERT_EQ(test_case.expected_error, ret) << test_case.name; + } else { + ASSERT_NE(0, ret) << test_case.name; + } } else { ASSERT_EQ(0, ret) << test_case.name; AssertGroupedResult(result, query_param, test_case); @@ -458,6 +488,21 @@ TEST_F(GroupByInterfaceTest, Dense) { HnswRabitqQuery(/*fetch_vector=*/false, /*is_linear=*/false, /*with_bf_pks=*/true), /*is_sparse=*/false, /*dimension=*/64}, + {"dense_ivf_rabitq_graph", DenseIvfRabitqParam(64), IvfRabitqQuery(), + /*is_sparse=*/false, /*dimension=*/64}, + {"dense_ivf_rabitq_linear", DenseIvfRabitqParam(64), + IvfRabitqQuery(/*is_linear=*/true), + /*is_sparse=*/false, /*dimension=*/64}, + {"dense_ivf_rabitq_bf_pks", DenseIvfRabitqParam(64), + IvfRabitqQuery(/*is_linear=*/false, /*with_bf_pks=*/true), + /*is_sparse=*/false, /*dimension=*/64}, + {"dense_ivf_rabitq_large_nprobe", DenseIvfRabitqParam(64), + [] { + auto query = IvfRabitqQuery(); + std::dynamic_pointer_cast(query)->nprobe = 1025; + return query; + }(), + /*is_sparse=*/false, /*dimension=*/64}, // Note: fetch_vector is not supported for RabitQ because the entity // stores quantized binary data (not original float vectors), and // RabitqReformer does not implement revert(). @@ -521,6 +566,20 @@ TEST_F(GroupByInterfaceTest, UnsupportedIndexTypes) { /*is_sparse=*/false, /*dimension=*/kDimension, /*with_refiner=*/true}, +#if RABITQ_SUPPORTED + {"unsupported_ivf_rabitq_fetch_vector", DenseIvfRabitqParam(64), + IvfRabitqQuery(/*is_linear=*/false, /*with_bf_pks=*/false, + /*fetch_vector=*/true), + /*is_sparse=*/false, /*dimension=*/64, /*with_refiner=*/false, + /*expected_error=*/zvec::core::IndexError_Unsupported}, + {"unsupported_ivf_rabitq_zero_nprobe", DenseIvfRabitqParam(64), + [] { + auto query = IvfRabitqQuery(); + std::dynamic_pointer_cast(query)->nprobe = 0; + return query; + }(), + /*is_sparse=*/false, /*dimension=*/64}, +#endif #if DISKANN_SUPPORTED {"unsupported_diskann_graph", DenseDiskAnnParam(), DiskAnnQuery()}, {"unsupported_diskann_linear", DenseDiskAnnParam(), diff --git a/tests/core/interface/index_interface_test.cc b/tests/core/interface/index_interface_test.cc index 4453756da..3db73a3a4 100644 --- a/tests/core/interface/index_interface_test.cc +++ b/tests/core/interface/index_interface_test.cc @@ -26,6 +26,7 @@ #include "zvec/core/framework/index_provider.h" #endif #include +#include "zvec/core/framework/index_error.h" #include "zvec/core/interface/index.h" #include "zvec/core/interface/index_factory.h" #include "zvec/core/interface/index_param.h" @@ -38,6 +39,63 @@ using namespace zvec::core_interface; +TEST(IndexInterface, IndexTypeKeepsExistingValues) { + EXPECT_EQ(5, static_cast(IndexType::kDiskAnn)); + EXPECT_EQ(6, static_cast(IndexType::kVamana)); + EXPECT_EQ(7, static_cast(IndexType::kIVFRabitq)); +} + +#if RABITQ_SUPPORTED +TEST(IndexInterface, IvfRabitqValidatesBuildParams) { + auto make_param = [](int nlist, int sample_count) { + return IVFRabitqIndexParamBuilder() + .WithMetricType(MetricType::kInnerProduct) + .WithDataType(DataType::DT_FP32) + .WithDimension(64) + .WithNlist(nlist) + .WithTotalBits(7) + .WithSampleCount(sample_count) + .Build(); + }; + + EXPECT_EQ(nullptr, IndexFactory::CreateAndInitIndex(*make_param(0, 0))); + EXPECT_EQ(nullptr, IndexFactory::CreateAndInitIndex(*make_param(-1, 0))); + EXPECT_EQ(nullptr, IndexFactory::CreateAndInitIndex(*make_param(32, -1))); + EXPECT_NE(nullptr, IndexFactory::CreateAndInitIndex(*make_param(1, 0))); + EXPECT_NE(nullptr, IndexFactory::CreateAndInitIndex(*make_param(1024, 1))); + EXPECT_NE(nullptr, IndexFactory::CreateAndInitIndex(*make_param(1025, 0))); +} + +TEST(IndexInterface, IvfRabitqFetchUnsupported) { + constexpr uint32_t kDimension = 64; + const std::string index_name{"ivf_rabitq_fetch.index"}; + zvec::test_util::RemoveTestFiles(index_name); + + auto param = IVFRabitqIndexParamBuilder() + .WithMetricType(MetricType::kL2sq) + .WithDataType(DataType::DT_FP32) + .WithDimension(kDimension) + .WithNlist(1) + .Build(); + auto index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, index); + ASSERT_EQ( + 0, index->Open(index_name, {StorageOptions::StorageType::kMMAP, true})); + + std::vector vector(kDimension, 1.0f); + VectorData vector_data{DenseVector{vector.data()}}; + ASSERT_EQ(0, index->Add(vector_data, 0)); + + VectorDataBuffer fetched; + EXPECT_EQ(zvec::core::IndexError_Unsupported, index->Fetch(0, &fetched)); + ASSERT_EQ(0, index->Train()); + EXPECT_EQ(zvec::core::IndexError_Unsupported, index->Fetch(0, &fetched)); + + ASSERT_EQ(0, index->Close()); + zvec::test_util::RemoveTestFiles(index_name); +} +#endif + TEST(IndexInterface, General) { constexpr uint32_t kDimension = 64; const std::string index_name{"test.index"}; @@ -2303,8 +2361,8 @@ TEST(IndexInterface, ExternalVectorFastSearchRecallRegression) { exact_results.reserve(kNumVectors); for (uint32_t i = 0; i < kNumVectors; ++i) { const float *vector = all_vectors.data() + i * kDimension; - const float score = std::inner_product( - query_vector.begin(), query_vector.end(), vector, 0.0f); + const float score = std::inner_product(query_vector.begin(), + query_vector.end(), vector, 0.0f); exact_results.emplace_back(score, i); } std::sort(exact_results.begin(), exact_results.end(), @@ -2329,13 +2387,11 @@ TEST(IndexInterface, ExternalVectorFastSearchRecallRegression) { .Build(); auto index = IndexFactory::CreateAndInitIndex(*param); ASSERT_NE(nullptr, index); - ASSERT_EQ(0, - index->Open(index_name, - {StorageOptions::StorageType::kMMAP, true})); + ASSERT_EQ( + 0, index->Open(index_name, {StorageOptions::StorageType::kMMAP, true})); for (uint32_t i = 0; i < kNumVectors; ++i) { - VectorData vector_data{ - DenseVector{all_vectors.data() + i * kDimension}}; + VectorData vector_data{DenseVector{all_vectors.data() + i * kDimension}}; ASSERT_EQ(0, index->AddWithSource(vector_data, i, source)); } 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 d837bb538..22b6f8481 100644 --- a/tests/db/collection_test.cc +++ b/tests/db/collection_test.cc @@ -3097,6 +3097,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 } } @@ -5592,6 +5595,76 @@ 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) { + ASSERT_EQ(*doc, expect_doc) + << "doc: " << doc->to_detail_string() + << "\nexpect_doc: " << expect_doc.to_detail_string(); + } + } + }; + + check_doc(); + + 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}); + ASSERT_TRUE(s.ok()) << s.message(); + + stats = collection->Stats().value(); + ASSERT_EQ(stats.doc_count, doc_count); + ASSERT_EQ(stats.index_completeness["dense_fp32"], 1); + + check_doc(); + + collection.reset(); + auto result = Collection::Open(col_path, options); + ASSERT_TRUE(result.has_value()); + collection = std::move(result.value()); + + check_doc(); + }; + + 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/CMakeLists.txt b/tests/db/index/CMakeLists.txt index 03da6ffd5..a700beef6 100644 --- a/tests/db/index/CMakeLists.txt +++ b/tests/db/index/CMakeLists.txt @@ -45,7 +45,7 @@ foreach(CC_SRCS ${ALL_TEST_SRCS}) core_quantizer_static core_knn_hnsw core_knn_hnsw_sparse sparsehash core_knn_flat core_knn_flat_sparse core_knn_ivf - core_knn_hnsw_rabitq core_mix_reducer + core_knn_hnsw_rabitq core_knn_ivf_rabitq core_mix_reducer Arrow::arrow_dataset ${CMAKE_THREAD_LIBS_INIT} ${CMAKE_DL_LIBS} diff --git a/tests/db/index/column/vector_column_group_by_test.cc b/tests/db/index/column/vector_column_group_by_test.cc index 4fedf3dcc..c0b41b374 100644 --- a/tests/db/index/column/vector_column_group_by_test.cc +++ b/tests/db/index/column/vector_column_group_by_test.cc @@ -72,6 +72,27 @@ std::vector AllPks() { return pks; } +class CapturingVectorColumnIndexer : public VectorColumnIndexer { + public: + Result Search( + const vector_column_params::VectorData &, + const vector_column_params::QueryParams &query_params) override { + search_called = true; + if (query_params.refiner_param) { + refiner_scale_factor = query_params.refiner_param->scale_factor_; + reference_indexer = query_params.refiner_param->reference_indexer; + } + IndexResults::Ptr result = std::make_shared( + false, core::IndexDocumentList{}, std::vector{}, + std::vector{}); + return result; + } + + bool search_called{false}; + float refiner_scale_factor{0.0f}; + VectorColumnIndexer::Ptr reference_indexer; +}; + class GroupByIndexerTest : public ::testing::Test { protected: void RunOk(const GroupByCase &tc) { @@ -271,6 +292,17 @@ TEST_F(GroupByIndexerTest, Dense) { std::make_shared(300), /*is_sparse=*/false, /*dimension=*/kGbDimension, /*optional=*/false, /*with_bf_pks=*/false, /*fetch_vector=*/true}, +#if RABITQ_SUPPORTED + {"dense_ivf_rabitq_graph", + std::make_shared(MetricType::IP, 4, 7, 0), + std::make_shared(4), + /*is_sparse=*/false, /*dimension=*/64}, + {"dense_ivf_rabitq_bf_pks", + std::make_shared(MetricType::IP, 4, 7, 0), + std::make_shared(4), + /*is_sparse=*/false, /*dimension=*/64, + /*optional=*/false, /*with_bf_pks=*/true}, +#endif }; for (const auto &tc : cases) { @@ -341,6 +373,35 @@ TEST_F(GroupByIndexerTest, CombinedSortsGroupsBeforeTruncating) { zvec::test_util::RemoveTestFiles(block1_path); } +TEST_F(GroupByIndexerTest, CombinedSupportsIvfRabitqRefiner) { + auto quantized_indexer = std::make_shared(); + auto normal_indexer = std::make_shared(); + auto index_params = + std::make_shared(MetricType::IP, 4, 7, 0); + FieldSchema schema("test", DataType::VECTOR_FP32, 64, false, index_params); + std::vector blocks{ + BlockMeta(0, BlockType::VECTOR_INDEX_QUANTIZE, 0, 0, 1, {"test"})}; + SegmentMeta segment_meta; + CombinedVectorColumnIndexer combined({quantized_indexer}, {normal_indexer}, + schema, segment_meta, blocks, + MetricType::IP, true); + + std::vector query(64, 1.0f); + vector_column_params::QueryParams query_params; + query_params.topk = 1; + query_params.query_params = + std::make_shared(4, 0.0f, false, true, 3.5f); + + auto result = combined.Search( + vector_column_params::VectorData{ + vector_column_params::DenseVector{query.data()}}, + query_params); + ASSERT_TRUE(result.has_value()); + ASSERT_TRUE(quantized_indexer->search_called); + ASSERT_FLOAT_EQ(3.5f, quantized_indexer->refiner_scale_factor); + ASSERT_EQ(normal_indexer, quantized_indexer->reference_indexer); +} + TEST_F(GroupByIndexerTest, Sparse) { std::vector cases{ {"sparse_flat_graph", std::make_shared(MetricType::IP), 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/doc_test.cc b/tests/db/index/common/doc_test.cc index b1f7a5b7f..e47174a43 100644 --- a/tests/db/index/common/doc_test.cc +++ b/tests/db/index/common/doc_test.cc @@ -1433,6 +1433,38 @@ TEST(SearchQuery, ValidateAndSanitize) { EXPECT_TRUE(s.ok()); } + // IVF RaBitQ nprobe must be positive + { + SearchQuery query; + query.target_.field_name_ = "embedding"; + query.topk_ = 10; + std::vector query_vector(128, 1.0f); + query.target_.set_vector( + std::string(reinterpret_cast(query_vector.data()), + query_vector.size() * sizeof(float))); + FieldSchema schema = + FieldSchema("embedding", DataType::VECTOR_FP32, 128, false, + std::make_shared(MetricType::L2)); + + query.target_.query_params_ = std::make_shared(0); + auto s = query.validate(&schema, nullptr); + EXPECT_FALSE(s.ok()); + EXPECT_EQ(s.code(), StatusCode::INVALID_ARGUMENT); + + query.target_.query_params_ = std::make_shared(-1); + s = query.validate(&schema, nullptr); + EXPECT_FALSE(s.ok()); + EXPECT_EQ(s.code(), StatusCode::INVALID_ARGUMENT); + + query.target_.query_params_ = std::make_shared(1025); + s = query.validate(&schema, nullptr); + EXPECT_TRUE(s.ok()) << s.message(); + + query.target_.query_params_ = std::make_shared(1024); + s = query.validate(&schema, nullptr); + EXPECT_TRUE(s.ok()) << s.message(); + } + // FTS clause validation { auto fts_params = std::make_shared(); @@ -1590,4 +1622,4 @@ TEST_F(DocDetailedTest, FieldExistenceChecks) { auto type_mismatch_opt = doc.get("existent"); EXPECT_FALSE(type_mismatch_opt.has_value()); -} \ 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/query_params_test.cc b/tests/db/index/common/query_params_test.cc index 13e8bd00b..fdb79e0d7 100644 --- a/tests/db/index/common/query_params_test.cc +++ b/tests/db/index/common/query_params_test.cc @@ -45,6 +45,18 @@ TEST(QueryParamsTest, IVFQueryParams) { EXPECT_EQ(params.nprobe(), 75); } +TEST(QueryParamsTest, IvfRabitqQueryParams) { + IvfRabitqQueryParams params; + EXPECT_EQ(params.type(), IndexType::IVF_RABITQ); + EXPECT_EQ(params.nprobe(), 10); + EXPECT_FLOAT_EQ(params.scale_factor(), 10.0f); + + params.set_nprobe(75); + params.set_scale_factor(3.5f); + EXPECT_EQ(params.nprobe(), 75); + EXPECT_FLOAT_EQ(params.scale_factor(), 3.5f); +} + TEST(QueryParamsTest, Polymorphism) { // Test polymorphic behavior QueryParams::Ptr hnsw_ptr = std::make_shared(100); @@ -76,4 +88,4 @@ TEST(QueryParamsTest, VirtualDestructor) { // This should not cause memory issues delete hnsw_ptr; delete ivf_ptr; -} \ 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 6984b60b9..983816bd9 100644 --- a/tests/db/index/common/schema_test.cc +++ b/tests/db/index/common/schema_test.cc @@ -899,9 +899,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(); } @@ -926,9 +925,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(); } @@ -957,6 +955,103 @@ 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(); + } +} + +TEST(FieldSchemaTest, IvfRabitqIndexValidationParameters) { + auto validate = [](int nlist, int sample_count) { + auto index_params = std::make_shared( + MetricType::L2, nlist, 7, sample_count); + FieldSchema field("vector_field", DataType::VECTOR_FP32, 128, false, + index_params); + return field.validate(); + }; + + EXPECT_FALSE(validate(0, 0).ok()); + EXPECT_FALSE(validate(-1, 0).ok()); + EXPECT_FALSE(validate(32, -1).ok()); + EXPECT_TRUE(validate(1, 0).ok()); + EXPECT_TRUE(validate(1024, 1).ok()); + EXPECT_TRUE(validate(1025, 0).ok()); +} #endif TEST(FieldSchemaTest, HnswRabitqIndexValidation_UnsupportedDataTypes) { @@ -972,7 +1067,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(); @@ -988,7 +1083,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/tests/shared/shared_api_test.cc b/tests/shared/shared_api_test.cc index cc157da1f..19a4d4763 100644 --- a/tests/shared/shared_api_test.cc +++ b/tests/shared/shared_api_test.cc @@ -16,6 +16,8 @@ #include #include #include +#include +#include #include int main() { @@ -36,5 +38,18 @@ int main() { return 3; } + zvec::IvfRabitqIndexParams index_params(zvec::MetricType::L2, 32, 8, 100); + if (index_params.type() != zvec::IndexType::IVF_RABITQ || + index_params.nlist() != 32 || index_params.total_bits() != 8 || + index_params.sample_count() != 100) { + return 4; + } + + zvec::IvfRabitqQueryParams query_params(4); + if (query_params.type() != zvec::IndexType::IVF_RABITQ || + query_params.nprobe() != 4) { + return 5; + } + return 0; } 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 =