diff --git a/.github/cmake/sdk-smoke/smoke.c b/.github/cmake/sdk-smoke/smoke.c index c1c6ce1fe1..b95c3b20ca 100644 --- a/.github/cmake/sdk-smoke/smoke.c +++ b/.github/cmake/sdk-smoke/smoke.c @@ -13,8 +13,8 @@ int main(void) { int minor = zvec_get_version_minor(); int patch = zvec_get_version_patch(); if (!zvec_check_version(major, minor, patch)) { - fprintf(stderr, "c_smoke: zvec_check_version(%d.%d.%d) failed\n", - major, minor, patch); + fprintf(stderr, "c_smoke: zvec_check_version(%d.%d.%d) failed\n", major, + minor, patch); return 1; } printf("c_smoke: OK\n"); diff --git a/python/tests/test_ivf_pq.py b/python/tests/test_ivf_pq.py new file mode 100644 index 0000000000..866beecc11 --- /dev/null +++ b/python/tests/test_ivf_pq.py @@ -0,0 +1,309 @@ +"""Public HNSW/IVF PQ configuration, persistence and coarse/refine behavior.""" + +import pickle + +import numpy as np +import pytest +import zvec + + +def index_param(kind, **kwargs): + if kind == "ivf": + return zvec.IVFIndexParam(n_list=4, n_iters=3, **kwargs) + return zvec.HnswIndexParam(m=16, ef_construction=100, **kwargs) + + +def query_param(kind, **kwargs): + if kind == "ivf": + return zvec.IVFQueryParam(nprobe=4, **kwargs) + return zvec.HnswQueryParam(ef=2000, **kwargs) + + +@pytest.mark.parametrize( + "kind,bits,fast_scan", + [ + ("ivf", 8, False), + ("ivf", 4, False), + ("ivf", 4, True), + ("hnsw", 8, False), + ("hnsw", 4, False), + ], +) +@pytest.mark.parametrize("rotate", [False, True]) +def test_pq_parameter_round_trip(kind, bits, fast_scan, rotate): + qp = zvec.QuantizerParam( + enable_rotate=rotate, + num_chunk=4, + num_bits=bits, + fast_scan=fast_scan, + opq_iter=3, + opq_pq_iter=2, + ) + param = index_param( + kind, + metric_type=zvec.MetricType.L2, + quantize_type=zvec.QuantizeType.PQ, + quantizer_param=qp, + ) + assert pickle.loads(pickle.dumps(qp)) == qp + restored = pickle.loads(pickle.dumps(param)) + assert restored.quantizer_param == qp + assert restored.to_dict() == param.to_dict() + assert param.to_dict()["quantizer_param"] == qp.to_dict() + assert '"fast_scan"' in repr(param) + + +def test_legacy_quantizer_and_ivf_pickle(): + qp = zvec.QuantizerParam.__new__(zvec.QuantizerParam) + qp.__setstate__((True,)) + assert qp.enable_rotate + assert qp.num_chunk == 8 and qp.num_bits == 8 + assert not qp.fast_scan + for state in [ + (zvec.MetricType.L2, 4, 3, False, zvec.QuantizeType.INT8), + (zvec.MetricType.L2, 4, 3, False, zvec.QuantizeType.INT8, True), + ]: + param = zvec.IVFIndexParam.__new__(zvec.IVFIndexParam) + param.__setstate__(state) + assert param.quantizer_param.enable_rotate == (len(state) == 6) + assert param.quantizer_param.num_bits == 8 + + +@pytest.mark.parametrize( + "kind,bits,fast_scan", + [ + ("ivf", 8, False), + ("ivf", 4, False), + ("ivf", 4, True), + ("hnsw", 8, False), + ("hnsw", 4, False), + ], +) +@pytest.mark.parametrize("rotate", [False, True]) +@pytest.mark.parametrize( + "metric", [zvec.MetricType.L2, zvec.MetricType.IP, zvec.MetricType.COSINE] +) +def test_pq_collection_reopen_and_refine( + tmp_path, kind, bits, fast_scan, rotate, metric +): + rng = np.random.default_rng(903) + vectors = rng.normal(size=(1100, 16)).astype(np.float32) + query = rng.normal(size=16).astype(np.float32) + qp = zvec.QuantizerParam( + enable_rotate=rotate, + num_chunk=4, + num_bits=bits, + fast_scan=fast_scan, + opq_iter=2, + opq_pq_iter=2, + ) + schema = zvec.CollectionSchema( + name="ivf_pq", + vectors=[ + zvec.VectorSchema( + "vector", + zvec.DataType.VECTOR_FP32, + dimension=16, + index_param=index_param( + kind, + metric_type=metric, + quantize_type=zvec.QuantizeType.PQ, + quantizer_param=qp, + ), + ) + ], + ) + path = str(tmp_path / "index") + coll = zvec.create_and_open(path, schema) + try: + for start in range(0, len(vectors), 200): + statuses = coll.insert( + [ + zvec.Doc(id=str(i), vectors={"vector": vectors[i].tolist()}) + for i in range(start, min(start + 200, len(vectors))) + ] + ) + assert all(status.ok() for status in statuses) + # Untrained writer blocks remain searchable through raw Flat. + assert coll.query( + zvec.Query(field_name="vector", vector=query.tolist()), topk=5 + ) + coll.optimize() + before = coll.query( + zvec.Query( + field_name="vector", + vector=query.tolist(), + param=query_param(kind), + ), + topk=40, + output_fields=[], + ) + finally: + coll.close() + coll = zvec.open(path, zvec.CollectionOption(read_only=True, enable_mmap=True)) + try: + restored = coll.schema.vector("vector").index_param + assert restored.quantize_type == zvec.QuantizeType.PQ + assert restored.quantizer_param == qp + coarse = coll.query( + zvec.Query( + field_name="vector", + vector=query.tolist(), + param=query_param(kind), + ), + topk=40, + output_fields=[], + ) + assert [doc.id for doc in coarse] == [doc.id for doc in before] + np.testing.assert_allclose( + [doc.score for doc in coarse], [doc.score for doc in before] + ) + refined = coll.query( + zvec.Query( + field_name="vector", + vector=query.tolist(), + param=query_param(kind, is_using_refiner=True, scale_factor=4), + ), + topk=10, + output_fields=[], + ) + ids = np.array([int(doc.id) for doc in coarse]) + candidates = vectors[ids] + if metric == zvec.MetricType.L2: + scores = np.sum((candidates - query) ** 2, axis=1) + order = np.argsort(scores)[:10] + else: + scores = candidates @ query + if metric == zvec.MetricType.COSINE: + scores /= np.linalg.norm(candidates, axis=1) * np.linalg.norm(query) + order = np.argsort(-scores)[:10] + if metric == zvec.MetricType.COSINE: + scores = 1.0 - scores + # Verify optimize built compressed postings rather than leaving the raw Flat fallback. + assert np.max(np.abs(np.array([doc.score for doc in coarse]) - scores)) > 1e-4 + assert [int(doc.id) for doc in refined] == ids[order].tolist() + np.testing.assert_allclose( + [doc.score for doc in refined], scores[order], rtol=1e-4, atol=1e-5 + ) + finally: + coll.close() + + +def test_legacy_hnsw_pickle(): + for state in [ + (zvec.MetricType.L2, 16, 100, zvec.QuantizeType.INT8, False), + ( + zvec.MetricType.L2, + 16, + 100, + zvec.QuantizeType.INT8, + False, + True, + True, + zvec.DataType.VECTOR_FP16, + ), + ]: + param = zvec.HnswIndexParam.__new__(zvec.HnswIndexParam) + param.__setstate__(state) + assert param.quantizer_param.enable_rotate == (len(state) == 8) + assert param.quantizer_param.num_bits == 8 + + +@pytest.mark.parametrize("bits", [4, 8]) +@pytest.mark.parametrize("enable_mmap", [False, True]) +def test_hnsw_pq_incremental_writes_and_retraining(tmp_path, bits, enable_mmap): + zvec.init() + rng = np.random.default_rng(137) + vectors = rng.normal(size=(303, 16)).astype(np.float32) + qp = zvec.QuantizerParam(num_chunk=4, num_bits=bits) + param = zvec.HnswIndexParam( + metric_type=zvec.MetricType.L2, + m=16, + ef_construction=100, + quantize_type=zvec.QuantizeType.PQ, + quantizer_param=qp, + flat_data_type=zvec.DataType.VECTOR_FP16, + ) + schema = zvec.CollectionSchema( + name="pq_writes", + vectors=[ + zvec.VectorSchema( + "vector", zvec.DataType.VECTOR_FP32, 16, index_param=param + ) + ], + ) + path = str(tmp_path / "index") + option = zvec.CollectionOption(enable_mmap=enable_mmap) + coll = zvec.create_and_open(path, schema, option=option) + + def check(ids): + hits = coll.query( + zvec.Query( + "vector", + vector=vectors[302].tolist(), + param=zvec.HnswQueryParam(is_linear=True, is_using_refiner=True), + ), + topk=400, + ) + assert {hit.id for hit in hits} == ids + assert coll.schema.vector("vector").index_param.quantizer_param == qp + + try: + for start in range(0, 300, 100): + statuses = coll.insert( + [ + zvec.Doc(id=str(i), vectors={"vector": vectors[i].tolist()}) + for i in range(start, start + 100) + ] + ) + assert all(status.ok() for status in statuses) + coll.optimize() + coll.close() + coll = zvec.open(path, option) + assert coll.insert( + zvec.Doc(id="300", vectors={"vector": vectors[300].tolist()}) + ).ok() + assert coll.update( + zvec.Doc(id="1", vectors={"vector": vectors[301].tolist()}) + ).ok() + assert coll.delete("0").ok() + expected = {str(i) for i in range(1, 301)} + check(expected) + coll.optimize() + check(expected) + coll.close() + coll = zvec.open(path, option) + check(expected) + finally: + coll.close() + + +@pytest.mark.parametrize( + "kwargs", + [ + {"fast_scan": True, "num_bits": 4}, + {"num_chunk": 0}, + {"num_chunk": 17}, + {"num_bits": 6}, + {"num_chunk": 3, "num_bits": 4}, + {"enable_rotate": True, "opq_iter": 0}, + ], +) +def test_hnsw_pq_rejects_invalid_configuration(tmp_path, kwargs): + schema = zvec.CollectionSchema( + name="invalid_pq", + vectors=[ + zvec.VectorSchema( + "vector", + zvec.DataType.VECTOR_FP32, + 16, + index_param=zvec.HnswIndexParam( + quantize_type=zvec.QuantizeType.PQ, + quantizer_param=zvec.QuantizerParam(**kwargs), + ), + ) + ], + ) + with pytest.raises(ValueError, match="PQ"): + coll = zvec.create_and_open(str(tmp_path / "index"), schema) + coll.close() diff --git a/python/zvec/model/param/__init__.pyi b/python/zvec/model/param/__init__.pyi index a5bba3bc6c..cf5bcd738c 100644 --- a/python/zvec/model/param/__init__.pyi +++ b/python/zvec/model/param/__init__.pyi @@ -335,7 +335,7 @@ class HnswIndexParam(VectorIndexParam): neighbors during index construction. Larger values yield better graph quality at the cost of slower build time. Default is 500. quantize_type (QuantizeType): Optional quantization type for vector - compression (e.g., FP16, INT8). Default is `QuantizeType.UNDEFINED` to + compression (FP16, INT8, INT4, or PQ). Default is `QuantizeType.UNDEFINED` to disable quantization. use_contiguous_memory (bool): If True, the HNSW streamer allocates a single contiguous memory arena for all graph nodes, improving cache @@ -347,6 +347,10 @@ class HnswIndexParam(VectorIndexParam): reference index. Defaults to ``VECTOR_FP32``; ``VECTOR_FP16`` and ``VECTOR_UINT8`` are also supported. + PQ learns its codebook during Collection.optimize(). Before optimization, + writes remain searchable through Flat. Configure PQ4/PQ8 and optional OPQ + with QuantizerParam; fast_scan is supported only by IVF. + Examples: >>> from zvec.typing import MetricType, QuantizeType >>> params = HnswIndexParam( @@ -1401,7 +1405,7 @@ class QuantizerParam: Attributes: enable_rotate (bool): Whether to apply random rotation before INT8/INT4 quantization to reduce quantization error. - Only effective with quantize_type=INT8 or INT4. Defaults to False. + For PQ, enables learned OPQ. Defaults to False. Examples: >>> qp = QuantizerParam(enable_rotate=True) @@ -1410,13 +1414,26 @@ class QuantizerParam: """ def __getstate__(self) -> tuple: ... - def __init__(self, enable_rotate: bool = False) -> None: + def __init__( + self, + enable_rotate: bool = False, + num_chunk: int = 8, + num_bits: int = 8, + fast_scan: bool = False, + opq_iter: int = 5, + opq_pq_iter: int = 4, + ) -> None: """ Constructs a QuantizerParam instance. Args: + num_chunk (int): PQ subquantizers; default 8. + num_bits (int): PQ bits per code, 4 or 8; default 8. + fast_scan (bool): Packed IVF PQ4 scan; default False. + opq_iter (int): OPQ iterations; default 5. + opq_pq_iter (int): PQ iterations per OPQ iteration; default 4. enable_rotate (bool, optional): Whether to apply random rotation - before INT8/INT4 quantization. Defaults to False. + before INT8/INT4 quantization, or learned OPQ for PQ. Defaults to False. """ def __repr__(self) -> str: ... @@ -1430,9 +1447,20 @@ class QuantizerParam: @property def enable_rotate(self) -> bool: """ - bool: Whether random rotation is enabled before INT8/INT4 quantization. + bool: Enable random rotation for INT8/INT4 or learned OPQ for PQ. """ + @property + def num_chunk(self) -> int: ... + @property + def num_bits(self) -> int: ... + @property + def fast_scan(self) -> bool: ... + @property + def opq_iter(self) -> int: ... + @property + def opq_pq_iter(self) -> int: ... + class VectorIndexParam(IndexParam): """ diff --git a/python/zvec/typing/__init__.pyi b/python/zvec/typing/__init__.pyi index be2828873f..5996033d0f 100644 --- a/python/zvec/typing/__init__.pyi +++ b/python/zvec/typing/__init__.pyi @@ -314,8 +314,11 @@ class QuantizeType: UNIFORM_UINT8 UNIFORM_UINT4 + + PQ """ + PQ: typing.ClassVar[QuantizeType] # value = FP16: typing.ClassVar[QuantizeType] # value = INT4: typing.ClassVar[QuantizeType] # value = INT8: typing.ClassVar[QuantizeType] # value = @@ -332,7 +335,7 @@ class QuantizeType: UNDEFINED: typing.ClassVar[QuantizeType] # value = __members__: typing.ClassVar[ dict[str, QuantizeType] - ] # value = {'UNDEFINED': , 'FP16': , 'INT8': , 'INT4': , 'RABITQ': , 'UNIFORM_UINT7': , 'UNIFORM_UINT8': , 'UNIFORM_UINT4': } + ] # value = {'UNDEFINED': , 'FP16': , 'INT8': , 'INT4': , 'RABITQ': , 'UNIFORM_UINT7': , 'UNIFORM_UINT8': , 'UNIFORM_UINT4': , 'PQ': } 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 f9108dc533..d57169ae70 100644 --- a/src/binding/c/c_api.cc +++ b/src/binding/c/c_api.cc @@ -1539,6 +1539,40 @@ bool zvec_index_params_get_quantizer_enable_rotate( return false; } +zvec_error_code_t zvec_index_params_set_quantizer_pq( + zvec_index_params_t *params, const zvec_pq_quantizer_params_t *pq) { + auto *ivf = params ? dynamic_cast( + reinterpret_cast(params)) : nullptr; + if (!ivf || !pq) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "IVF index params and PQ params are required"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + if (pq->num_chunk <= 0 || (pq->num_bits != 4 && pq->num_bits != 8) || + (pq->fast_scan && pq->num_bits != 4) || + (pq->enable_rotate && (pq->opq_iter == 0 || pq->opq_pq_iter == 0))) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "Invalid PQ/OPQ/FastScan parameters"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + ivf->set_quantizer_param(zvec::QuantizerParam( + pq->enable_rotate, pq->num_chunk, pq->num_bits, pq->fast_scan, + pq->opq_iter, pq->opq_pq_iter)); + return ZVEC_OK; +} + +zvec_error_code_t zvec_index_params_get_quantizer_pq( + const zvec_index_params_t *params, zvec_pq_quantizer_params_t *pq) { + const auto *ivf = params ? dynamic_cast( + reinterpret_cast(params)) : nullptr; + if (!ivf || !pq) { + SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT, "IVF index params and PQ output are required"); + return ZVEC_ERROR_INVALID_ARGUMENT; + } + const auto &qp = ivf->quantizer_param(); + *pq = {qp.num_chunk(), qp.num_bits(), qp.enable_rotate(), qp.fast_scan(), + qp.opq_iter(), qp.opq_pq_iter()}; + return ZVEC_OK; +} + /** * @brief Get index type from index parameters * @param params Index parameters diff --git a/src/binding/c/exported_symbols.lds b/src/binding/c/exported_symbols.lds index b8fd17ad73..594b3edee4 100644 --- a/src/binding/c/exported_symbols.lds +++ b/src/binding/c/exported_symbols.lds @@ -216,6 +216,7 @@ zvec_index_params_get_metric_type; zvec_index_params_get_quantize_type; zvec_index_params_get_quantizer_enable_rotate; + zvec_index_params_get_quantizer_pq; zvec_index_params_get_type; zvec_index_params_get_vamana_params; zvec_index_params_get_vamana_two_pass_build; @@ -228,6 +229,7 @@ zvec_index_params_set_metric_type; zvec_index_params_set_quantize_type; zvec_index_params_set_quantizer_enable_rotate; + zvec_index_params_set_quantizer_pq; zvec_index_params_set_vamana_params; zvec_index_params_set_vamana_two_pass_build; zvec_index_type_to_string; diff --git a/src/binding/c/exported_symbols.txt b/src/binding/c/exported_symbols.txt index 6cf597feb4..c6c1dfdf43 100644 --- a/src/binding/c/exported_symbols.txt +++ b/src/binding/c/exported_symbols.txt @@ -215,6 +215,7 @@ _zvec_index_params_get_ivf_rabitq_params _zvec_index_params_get_metric_type _zvec_index_params_get_quantize_type _zvec_index_params_get_quantizer_enable_rotate +_zvec_index_params_get_quantizer_pq _zvec_index_params_get_type _zvec_index_params_get_vamana_params _zvec_index_params_get_vamana_two_pass_build @@ -227,6 +228,7 @@ _zvec_index_params_set_ivf_rabitq_params _zvec_index_params_set_metric_type _zvec_index_params_set_quantize_type _zvec_index_params_set_quantizer_enable_rotate +_zvec_index_params_set_quantizer_pq _zvec_index_params_set_vamana_params _zvec_index_params_set_vamana_two_pass_build _zvec_index_type_to_string diff --git a/src/binding/python/model/param/python_param.cc b/src/binding/python/model/param/python_param.cc index 79636f772b..eed208ee4f 100644 --- a/src/binding/python/model/param/python_param.cc +++ b/src/binding/python/model/param/python_param.cc @@ -60,10 +60,23 @@ static std::string metric_type_to_string(const MetricType type) { } } +static py::dict quantizer_param_to_dict(const QuantizerParam ¶ms) { + py::dict result; + result["enable_rotate"] = params.enable_rotate(); + result["num_chunk"] = params.num_chunk(); + result["num_bits"] = params.num_bits(); + result["fast_scan"] = params.fast_scan(); + result["opq_iter"] = params.opq_iter(); + result["opq_pq_iter"] = params.opq_pq_iter(); + return result; +} + static std::string quantize_type_to_string(const QuantizeType type) { switch (type) { case QuantizeType::UNDEFINED: return "UNDEFINED"; + case QuantizeType::PQ: + return "PQ"; case QuantizeType::INT8: return "INT8"; case QuantizeType::INT4: @@ -425,39 +438,56 @@ Constructs an FtsIndexParam instance. m, "QuantizerParam", R"pbdoc( Parameters for quantizer configuration. -Encapsulates quantization-related settings such as enable_rotate. -Designed for future extensibility. +For HNSW/IVF PQ, configure subquantizers, code width, FastScan and learned OPQ. Attributes: enable_rotate (bool): Whether to apply random rotation before INT8/INT4 quantization to reduce quantization error. - Only effective with quantize_type=INT8 or INT4. Defaults to False. + With quantize_type=PQ, enables learned OPQ. Defaults to False. + num_chunk (int): PQ subquantizers; default 8. PQ4 requires dimension divisible by this value. + num_bits (int): Bits per PQ code, 4 or 8; default 8. + fast_scan (bool): Enable packed IVF PQ4 scans; default False. + opq_iter (int): OPQ training iterations; default 5, positive when OPQ is enabled. + opq_pq_iter (int): PQ iterations per OPQ iteration; default 4, positive when OPQ is enabled. Examples: >>> qp = QuantizerParam(enable_rotate=True) >>> print(qp.enable_rotate) True )pbdoc"); - quantizer_param.def(py::init(), py::arg("enable_rotate") = false) + quantizer_param + .def(py::init(), + py::arg("enable_rotate") = false, py::arg("num_chunk") = 8, + py::arg("num_bits") = 8, py::arg("fast_scan") = false, + py::arg("opq_iter") = 5, py::arg("opq_pq_iter") = 4) + .def_property_readonly("num_chunk", &QuantizerParam::num_chunk) + .def_property_readonly("num_bits", &QuantizerParam::num_bits) + .def_property_readonly("fast_scan", &QuantizerParam::fast_scan) + .def_property_readonly("opq_iter", &QuantizerParam::opq_iter) + .def_property_readonly("opq_pq_iter", &QuantizerParam::opq_pq_iter) .def_property_readonly( "enable_rotate", [](const QuantizerParam &self) -> bool { return self.enable_rotate(); }, - "bool: Whether random rotation is enabled before INT8/INT4 " - "quantization.") + "bool: Enable random rotation for scalar quantizers or learned OPQ " + "for PQ.") .def( "to_dict", [](const QuantizerParam &self) -> py::dict { - py::dict dict; - dict["enable_rotate"] = self.enable_rotate(); - return dict; + return quantizer_param_to_dict(self); }, "Convert to dictionary with all fields") .def("__repr__", [](const QuantizerParam &self) -> std::string { return "{\"enable_rotate\":" + - std::string(self.enable_rotate() ? "true" : "false") + "}"; + std::string(self.enable_rotate() ? "true" : "false") + + ", \"num_chunk\":" + std::to_string(self.num_chunk()) + + ", \"num_bits\":" + std::to_string(self.num_bits()) + + ", \"fast_scan\":" + std::to_string(self.fast_scan()) + + ", \"opq_iter\":" + std::to_string(self.opq_iter()) + + ", \"opq_pq_iter\":" + std::to_string(self.opq_pq_iter()) + + "}"; }) .def( "__eq__", @@ -468,12 +498,19 @@ Designed for future extensibility. py::is_operator()) .def(py::pickle( [](const QuantizerParam &self) { - return py::make_tuple(self.enable_rotate()); + return py::make_tuple(self.enable_rotate(), self.num_chunk(), + self.num_bits(), self.fast_scan(), + self.opq_iter(), self.opq_pq_iter()); }, [](py::tuple t) { - if (t.size() != 1) + if (t.size() != 1 && t.size() != 6) throw std::runtime_error("Invalid state for QuantizerParam"); - return std::make_shared(t[0].cast()); + if (t.size() == 1) + return std::make_shared(t[0].cast()); + return std::make_shared( + t[0].cast(), t[1].cast(), t[2].cast(), + t[3].cast(), t[4].cast(), + t[5].cast()); })); // binding base vector index params @@ -516,9 +553,8 @@ Encapsulates common settings for all vector index types. dict["metric_type"] = metric_type_to_string(self.metric_type()); dict["quantize_type"] = quantize_type_to_string(self.quantize_type()); - py::dict qp_dict; - qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate(); - dict["quantizer_param"] = qp_dict; + dict["quantizer_param"] = + quantizer_param_to_dict(self.quantizer_param()); return dict; }, "Convert to dictionary with all fields") @@ -553,7 +589,7 @@ encapsulates its construction hyperparameters. neighbors during index construction. Larger values yield better graph quality at the cost of slower build time. Default is 500. quantize_type (QuantizeType): Optional quantization type for vector - compression (e.g., FP16, INT8). Default is `QuantizeType.UNDEFINED` to + compression (FP16, INT8, INT4, or PQ). Default is `QuantizeType.UNDEFINED` to disable quantization. use_flat_contiguous_memory (bool): Whether the Flat reference index used for refine should use contiguous vector memory. @@ -561,6 +597,10 @@ encapsulates its construction hyperparameters. index. Defaults to ``VECTOR_FP32``; ``VECTOR_FP16`` and ``VECTOR_UINT8`` are also supported. +PQ learns its codebook during Collection.optimize(). Before optimization, +writes remain searchable through Flat. Configure PQ4/PQ8 and optional OPQ +with QuantizerParam; fast_scan is supported only by IVF. + Examples: >>> from zvec.typing import MetricType, QuantizeType >>> params = HnswIndexParam( @@ -626,45 +666,48 @@ encapsulates its construction hyperparameters. dict["use_flat_contiguous_memory"] = self.use_flat_contiguous_memory(); dict["flat_data_type"] = data_type_to_string(self.flat_data_type()); - py::dict qp_dict; - qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate(); - dict["quantizer_param"] = qp_dict; + dict["quantizer_param"] = + quantizer_param_to_dict(self.quantizer_param()); return dict; }, "Convert to dictionary with all fields") - .def( - "__repr__", - [](const HnswIndexParams &self) -> std::string { - return "{" - "\"metric_type\":" + - metric_type_to_string(self.metric_type()) + - ", \"m\":" + std::to_string(self.m()) + - ", \"ef_construction\":" + - std::to_string(self.ef_construction()) + - ", \"quantize_type\":" + - quantize_type_to_string(self.quantize_type()) + - ", \"use_contiguous_memory\":" + - (self.use_contiguous_memory() ? "true" : "false") + - ", \"use_flat_contiguous_memory\":" + - (self.use_flat_contiguous_memory() ? "true" : "false") + - ", \"flat_data_type\":\"" + - data_type_to_string(self.flat_data_type()) + "\"" + - ", \"quantizer_param\":{" + "\"enable_rotate\":" + - (self.quantizer_param().enable_rotate() ? "true" : "false") + - "}}"; - }) + .def("__repr__", + [](const HnswIndexParams &self) -> std::string { + return "{" + "\"metric_type\":" + + metric_type_to_string(self.metric_type()) + + ", \"m\":" + std::to_string(self.m()) + + ", \"ef_construction\":" + + std::to_string(self.ef_construction()) + + ", \"quantize_type\":" + + quantize_type_to_string(self.quantize_type()) + + ", \"use_contiguous_memory\":" + + (self.use_contiguous_memory() ? "true" : "false") + + ", \"use_flat_contiguous_memory\":" + + (self.use_flat_contiguous_memory() ? "true" : "false") + + ", \"flat_data_type\":\"" + + data_type_to_string(self.flat_data_type()) + "\"" + + ", \"quantizer_param\":" + + py::str(py::cast(self.quantizer_param())) + .cast() + + "}"; + }) .def(py::pickle( [](const HnswIndexParams &self) { return py::make_tuple( self.metric_type(), self.m(), self.ef_construction(), self.quantize_type(), self.use_contiguous_memory(), self.quantizer_param().enable_rotate(), - self.use_flat_contiguous_memory(), self.flat_data_type()); + self.use_flat_contiguous_memory(), self.flat_data_type(), + self.quantizer_param()); }, [](py::tuple t) { - if (t.size() < 5 || t.size() > 8) + if (t.size() < 5 || t.size() > 9) throw std::runtime_error("Invalid state for HnswIndexParams"); - QuantizerParam qp(t.size() >= 6 ? t[5].cast() : false); + QuantizerParam qp = + t.size() == 9 + ? t[8].cast() + : QuantizerParam(t.size() >= 6 ? t[5].cast() : false); return std::make_shared( t[0].cast(), t[1].cast(), t[2].cast(), t[3].cast(), t[4].cast(), qp, @@ -988,9 +1031,8 @@ its construction hyperparameters. dict["flat_data_type"] = data_type_to_string(self.flat_data_type()); dict["quantize_type"] = quantize_type_to_string(self.quantize_type()); - py::dict qp_dict; - qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate(); - dict["quantizer_param"] = qp_dict; + dict["quantizer_param"] = + quantizer_param_to_dict(self.quantizer_param()); return dict; }, "Convert to dictionary with all fields") @@ -1109,9 +1151,8 @@ Constructs a FlatIndexParam instance. dict["metric_type"] = metric_type_to_string(self.metric_type()); dict["quantize_type"] = quantize_type_to_string(self.quantize_type()); - py::dict qp_dict; - qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate(); - dict["quantizer_param"] = qp_dict; + dict["quantizer_param"] = + quantizer_param_to_dict(self.quantizer_param()); dict["use_contiguous_memory"] = self.use_contiguous_memory(); return dict; }, @@ -1222,38 +1263,40 @@ Constructs an IVFIndexParam instance. dict["use_soar"] = self.use_soar(); dict["quantize_type"] = quantize_type_to_string(self.quantize_type()); - py::dict qp_dict; - qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate(); - dict["quantizer_param"] = qp_dict; + dict["quantizer_param"] = + quantizer_param_to_dict(self.quantizer_param()); return dict; }, "Convert to dictionary with all fields") - .def( - "__repr__", - [](const IVFIndexParams &self) { - return "{" - "\"metric_type\":" + - metric_type_to_string(self.metric_type()) + - ", \"n_list\":" + std::to_string(self.n_list()) + - ", \"n_iters\":" + std::to_string(self.n_iters()) + - ", \"use_soar\":" + std::to_string(self.use_soar()) + - ", \"quantize_type\":" + - quantize_type_to_string(self.quantize_type()) + - ", \"quantizer_param\":{" + "\"enable_rotate\":" + - (self.quantizer_param().enable_rotate() ? "true" : "false") + - "}}"; - }) + .def("__repr__", + [](const IVFIndexParams &self) { + return "{" + "\"metric_type\":" + + metric_type_to_string(self.metric_type()) + + ", \"n_list\":" + std::to_string(self.n_list()) + + ", \"n_iters\":" + std::to_string(self.n_iters()) + + ", \"use_soar\":" + std::to_string(self.use_soar()) + + ", \"quantize_type\":" + + quantize_type_to_string(self.quantize_type()) + + ", \"quantizer_param\":" + + py::str(py::cast(self.quantizer_param())) + .cast() + + "}"; + }) .def(py::pickle( [](const IVFIndexParams &self) { - return py::make_tuple(self.metric_type(), self.n_list(), - self.n_iters(), self.use_soar(), - self.quantize_type(), - self.quantizer_param().enable_rotate()); + return py::make_tuple( + self.metric_type(), self.n_list(), self.n_iters(), + self.use_soar(), self.quantize_type(), + self.quantizer_param().enable_rotate(), self.quantizer_param()); }, [](py::tuple t) { - if (t.size() != 5 && t.size() != 6) + if (t.size() != 5 && t.size() != 6 && t.size() != 7) throw std::runtime_error("Invalid state for IVFIndexParams"); - QuantizerParam qp(t.size() >= 6 ? t[5].cast() : false); + QuantizerParam qp = + t.size() == 7 + ? t[6].cast() + : QuantizerParam(t.size() >= 6 ? t[5].cast() : false); return std::make_shared( t[0].cast(), t[1].cast(), t[2].cast(), t[3].cast(), t[4].cast(), qp); @@ -1345,9 +1388,8 @@ Constructs an DiskAnnIndexParams instance. dict["pq_chunk_num"] = self.pq_chunk_num(); dict["quantize_type"] = quantize_type_to_string(self.quantize_type()); - py::dict qp_dict; - qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate(); - dict["quantizer_param"] = qp_dict; + dict["quantizer_param"] = + quantizer_param_to_dict(self.quantizer_param()); return dict; }, "Convert to dictionary with all fields") diff --git a/src/binding/python/typing/python_type.cc b/src/binding/python/typing/python_type.cc index 747b233af8..16bf56ba16 100644 --- a/src/binding/python/typing/python_type.cc +++ b/src/binding/python/typing/python_type.cc @@ -142,7 +142,8 @@ Enumeration of supported quantization types for vector compression. .value("RABITQ", QuantizeType::RABITQ) .value("UNIFORM_UINT7", QuantizeType::UNIFORM_UINT7) .value("UNIFORM_UINT8", QuantizeType::UNIFORM_UINT8) - .value("UNIFORM_UINT4", QuantizeType::UNIFORM_UINT4); + .value("UNIFORM_UINT4", QuantizeType::UNIFORM_UINT4) + .value("PQ", QuantizeType::PQ); } void ZVecPyTyping::bind_io_backend_types(py::module_ &m) { diff --git a/src/core/algorithm/hnsw/hnsw_streamer.cc b/src/core/algorithm/hnsw/hnsw_streamer.cc index 5f625346f8..23a4b59429 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer.cc +++ b/src/core/algorithm/hnsw/hnsw_streamer.cc @@ -16,6 +16,7 @@ #include #include #include +#include #include #include "utility/sparse_utility.h" #include "hnsw_algorithm.h" @@ -45,8 +46,105 @@ void BindTurboDistances(TurboQuantizer::Pointer quantizer, }; } +const std::string HNSW_TURBO_QUANTIZER_SEG_ID = "hnsw.turbo_quantizer"; + +int DumpTurboQuantizerSegment(const IndexStorage::Pointer &storage, + const zvec::turbo::Quantizer &quantizer) { + if (!quantizer.require_train()) return 0; + std::string state; + int ret = quantizer.serialize(&state); + if (ret != 0) return ret; + if (state.empty()) return IndexError_InvalidFormat; + ret = storage->append(HNSW_TURBO_QUANTIZER_SEG_ID, + (state.size() + 31u) & ~size_t(31u)); + if (ret != 0) return ret; + auto segment = storage->get(HNSW_TURBO_QUANTIZER_SEG_ID); + if (!segment || + segment->write(0, state.data(), state.size()) != state.size() || + segment->resize(state.size()) != state.size()) { + return IndexError_WriteData; + } + segment->update_data_crc(ailego::Crc32c::Hash(state.data(), state.size(), 0)); + return 0; +} + } // namespace +int HnswStreamer::RestoreQuantizer( + const IndexStorage::Pointer &storage, const IndexMeta &quantized_meta, + std::shared_ptr *out_quantizer) { + if (!storage || out_quantizer == nullptr || + quantized_meta.quantizer_name().empty()) { + return IndexError_NoExist; + } + std::shared_ptr quantizer = + *out_quantizer + ? *out_quantizer + : IndexFactory::CreateQuantizer(quantized_meta.quantizer_name()); + if (!quantizer) { + LOG_ERROR("Failed to create turbo quantizer %s", + quantized_meta.quantizer_name().c_str()); + return IndexError_NoExist; + } + + // The quantized meta records the raw input layout next to the quantizer + // params, so init() can rebuild the code layout before the state is + // deserialized on top. + const ailego::Params ¶ms = quantized_meta.quantizer_params(); + int64_t input_dim = 0; + int64_t input_type = IndexMeta::DataType::DT_FP32; + if (!params.get("input_dimension", &input_dim) || input_dim <= 0) { + LOG_ERROR("Quantizer %s misses the recorded input dimension", + quantized_meta.quantizer_name().c_str()); + return IndexError_InvalidFormat; + } + params.get("input_data_type", &input_type); + IndexMeta raw_meta; + raw_meta.set_meta(static_cast(input_type), + static_cast(input_dim)); + raw_meta.set_metric(quantized_meta.metric_name(), 0, + quantized_meta.metric_params()); + int ret = quantizer->init(raw_meta, params); + if (ret != 0) { + LOG_ERROR("Failed to init turbo quantizer %s, ret=%d", + quantized_meta.quantizer_name().c_str(), ret); + return ret; + } + + auto segment = storage->get(HNSW_TURBO_QUANTIZER_SEG_ID); + if (quantizer->require_train() && (!segment || !segment->data_size())) { + LOG_ERROR("Missing trained HNSW quantizer state"); + return IndexError_InvalidFormat; + } + if (segment && segment->data_size()) { + IndexStorage::MemoryBlock block; + const size_t size = segment->data_size(); + if (segment->read(0, block, size) != size) { + return IndexError_ReadData; + } + if (segment->data_crc() && + ailego::Crc32c::Hash(block.data(), size, 0) != segment->data_crc()) { + return IndexError_InvalidChecksum; + } + ret = quantizer->deserialize(block.data(), size); + if (ret != 0) return ret; + } + // Deserialization may change the layout; validate the restored state. + const IndexMeta &rebuilt = quantizer->meta(); + if (rebuilt.data_type() != quantized_meta.data_type() || + rebuilt.dimension() != quantized_meta.dimension() || + rebuilt.element_size() != quantized_meta.element_size() || + rebuilt.extra_meta_size() != quantized_meta.extra_meta_size() || + quantizer->quantized_datapoint_vector_length() != + quantized_meta.element_size() || + quantizer->dim() != input_dim) { + return IndexError_Mismatch; + } + + *out_quantizer = quantizer; + return 0; +} + HnswStreamer::HnswStreamer() = default; HnswStreamer::~HnswStreamer() { @@ -68,6 +166,31 @@ void HnswStreamer::merge_trained_meta(const IndexMeta &trained_meta) { } } +int HnswStreamer::check_params(const void *query, const IndexQueryMeta &qmeta, + bool search) const { + if (ailego_unlikely(!query)) { + LOG_ERROR("null query"); + return IndexError_InvalidArgument; + } + if (ailego_unlikely(qmeta.dimension() != meta_.dimension() || + qmeta.data_type() != meta_.data_type() || + qmeta.element_size() != meta_.element_size())) { + // A query prepared by the active turbo quantizer may keep the + // quantizer's raw input layout instead of the stored code layout + // (e.g. a PQ lookup table); identify it by the quantizer stamp. + const bool quantizer_query = + search && quantizer_ != nullptr && quantizer_->dim() > 0 && + qmeta.data_type() == IndexMeta::DT_FP32 && + qmeta.quantize_type() == static_cast(quantizer_->type()) && + qmeta.dimension() == static_cast(quantizer_->dim()); + if (ailego_unlikely(!quantizer_query)) { + LOG_ERROR("Unsupported query meta"); + return IndexError_Mismatch; + } + } + return 0; +} + int HnswStreamer::init(const IndexMeta &imeta, const ailego::Params ¶ms) { meta_ = imeta; meta_.set_streamer("HnswStreamer", HnswEntity::kRevision, params); @@ -219,8 +342,21 @@ int HnswStreamer::init( return this->init(imeta, params); } + if (dynamic_cast(quantizer.get()) || + (quantizer->require_train() && imeta.quantizer_name().empty())) { + return IndexError_Unsupported; + } quantizer_ = quantizer; - int ret = this->init(imeta, params); + IndexMeta meta = imeta; + auto qparams = meta.quantizer_params(); + qparams.set("input_dimension", static_cast(quantizer->dim())); + qparams.set("input_data_type", + static_cast(quantizer->input_data_type() == + turbo::DataType::kFp16 + ? IndexMeta::DT_FP16 + : IndexMeta::DT_FP32)); + meta.set_quantizer(meta.quantizer_name(), meta.quantizer_revision(), qparams); + int ret = this->init(meta, params); if (ret != 0) { quantizer_.reset(); } @@ -302,6 +438,10 @@ int HnswStreamer::setup_entity() { } int HnswStreamer::open(IndexStorage::Pointer stg) { + return open(std::move(stg), false); +} + +int HnswStreamer::open(IndexStorage::Pointer stg, bool read_only) { LOG_INFO("HnswStreamer open"); if (ailego_unlikely(state_ != STATE_INITED)) { @@ -334,6 +474,9 @@ int HnswStreamer::open(IndexStorage::Pointer stg) { return ret; } + // The entity takes ownership of the storage handle; keep a reference for + // the quantizer state restore below. + IndexStorage::Pointer stg_ref = stg; ret = entity_->open(std::move(stg), max_index_size_, check_crc_enabled_); if (ret != 0) { return ret; @@ -341,6 +484,14 @@ int HnswStreamer::open(IndexStorage::Pointer stg) { IndexMeta index_meta; ret = entity_->get_index_meta(&index_meta); if (ret == IndexError_NoExist) { + if (!quantizer_ && !meta_.quantizer_name().empty()) { + ret = RestoreQuantizer(stg_ref, meta_, &quantizer_); + if (ret != 0) return ret; + } + if (quantizer_) { + ret = DumpTurboQuantizerSegment(stg_ref, *quantizer_); + if (ret != 0) return ret; + } // Set IndexMeta for the new index ret = entity_->set_index_meta(meta_); if (ret != 0) { @@ -359,6 +510,15 @@ int HnswStreamer::open(IndexStorage::Pointer stg) { LOG_ERROR("IndexMeta mismatch from the previous in index"); return IndexError_Mismatch; } + if (!index_meta.quantizer_name().empty()) { + meta_.set_quantizer(index_meta.quantizer_name(), + index_meta.quantizer_revision(), + index_meta.quantizer_params()); + if (!quantizer_ || quantizer_->require_train()) { + ret = RestoreQuantizer(stg_ref, meta_, &quantizer_); + if (ret != 0) return ret; + } + } // The IndexMetric Params may be updated like MipsSquaredEuclidean auto metric_params = index_meta.metric_params(); metric_params.merge(meta_.metric_params()); @@ -382,6 +542,15 @@ int HnswStreamer::open(IndexStorage::Pointer stg) { }; if (quantizer_) { + if (dynamic_cast(quantizer_.get())) { + return IndexError_Unsupported; + } + // Read-only queries use ADC; the centroid-pair SDC table is only needed + // for graph construction. Prepare it before admitting concurrent inserts. + if (!read_only) { + ret = quantizer_->prepare_build(); + if (ret != 0) return ret; + } if (use_external_vector_) { // External records and build queries are raw; search queries are encoded. BindTurboDistances<&TurboQuantizer::calc_distance_input_input, @@ -573,6 +742,7 @@ int HnswStreamer::open(IndexStorage::Pointer stg) { return ret; } + read_only_ = read_only; state_ = STATE_OPENED; magic_ = IndexContext::GenerateMagic(); @@ -620,6 +790,21 @@ int HnswStreamer::dump(const IndexDumper::Pointer &dumper) { LOG_ERROR("Failed to serialize meta into dumper."); return ret; } + if (quantizer_ && quantizer_->require_train()) { + std::string state; + ret = quantizer_->serialize(&state); + if (ret != 0) return ret; + if (state.empty()) return IndexError_InvalidFormat; + const size_t size = state.size(); + const uint32_t crc = ailego::Crc32c::Hash(state.data(), size, 0); + state.resize((size + 31u) & ~size_t(31u)); + if (dumper->write(state.data(), state.size()) != state.size()) { + return IndexError_WriteData; + } + ret = dumper->append(HNSW_TURBO_QUANTIZER_SEG_ID, size, state.size() - size, + crc); + if (ret != 0) return ret; + } return entity_->dump(dumper); } @@ -718,6 +903,7 @@ void HnswStreamer::bind_search_dist_space(HnswContext *ctx) const { int HnswStreamer::add_with_id_impl(uint32_t id, const void *query, const IndexQueryMeta &qmeta, IndexStreamer::Context::Pointer &context) { + if (read_only_) return IndexError_Unsupported; int ret = check_params(query, qmeta); if (ailego_unlikely(ret != 0)) { return ret; @@ -820,6 +1006,7 @@ int HnswStreamer::add_with_id_impl(uint32_t id, const void *query, int HnswStreamer::add_impl(uint64_t pkey, const void *query, const IndexQueryMeta &qmeta, IndexStreamer::Context::Pointer &context) { + if (read_only_) return IndexError_Unsupported; int ret = check_params(query, qmeta); if (ailego_unlikely(ret != 0)) { return ret; @@ -928,7 +1115,7 @@ int HnswStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, int HnswStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, uint32_t count, IndexStreamer::Context::Pointer &context) const { - int ret = check_params(query, qmeta); + int ret = check_params(query, qmeta, true); if (ailego_unlikely(ret != 0)) { return ret; } @@ -971,7 +1158,9 @@ int HnswStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, return ret; } ctx->topk_to_result(q); - query = static_cast(query) + qmeta.element_size(); + query = static_cast(query) + + (quantizer_ ? quantizer_->quantized_query_vector_length() + : qmeta.element_size()); } if (ailego_unlikely(ctx->error())) { @@ -987,7 +1176,7 @@ int HnswStreamer::search_candidates_impl( std::vector *scores) const { keys.clear(); if (scores) scores->clear(); - int ret = check_params(query, qmeta); + int ret = check_params(query, qmeta, true); if (ailego_unlikely(ret != 0)) { return ret; } @@ -1067,7 +1256,7 @@ int HnswStreamer::search_bf_impl( int HnswStreamer::search_bf_impl( const void *query, const IndexQueryMeta &qmeta, uint32_t count, IndexStreamer::Context::Pointer &context) const { - int ret = check_params(query, qmeta); + int ret = check_params(query, qmeta, true); if (ailego_unlikely(ret != 0)) { return ret; } @@ -1127,7 +1316,9 @@ int HnswStreamer::search_bf_impl( } } ctx->topk_to_result(q); - query = static_cast(query) + qmeta.element_size(); + query = static_cast(query) + + (quantizer_ ? quantizer_->quantized_query_vector_length() + : qmeta.element_size()); } } else { auto &filter = ctx->filter(); @@ -1151,7 +1342,9 @@ int HnswStreamer::search_bf_impl( } } ctx->topk_to_result(q); - query = static_cast(query) + qmeta.element_size(); + query = static_cast(query) + + (quantizer_ ? quantizer_->quantized_query_vector_length() + : qmeta.element_size()); } } @@ -1166,7 +1359,7 @@ int HnswStreamer::search_bf_by_p_keys_impl( const void *query, const std::vector> &p_keys, const IndexQueryMeta &qmeta, uint32_t count, Context::Pointer &context) const { - int ret = check_params(query, qmeta); + int ret = check_params(query, qmeta, true); if (ailego_unlikely(ret != 0)) { return ret; } @@ -1231,7 +1424,9 @@ int HnswStreamer::search_bf_by_p_keys_impl( } } ctx->topk_to_result(q); - query = static_cast(query) + qmeta.element_size(); + query = static_cast(query) + + (quantizer_ ? quantizer_->quantized_query_vector_length() + : qmeta.element_size()); } } else { auto &filter = ctx->filter(); @@ -1255,7 +1450,9 @@ int HnswStreamer::search_bf_by_p_keys_impl( } } ctx->topk_to_result(q); - query = static_cast(query) + qmeta.element_size(); + query = static_cast(query) + + (quantizer_ ? quantizer_->quantized_query_vector_length() + : qmeta.element_size()); } } diff --git a/src/core/algorithm/hnsw/hnsw_streamer.h b/src/core/algorithm/hnsw/hnsw_streamer.h index c78b7de931..82d6d198e3 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer.h +++ b/src/core/algorithm/hnsw/hnsw_streamer.h @@ -95,6 +95,11 @@ class HnswStreamer : public IndexStreamer { return quantizer_ != nullptr; } + //! Reconstruct the persisted quantizer, sharing it with the interface layer. + static int RestoreQuantizer( + const IndexStorage::Pointer &storage, const IndexMeta &meta, + std::shared_ptr *quantizer); + //! Whether graph construction is backed by a turbo quantizer. Compatible //! FP32 providers use a dedicated quantizer; other providers use a metric. bool uses_turbo_build_distance() const { @@ -185,6 +190,7 @@ class HnswStreamer : public IndexStreamer { //! Open index from file path int open(IndexStorage::Pointer stg) override; + int open(IndexStorage::Pointer stg, bool read_only) override; //! Close file int close() override; @@ -208,20 +214,8 @@ class HnswStreamer : public IndexStreamer { void print_debug_info() override; private: - inline int check_params(const void *query, - const IndexQueryMeta &qmeta) const { - if (ailego_unlikely(!query)) { - LOG_ERROR("null query"); - return IndexError_InvalidArgument; - } - if (ailego_unlikely(qmeta.dimension() != meta_.dimension() || - qmeta.data_type() != meta_.data_type() || - qmeta.element_size() != meta_.element_size())) { - LOG_ERROR("Unsupported query meta"); - return IndexError_Mismatch; - } - return 0; - } + int check_params(const void *query, const IndexQueryMeta &qmeta, + bool search = false) const; inline int check_sparse_count_is_zero(const uint32_t *sparse_count, uint32_t count) const { @@ -309,6 +303,7 @@ class HnswStreamer : public IndexStreamer { uint32_t magic_{0U}; State state_{STATE_INIT}; + bool read_only_{false}; bool bf_enabled_{false}; bool check_crc_enabled_{false}; bool filter_same_key_{false}; diff --git a/src/core/algorithm/ivf/ivf_builder.cc b/src/core/algorithm/ivf/ivf_builder.cc index 9eafc6fcb5..cccb4c8880 100644 --- a/src/core/algorithm/ivf/ivf_builder.cc +++ b/src/core/algorithm/ivf/ivf_builder.cc @@ -182,10 +182,6 @@ int IVFBuilder::init(const IndexMeta &meta, const ailego::Params ¶ms, LOG_ERROR("IVFBuilder state wrong. state=%d", state_); return IndexError_Logic; } - if (dynamic_cast(quantizer.get())) { - LOG_ERROR("Turbo IVF postings do not support packed-code quantizers"); - return IndexError_Unsupported; - } if (meta.major_order() == IndexMeta::MO_COLUMN || quantizer->meta().major_order() == IndexMeta::MO_COLUMN) { LOG_ERROR("Turbo IVF postings require row major order"); @@ -227,7 +223,12 @@ int IVFBuilder::init(const IndexMeta &meta, const ailego::Params ¶ms, } turbo_quantizer_ = quantizer; - int ret = this->init(raw_meta, params); + ailego::Params posting_params = params; + if (dynamic_cast(quantizer.get())) { + posting_params.set(PARAM_IVF_BUILDER_BLOCK_VECTOR_COUNT, uint32_t{32}); + posting_params.set(PARAM_IVF_BUILDER_STORE_ORIGINAL_FEATURES, true); + } + int ret = this->init(raw_meta, posting_params); if (ret != 0) { turbo_quantizer_.reset(); return ret; @@ -860,7 +861,7 @@ int IVFBuilder::dump_index(const IndexDumper::Pointer &dumper) { } IVFDumper::Pointer ivf_dumper = std::make_shared( posting_meta, dumper, centroid_index_->centroids_count(), - block_vector_count_); + block_vector_count_, turbo_quantizer_); if (!ivf_dumper) { LOG_ERROR("Alloc IVFDumper failed"); return IndexError_NoMemory; diff --git a/src/core/algorithm/ivf/ivf_dumper.cc b/src/core/algorithm/ivf/ivf_dumper.cc index 44c3ffdd2c..4977efb734 100644 --- a/src/core/algorithm/ivf/ivf_dumper.cc +++ b/src/core/algorithm/ivf/ivf_dumper.cc @@ -37,7 +37,8 @@ int IVFDumper::dump_inverted_block(uint32_t inverted_list_id, int ret = this->check_dump_inverted_list(inverted_list_id); ivf_check_error_code(ret); - if (block_.match_order(column_major ? IndexMeta::MajorOrder::MO_COLUMN + if (!dynamic_cast(quantizer_.get()) && + block_.match_order(column_major ? IndexMeta::MajorOrder::MO_COLUMN : IndexMeta::MajorOrder::MO_ROW) && vector_count == block_.capacity()) { // Dump the block directly @@ -409,7 +410,19 @@ int IVFDumper::dump_block() { } size_t size = ailego_align(block_.bytes(), 32); - if (dumper_->write(block_.data(), size) != size) { + const void *data = block_.data(); + std::string packed; + if (auto *packer = + dynamic_cast(quantizer_.get())) { + // Packed blocks have a fixed 32-code stride, including the final block. + size = block_.block_size(); + packed.resize(size, '\0'); + int ret = packer->pack_codes(data, block_.size(), block_.element_size(), + &packed[0]); + ivf_check_error_code(ret); + data = packed.data(); + } + if (dumper_->write(data, size) != size) { LOG_ERROR("Failed to write data into dumper %s", dumper_->name().c_str()); return IndexError_WriteData; } diff --git a/src/core/algorithm/ivf/ivf_dumper.h b/src/core/algorithm/ivf/ivf_dumper.h index 79acdc35bb..72b2218e49 100644 --- a/src/core/algorithm/ivf/ivf_dumper.h +++ b/src/core/algorithm/ivf/ivf_dumper.h @@ -16,6 +16,7 @@ #include #include #include "metric/metric_params.h" +#include "turbo/quantizer/common/pq_quantizer/packed_code_quantizer.h" #include "turbo/quantizer/quantizer.h" #include "ivf_index_format.h" #include "ivf_params.h" @@ -172,9 +173,11 @@ class IVFDumper { //! Constructor IVFDumper(const IndexMeta &meta, const IndexDumper::Pointer &dumper, - size_t inverted_list_count, size_t block_vector_count) + size_t inverted_list_count, size_t block_vector_count, + turbo::Quantizer::Pointer quantizer = nullptr) : meta_(meta), dumper_(dumper), + quantizer_(std::move(quantizer)), block_vector_count_(block_vector_count), inverted_lists_meta_(inverted_list_count) { block_.init(meta, block_vector_count_); @@ -259,6 +262,7 @@ class IVFDumper { Block block_{}; // vectors grouped in block const IndexMeta meta_{}; // IndexMeta of the inverted index const IndexDumper::Pointer dumper_{}; + turbo::Quantizer::Pointer quantizer_{}; size_t block_vector_count_{kDefaultBlockCount}; std::vector inverted_lists_meta_{}; std::vector keys_{}; diff --git a/src/core/algorithm/ivf/ivf_entity.cc b/src/core/algorithm/ivf/ivf_entity.cc index c80b413e54..3028705cbd 100644 --- a/src/core/algorithm/ivf/ivf_entity.cc +++ b/src/core/algorithm/ivf/ivf_entity.cc @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include "metric/metric_params.h" @@ -95,7 +96,8 @@ template bool calculate_row_major_block(ScatterCursor *cursor, const Distance &distance, size_t vector_count, size_t element_size, size_t serialized_size, uint8_t *scratch, - float *distances, bool calculate) { + float *distances, bool calculate, + bool packed = false) { if (element_size == 0 || vector_count > std::numeric_limits::max() / element_size) { return false; @@ -108,6 +110,15 @@ bool calculate_row_major_block(ScatterCursor *cursor, const Distance &distance, return cursor->skip(serialized_size); } + if (packed) { + if (cursor->contiguous_size() >= serialized_size) { + distance(cursor->data(), vector_count, distances); + return cursor->skip(serialized_size); + } + if (!scratch || !cursor->copy(scratch, serialized_size)) return false; + distance(scratch, vector_count, distances); + return true; + } size_t processed = 0; while (processed < vector_count) { const size_t contiguous = cursor->contiguous_size(); @@ -606,6 +617,9 @@ int IVFEntity::load_header(const IndexStorage::Pointer &container) { return IndexError_InvalidFormat; } + distance_quantizer_.reset(); + packed_quantizer_ = nullptr; + packed_query_.clear(); const auto &attributes = meta_.attributes(); if (attributes.has(IVF_TURBO_FORMAT_VERSION_ATTRIBUTE)) { uint32_t version = 0; @@ -629,6 +643,25 @@ int IVFEntity::load_header(const IndexStorage::Pointer &container) { int ret = reformer_.init(meta_); ivf_check_error_code(ret); + //! Route plain fp32/fp16 postings through the Turbo batch distance kernels + //! when the stored layout matches the distance quantizer contract: dense + //! row-major rows with no legacy reformer (e.g. int8/int4 scales, OPQ + //! rotation). Legacy column-major and reformed indexes keep the metric + //! calculator below. + if (meta_.major_order() == IndexMeta::MajorOrder::MO_ROW && + meta_.reformer_name().empty()) { + auto distance_quantizer = CreateDistanceQuantizer(meta_); + if (distance_quantizer) { + distance_quantizer_ = std::move(distance_quantizer); + calculator_ = std::make_shared( + meta_, nullptr, header_.block_vector_count); + if (!calculator_) { + return IndexError_NoMemory; + } + return 0; + } + } + //! Create the distance calculator auto metric = IndexFactory::CreateMetric(meta_.metric_name()); if (!metric) { @@ -669,10 +702,6 @@ int IVFEntity::load_quantizer(const IndexStorage::Pointer &container) { meta_.quantizer_name().c_str()); return IndexError_NoExist; } - if (dynamic_cast(quantizer.get())) { - LOG_ERROR("Packed Turbo codes are incompatible with row-major IVF"); - return IndexError_Unsupported; - } ret = quantizer->init(raw_meta, meta_.quantizer_params()); ivf_check_with_msg(ret, "Failed to initialize Turbo IVF quantizer"); const bool matched_input = @@ -711,7 +740,17 @@ int IVFEntity::load_quantizer(const IndexStorage::Pointer &container) { return IndexError_InvalidFormat; } + packed_quantizer_ = + dynamic_cast(quantizer.get()); + if (packed_quantizer_ && + (header_.block_vector_count != 32 || + header_.block_size != 32 * meta_.element_size() || + header_.inverted_body_size != + uint64_t{header_.block_count} * header_.block_size)) { + return IndexError_InvalidFormat; + } quantizer_ = std::move(quantizer); + packed_query_.clear(); query_distance_ = turbo::DistanceImpl{}; calculator_ = std::make_shared( meta_, nullptr, header_.block_vector_count); @@ -719,10 +758,17 @@ int IVFEntity::load_quantizer(const IndexStorage::Pointer &container) { } int IVFEntity::bind_query(const void *query, const IndexQueryMeta &qmeta) { + if (distance_quantizer_) { + // The legacy pipeline has already converted the query to the posting + // layout. + query_distance_ = distance_quantizer_->distance(query, qmeta); + return query_distance_.valid() ? 0 : IndexError_Unsupported; + } if (!quantizer_) { return 0; } query_distance_ = turbo::DistanceImpl{}; + packed_query_.clear(); std::string encoded; IndexQueryMeta encoded_meta; int ret = quantizer_->quantize(query, qmeta, &encoded, &encoded_meta); @@ -731,6 +777,10 @@ int IVFEntity::bind_query(const void *query, const IndexQueryMeta &qmeta) { LOG_ERROR("Unexpected Turbo IVF query code size"); return IndexError_InvalidArgument; } + if (packed_quantizer_) { + packed_query_ = std::move(encoded); + return 0; + } query_distance_ = quantizer_->distance(encoded.data(), encoded_meta); if (!query_distance_.valid()) { LOG_ERROR("Turbo quantizer does not support IVF posting distances"); @@ -824,6 +874,17 @@ int IVFEntity::load(const IndexStorage::Pointer &container) { } } + if (packed_quantizer_ && + (!features_ || + features_->data_size() != + vector_count() * quantizer_->dim() * + (quantizer_->input_data_type() == turbo::DataType::kFp16 + ? 2u + : 4u))) { + LOG_ERROR("Packed IVF requires original vectors for fetch and merge"); + return IndexError_InvalidFormat; + } + LOG_DEBUG( "Load inverted index done, docs=%u invertedListCnt=%u " "elementSize=%u metric=%s reformer=%s", @@ -837,7 +898,8 @@ int IVFEntity::search(size_t inverted_list_id, const void *query, const IndexFilter &filter, uint32_t *scan_count, IndexDocumentHeap *heap, IndexContext::Stats *context_stats) const { - if (quantizer_ && !query_distance_.valid()) { + if ((quantizer_ || distance_quantizer_) && !query_distance_.valid() && + packed_query_.empty()) { return IndexError_InvalidArgument; } ailego_assert_with(inverted_list_id < header_.inverted_list_count, @@ -858,7 +920,7 @@ int IVFEntity::search(size_t inverted_list_id, const void *query, }; const bool row_major = meta_.major_order() == IndexMeta::MO_ROW; if (row_major) { - scatter_vector_.resize(element_size); + scatter_vector_.resize(packed_quantizer_ ? block_size : element_size); } const auto norm_val = this->inverted_list_normalize_value(inverted_list_id); for (size_t i = 0; i < list_meta->block_count; i += batch_size) { @@ -913,10 +975,11 @@ int IVFEntity::search(size_t inverted_list_id, const void *query, } const size_t serialized_size = std::min(block_size, size - block_offset); if (keeps == 0) { - if (row_major && !calculate_row_major_block( - &scatter_cursor, distance, vecs_count, - element_size, serialized_size, - scatter_vector_.data(), distances.data(), false)) { + if (row_major && + !calculate_row_major_block(&scatter_cursor, distance, vecs_count, + element_size, serialized_size, + scatter_vector_.data(), distances.data(), + false, packed_quantizer_ != nullptr)) { LOG_ERROR("Invalid scatter-read IVF block, off=%zu, size=%zu", off + block_offset, serialized_size); return IndexError_ReadData; @@ -928,7 +991,7 @@ int IVFEntity::search(size_t inverted_list_id, const void *query, if (!calculate_row_major_block(&scatter_cursor, distance, vecs_count, element_size, serialized_size, scatter_vector_.data(), distances.data(), - true)) { + true, packed_quantizer_ != nullptr)) { LOG_ERROR("Invalid scatter-read IVF block, off=%zu, size=%zu", off + block_offset, serialized_size); return IndexError_ReadData; @@ -961,7 +1024,8 @@ int IVFEntity::search(size_t inverted_list_id, const void *query, int IVFEntity::search(size_t inverted_list_id, const void *query, uint32_t *scan_count, IndexDocumentHeap *heap, IndexContext::Stats *context_stats) const { - if (quantizer_ && !query_distance_.valid()) { + if ((quantizer_ || distance_quantizer_) && !query_distance_.valid() && + packed_query_.empty()) { return IndexError_InvalidArgument; } ailego_assert_with(inverted_list_id < header_.inverted_list_count, @@ -982,7 +1046,7 @@ int IVFEntity::search(size_t inverted_list_id, const void *query, }; const bool row_major = meta_.major_order() == IndexMeta::MO_ROW; if (row_major) { - scatter_vector_.resize(element_size); + scatter_vector_.resize(packed_quantizer_ ? block_size : element_size); } const auto norm_val = this->inverted_list_normalize_value(inverted_list_id); for (size_t i = 0; i < list_meta->block_count; i += batch_size) { @@ -1030,7 +1094,7 @@ int IVFEntity::search(size_t inverted_list_id, const void *query, if (!calculate_row_major_block(&scatter_cursor, distance, vecs_count, element_size, serialized_size, scatter_vector_.data(), distances.data(), - true)) { + true, packed_quantizer_ != nullptr)) { LOG_ERROR("Invalid scatter-read IVF block, off=%zu, size=%zu", off + block_offset, serialized_size); return IndexError_ReadData; @@ -1055,6 +1119,64 @@ int IVFEntity::search(size_t inverted_list_id, const void *query, return 0; } +int IVFEntity::search_by_keys(const std::vector &keys, + const IndexFilter &filter, + IndexDocumentHeap *heap, + IndexContext::Stats *stats) const { + if (!quantizer_) return IndexError_Unsupported; + if (!query_distance_.valid() && packed_query_.empty()) { + return IndexError_InvalidArgument; + } + std::vector distances(header_.block_vector_count); + // Consecutive candidates in the same packed block reuse its distance scan. + uint64_t cached_block = std::numeric_limits::max(); + for (uint64_t key : keys) { + if (key == kInvalidKey) continue; + if (filter.is_valid() && filter(key)) { + ++(*stats->mutable_filtered_count()); + continue; + } + uint32_t id; + int ret = key_to_id(key, &id); + ivf_check_error_code(ret); + if (id >= vector_count()) continue; + InvertedVecLocation loc(0, false); + if (offsets_->fetch(id * sizeof(loc), &loc, sizeof(loc)) != sizeof(loc)) { + return IndexError_ReadData; + } + if (loc.column_major) return IndexError_InvalidFormat; + float score = 0; + IndexStorage::MemoryBlock data; + if (packed_quantizer_) { + // Packed offsets retain logical row positions for locating block/lane. + const uint64_t block = + loc.offset / header_.block_size * header_.block_size; + const size_t lane = (loc.offset - block) / meta_.element_size(); + if (lane >= header_.block_vector_count) return IndexError_InvalidFormat; + if (block != cached_block) { + if (inverted_->read(block, data, header_.block_size) != + header_.block_size) { + return IndexError_ReadData; + } + query_features_distance(nullptr, data.data(), + header_.block_vector_count, distances.data()); + *stats->mutable_dist_calced_count() += header_.block_vector_count; + cached_block = block; + } + score = distances[lane]; + } else { + if (inverted_->read(loc.offset, data, meta_.element_size()) != + meta_.element_size()) { + return IndexError_ReadData; + } + query_features_distance(nullptr, data.data(), 1, &score); + ++(*stats->mutable_dist_calced_count()); + } + heap->emplace(key, score, id); + } + return 0; +} + //! search all inverted list with filter int IVFEntity::search(const void *query, const IndexFilter &filter, IndexDocumentHeap *heap, @@ -1089,7 +1211,11 @@ int IVFEntity::search_by_p_keys(const void *query, const IndexFilter &filter, IndexDocumentHeap *heap, IndexContext::Stats *context_stats) const { - if (quantizer_ && !query_distance_.valid()) { + if (quantizer_) { + // Packed PQ codes must be scored as blocks, not read as individual rows. + return search_by_keys(p_keys, filter, heap, context_stats); + } + if (distance_quantizer_ && !query_distance_.valid()) { return IndexError_InvalidArgument; } for (uint64_t key : p_keys) { @@ -1113,7 +1239,7 @@ int IVFEntity::search_by_p_keys(const void *query, ret = read_posting_vector(id, vector); ivf_check_error_code(ret); float distance; - if (quantizer_) { + if (distance_quantizer_) { calculator_->query_features_distance(query_distance_, vector.data(), 1, &distance); } else { @@ -1332,7 +1458,10 @@ IVFEntity::Pointer IVFEntity::clone(const IVFEntity::Pointer &entity) const { entity->reformer_ = this->reformer_; entity->calculator_ = this->calculator_; entity->quantizer_ = this->quantizer_; + entity->distance_quantizer_ = this->distance_quantizer_; entity->query_distance_ = turbo::DistanceImpl{}; + entity->packed_quantizer_ = packed_quantizer_; + entity->packed_query_.clear(); entity->header_ = this->header_; entity->container_ = this->container_; diff --git a/src/core/algorithm/ivf/ivf_entity.h b/src/core/algorithm/ivf/ivf_entity.h index 09087b5101..8568948cf2 100644 --- a/src/core/algorithm/ivf/ivf_entity.h +++ b/src/core/algorithm/ivf/ivf_entity.h @@ -13,6 +13,7 @@ // limitations under the License. #pragma once +#include #include #include #include "ivf_distance_calculator.h" @@ -67,6 +68,11 @@ class IVFEntity { int search(const void *query, IndexDocumentHeap *heap, IndexContext::Stats *context_stats) const; + //! Score only the requested IDs in this index's posting score space. + int search_by_keys(const std::vector &keys, + const IndexFilter &filter, IndexDocumentHeap *heap, + IndexContext::Stats *stats) const; + //! Search only the specified primary keys, using posting-space distances. int search_by_p_keys(const void *query, const std::vector &p_keys, const IndexFilter &filter, IndexDocumentHeap *heap, @@ -134,7 +140,8 @@ class IVFEntity { ailego_assert_with(*vecs_count <= header_.block_vector_count, "invalid vecs"); const size_t off = iv_meta->offset + local_block_id * header_.block_size; - const size_t size = *vecs_count * meta_.element_size(); + const size_t size = packed_quantizer_ ? header_.block_size + : *vecs_count * meta_.element_size(); if (inverted_->read(off, data, size) != size) { LOG_ERROR("Failed to read block off=%zu size=%zu", off, size); return nullptr; @@ -376,7 +383,10 @@ class IVFEntity { void query_features_distance(const void *query, const void *features, size_t count, float *out) const { - if (quantizer_) { + if (packed_quantizer_) { + packed_quantizer_->calc_distance_packed_block(features, count, + packed_query_.data(), out); + } else if (quantizer_ || distance_quantizer_) { calculator_->query_features_distance(query_distance_, features, count, out); } else { @@ -399,8 +409,13 @@ class IVFEntity { mutable IVFReformerWrapper reformer_{}; IVFDistanceCalculator::Pointer calculator_{}; turbo::Quantizer::Pointer quantizer_{}; + // Legacy rows are already encoded. This helper only supplies distance + // kernels; it must not change query transformation or provider decoding. + turbo::Quantizer::Pointer distance_quantizer_{}; // Each context owns its entity clone and therefore its query buffer/LUT. turbo::DistanceImpl query_distance_{}; + turbo::PackedCodeQuantizer *packed_quantizer_{nullptr}; + std::string packed_query_{}; IndexStorage::Pointer container_{}; IndexStorage::Segment::Pointer inverted_{}; IndexStorage::Segment::Pointer inverted_meta_{}; diff --git a/src/core/algorithm/ivf/ivf_searcher_context.h b/src/core/algorithm/ivf/ivf_searcher_context.h index fe8ea8c7a5..7cc6057e46 100644 --- a/src/core/algorithm/ivf/ivf_searcher_context.h +++ b/src/core/algorithm/ivf/ivf_searcher_context.h @@ -72,8 +72,9 @@ class IVFSearcherContext : public IndexSearcher::Context { std::round(entity_->inverted_list_count() * scan_ratio_)), 1u); - uint32_t nprobe = 0; - params.get(PARAM_IVF_SEARCHER_NPROBE, &nprobe); + requested_nprobe_ = 0; + params.get(PARAM_IVF_SEARCHER_NPROBE, &requested_nprobe_); + uint32_t nprobe = requested_nprobe_; if (nprobe > 0) { nprobe = std::min(nprobe, static_cast(entity_->inverted_list_count())); @@ -166,7 +167,11 @@ class IVFSearcherContext : public IndexSearcher::Context { const ailego::Params ¶ms, uint32_t magic_num) { entity_ = new_entity; centroid_searcher_ctx_ = std::move(centroid_ctx); - int ret = this->update(params); + // Public indexes share a per-thread context. Preserve the request made + // before rebinding, and clamp it against the new entity's list count. + ailego::Params rebound_params = params; + rebound_params.set(PARAM_IVF_SEARCHER_NPROBE, requested_nprobe_); + int ret = this->update(rebound_params); ivf_check_error_code(ret); magic_ = magic_num; @@ -230,6 +235,7 @@ class IVFSearcherContext : public IndexSearcher::Context { bool fetch_vector_{false}; uint32_t topk_{0}; uint32_t magic_{0}; + uint32_t requested_nprobe_{0}; float scan_ratio_{kDefaultScanRatio}; uint32_t max_scan_count_{0}; uint32_t bruteforce_threshold_{kDefaultBfThreshold}; diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc index e80bce6030..b0d4ca0d6a 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -231,6 +232,7 @@ core::IndexContext::Pointer &Index::acquire_context() { } int Index::train() { + if (pending_quantizer_training_) return core::IndexError_NoReady; is_trained_ = true; return 0; } @@ -244,6 +246,7 @@ bool Index::is_trained() const { } uint32_t Index::get_doc_count() const { + if (pending_quantizer_training_) return 0; if (streamer_ == nullptr) { return -1; } @@ -258,6 +261,7 @@ core::IndexStreamer::Pointer Index::index_searcher() { } core::IndexProvider::Pointer Index::create_index_provider() const { + if (pending_quantizer_training_) return nullptr; return streamer_->create_provider(); } @@ -425,6 +429,32 @@ int Index::create_and_init_converter_reformer( return init_converter_reformer(converter_name, converter_params); } +int Index::init_pq_quantizer(const QuantizerParam ¶m, + const BaseIndexParam &index_param, + std::shared_ptr *quantizer) { + const auto *pq = dynamic_cast(¶m); + if (!pq || pq->num_chunk <= 0 || pq->num_chunk > index_param.dimension || + (pq->num_bits != 4 && pq->num_bits != 8) || + (pq->fast_scan && + (pq->num_bits != 4 || index_param.index_type != IndexType::kIVF)) || + (pq->num_bits == 4 && index_param.dimension % pq->num_chunk != 0) || + (pq->enable_rotate && (pq->opq_iter == 0 || pq->opq_pq_iter == 0))) { + return core::IndexError_InvalidArgument; + } + const char *name = pq->fast_scan ? "PqFastQuantizer" + : (pq->num_bits == 4 ? "PqInt4Quantizer" + : "PqInt8Quantizer"); + ailego::Params params; + params.set("num_chunk", static_cast(pq->num_chunk)); + params.set("rotate_type", std::string(pq->enable_rotate ? "opq" : "none")); + params.set("opq_iter", pq->opq_iter); + params.set("opq_pq_iter", pq->opq_pq_iter); + proxima_index_meta_.set_quantizer(name, 0, params); + *quantizer = core::IndexFactory::CreateQuantizer(name); + return *quantizer ? (*quantizer)->init(proxima_index_meta_, params) + : core::IndexError_NoExist; +} + int Index::init_converter_reformer(const std::string &converter_name, const ailego::Params &converter_params) { proxima_index_meta_.set_converter(converter_name, 0, converter_params); @@ -571,7 +601,9 @@ int Index::open(const std::string &file_path, StorageOptions storage_options) { if (ret != 0) { return ret; } - if (streamer_ == nullptr || streamer_->open(storage_) != 0) { + if (!pending_quantizer_training_ && + (streamer_ == nullptr || + streamer_->open(storage_, storage_options.read_only) != 0)) { LOG_ERROR("Failed to open streamer, path: %s", file_path.c_str()); return core::IndexError_Runtime; } @@ -624,7 +656,7 @@ int Index::open(const std::string &file_path, StorageOptions storage_options) { } // TODO: context pool - if (!init_context()) { // to validate if any error, will be overwritten + if (!pending_quantizer_training_ && !init_context()) { LOG_ERROR("Failed to init context"); return core::IndexError_Runtime; } @@ -640,7 +672,7 @@ int Index::close() { return core::IndexError_Runtime; } - if (!is_read_only_) { + if (!is_read_only_ && !pending_quantizer_training_) { if (ailego_unlikely(flush() != 0)) { LOG_ERROR("Failed to cleanup streamer"); return core::IndexError_Runtime; @@ -667,6 +699,7 @@ int Index::close() { } int Index::flush() { + if (pending_quantizer_training_) return core::IndexError_NoReady; if (!is_open_) { LOG_ERROR("Index is not open"); return core::IndexError_Runtime; @@ -695,6 +728,7 @@ bool Index::is_dirty() const { } int Index::fetch(const uint32_t doc_id, VectorDataBuffer *vector_data_buffer) { + if (pending_quantizer_training_) return core::IndexError_NoReady; if (!is_open_) { LOG_ERROR("Index is not open"); return core::IndexError_Runtime; @@ -706,6 +740,7 @@ int Index::fetch(const uint32_t doc_id, VectorDataBuffer *vector_data_buffer) { } int Index::add(const VectorData &vector_data, const uint32_t doc_id) { + if (pending_quantizer_training_) return core::IndexError_NoReady; if (!is_open_) { LOG_ERROR("Index is not open"); return core::IndexError_Runtime; @@ -802,9 +837,9 @@ int Index::search(const VectorData &vector_data, context->reset(); return core::IndexError_Runtime; } - // TODO: tackle query_param's type info loss to loosen the constraint - if (reference_index->param_.index_type != IndexType::kFlat) { - LOG_ERROR("Reference index is not flat"); + if (reference_index->param_.index_type != IndexType::kFlat && + reference_index->param_.index_type != IndexType::kIVF) { + LOG_ERROR("Reference index must be Flat or IVF"); context->reset(); return core::IndexError_Runtime; } @@ -814,32 +849,44 @@ int Index::search(const VectorData &vector_data, context->reset(); return coarse_topk; } - context->set_topk(coarse_topk); - context->set_fetch_vector(false); // no need to fetch vector - std::string transformed_vector; - const void *query = nullptr; - core::IndexQueryMeta query_meta; - ret = _prepare_dense_query(vector_data, &transformed_vector, &query, - &query_meta); - if (ret != 0) { - context->reset(); - return ret; - } - // Refine replaces the coarse scores; only candidate membership and order - // cross this boundary, not document/vector materialization or - // normalization. - std::vector> keys(1); - if (_execute_dense_search(query, query_meta, search_param, context, - &keys[0]) != 0) { - LOG_ERROR("Failed to search"); + try { + context->set_topk(coarse_topk); + context->set_fetch_vector(false); // no need to fetch vector + std::string transformed_vector; + const void *query = nullptr; + core::IndexQueryMeta query_meta; + ret = _prepare_dense_query(vector_data, &transformed_vector, &query, + &query_meta); + if (ret != 0) { + context->reset(); + return ret; + } + // Refine replaces the coarse scores; only candidate membership and order + // cross this boundary, not document/vector materialization or + // normalization. + std::vector> keys(1); + if (_execute_dense_search(query, query_meta, search_param, context, + &keys[0]) != 0) { + LOG_ERROR("Failed to search"); + if (context) { + context->reset(); + } + return core::IndexError_Runtime; + } + result->reverted_vector_list_.clear(); + ret = reference_index->_refine_dense_candidates(vector_data, search_param, + keys, result); + } catch (const std::bad_alloc &) { + // A huge scale factor (topk * SF up to INT_MAX) passes validation but + // the coarse result heap cannot be allocated; surface it as an error + // instead of terminating the process. + LOG_ERROR("Failed to allocate coarse result heap for %d candidates", + coarse_topk); if (context) { context->reset(); } - return core::IndexError_Runtime; + return core::IndexError_NoMemory; } - result->reverted_vector_list_.clear(); - ret = reference_index->_refine_dense_candidates(vector_data, search_param, - keys, result); } if (context) { context->reset(); @@ -920,15 +967,11 @@ int Index::_dense_add(const VectorData &vector_data, const uint32_t doc_id, } const DenseVector &dense_vector = std::get(vector_data.vector); if (turbo_quantizer_ != nullptr) { - core::IndexQueryMeta new_meta; auto *new_vector = context->mutable_features(); - if (turbo_quantizer_->quantize(dense_vector.data, input_vector_meta_, - new_vector, &new_meta) != 0) { - LOG_ERROR("Failed to quantize vector with turbo quantizer"); - return core::IndexError_Runtime; - } - if (streamer_->add_with_id_impl(doc_id, new_vector->data(), new_meta, - context) != 0) { + new_vector->resize(turbo_quantizer_->quantized_datapoint_vector_length()); + turbo_quantizer_->quantize_data(dense_vector.data, new_vector->data()); + if (streamer_->add_with_id_impl(doc_id, new_vector->data(), + streamer_vector_meta_, context) != 0) { LOG_ERROR("Failed to add vector"); return core::IndexError_Runtime; } @@ -1478,7 +1521,8 @@ int Index::merge(const std::vector &indexes, LOG_ERROR("Failed to init reducer"); return core::IndexError_Runtime; } - if (converter_ != nullptr && reformer_ == nullptr) { + if (pending_quantizer_training_ || + (converter_ != nullptr && reformer_ == nullptr)) { std::vector sources; sources.reserve(indexes.size()); for (const auto &index : indexes) { @@ -1513,19 +1557,36 @@ int Index::merge(const std::vector &indexes, input_vector_meta_); if (holder->error() != 0) return holder->error(); if (holder->count() == 0) { - LOG_ERROR("No vectors available to train converter"); + LOG_ERROR("No vectors available to train quantization"); return core::IndexError_InvalidArgument; } - int ret = converter_->train(holder); + int ret = pending_quantizer_training_ ? turbo_quantizer_->train(holder) + : converter_->train(holder); if (holder->error() != 0) return holder->error(); if (ret != 0) { - LOG_ERROR("Failed to train converter from merge sources"); + LOG_ERROR("Failed to train quantization from merge sources"); return ret; } - ret = CreateReformerFromConverterMeta(converter_, &proxima_index_meta_, - &streamer_vector_meta_, &reformer_); - if (ret != 0) return ret; - streamer_->merge_trained_meta(proxima_index_meta_); + if (pending_quantizer_training_) { + // Training may change the code layout. Persist the codebook before any + // graph nodes are inserted, using the normal streamer open path. + proxima_index_meta_ = turbo_quantizer_->meta(); + streamer_vector_meta_.set_meta( + proxima_index_meta_.data_type(), proxima_index_meta_.dimension(), + static_cast(turbo_quantizer_->type()), + proxima_index_meta_.extra_meta_size()); + ret = streamer_->init(proxima_index_meta_, proxima_index_params_, + turbo_quantizer_); + if (ret != 0) return ret; + ret = streamer_->open(storage_); + if (ret != 0) return ret; + pending_quantizer_training_ = false; + } else { + ret = CreateReformerFromConverterMeta(converter_, &proxima_index_meta_, + &streamer_vector_meta_, &reformer_); + if (ret != 0) return ret; + streamer_->merge_trained_meta(proxima_index_meta_); + } } if (reducer->set_target_streamer_wiht_info(builder_, streamer_, converter_, reformer_, input_vector_meta_, diff --git a/src/core/interface/index_param.cc b/src/core/interface/index_param.cc index 0a3c135c0a..73454d4920 100644 --- a/src/core/interface/index_param.cc +++ b/src/core/interface/index_param.cc @@ -509,6 +509,9 @@ ailego::JsonObject PqQuantizerParam::serialize_to_json_object( auto json_obj = QuantizerParam::serialize_to_json_object(omit_empty_value); json_obj.set("num_chunk", ailego::JsonValue(num_chunk)); json_obj.set("num_bits", ailego::JsonValue(num_bits)); + json_obj.set("fast_scan", ailego::JsonValue(fast_scan)); + json_obj.set("opq_iter", ailego::JsonValue(opq_iter)); + json_obj.set("opq_pq_iter", ailego::JsonValue(opq_pq_iter)); return json_obj; } @@ -519,6 +522,9 @@ bool PqQuantizerParam::deserialize_from_json_object( } DESERIALIZE_VALUE_FIELD(json_obj, num_chunk); DESERIALIZE_VALUE_FIELD(json_obj, num_bits); + DESERIALIZE_VALUE_FIELD(json_obj, fast_scan); + DESERIALIZE_VALUE_FIELD(json_obj, opq_iter); + DESERIALIZE_VALUE_FIELD(json_obj, opq_pq_iter); return true; } diff --git a/src/core/interface/indexes/hnsw_index.cc b/src/core/interface/indexes/hnsw_index.cc index 3dd81f5197..7a92e89f37 100644 --- a/src/core/interface/indexes/hnsw_index.cc +++ b/src/core/interface/indexes/hnsw_index.cc @@ -80,29 +80,108 @@ const char *ResolveTurboQuantizerName(const QuantizerParam &quantizer_param, } // namespace int HNSWIndex::prepare_streamer_open(const StorageOptions &options) { - if (!turbo_quantizer_ || options.create_new) { + if (options.create_new) { return 0; } core::IndexMeta persisted_meta; if (core::IndexHelper::DeserializeFromStorage(storage_.get(), - &persisted_meta) != 0 || - !persisted_meta.quantizer_name().empty()) { + &persisted_meta) != 0) { + return pending_quantizer_training_ ? core::IndexError_InvalidFormat : 0; + } + if (persisted_meta.quantizer_name().empty()) { + if (pending_quantizer_training_) return core::IndexError_InvalidFormat; + if (!turbo_quantizer_) { + return 0; + } + // Reuse normal initialization for old indexes, without selecting Turbo + // again. + turbo_quantizer_.reset(); + streamer_.reset(); + converter_.reset(); + reformer_.reset(); + metric_.reset(); + proxima_index_meta_.clear(); + use_legacy_pipeline_ = true; + return Index::init(param_); + } + // The persisted descriptor decides the search pipeline, including when a + // caller reopens an index with a different quantizer configuration: a + // trained quantizer (e.g. PQ) is restored from its serialized state, and + // the streamer is re-initialized on the persisted code layout. + bool matches_config = false; + if (turbo_quantizer_) { + const core::IndexMeta &config_meta = turbo_quantizer_->meta(); + matches_config = + !turbo_quantizer_->require_train() && + proxima_index_meta_.quantizer_name() == + persisted_meta.quantizer_name() && + config_meta.data_type() == persisted_meta.data_type() && + config_meta.dimension() == persisted_meta.dimension() && + config_meta.element_size() == persisted_meta.element_size(); + } + if (matches_config) { return 0; } - // Reuse normal initialization for old indexes, without selecting Turbo again. - turbo_quantizer_.reset(); - streamer_.reset(); + + std::shared_ptr restored; + int ret = + core::HnswStreamer::RestoreQuantizer(storage_, persisted_meta, &restored); + if (ret != 0) { + LOG_ERROR("Failed to restore turbo quantizer %s, ret=%d", + persisted_meta.quantizer_name().c_str(), ret); + return ret; + } + if (!restored) { + LOG_ERROR("Restored turbo quantizer %s is null", + persisted_meta.quantizer_name().c_str()); + return core::IndexError_Runtime; + } + + pending_quantizer_training_ = false; + turbo_quantizer_ = restored; converter_.reset(); reformer_.reset(); metric_.reset(); - proxima_index_meta_.clear(); - use_legacy_pipeline_ = true; - return Index::init(param_); + proxima_index_meta_ = persisted_meta; + streamer_vector_meta_.set_meta( + persisted_meta.data_type(), persisted_meta.dimension(), + static_cast(turbo_quantizer_->type()), + persisted_meta.extra_meta_size()); + streamer_vector_meta_.set_meta_type(persisted_meta.meta_type()); + LOG_INFO("Restored turbo quantizer %s from the persisted index meta", + persisted_meta.quantizer_name().c_str()); + ret = streamer_->init(proxima_index_meta_, proxima_index_params_, + turbo_quantizer_); + if (ret != 0) { + LOG_ERROR("Failed to re-init streamer with restored turbo quantizer %s", + persisted_meta.quantizer_name().c_str()); + return ret; + } + return 0; } int HNSWIndex::create_and_init_converter_reformer( const QuantizerParam &quantizer_param, const BaseIndexParam &index_param) { const auto &hnsw_param = dynamic_cast(index_param); + if (quantizer_param.type == QuantizerType::kPQ) { + if (hnsw_param.is_sparse || hnsw_param.data_type != DataType::DT_FP32 || + (hnsw_param.metric_type != MetricType::kL2sq && + hnsw_param.metric_type != MetricType::kInnerProduct && + hnsw_param.metric_type != MetricType::kCosine) || + hnsw_param.use_external_vector || hnsw_param.provider) { + return core::IndexError_Unsupported; + } + int ret = + init_pq_quantizer(quantizer_param, index_param, &turbo_quantizer_); + if (ret != 0) return ret; + pending_quantizer_training_ = true; + proxima_index_meta_ = turbo_quantizer_->meta(); + streamer_vector_meta_.set_meta( + proxima_index_meta_.data_type(), proxima_index_meta_.dimension(), + static_cast(turbo_quantizer_->type()), + proxima_index_meta_.extra_meta_size()); + return 0; + } const char *quantizer_name = use_legacy_pipeline_ ? nullptr diff --git a/src/core/interface/indexes/ivf_index.cc b/src/core/interface/indexes/ivf_index.cc index 380c5fc521..81e17cd642 100644 --- a/src/core/interface/indexes/ivf_index.cc +++ b/src/core/interface/indexes/ivf_index.cc @@ -29,7 +29,8 @@ int IVFIndex::create_and_init_converter_reformer( const QuantizerParam ¶m, const BaseIndexParam &index_param) { // Clustering and centroid selection use the input vectors. Only the // inverted lists are encoded, after their centroid assignments are known. - if (index_param.is_sparse || param.enable_rotate || + if (index_param.is_sparse || + (param.enable_rotate && param.type != QuantizerType::kPQ) || index_param.data_type != DataType::DT_FP32 || (index_param.metric_type != MetricType::kL2sq && index_param.metric_type != MetricType::kInnerProduct && @@ -37,7 +38,10 @@ int IVFIndex::create_and_init_converter_reformer( return Index::create_and_init_converter_reformer(param, index_param); } const char *name = nullptr; + ailego::Params quantizer_params; switch (param.type) { + case QuantizerType::kPQ: + return init_pq_quantizer(param, index_param, &ivf_quantizer_); case QuantizerType::kNone: name = "Fp32Quantizer"; break; @@ -53,12 +57,12 @@ int IVFIndex::create_and_init_converter_reformer( default: return Index::create_and_init_converter_reformer(param, index_param); } - proxima_index_meta_.set_quantizer(name, 0, ailego::Params{}); + proxima_index_meta_.set_quantizer(name, 0, quantizer_params); ivf_quantizer_ = core::IndexFactory::CreateQuantizer(name); if (!ivf_quantizer_) { return core::IndexError_NoExist; } - return ivf_quantizer_->init(proxima_index_meta_, ailego::Params{}); + return ivf_quantizer_->init(proxima_index_meta_, quantizer_params); } int IVFIndex::create_and_init_streamer(const BaseIndexParam ¶m) { @@ -404,11 +408,12 @@ int IVFIndex::_dense_fetch(const uint32_t doc_id, } } -int IVFIndex::_dense_search(const VectorData &query, - const BaseIndexQueryParam::Pointer &search_param, - SearchResult *result, - core::IndexContext::Pointer &context) { - int ret = Index::_dense_search(query, search_param, result, context); +int IVFIndex::_collect_dense_result( + const VectorData &query, const core::IndexQueryMeta &query_meta, + const BaseIndexQueryParam::Pointer &search_param, SearchResult *result, + core::IndexContext::Pointer &context) { + int ret = Index::_collect_dense_result(query, query_meta, search_param, + result, context); if (ret != 0 || !ivf_quantizer_ || !context->fetch_vector()) return ret; auto provider = streamer_->create_provider(); if (!provider) return core::IndexError_NoReady; @@ -435,6 +440,7 @@ int IVFIndex::_prepare_for_search( return core::IndexError_Unsupported; } + context->reset_threshold(); context->set_topk(ivf_search_param->topk); context->set_fetch_vector(ivf_search_param->fetch_vector); if (ivf_search_param->filter && ivf_search_param->filter->is_valid()) { @@ -446,12 +452,9 @@ int IVFIndex::_prepare_for_search( context->set_threshold(ivf_search_param->radius); } - if (ivf_search_param->nprobe > 0) { - ailego::Params params; - params.set(core::PARAM_IVF_SEARCHER_NPROBE, ivf_search_param->nprobe); - context->update(params); - } - return 0; + ailego::Params params; + params.set(core::PARAM_IVF_SEARCHER_NPROBE, ivf_search_param->nprobe); + return context->update(params); } int IVFIndex::merge(const std::vector &indexes, diff --git a/src/core/mixed_reducer/mixed_streamer_reducer.cc b/src/core/mixed_reducer/mixed_streamer_reducer.cc index 754632fafc..3451880c6c 100644 --- a/src/core/mixed_reducer/mixed_streamer_reducer.cc +++ b/src/core/mixed_reducer/mixed_streamer_reducer.cc @@ -373,7 +373,7 @@ int MixedStreamerReducer::read_vec(size_t source_streamer_index, const auto &source_meta = streamer->meta(); const auto &target_meta = target_streamer_->meta(); const bool same_layout = - target_builder_ == nullptr && + target_builder_ == nullptr && !quantizer->require_train() && turbo::QuantizerStorageDataTypeMatches(target_meta, source_meta) && target_meta.quantizer_name() == source_meta.quantizer_name() && target_meta.metric_name() == source_meta.metric_name() && @@ -507,17 +507,15 @@ void MixedStreamerReducer::add_vec(int *result) { // the target layout. A target without quantizer or reformer stores // the original format directly, so no encoding is needed. if (target_streamer_quantizer_ != nullptr) { - IndexQueryMeta quantized_meta; - if (target_streamer_quantizer_->quantize(vector, original_query_meta_, - &new_vector, - &quantized_meta) != 0) { - LOG_ERROR("Failed to quantize vector. pkey[%zu]", - (size_t)vector_item.pkey_); - *result = IndexError_Runtime; - return; - } + new_vector.resize( + target_streamer_quantizer_->quantized_datapoint_vector_length()); + target_streamer_quantizer_->quantize_data(vector, new_vector.data()); vector = new_vector.data(); - add_meta = quantized_meta; + add_meta.set_meta( + target_streamer_query_meta.data_type(), + target_streamer_query_meta.dimension(), + static_cast(target_streamer_quantizer_->type()), + target_streamer_query_meta.extra_meta_size()); } else if (target_streamer_reformer_ != nullptr) { IndexQueryMeta new_meta; if (target_streamer_reformer_->convert(vector, original_query_meta_, diff --git a/src/db/index/column/vector_column/engine_helper.hpp b/src/db/index/column/vector_column/engine_helper.hpp index efef2b4481..97b487d439 100644 --- a/src/db/index/column/vector_column/engine_helper.hpp +++ b/src/db/index/column/vector_column/engine_helper.hpp @@ -453,6 +453,8 @@ class ProximaEngineHelper { return core_interface::QuantizerType::kFP16; case QuantizeType::INT8: return core_interface::QuantizerType::kInt8; + case QuantizeType::PQ: + return core_interface::QuantizerType::kPQ; case QuantizeType::INT4: return core_interface::QuantizerType::kInt4; case QuantizeType::RABITQ: @@ -527,8 +529,19 @@ class ProximaEngineHelper { if (auto quantize_type = convert_to_engine_quantize_type(db_index_params->quantize_type()); quantize_type.has_value()) { - index_param_builder->with_quantizer_param( - core_interface::QuantizerParam::Create(quantize_type.value())); + auto quantizer = + core_interface::QuantizerParam::Create(quantize_type.value()); + if (quantize_type.value() == core_interface::QuantizerType::kPQ) { + auto pq = std::static_pointer_cast( + quantizer); + const auto ¶ms = db_index_params->quantizer_param(); + pq->num_chunk = params.num_chunk(); + pq->num_bits = params.num_bits(); + pq->fast_scan = params.fast_scan(); + pq->opq_iter = params.opq_iter(); + pq->opq_pq_iter = params.opq_pq_iter(); + } + index_param_builder->with_quantizer_param(quantizer); } else { return tl::make_unexpected( Status::InvalidArgument("unsupported quantize type")); diff --git a/src/db/index/common/index_params.cc b/src/db/index/common/index_params.cc index 0b696956a0..b9ea241f0d 100644 --- a/src/db/index/common/index_params.cc +++ b/src/db/index/common/index_params.cc @@ -41,6 +41,13 @@ std::string VectorIndexParams::vector_index_params_to_string( oss << class_name << "{" << "metric:" << MetricTypeCodeBook::AsString(metric_type) << ",quantize:" << QuantizeTypeCodeBook::AsString(quantize_type); + if (quantize_type == QuantizeType::PQ) { + oss << ",num_chunk:" << quantizer_param_.num_chunk() + << ",num_bits:" << quantizer_param_.num_bits() + << ",fast_scan:" << quantizer_param_.fast_scan() + << ",opq_iter:" << quantizer_param_.opq_iter() + << ",opq_pq_iter:" << quantizer_param_.opq_pq_iter(); + } return oss.str(); } @@ -133,4 +140,4 @@ Result> AcquireFtsPipeline( } // namespace detail -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/db/index/common/manifest/manifest_codec.cc b/src/db/index/common/manifest/manifest_codec.cc index f5b16ea492..036ce0e6ad 100644 --- a/src/db/index/common/manifest/manifest_codec.cc +++ b/src/db/index/common/manifest/manifest_codec.cc @@ -33,7 +33,12 @@ using pbwire::Writer; // breaks compatibility with existing manifest files. namespace f_quantizer { constexpr uint32_t kEnableRotate = 1; -} +constexpr uint32_t kNumChunk = 2; +constexpr uint32_t kNumBits = 3; +constexpr uint32_t kFastScan = 4; +constexpr uint32_t kOpqIter = 5; +constexpr uint32_t kOpqPqIter = 6; +} // namespace f_quantizer namespace f_base { constexpr uint32_t kMetricType = 1; constexpr uint32_t kQuantizeType = 2; @@ -158,7 +163,7 @@ struct BaseParams { // HNSW_RABITQ, so presence must be reproduced exactly to keep the encoded // bytes identical. bool has_quantizer_param{false}; - bool enable_rotate{false}; + QuantizerParam quantizer_param{}; }; void EncodeBase(const BaseParams &base, std::string *out) { @@ -172,7 +177,20 @@ void EncodeBase(const BaseParams &base, std::string *out) { if (base.has_quantizer_param) { std::string quantizer; Writer qw(&quantizer); - qw.put_bool(f_quantizer::kEnableRotate, base.enable_rotate); + qw.put_bool(f_quantizer::kEnableRotate, + base.quantizer_param.enable_rotate()); + if (base.quantize_type == QuantizeType::PQ) { + qw.put_varint(f_quantizer::kNumChunk, base.quantizer_param.num_chunk(), + false); + qw.put_varint(f_quantizer::kNumBits, base.quantizer_param.num_bits(), + false); + qw.put_varint(f_quantizer::kFastScan, base.quantizer_param.fast_scan(), + false); + qw.put_varint(f_quantizer::kOpqIter, base.quantizer_param.opq_iter(), + false); + qw.put_varint(f_quantizer::kOpqPqIter, base.quantizer_param.opq_pq_iter(), + false); + } w.put_message(f_base::kQuantizerParam, quantizer); } } @@ -194,8 +212,27 @@ BaseParams DecodeBase(std::string_view buf) { base.has_quantizer_param = true; Reader qr(r.bytes()); while (qr.next()) { - if (qr.field() == f_quantizer::kEnableRotate) { - base.enable_rotate = qr.bool_value(); + switch (qr.field()) { + case f_quantizer::kEnableRotate: + base.quantizer_param.set_enable_rotate(qr.bool_value()); + break; + case f_quantizer::kNumChunk: + base.quantizer_param.set_num_chunk(qr.int32_value()); + break; + case f_quantizer::kNumBits: + base.quantizer_param.set_num_bits(qr.int32_value()); + break; + case f_quantizer::kFastScan: + base.quantizer_param.set_fast_scan(qr.bool_value()); + break; + case f_quantizer::kOpqIter: + base.quantizer_param.set_opq_iter(qr.uint32_value()); + break; + case f_quantizer::kOpqPqIter: + base.quantizer_param.set_opq_pq_iter(qr.uint32_value()); + break; + default: + break; } } break; @@ -214,7 +251,7 @@ BaseParams MakeBase(const Params *params) { base.metric_type = params->metric_type(); base.quantize_type = params->quantize_type(); base.has_quantizer_param = true; - base.enable_rotate = params->quantizer_param().enable_rotate(); + base.quantizer_param = params->quantizer_param(); return base; } @@ -271,8 +308,8 @@ HnswIndexParams::OPtr DecodeHnsw(std::string_view buf) { } return std::make_shared( base.metric_type, m, ef_construction, base.quantize_type, - use_contiguous_memory, QuantizerParam(base.enable_rotate), - use_flat_contiguous_memory, flat_data_type); + use_contiguous_memory, base.quantizer_param, use_flat_contiguous_memory, + flat_data_type); } void EncodeHnswRabitq(const HnswRabitqIndexParams *params, std::string *out) { @@ -417,7 +454,7 @@ FlatIndexParams::OPtr DecodeFlat(std::string_view buf) { } } return std::make_shared( - base.metric_type, base.quantize_type, QuantizerParam(base.enable_rotate), + base.metric_type, base.quantize_type, base.quantizer_param, use_contiguous_memory, storage_data_type); } @@ -457,7 +494,7 @@ IVFIndexParams::OPtr DecodeIvf(std::string_view buf) { } return std::make_shared(base.metric_type, n_list, n_iters, use_soar, base.quantize_type, - QuantizerParam(base.enable_rotate)); + base.quantizer_param); } void EncodeDiskAnn(const DiskAnnIndexParams *params, std::string *out) { @@ -499,7 +536,7 @@ DiskAnnIndexParams::OPtr DecodeDiskAnn(std::string_view buf) { } return std::make_shared( base.metric_type, max_degree, list_size, pq_chunk_num, base.quantize_type, - QuantizerParam(base.enable_rotate)); + base.quantizer_param); } void EncodeVamana(const VamanaIndexParams *params, std::string *out) { @@ -577,8 +614,8 @@ VamanaIndexParams::OPtr DecodeVamana(std::string_view buf) { return std::make_shared( base.metric_type, max_degree, search_list_size, alpha, saturate_graph, use_contiguous_memory, use_id_map, base.quantize_type, - QuantizerParam(base.enable_rotate), two_pass_build, - use_flat_contiguous_memory, flat_data_type); + base.quantizer_param, two_pass_build, use_flat_contiguous_memory, + flat_data_type); } void EncodeInvert(const InvertIndexParams *params, std::string *out) { diff --git a/src/db/index/common/manifest/pb_wire.h b/src/db/index/common/manifest/pb_wire.h index aa26e235c8..fa53a4c864 100644 --- a/src/db/index/common/manifest/pb_wire.h +++ b/src/db/index/common/manifest/pb_wire.h @@ -59,8 +59,9 @@ class Writer { //! Writes a varint field. Skipped when the value is zero, matching proto3 //! semantics where default-valued singular fields are not serialized. - void put_varint(uint32_t field, uint64_t value) { - if (value == 0) { + // Set omit_zero=false for explicitly present fields with nonzero defaults. + void put_varint(uint32_t field, uint64_t value, bool omit_zero = true) { + if (omit_zero && value == 0) { return; } put_varint_always(field, value); diff --git a/src/db/index/common/manifest_enum.h b/src/db/index/common/manifest_enum.h index 33f2aabfb2..3b95e1526e 100644 --- a/src/db/index/common/manifest_enum.h +++ b/src/db/index/common/manifest_enum.h @@ -86,6 +86,7 @@ enum class QuantizeType : int32_t { QT_UNIFORM_UINT7 = 5, QT_UNIFORM_UINT8 = 6, QT_UNIFORM_UINT4 = 7, + QT_PQ = 8, }; //! Mirrors proto enum MetricType. diff --git a/src/db/index/common/schema.cc b/src/db/index/common/schema.cc index bea3f86328..fc2e614f05 100644 --- a/src/db/index/common/schema.cc +++ b/src/db/index/common/schema.cc @@ -35,7 +35,8 @@ std::unordered_map> quantize_type_map = { {DataType::VECTOR_FP32, {QuantizeType::FP16, QuantizeType::INT4, QuantizeType::INT8, QuantizeType::RABITQ, QuantizeType::UNIFORM_UINT7, - QuantizeType::UNIFORM_UINT8, QuantizeType::UNIFORM_UINT4}}, + QuantizeType::UNIFORM_UINT8, QuantizeType::UNIFORM_UINT4, + QuantizeType::PQ}}, // {DataType::VECTOR_FP64, {QuantizeType::FP16}}, {DataType::SPARSE_VECTOR_FP32, {QuantizeType::FP16}}, }; @@ -263,6 +264,30 @@ Status FieldSchema::validate() const { if (vector_index_params->quantize_type() != QuantizeType::UNDEFINED) { const auto quantize_type = vector_index_params->quantize_type(); + if (quantize_type == QuantizeType::PQ) { + const auto &pq = vector_index_params->quantizer_param(); + const auto metric = vector_index_params->metric_type(); + if ((index_params_->type() != IndexType::IVF && + index_params_->type() != IndexType::HNSW) || + data_type_ != DataType::VECTOR_FP32 || + (metric != MetricType::L2 && metric != MetricType::IP && + metric != MetricType::COSINE)) { + return Status::InvalidArgument( + "PQ requires an IVF or HNSW FP32 field with L2, IP or COSINE " + "metric"); + } + if (pq.num_chunk() <= 0 || + static_cast(pq.num_chunk()) > dimension_ || + (pq.num_bits() != 4 && pq.num_bits() != 8) || + (pq.fast_scan() && (pq.num_bits() != 4 || + index_params_->type() != IndexType::IVF)) || + (pq.num_bits() == 4 && dimension_ % pq.num_chunk() != 0) || + (pq.enable_rotate() && + (pq.opq_iter() == 0 || pq.opq_pq_iter() == 0))) { + return Status::InvalidArgument( + "Invalid PQ/OPQ parameters; FastScan requires IVF PQ4"); + } + } const bool is_uniform = quantize_type == QuantizeType::UNIFORM_UINT7 || quantize_type == QuantizeType::UNIFORM_UINT8 || quantize_type == QuantizeType::UNIFORM_UINT4; diff --git a/src/db/index/common/type_helper.h b/src/db/index/common/type_helper.h index 38fa691aae..bc528dd6b5 100644 --- a/src/db/index/common/type_helper.h +++ b/src/db/index/common/type_helper.h @@ -480,6 +480,8 @@ struct QuantizeTypeCodeBook { return QuantizeType::UNIFORM_UINT7; case wire::QuantizeType::QT_UNIFORM_UINT8: return QuantizeType::UNIFORM_UINT8; + case wire::QuantizeType::QT_PQ: + return QuantizeType::PQ; case wire::QuantizeType::QT_UNIFORM_UINT4: return QuantizeType::UNIFORM_UINT4; default: @@ -501,6 +503,8 @@ struct QuantizeTypeCodeBook { return wire::QuantizeType::QT_UNIFORM_UINT7; case QuantizeType::UNIFORM_UINT8: return wire::QuantizeType::QT_UNIFORM_UINT8; + case QuantizeType::PQ: + return wire::QuantizeType::QT_PQ; case QuantizeType::UNIFORM_UINT4: return wire::QuantizeType::QT_UNIFORM_UINT4; default: @@ -522,6 +526,8 @@ struct QuantizeTypeCodeBook { return "UNIFORM_UINT7"; case QuantizeType::UNIFORM_UINT8: return "UNIFORM_UINT8"; + case QuantizeType::PQ: + return "PQ"; case QuantizeType::UNIFORM_UINT4: return "UNIFORM_UINT4"; default: diff --git a/src/db/index/segment/segment.cc b/src/db/index/segment/segment.cc index 11879ebb7d..4b71eaa74a 100644 --- a/src/db/index/segment/segment.cc +++ b/src/db/index/segment/segment.cc @@ -860,7 +860,8 @@ Status SegmentImpl::insert_vector_indexer(Doc &doc) { vector_index_params->quantize_type() != QuantizeType::UNIFORM_UINT8 && vector_index_params->quantize_type() != - QuantizeType::UNIFORM_UINT4) { + QuantizeType::UNIFORM_UINT4 && + vector_index_params->quantize_type() != QuantizeType::PQ) { LOG_ERROR("quant vector indexer not found for field %s", field->name().c_str()); return Status::InternalError( @@ -4257,11 +4258,12 @@ Status SegmentImpl::init_memory_components() { } memory_vector_indexers_.insert({field->name(), vector_indexer}); - // Uniform quantizers need the full dataset to train their scale/bias. + // Uniform and PQ quantizers require dataset-level training. // Until optimize creates that index, retain only the raw writer Flat. if (index_params->quantize_type() == QuantizeType::UNIFORM_UINT7 || index_params->quantize_type() == QuantizeType::UNIFORM_UINT8 || - index_params->quantize_type() == QuantizeType::UNIFORM_UINT4) { + index_params->quantize_type() == QuantizeType::UNIFORM_UINT4 || + index_params->quantize_type() == QuantizeType::PQ) { continue; } diff --git a/src/db/index/segment/segment_helper.cc b/src/db/index/segment/segment_helper.cc index c8fa1c396e..8382a8876f 100644 --- a/src/db/index/segment/segment_helper.cc +++ b/src/db/index/segment/segment_helper.cc @@ -791,6 +791,15 @@ bool CanReuseFirstIndexer(const std::vector &indexers, std::dynamic_pointer_cast(f.index_params()); return params ? params->quantize_type() : QuantizeType::UNDEFINED; }; + if (quantize_type_of(output_field) == QuantizeType::PQ) { + const auto &first = + static_cast(*first_field.index_params()); + const auto &output = + static_cast(*output_field.index_params()); + if (first.metric_type() != output.metric_type() || + first.quantizer_param() != output.quantizer_param()) + return false; + } return quantize_type_of(first_field) == quantize_type_of(output_field); } diff --git a/src/include/zvec/c_api.h b/src/include/zvec/c_api.h index b8fa542aae..627ee3c231 100644 --- a/src/include/zvec/c_api.h +++ b/src/include/zvec/c_api.h @@ -883,6 +883,7 @@ typedef uint32_t zvec_quantize_type_t; #define ZVEC_QUANTIZE_TYPE_INT8 2 #define ZVEC_QUANTIZE_TYPE_INT4 3 #define ZVEC_QUANTIZE_TYPE_RABITQ 4 +#define ZVEC_QUANTIZE_TYPE_PQ 8 // ============================================================================= // Collection Structures (Opaque Pointer Pattern) @@ -989,15 +990,16 @@ ZVEC_EXPORT zvec_quantize_type_t ZVEC_CALL zvec_index_params_get_quantize_type(const zvec_index_params_t *params); /** - * @brief Set enable_rotate for quantizer (only effective with INT8/INT4 - * quantize type) + * @brief Set enable_rotate for quantizer (random rotation for INT8/INT4; OPQ + * for PQ, quantize type) * * When enabled, vectors are randomly rotated before INT8/INT4 quantization to - * reduce quantization error. The rotation matrix is stored with the index + * reduce quantization error. With PQ, learned OPQ is used instead. + * The rotation matrix is stored with the index * and automatically applied to query vectors at search time. * * @param params Index parameters (must be vector index type) - * @param enable_rotate Whether to enable random rotation before quantization + * @param enable_rotate Whether to enable rotation (learned OPQ with PQ) * @return ZVEC_OK on success, error code on failure */ ZVEC_EXPORT zvec_error_code_t ZVEC_CALL @@ -1012,6 +1014,34 @@ zvec_index_params_set_quantizer_enable_rotate(zvec_index_params_t *params, ZVEC_EXPORT bool ZVEC_CALL zvec_index_params_get_quantizer_enable_rotate( const zvec_index_params_t *params); +/** IVF PQ configuration. Set the index quantize type to ZVEC_QUANTIZE_TYPE_PQ. + * PQ supports FP32 IVF fields with L2, IP or COSINE metrics. num_bits is 4 or + * 8; FastScan requires 4 bits and dimension divisible by num_chunk. + * enable_rotate selects learned OPQ; its iteration counts must then be + * positive. + */ +typedef struct { + int num_chunk; + int num_bits; + bool enable_rotate; + bool fast_scan; + uint32_t opq_iter; + uint32_t opq_pq_iter; +} zvec_pq_quantizer_params_t; + +#define ZVEC_PQ_QUANTIZER_PARAMS_DEFAULT {8, 8, false, false, 5, 4} + +/** Set IVF PQ configuration. Copies values; schema validation checks dimension. + */ +ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_index_params_set_quantizer_pq( + zvec_index_params_t *params, const zvec_pq_quantizer_params_t *pq); + +/** Read IVF PQ configuration. Returns INVALID_ARGUMENT for null/non-IVF params. + */ +ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_index_params_get_quantizer_pq( + const zvec_index_params_t *params, zvec_pq_quantizer_params_t *pq); + + /** * @brief Set HNSW specific parameters * @param params Index parameters (must be HNSW type) diff --git a/src/include/zvec/core/framework/index_streamer.h b/src/include/zvec/core/framework/index_streamer.h index 6f300be466..90d75d5d57 100644 --- a/src/include/zvec/core/framework/index_streamer.h +++ b/src/include/zvec/core/framework/index_streamer.h @@ -59,6 +59,12 @@ class IndexStreamer : public IndexRunner { //! Open a index from storage virtual int open(IndexStorage::Pointer stg) = 0; + //! Open with explicit access intent. Legacy streamers retain their open + //! behavior; implementations may skip build-only state for read-only use. + virtual int open(IndexStorage::Pointer stg, bool /*read_only*/) { + return open(stg); + } + //! Flush index virtual int flush(uint64_t check_point) = 0; diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index c695140297..a11724d853 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -209,11 +209,10 @@ class ZVEC_CORE_API Index { int _search_refine_internal_ids( const VectorData &query, const BaseIndexQueryParam::Pointer &search_param, int64_t *output_ids, float *output_scores); - int _collect_dense_result(const VectorData &query, - const core::IndexQueryMeta &query_meta, - const BaseIndexQueryParam::Pointer &search_param, - SearchResult *result, - core::IndexContext::Pointer &context); + virtual int _collect_dense_result( + const VectorData &query, const core::IndexQueryMeta &query_meta, + const BaseIndexQueryParam::Pointer &search_param, SearchResult *result, + core::IndexContext::Pointer &context); int _refine_dense_candidates(const VectorData &query, const BaseIndexQueryParam::Pointer &search_param, const std::vector> &keys, @@ -242,6 +241,9 @@ class ZVEC_CORE_API Index { int create_and_init_metric(const BaseIndexParam ¶m); virtual int create_and_init_converter_reformer( const QuantizerParam ¶m, const BaseIndexParam &index_param); + int init_pq_quantizer(const QuantizerParam ¶m, + const BaseIndexParam &index_param, + std::shared_ptr *quantizer); int init_converter_reformer(const std::string &converter_name, const ailego::Params &converter_params = {}); virtual int create_and_init_streamer(const BaseIndexParam ¶m) = 0; @@ -257,6 +259,8 @@ class ZVEC_CORE_API Index { protected: bool is_trained_{false}; + // Trained streaming quantizers open their empty target after merge training. + bool pending_quantizer_training_{false}; BaseIndexParam param_; ailego::Params proxima_index_params_{}; @@ -356,10 +360,11 @@ class ZVEC_CORE_API IVFIndex : public Index { const QuantizerParam ¶m, const BaseIndexParam &index_param) override; int create_and_init_streamer(const BaseIndexParam ¶m) override; - int _dense_search(const VectorData &query, - const BaseIndexQueryParam::Pointer &search_param, - SearchResult *result, - core::IndexContext::Pointer &context) override; + int _collect_dense_result(const VectorData &query, + const core::IndexQueryMeta &query_meta, + const BaseIndexQueryParam::Pointer &search_param, + SearchResult *result, + core::IndexContext::Pointer &context) override; int _prepare_for_search(const VectorData &query, const BaseIndexQueryParam::Pointer &search_param, diff --git a/src/include/zvec/core/interface/index_param.h b/src/include/zvec/core/interface/index_param.h index e40724c565..3750b63bfc 100644 --- a/src/include/zvec/core/interface/index_param.h +++ b/src/include/zvec/core/interface/index_param.h @@ -161,8 +161,11 @@ struct ZVEC_CORE_API QuantizerParam : public SerializableBase { //! Product-Quantization specific params struct PqQuantizerParam : public QuantizerParam { - int num_chunk = 8; // M: number of chunks - int num_bits = 8; // bits per chunk + int num_chunk = 8; // M: number of chunks + int num_bits = 8; // bits per chunk + bool fast_scan = false; // IVF packed PQ4 scan; requires num_bits == 4 + uint32_t opq_iter = 5; + uint32_t opq_pq_iter = 4; // Constructors PqQuantizerParam(int chunks = 8, int bits = 8, bool rotate = false) diff --git a/src/include/zvec/db/index_params.h b/src/include/zvec/db/index_params.h index 15b9292482..f0e49ca737 100644 --- a/src/include/zvec/db/index_params.h +++ b/src/include/zvec/db/index_params.h @@ -122,13 +122,21 @@ class ZVEC_API InvertIndexParams : public IndexParams { /* * Quantizer parameters for vector indexes. - * Encapsulates quantization-related settings such as enable_rotate. - * Designed for future extensibility (e.g., num_bits, calibration_size). + * Rotation settings and IVF product-quantization configuration. + * PQ fields apply when quantize_type is PQ; enable_rotate selects learned OPQ. */ class QuantizerParam { public: QuantizerParam() = default; - explicit QuantizerParam(bool enable_rotate) : enable_rotate_(enable_rotate) {} + explicit QuantizerParam(bool enable_rotate, int num_chunk = 8, + int num_bits = 8, bool fast_scan = false, + uint32_t opq_iter = 5, uint32_t opq_pq_iter = 4) + : enable_rotate_(enable_rotate), + num_chunk_(num_chunk), + num_bits_(num_bits), + fast_scan_(fast_scan), + opq_iter_(opq_iter), + opq_pq_iter_(opq_pq_iter) {} bool enable_rotate() const { return enable_rotate_; @@ -138,8 +146,46 @@ class QuantizerParam { enable_rotate_ = v; } + int num_chunk() const { + return num_chunk_; + } + void set_num_chunk(int value) { + num_chunk_ = value; + } + + int num_bits() const { + return num_bits_; + } + void set_num_bits(int value) { + num_bits_ = value; + } + + bool fast_scan() const { + return fast_scan_; + } + void set_fast_scan(bool value) { + fast_scan_ = value; + } + + uint32_t opq_iter() const { + return opq_iter_; + } + void set_opq_iter(uint32_t value) { + opq_iter_ = value; + } + + uint32_t opq_pq_iter() const { + return opq_pq_iter_; + } + void set_opq_pq_iter(uint32_t value) { + opq_pq_iter_ = value; + } + bool operator==(const QuantizerParam &other) const { - return enable_rotate_ == other.enable_rotate_; + return enable_rotate_ == other.enable_rotate_ && + num_chunk_ == other.num_chunk_ && num_bits_ == other.num_bits_ && + fast_scan_ == other.fast_scan_ && opq_iter_ == other.opq_iter_ && + opq_pq_iter_ == other.opq_pq_iter_; } bool operator!=(const QuantizerParam &other) const { @@ -147,9 +193,13 @@ class QuantizerParam { } private: - // When enabled, vectors are rotated before INT8 quantization to reduce - // quantization error. Only effective with quantize_type=INT8. + // Random rotation for scalar quantizers; learned OPQ for PQ. bool enable_rotate_{false}; + int num_chunk_{8}; + int num_bits_{8}; + bool fast_scan_{false}; + uint32_t opq_iter_{5}; + uint32_t opq_pq_iter_{4}; }; /* diff --git a/src/include/zvec/db/type.h b/src/include/zvec/db/type.h index e41e7557bd..9d187836c4 100644 --- a/src/include/zvec/db/type.h +++ b/src/include/zvec/db/type.h @@ -84,6 +84,7 @@ enum class QuantizeType : uint32_t { UNIFORM_UINT7 = 5, UNIFORM_UINT8 = 6, UNIFORM_UINT4 = 7, + PQ = 8, }; enum class MetricType : uint32_t { diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h index 7f8ea028e4..9fdda2a01d 100644 --- a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h @@ -14,11 +14,11 @@ #pragma once +#include #include #include #include #include -#include "quantizer/quantizer.h" namespace zvec { namespace turbo { diff --git a/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.cc b/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.cc index d4c6dae675..0b15ed3964 100644 --- a/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.cc +++ b/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.cc @@ -655,7 +655,7 @@ int PqInt4Quantizer::quantize(const void *query, const IndexQueryMeta &qmeta, quantize_query(query, &(*out)[0]); *ometa = qmeta; - ometa->set_meta(IndexMeta::DataType::DT_INT4, num_chunk_, + ometa->set_meta(IndexMeta::DataType::DT_FP32, original_dim_, static_cast(type_), 0); return 0; } @@ -1109,8 +1109,10 @@ int PqInt4Quantizer::deserialize(const void *data, size_t len) { payload.rotate_type, original_dim_, ptr, len - static_cast(ptr - reinterpret_cast(data))); if (opq_ret != 0) return opq_ret; - // dist_table_ is intentionally not restored: SDC is only needed during - // offline build, not after deserialization (search uses ADC). + // ADC search only needs the codebook and a query LUT. Leave the SDC table + // empty; prepare_build() restores it before graph construction or inserts + // into a reopened index. + dist_table_.clear(); // Re-dispatch kernels and batch distance functions. if (setup_functions() != 0) { diff --git a/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.h b/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.h index 16c472697a..63e5945ca6 100644 --- a/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.h +++ b/src/turbo/quantizer/pq_int4_quantizer/pq_int4_quantizer.h @@ -71,6 +71,13 @@ class PqInt4Quantizer : public Quantizer, public PrecomputeTableQuantizer { int train(IndexHolder::Pointer holder) override; + // Rebuild the centroid-pair table for SDC before graph inserts. + // Read-only ADC queries use their own LUT and do not need this table. + int prepare_build() override { + if (dist_table_.empty()) compute_dist_table(); + return 0; + } + size_t quantized_datapoint_vector_length() const override { return packed_code_length() + extra_meta_size_; } diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc index a7226a9c60..0b349b3eba 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.cc @@ -1202,8 +1202,9 @@ int PqInt8Quantizer::deserialize(const void *data, size_t len) { int opq_ret = opq_.deserialize(payload.rotate_type, original_dim_, ptr, static_cast(preprocessor_bytes)); if (opq_ret != 0) return opq_ret; - // dist_table_ is intentionally not restored: SDC is only needed during - // offline build, not after deserialization (search uses ADC). + // ADC search only needs the codebook and a query LUT. Release the SDC table; + // prepare_build() restores it before graph construction or inserts into a + // reopened writable index. std::vector().swap(dist_table_); // Re-dispatch kernels and batch distance functions. diff --git a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h index a607eb2840..2969919ddc 100644 --- a/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h +++ b/src/turbo/quantizer/pq_int8_quantizer/pq_int8_quantizer.h @@ -67,6 +67,13 @@ class PqInt8Quantizer : public Quantizer, public PrecomputeTableQuantizer { int train(IndexHolder::Pointer holder) override; + // Rebuild the centroid-pair table for SDC before graph inserts. + // Read-only ADC queries use their own LUT and do not need this table. + int prepare_build() override { + if (dist_table_.empty()) compute_dist_table(); + return 0; + } + size_t quantized_datapoint_vector_length() const override { return num_chunk_ + extra_meta_size_; } diff --git a/src/turbo/quantizer/quantizer.h b/src/turbo/quantizer/quantizer.h index 20d8d9e454..40e551e985 100644 --- a/src/turbo/quantizer/quantizer.h +++ b/src/turbo/quantizer/quantizer.h @@ -121,6 +121,11 @@ class Quantizer { return train(holder); } + //! Prepare derived state for graph construction before concurrent inserts. + virtual int prepare_build() { + return 0; + } + //! Byte length of a quantized datapoint vector virtual size_t quantized_datapoint_vector_length() const = 0; diff --git a/tests/ailego/buffer/vector_page_table_test.cc b/tests/ailego/buffer/vector_page_table_test.cc index 9f422fe112..871b2560ed 100644 --- a/tests/ailego/buffer/vector_page_table_test.cc +++ b/tests/ailego/buffer/vector_page_table_test.cc @@ -271,13 +271,12 @@ version_t FindLiveVersion(SizedExternalCache &cache, eviction_key_t owner_key) { TEST_F(BufferPoolTest, HeapAllocatedOwnersPreserveCacheLineAlignment) { auto table = std::make_unique(); auto shared_table = std::make_shared(); - EXPECT_EQ(0U, reinterpret_cast(table.get()) % - AILEGO_CACHE_LINE_SIZE); + EXPECT_EQ(0U, + reinterpret_cast(table.get()) % AILEGO_CACHE_LINE_SIZE); EXPECT_EQ(0U, reinterpret_cast(shared_table.get()) % AILEGO_CACHE_LINE_SIZE); - EXPECT_EQ(0U, - reinterpret_cast(&MemoryLimitPool::get_instance()) % - AILEGO_CACHE_LINE_SIZE); + EXPECT_EQ(0U, reinterpret_cast(&MemoryLimitPool::get_instance()) % + AILEGO_CACHE_LINE_SIZE); } TEST_F(BufferPoolTest, AdmissionControlRejectsFirstColdMissAfterPressure) { diff --git a/tests/android_gmock_main.cc b/tests/android_gmock_main.cc index 20cd45da7d..1538da6b95 100644 --- a/tests/android_gmock_main.cc +++ b/tests/android_gmock_main.cc @@ -5,10 +5,10 @@ // (exit code 139) or aborts (134/135) during teardown even when all tests // passed. Using _exit() skips the static destructor phase entirely. +#include #include #include #include -#include #include "gmock/gmock.h" GTEST_API_ int main(int argc, char **argv) { diff --git a/tests/android_gtest_main.cc b/tests/android_gtest_main.cc index 372e2b4632..81ee6e36ec 100644 --- a/tests/android_gtest_main.cc +++ b/tests/android_gtest_main.cc @@ -5,10 +5,10 @@ // (exit code 139) or aborts (134/135) during teardown even when all tests // passed. Using _exit() skips the static destructor phase entirely. +#include #include #include #include -#include #include "gtest/gtest.h" GTEST_API_ int main(int argc, char **argv) { diff --git a/tests/c/c_api_test.c b/tests/c/c_api_test.c index e8fba3d862..dc0c19bbb6 100644 --- a/tests/c/c_api_test.c +++ b/tests/c/c_api_test.c @@ -3845,6 +3845,42 @@ void test_index_params_functions(void) { TEST_END(); } +void test_ivf_pq_params(void) { + TEST_START(); + zvec_index_params_t *ivf = zvec_index_params_create(ZVEC_INDEX_TYPE_IVF); + TEST_ASSERT(ivf != NULL); + TEST_ASSERT(zvec_index_params_set_quantize_type(ivf, ZVEC_QUANTIZE_TYPE_PQ) == + ZVEC_OK); + zvec_pq_quantizer_params_t pq = ZVEC_PQ_QUANTIZER_PARAMS_DEFAULT; + pq.num_chunk = 4; + pq.num_bits = 4; + pq.fast_scan = true; + pq.enable_rotate = true; + pq.opq_iter = 3; + pq.opq_pq_iter = 2; + TEST_ASSERT(zvec_index_params_set_quantizer_pq(ivf, &pq) == ZVEC_OK); + zvec_pq_quantizer_params_t restored = ZVEC_PQ_QUANTIZER_PARAMS_DEFAULT; + TEST_ASSERT(zvec_index_params_get_quantizer_pq(ivf, &restored) == ZVEC_OK); + TEST_ASSERT(restored.num_chunk == 4 && restored.num_bits == 4); + TEST_ASSERT(restored.fast_scan && restored.enable_rotate); + TEST_ASSERT(restored.opq_iter == 3 && restored.opq_pq_iter == 2); + TEST_ASSERT(zvec_index_params_get_quantize_type(ivf) == + ZVEC_QUANTIZE_TYPE_PQ); + pq.num_bits = 8; + TEST_ASSERT(zvec_index_params_set_quantizer_pq(ivf, &pq) == + ZVEC_ERROR_INVALID_ARGUMENT); + TEST_ASSERT(zvec_index_params_get_quantizer_pq(ivf, NULL) == + ZVEC_ERROR_INVALID_ARGUMENT); + TEST_ASSERT(zvec_index_params_set_quantizer_pq(NULL, &pq) == + ZVEC_ERROR_INVALID_ARGUMENT); + zvec_index_params_t *flat = zvec_index_params_create(ZVEC_INDEX_TYPE_FLAT); + TEST_ASSERT(zvec_index_params_set_quantizer_pq(flat, &pq) == + ZVEC_ERROR_INVALID_ARGUMENT); + zvec_index_params_destroy(flat); + zvec_index_params_destroy(ivf); + TEST_END(); +} + void test_quantizer_enable_rotate(void) { TEST_START(); @@ -7400,6 +7436,7 @@ int main(void) { test_index_params(); test_index_params_functions(); test_quantizer_enable_rotate(); + test_ivf_pq_params(); test_int8_rotate_e2e(); test_index_params_api_functions(); test_index_creation_and_management(); diff --git a/tests/core/algorithm/diskann/diskann_file_reader_windows_test.cc b/tests/core/algorithm/diskann/diskann_file_reader_windows_test.cc index dbc7e44b58..c459398b47 100644 --- a/tests/core/algorithm/diskann/diskann_file_reader_windows_test.cc +++ b/tests/core/algorithm/diskann/diskann_file_reader_windows_test.cc @@ -150,12 +150,11 @@ bool verify_page(const uint8_t *buffer, size_t page, uint8_t bias = 0) { } DWORD replace_open_file_atomically(const wchar_t *replacement_path, - const wchar_t *target_path) { + const wchar_t *target_path) { // diskann_file_reader.h targets Windows Vista, whose SDK view predates the // extended rename declarations. These values and this layout are the Win10 // FILE_RENAME_INFO_EX ABI used by FileRenameInfoEx. - constexpr auto kFileRenameInfoEx = - static_cast(22); + constexpr auto kFileRenameInfoEx = static_cast(22); constexpr DWORD kReplaceIfExists = 0x00000001; constexpr DWORD kPosixSemantics = 0x00000002; struct ExtendedFileRenameInfo { @@ -171,10 +170,10 @@ DWORD replace_open_file_atomically(const wchar_t *replacement_path, static_assert(offsetof(ExtendedFileRenameInfo, file_name) == offsetof(FILE_RENAME_INFO, FileName)); - HANDLE replacement_handle = ::CreateFileW( - replacement_path, DELETE | SYNCHRONIZE, - FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, nullptr, - OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr); + HANDLE replacement_handle = + ::CreateFileW(replacement_path, DELETE | SYNCHRONIZE, + FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, + nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr); if (replacement_handle == INVALID_HANDLE_VALUE) { return ::GetLastError(); } @@ -185,8 +184,7 @@ DWORD replace_open_file_atomically(const wchar_t *replacement_path, return ERROR_FILENAME_EXCED_RANGE; } - const size_t target_path_bytes = - target_path_length * sizeof(target_path[0]); + const size_t target_path_bytes = target_path_length * sizeof(target_path[0]); ExtendedFileRenameInfo rename_info{}; rename_info.flags = kReplaceIfExists | kPosixSemantics; rename_info.root_directory = nullptr; @@ -225,8 +223,7 @@ class ScopedCurrentDirectory { return; } original_.resize(capacity, L'\0'); - const DWORD length = - ::GetCurrentDirectoryW(capacity, original_.data()); + const DWORD length = ::GetCurrentDirectoryW(capacity, original_.data()); if (length == 0 || length >= capacity) { original_.clear(); return; @@ -399,8 +396,7 @@ TEST(DiskAnnFileReaderWindowsTest, OpenSupportsUtf8Path) { EXPECT_EQ(ctx, nullptr); } -TEST(DiskAnnFileReaderWindowsTest, - RelativePathSurvivesWorkingDirectoryChange) { +TEST(DiskAnnFileReaderWindowsTest, RelativePathSurvivesWorkingDirectoryChange) { TemporaryFile file; ASSERT_TRUE(file.valid()); ASSERT_TRUE(file.write_pages()); @@ -623,8 +619,7 @@ TEST(DiskAnnFileReaderWindowsTest, DestroyContextDrainsOutstandingBatch) { // reader does not opt into FILE_SKIP_COMPLETION_PORT_ON_SUCCESS, teardown // must still dequeue its completion packet along with pending reads. const size_t length = i == 0 ? 0 : kPageSize; - requests.emplace_back(i * kPageSize, length, - output.get() + i * kPageSize); + requests.emplace_back(i * kPageSize, length, output.get() + i * kPageSize); } PendingBatch batch; @@ -727,9 +722,9 @@ TEST(DiskAnnFileReaderWindowsTest, ERROR_SUCCESS); WindowsAlignedFileReader original_reader; - ASSERT_EQ(original_reader.open_from_handle(original.path(), - source.native_handle()), - 0); + ASSERT_EQ( + original_reader.open_from_handle(original.path(), source.native_handle()), + 0); HANDLE stable_handle = WindowsAlignedFileReaderTestPeer::stable_file_handle(original_reader); ASSERT_NE(stable_handle, INVALID_HANDLE_VALUE); diff --git a/tests/core/algorithm/flat/flat_streamer_test.cc b/tests/core/algorithm/flat/flat_streamer_test.cc index a1df97c31f..d473bb47dd 100644 --- a/tests/core/algorithm/flat/flat_streamer_test.cc +++ b/tests/core/algorithm/flat/flat_streamer_test.cc @@ -87,8 +87,8 @@ class FlatQueryCopyTestMetric : public IndexMetric { size_t dimension, float *distances, const void **) { const auto *prepared_query = static_cast(query); for (size_t i = 0; i < count; ++i) { - EXPECT_EQ(0U, reinterpret_cast(rows[i]) % - AILEGO_CACHE_LINE_SIZE); + EXPECT_EQ( + 0U, reinterpret_cast(rows[i]) % AILEGO_CACHE_LINE_SIZE); const auto *row = static_cast(rows[i]); double sum = 0; for (size_t d = 0; d < dimension; ++d) { diff --git a/tests/core/algorithm/hnsw/hnsw_streamer_test.cc b/tests/core/algorithm/hnsw/hnsw_streamer_test.cc index da5b6323aa..8ce5becbb7 100644 --- a/tests/core/algorithm/hnsw/hnsw_streamer_test.cc +++ b/tests/core/algorithm/hnsw/hnsw_streamer_test.cc @@ -5120,6 +5120,133 @@ TEST_F(HnswStreamerTest, TestTurboInt8QuantizerDistance) { ASSERT_EQ(0, storage->close()); } +TEST_F(HnswStreamerTest, PqPersistenceAndIncrementalBuild) { + ASSERT_EQ(0, MemoryLimitPool::get_instance().init(256UL * 1024UL * 1024UL)); + auto holder = std::make_shared>(dim); + std::mt19937 rng(42); + std::uniform_real_distribution dist(-1, 1); + std::vector> vectors; + for (size_t i = 0; i < 256; ++i) { + NumericalVector v(dim); + for (size_t j = 0; j < dim; ++j) v[j] = dist(rng); + holder->emplace(i, v); + vectors.push_back(v); + } + for (const auto *name : {"PqInt8Quantizer", "PqInt4Quantizer"}) { + SCOPED_TRACE(name); + ailego::Params params; + params.set("num_chunk", 4U); + auto quantizer = IndexFactory::CreateQuantizer(name); + ASSERT_NE(nullptr, quantizer); + ASSERT_EQ(0, quantizer->init(*index_meta_ptr_, params)); + ASSERT_EQ(0, quantizer->train(holder)); + auto meta = quantizer->meta(); + meta.set_quantizer(name, 0, params); + IndexQueryMeta codes; + codes.set_meta(meta.data_type(), meta.dimension(), + static_cast(quantizer->type()), + meta.extra_meta_size()); + IndexQueryMeta raw(IndexMeta::DT_FP32, dim), queries; + std::string lut, second; + ASSERT_EQ(0, quantizer->quantize(vectors[3].data(), raw, &lut, &queries)); + ASSERT_EQ(0, + quantizer->quantize(vectors[7].data(), raw, &second, &queries)); + lut += second; + std::string code(quantizer->quantized_datapoint_vector_length(), '\0'); + for (const auto *storage_name : {"MMapFileStorage", "BufferStorage"}) { + SCOPED_TRACE(storage_name); + const auto path = dir_ + name + storage_name; + auto storage = IndexFactory::CreateStorage(storage_name); + ASSERT_EQ(0, storage->init(ailego::Params())); + ASSERT_EQ(0, storage->open(path, true)); + IndexStreamer::Pointer streamer = std::make_shared(); + ASSERT_EQ(0, streamer->init(meta, ailego::Params(), quantizer)); + ASSERT_EQ(0, streamer->open(storage)); + auto ctx = streamer->create_context(); + for (size_t i = 0; i < 32; ++i) { + quantizer->quantize_data(vectors[i].data(), code.data()); + ASSERT_EQ(0, streamer->add_impl(i, code.data(), codes, ctx)); + } + // LUT metadata is valid only on the query path. + EXPECT_EQ(IndexError_Mismatch, + streamer->add_impl(99, lut.data(), queries, ctx)); + ctx->set_topk(4); + ASSERT_EQ(0, streamer->search_bf_impl(lut.data(), queries, 2, ctx)); + auto first_result = ctx->result(0); + auto second_result = ctx->result(1); + ASSERT_EQ(4U, first_result.size()); + ASSERT_EQ(4U, second_result.size()); + meta = streamer->meta(); + ASSERT_EQ(0, streamer->close()); + ASSERT_EQ(0, storage->close()); + + ASSERT_EQ(0, storage->open(path, true)); + streamer = std::make_shared(); + ASSERT_EQ(0, streamer->init(meta, ailego::Params())); + ASSERT_EQ(0, streamer->open(storage)); + ctx = streamer->create_context(); + ctx->set_topk(4); + ASSERT_EQ(0, streamer->search_bf_impl(lut.data(), queries, 2, ctx)); + for (size_t i = 0; i < 4; ++i) { + EXPECT_EQ(first_result[i].key(), ctx->result(0)[i].key()); + EXPECT_FLOAT_EQ(first_result[i].score(), ctx->result(0)[i].score()); + EXPECT_EQ(second_result[i].key(), ctx->result(1)[i].key()); + EXPECT_FLOAT_EQ(second_result[i].score(), ctx->result(1)[i].score()); + } + // Compare the second batch result with an independent single query. + ASSERT_EQ(0, streamer->search_bf_impl(second.data(), queries, ctx)); + for (size_t i = 0; i < 4; ++i) { + EXPECT_EQ(second_result[i].key(), ctx->result()[i].key()); + } + quantizer->quantize_data(vectors[32].data(), code.data()); + ASSERT_EQ(0, streamer->add_impl(32, code.data(), codes, ctx)); + // Exercise graph traversal as well as brute force after reopening. + dynamic_cast(ctx.get())->set_bruteforce_threshold(0); + ASSERT_EQ(0, streamer->search_impl(second.data(), queries, ctx)); + ASSERT_FALSE(ctx->result().empty()); + + auto dumper = IndexFactory::CreateDumper("FileDumper"); + ASSERT_EQ(0, dumper->create(path + ".dump")); + ASSERT_EQ(0, streamer->dump(dumper)); + ASSERT_EQ(0, dumper->close()); + auto dumped = IndexFactory::CreateStorage("MMapFileReadStorage"); + ASSERT_EQ(0, dumped->init(ailego::Params())); + ASSERT_EQ(0, dumped->open(path + ".dump", false)); + std::shared_ptr restored; + ASSERT_EQ(0, HnswStreamer::RestoreQuantizer(dumped, meta, &restored)); + EXPECT_FLOAT_EQ( + quantizer->calc_distance_dp_query(code.data(), second.data()), + restored->calc_distance_dp_query(code.data(), second.data())); + ASSERT_EQ(0, dumped->close()); + ASSERT_EQ(0, streamer->close()); + ASSERT_EQ(0, storage->close()); + + // Truncating a trained quantizer must fail, never fall back to a metric. + ASSERT_EQ(0, storage->open(path, true)); + auto segment = storage->get("hnsw.turbo_quantizer"); + ASSERT_NE(nullptr, segment); + char damaged = 0; + ASSERT_EQ(1U, segment->write(0, &damaged, 1)); + restored.reset(); + EXPECT_EQ(IndexError_InvalidChecksum, + HnswStreamer::RestoreQuantizer(storage, meta, &restored)); + segment->resize(0); + restored.reset(); + EXPECT_EQ(IndexError_InvalidFormat, + HnswStreamer::RestoreQuantizer(storage, meta, &restored)); + streamer = std::make_shared(); + ASSERT_EQ(0, streamer->init(meta, ailego::Params())); + EXPECT_EQ(IndexError_InvalidFormat, streamer->open(storage)); + streamer.reset(); + ASSERT_EQ(0, storage->close()); + ASSERT_EQ(0, storage->open(path + ".missing", true)); + EXPECT_EQ(IndexError_InvalidFormat, + HnswStreamer::RestoreQuantizer(storage, meta, &restored)); + ASSERT_EQ(0, storage->close()); + } + } +} + } // namespace core } // namespace zvec diff --git a/tests/core/algorithm/ivf/ivf_searcher_test.cc b/tests/core/algorithm/ivf/ivf_searcher_test.cc index df6e8aa5f7..c2956f44fb 100644 --- a/tests/core/algorithm/ivf/ivf_searcher_test.cc +++ b/tests/core/algorithm/ivf/ivf_searcher_test.cc @@ -1888,14 +1888,14 @@ TEST_F(IVFSearcherTest, TestRowMajorFp16WithBuildMemory) { IVFBuilder builder; index_meta_.set_meta(IndexMeta::DataType::DT_FP16, dimension_); int ret = builder.init(index_meta_, params_); - EXPECT_EQ(0, ret); + ASSERT_EQ(0, ret); ret = builder.train(threads_, holder_); - EXPECT_EQ(0, ret); + ASSERT_EQ(0, ret); ret = builder.build(threads_, holder_); - EXPECT_EQ(0, ret); + ASSERT_EQ(0, ret); IndexDumper::Pointer dumper = IndexFactory::CreateDumper("FileDumper"); ret = dumper->create(index_path_); - EXPECT_EQ(0, ret); + ASSERT_EQ(0, ret); ret = builder.dump(dumper); EXPECT_EQ((size_t)1000, builder.stats().built_count()); @@ -1909,7 +1909,7 @@ TEST_F(IVFSearcherTest, TestRowMajorFp16WithBuildMemory) { params.set(PARAM_IVF_SEARCHER_BRUTE_FORCE_THRESHOLD, 1); ret = searcher.init(params); - EXPECT_EQ(0, ret); + ASSERT_EQ(0, ret); IndexStorage::Pointer container = IndexFactory::CreateStorage("MMapFileReadStorage"); @@ -1919,10 +1919,38 @@ TEST_F(IVFSearcherTest, TestRowMajorFp16WithBuildMemory) { container_params.set("proxima.mmap_file.container.memory_warmup", true); container->init(container_params); ret = container->open(index_path_, false); - EXPECT_EQ(0, ret); + ASSERT_EQ(0, ret); ret = searcher.load(container, IndexMetric::Pointer()); - EXPECT_EQ(0, ret); + ASSERT_EQ(0, ret); + + // A distance-only helper must leave legacy providers in the FP16 domain. + auto check_provider = [&](const IndexProvider::Pointer &provider) { + ASSERT_NE(nullptr, provider); + ASSERT_EQ(IndexMeta::DT_FP16, provider->data_type()); + ASSERT_EQ(dimension_, provider->dimension()); + ASSERT_EQ(dimension_ * sizeof(Float16), provider->element_size()); + ASSERT_EQ(1000u, provider->count()); + for (size_t key = 0; key < provider->count(); ++key) { + std::vector expected(dimension_, Float16(0.01f * key)); + const void *data = provider->get_vector(key); + ASSERT_NE(nullptr, data); + EXPECT_EQ(0, std::memcmp(expected.data(), data, provider->element_size())) + << "key=" << key; + } + auto iter = provider->create_iterator(); + ASSERT_NE(nullptr, iter); + size_t count = 0; + for (; iter->is_valid(); iter->next(), ++count) { + std::vector expected(dimension_, Float16(0.01f * iter->key())); + ASSERT_NE(nullptr, iter->data()); + EXPECT_EQ(0, std::memcmp(expected.data(), iter->data(), + provider->element_size())) + << "key=" << iter->key(); + } + EXPECT_EQ(provider->count(), count); + }; + check_provider(searcher.create_provider()); std::vector query; for (size_t i = 0; i < dimension_; ++i) { @@ -1954,10 +1982,10 @@ TEST_F(IVFSearcherTest, TestRowMajorFp16WithBuildMemory) { size_t topk = 1000; context->set_topk(topk); ret = searcher.search_bf_impl(query_buf.data(), qmeta, context); - EXPECT_EQ(0, ret); + ASSERT_EQ(0, ret); const IndexDocumentList &result = context->result(0); - EXPECT_EQ((size_t)topk, result.size()); + ASSERT_EQ((size_t)topk, result.size()); for (size_t i = 0; i < topk; ++i) { EXPECT_EQ((uint64_t)i, result[i].key()); EXPECT_NEAR((float)(0.01f * i + 0.1) * (0.01f * i + 0.1) * dimension_ / @@ -1971,11 +1999,11 @@ TEST_F(IVFSearcherTest, TestRowMajorFp16WithBuildMemory) { size_t topk = 1; context->set_topk(topk); ret = searcher.search_bf_impl(query1_buf.data(), qmeta, qnum, context); - EXPECT_EQ(0, ret); + ASSERT_EQ(0, ret); for (size_t q = 0; q < qnum; ++q) { const IndexDocumentList &result = context->result(q); - EXPECT_EQ((size_t)topk, result.size()); + ASSERT_EQ((size_t)topk, result.size()); EXPECT_EQ((uint64_t)q, result[0].key()); EXPECT_FLOAT_EQ((float)0, result[0].score()); } @@ -1986,10 +2014,10 @@ TEST_F(IVFSearcherTest, TestRowMajorFp16WithBuildMemory) { size_t topk = 100; context->set_topk(topk); ret = searcher.search_impl(query_buf.data(), qmeta, context); - EXPECT_EQ(0, ret); + ASSERT_EQ(0, ret); const IndexDocumentList &result = context->result(0); - EXPECT_EQ((size_t)topk, result.size()); + ASSERT_EQ((size_t)topk, result.size()); for (size_t i = 0; i < topk; ++i) { EXPECT_EQ((uint64_t)i, result[i].key()); EXPECT_NEAR((float)(0.01f * i + 0.1) * (0.01f * i + 0.1) * dimension_ / @@ -2003,18 +2031,41 @@ TEST_F(IVFSearcherTest, TestRowMajorFp16WithBuildMemory) { size_t topk = 1; context->set_topk(topk); ret = searcher.search_impl(query1_buf.data(), qmeta, qnum, context); - EXPECT_EQ(0, ret); + ASSERT_EQ(0, ret); for (size_t q = 0; q < qnum; ++q) { const IndexDocumentList &result = context->result(q); - EXPECT_EQ((size_t)topk, result.size()); + ASSERT_EQ((size_t)topk, result.size()); EXPECT_EQ((uint64_t)q, result[0].key()); EXPECT_FLOAT_EQ((float)0, result[0].score()); } } + IVFStreamer streamer; + ASSERT_EQ(0, streamer.init(index_meta_, params)); + ASSERT_EQ(0, streamer.open(container)); + check_provider(streamer.create_provider()); + auto streamer_context = streamer.create_context(); + ASSERT_NE(nullptr, streamer_context); + streamer_context->set_topk(1); + for (bool brute_force : {false, true}) { + SCOPED_TRACE(brute_force); + ret = brute_force ? streamer.search_bf_impl(query1_buf.data(), qmeta, qnum, + streamer_context) + : streamer.search_impl(query1_buf.data(), qmeta, qnum, + streamer_context); + ASSERT_EQ(0, ret); + for (size_t q = 0; q < qnum; ++q) { + const auto &result = streamer_context->result(q); + ASSERT_EQ(1u, result.size()); + EXPECT_EQ(q, result[0].key()); + EXPECT_FLOAT_EQ(0.0f, result[0].score()); + } + } + ASSERT_EQ(0, streamer.close()); + ret = searcher.unload(); - EXPECT_EQ(0, ret); + ASSERT_EQ(0, ret); } TEST_F(IVFSearcherTest, TestColumnMajorFloatWithHnswGraphType) { diff --git a/tests/core/algorithm/vamana/vamana_streamer_test.cc b/tests/core/algorithm/vamana/vamana_streamer_test.cc index 79e00e4d8e..592587410f 100644 --- a/tests/core/algorithm/vamana/vamana_streamer_test.cc +++ b/tests/core/algorithm/vamana/vamana_streamer_test.cc @@ -112,7 +112,7 @@ TEST(VamanaQueryPrefetchTest, ResolvesSharedDefaultsFromStoredVectorSchema) { EXPECT_EQ( std::make_pair(std::min(64U, kPrefetchBudgetLines / lines), lines), VamanaContext::resolve_query_prefetch(bytes, 64, default_offset, - default_lines)); + default_lines)); } } @@ -124,21 +124,21 @@ TEST(VamanaQueryPrefetchTest, CacheLineBoundaryAndByteBudget) { for (size_t bytes : {line - 1, line}) { EXPECT_EQ(std::make_pair(kPrefetchBudgetLines, 1U), VamanaContext::resolve_query_prefetch(bytes, 256, default_offset, - default_lines)); + default_lines)); } for (size_t bytes : {line + 1, 2 * line}) { EXPECT_EQ(std::make_pair(kPrefetchBudgetLines / 2, 2U), VamanaContext::resolve_query_prefetch(bytes, 256, default_offset, - default_lines)); + default_lines)); } // An explicit line request is clamped to the stored body, including a // partial final line. The automatic prefix still obeys the 6 KiB budget. EXPECT_EQ(std::make_pair(kPrefetchBudgetLines / 3, 3U), VamanaContext::resolve_query_prefetch(2 * line + 1, 256, - default_offset, 4)); + default_offset, 4)); EXPECT_EQ(std::make_pair(0U, 0U), VamanaContext::resolve_query_prefetch(line, 0, default_offset, - default_lines)); + default_lines)); } TEST(VamanaQueryPrefetchTest, ManualFieldsOverrideDefaultsIndependently) { @@ -1001,7 +1001,8 @@ TEST_F(VamanaStreamerTest, TestContiguousPackedGraphAndExtraValuesLayout) { constexpr size_t kCount = 192; constexpr uint64_t kKeyBase = 10000; constexpr uint32_t kExpectedLines = - (kOriginalDimension + AILEGO_CACHE_LINE_SIZE - 1) / AILEGO_CACHE_LINE_SIZE; + (kOriginalDimension + AILEGO_CACHE_LINE_SIZE - 1) / + AILEGO_CACHE_LINE_SIZE; constexpr uint32_t kExpectedOffset = std::min(32U, kPrefetchBudgetLines / kExpectedLines); @@ -1202,14 +1203,13 @@ TEST_F(VamanaStreamerTest, TestContiguousKeepsInt8RecordTailInline) { // the inline quantizer tail must remain byte-for-byte unchanged. const void *stored = contiguous_entity->get_vector_ptr(probe); ASSERT_NE(nullptr, stored); - EXPECT_EQ(0U, - reinterpret_cast(stored) % AILEGO_CACHE_LINE_SIZE); + EXPECT_EQ(0U, reinterpret_cast(stored) % AILEGO_CACHE_LINE_SIZE); std::string expected_record; IndexQueryMeta encoded_meta; ASSERT_EQ(0, reformer->convert(vectors[probe].data(), raw_query_meta, &expected_record, &encoded_meta)); - EXPECT_EQ(0, std::memcmp(stored, expected_record.data(), - expected_record.size())); + EXPECT_EQ( + 0, std::memcmp(stored, expected_record.data(), expected_record.size())); std::string query; IndexQueryMeta query_meta; diff --git a/tests/core/interface/hnsw_turbo_index_test.cc b/tests/core/interface/hnsw_turbo_index_test.cc index 35c00a92f5..116aea9dd4 100644 --- a/tests/core/interface/hnsw_turbo_index_test.cc +++ b/tests/core/interface/hnsw_turbo_index_test.cc @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -1089,3 +1090,229 @@ INSTANTIATE_TEST_SUITE_P(AddApis, HnswExternalCoreCompatibilityTest, } // namespace } // namespace core } // namespace zvec + +TEST(HnswPqIndexTest, ReadOnlyOpenSkipsBuildPreparation) { + using namespace zvec; + class CountingQuantizer : public turbo::Fp32Quantizer { + public: + int prepare_build() override { + ++prepare_count; + return 0; + } + int prepare_count{0}; + }; + + const std::string path = "hnsw_read_only_prepare.index"; + test_util::RemoveTestFiles(path); + core::IndexMeta meta(core::IndexMeta::DT_FP32, kDimension); + meta.set_metric("SquaredEuclidean", 0, ailego::Params()); + auto quantizer = std::make_shared(); + ASSERT_EQ(0, quantizer->init(meta, ailego::Params())); + auto param = MakeParam(MetricType::kL2sq, QuantizerType::kNone); + auto seed = IndexFactory::CreateAndInitIndex(*param); + core::IndexStreamer::Pointer streamer = + std::make_shared(); + ASSERT_EQ( + 0, streamer->init(meta, seed->index_searcher()->meta().streamer_params(), + quantizer)); + auto storage = core::IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_EQ(0, storage->init(ailego::Params())); + ASSERT_EQ(0, storage->open(path, true)); + ASSERT_EQ(0, streamer->open(storage)); + EXPECT_EQ(1, quantizer->prepare_count); + ASSERT_EQ(0, streamer->close()); + + quantizer->prepare_count = 0; + ASSERT_EQ(0, streamer->open(storage, true)); + EXPECT_EQ(0, quantizer->prepare_count); + auto context = streamer->create_context(); + std::vector vector(kDimension, 1.0f); + core::IndexQueryMeta qmeta; + qmeta.set_meta(meta.data_type(), meta.dimension()); + EXPECT_NE(0, streamer->add_impl(0, vector.data(), qmeta, context)); + EXPECT_NE(0, streamer->add_with_id_impl(0, vector.data(), qmeta, context)); + ASSERT_EQ(0, streamer->close()); + + // The legacy open entry point still means writable, including after a + // read-only open on the same streamer. + ASSERT_EQ(0, streamer->open(storage)); + EXPECT_EQ(1, quantizer->prepare_count); + context = streamer->create_context(); + ASSERT_EQ(0, streamer->add_with_id_impl(0, vector.data(), qmeta, context)); + ASSERT_EQ(0, streamer->close()); + ASSERT_EQ(0, storage->close()); + test_util::RemoveTestFiles(path); +} + +TEST(HnswPqIndexTest, ReopenAndAddUsesDataCodes) { + using namespace zvec; + const std::string path = "hnsw_pq_reopen_test.index"; + test_util::RemoveTestFiles(path); + const auto vectors = RandomVectors(256); + auto holder = + std::make_shared>( + kDimension); + for (size_t i = 0; i < vectors.size(); ++i) { + ailego::NumericalVector v(kDimension); + std::copy(vectors[i].begin(), vectors[i].end(), v.data()); + holder->emplace(i, v); + } + core::IndexMeta raw(core::IndexMeta::DT_FP32, kDimension); + raw.set_metric("SquaredEuclidean", 0, ailego::Params()); + ailego::Params params; + params.set("num_chunk", 4U); + auto quantizer = core::IndexFactory::CreateQuantizer("PqInt8Quantizer"); + ASSERT_EQ(0, quantizer->init(raw, params)); + ASSERT_EQ(0, quantizer->train(holder)); + auto meta = quantizer->meta(); + meta.set_quantizer("PqInt8Quantizer", 0, params); + auto param = MakeParam(MetricType::kL2sq, QuantizerType::kNone); + auto seed = IndexFactory::CreateAndInitIndex(*param); + core::IndexStreamer::Pointer streamer = + std::make_shared(); + ASSERT_EQ( + 0, streamer->init(meta, seed->index_searcher()->meta().streamer_params(), + quantizer)); + auto storage = core::IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_EQ(0, storage->init(ailego::Params())); + ASSERT_EQ(0, storage->open(path, true)); + ASSERT_EQ(0, streamer->open(storage)); + auto ctx = streamer->create_context(); + core::IndexQueryMeta codes; + codes.set_meta(meta.data_type(), meta.dimension(), + static_cast(quantizer->type()), + meta.extra_meta_size()); + std::string code(quantizer->quantized_datapoint_vector_length(), '\0'); + for (size_t i = 0; i < 32; ++i) { + quantizer->quantize_data(vectors[i].data(), code.data()); + ASSERT_EQ(0, streamer->add_with_id_impl(i, code.data(), codes, ctx)); + } + ASSERT_EQ(0, streamer->close()); + ASSERT_EQ(0, storage->close()); + auto index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_EQ(0, index->open(path, {StorageOptions::StorageType::kMMAP, false})); + ASSERT_EQ(0, index->add(VectorData{DenseVector{vectors[32].data()}}, 32)); + quantizer->quantize_data(vectors[32].data(), code.data()); + std::string expected; + ASSERT_EQ(0, quantizer->dequantize(code.data(), codes, &expected)); + VectorDataBuffer fetched; + ASSERT_EQ(0, index->fetch(32, &fetched)); + EXPECT_EQ(expected, std::get(fetched.vector_buffer).data); + auto rows = SearchRows(index.get(), vectors[32], true); + ASSERT_FALSE(rows.empty()); + EXPECT_EQ(32U, rows.front().first); + ASSERT_EQ(0, index->close()); + test_util::RemoveTestFiles(path); +} + +class HnswPqLifecycleTest + : public testing::TestWithParam< + std::tuple> {}; + +TEST_P(HnswPqLifecycleTest, TrainMergeReopenAndAppend) { + using namespace zvec; + const auto [bits, rotate, metric, storage_type] = GetParam(); + const std::string raw_path = "hnsw_pq_raw.index"; + const std::string pq_path = "hnsw_pq_target.index"; + const std::string merge_path = "hnsw_pq_merge.index"; + for (auto &path : {raw_path, pq_path, merge_path}) + test_util::RemoveTestFiles(path); + ailego::MemoryLimitPool::get_instance().init(256 * 1024 * 1024); + auto vectors = RandomVectors(300); + auto raw_param = FlatIndexParamBuilder() + .with_dimension(kDimension) + .with_data_type(DataType::DT_FP32) + .with_metric_type(metric) + .build(); + auto raw = IndexFactory::CreateAndInitIndex(*raw_param); + ASSERT_NE(nullptr, raw); + ASSERT_EQ(0, raw->open(raw_path, {storage_type, true})); + AddVectors(raw.get(), vectors); + PqQuantizerParam pq(bits == 8 ? 5 : 4, bits, rotate); + pq.opq_iter = 2; + pq.opq_pq_iter = 2; + auto param = MakeDefaultParam(metric); + param->quantizer_param = pq.clone(); + auto index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, index); + ASSERT_EQ(0, index->open(pq_path, {storage_type, true})); + EXPECT_EQ(0U, index->get_doc_count()); + EXPECT_NE(0, index->add(VectorData{DenseVector{vectors[0].data()}}, 0)); + EXPECT_NE(0, index->train()); + ASSERT_EQ(0, index->merge({raw}, IndexFilter())); + EXPECT_EQ(vectors.size(), index->get_doc_count()); + auto before = SearchRows(index.get(), vectors[11], true); + auto graph_before = SearchRows(index.get(), vectors[11], false); + ASSERT_FALSE(before.empty()); + ASSERT_EQ(0, index->close()); + index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_EQ(0, index->open(pq_path, {storage_type, false, true})); + EXPECT_EQ(before, SearchRows(index.get(), vectors[11], true)); + EXPECT_EQ(graph_before, SearchRows(index.get(), vectors[11], false)); + EXPECT_NE(0, index->add(VectorData{DenseVector{vectors[0].data()}}, 300)); + ASSERT_EQ(0, index->close()); + // A writable reopen must restore the SDC table before adding vectors. + index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_EQ(0, index->open(pq_path, {storage_type, false})); + EXPECT_EQ(before, SearchRows(index.get(), vectors[11], true)); + ASSERT_EQ(0, index->add(VectorData{DenseVector{vectors[0].data()}}, 300)); + EXPECT_EQ(301U, index->get_doc_count()); + + // Candidate extraction accepts the PQ query layout and preserves scores. + auto query_param = + HNSWQueryParamBuilder().with_topk(kTopK).with_ef_search(100).build(); + const VectorData query{DenseVector{vectors[11].data()}}; + SearchResult expected_rows; + ASSERT_EQ(0, index->search(query, query_param, &expected_rows)); + std::vector ids(kTopK); + std::vector scores(kTopK); + ASSERT_EQ(0, index->search_internal_ids(query, query_param, ids.data(), + scores.data())); + ASSERT_EQ(kTopK, expected_rows.doc_list_.size()); + for (size_t i = 0; i < kTopK; ++i) { + EXPECT_EQ(expected_rows.doc_list_[i].key(), static_cast(ids[i])); + EXPECT_FLOAT_EQ(expected_rows.doc_list_[i].score(), scores[i]); + } + + // Train a separate codebook from reconstructed source vectors. Codes from + // the first index must be decoded and encoded again, even with equal layout. + auto merged = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, merged); + ASSERT_EQ(0, merged->open(merge_path, {storage_type, true})); + ASSERT_EQ(0, merged->merge({index}, IndexFilter())); + ASSERT_EQ(0, merged->flush()); + auto storage = core::IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_EQ(0, storage->init(ailego::Params())); + ASSERT_EQ(0, storage->open(merge_path, false)); + std::shared_ptr quantizer; + ASSERT_EQ(0, core::HnswStreamer::RestoreQuantizer( + storage, merged->index_searcher()->meta(), &quantizer)); + std::string code(quantizer->quantized_datapoint_vector_length(), '\0'); + for (uint32_t id : {0U, 11U, 299U, 300U}) { + VectorDataBuffer source, actual; + ASSERT_EQ(0, index->fetch(id, &source)); + const auto &decoded = + std::get(source.vector_buffer).data; + quantizer->quantize_data(decoded.data(), code.data()); + std::string expected; + ASSERT_EQ(0, quantizer->dequantize(code.data(), core::IndexQueryMeta(), + &expected)); + ASSERT_EQ(0, merged->fetch(id, &actual)); + EXPECT_EQ(expected, std::get(actual.vector_buffer).data); + } + ASSERT_EQ(0, storage->close()); + ASSERT_EQ(0, merged->close()); + ASSERT_EQ(0, index->close()); + ASSERT_EQ(0, raw->close()); + for (auto &path : {raw_path, pq_path, merge_path}) + test_util::RemoveTestFiles(path); +} + +INSTANTIATE_TEST_SUITE_P( + Pq, HnswPqLifecycleTest, + testing::Combine( + testing::Values(4, 8), testing::Bool(), + testing::Values(MetricType::kL2sq, MetricType::kInnerProduct, + MetricType::kCosine), + testing::Values(StorageOptions::StorageType::kMMAP, + StorageOptions::StorageType::kBufferPool))); diff --git a/tests/core/interface/index_interface_test.cc b/tests/core/interface/index_interface_test.cc index f3950b4610..047e0569f3 100644 --- a/tests/core/interface/index_interface_test.cc +++ b/tests/core/interface/index_interface_test.cc @@ -5100,6 +5100,94 @@ TEST(IndexInterface, QuantizerParamLegacyJsonCompat) { param->quantizer_param)); } +TEST(IndexInterface, RefineOversizedScaleFactorDoesNotCrash) { + // Regression: floor(topk * SF) below INT_MAX passes validation, but the + // coarse result heap allocation throws std::bad_alloc (e.g. topk=50 with + // SF=10M reserves a ~40GB heap). The query must return an error code + // instead of terminating the process. On hosts with enough RAM the + // allocation succeeds and the query legitimately returns results; both + // outcomes are acceptable, a crash is not. + constexpr uint32_t kDimension = 17; + constexpr uint32_t kVectorCount = 64; + constexpr uint32_t kTopk = 50; + const std::string coarse_path{"refine_huge_sf_coarse.index"}; + const std::string fine_path{"refine_huge_sf_fine.index"}; + zvec::test_util::RemoveTestFiles(coarse_path); + zvec::test_util::RemoveTestFiles(fine_path); + + std::vector> vectors(kVectorCount, + std::vector(kDimension)); + for (uint32_t i = 0; i < kVectorCount; ++i) { + for (uint32_t d = 0; d < kDimension; ++d) { + vectors[i][d] = static_cast((i * 47 + d * 29) % 311) - 23.25F; + } + } + + auto coarse_param = IVFIndexParamBuilder() + .with_metric_type(MetricType::kL2sq) + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .with_is_sparse(false) + .with_n_list(4) + .build(); + auto coarse = IndexFactory::CreateAndInitIndex(*coarse_param); + ASSERT_NE(nullptr, coarse); + ASSERT_EQ( + 0, coarse->open(coarse_path, {StorageOptions::StorageType::kMMAP, true})); + for (uint32_t i = 0; i < kVectorCount; ++i) { + ASSERT_EQ(0, coarse->add(VectorData{DenseVector{vectors[i].data()}}, i)); + } + ASSERT_EQ(0, coarse->train()); + + auto fine_param = FlatIndexParamBuilder() + .with_metric_type(MetricType::kL2sq) + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .with_is_sparse(false) + .build(); + auto fine = IndexFactory::CreateAndInitIndex(*fine_param); + ASSERT_NE(nullptr, fine); + ASSERT_EQ(0, + fine->open(fine_path, {StorageOptions::StorageType::kMMAP, true})); + for (uint32_t i = 0; i < kVectorCount; ++i) { + ASSERT_EQ(0, fine->add(VectorData{DenseVector{vectors[i].data()}}, i)); + } + ASSERT_EQ(0, fine->train()); + + auto refiner = std::make_shared(); + refiner->scale_factor_ = 10000000.0f; // topk * SF = 500M candidates + refiner->reference_index = fine; + auto huge_param = IVFQueryParamBuilder() + .with_nprobe(4) + .with_topk(kTopk) + .with_refiner_param(refiner) + .build(); + SearchResult result; + const int rc = coarse->search(VectorData{DenseVector{vectors[0].data()}}, + huge_param, &result); + if (rc != 0) { + EXPECT_EQ(static_cast(zvec::core::IndexError_NoMemory), rc); + } else { + EXPECT_EQ(kTopk, result.doc_list_.size()); + } + + // The index must stay usable after a failed oversized allocation. + refiner->scale_factor_ = 4.0f; + auto valid_param = IVFQueryParamBuilder() + .with_nprobe(4) + .with_topk(kTopk) + .with_refiner_param(refiner) + .build(); + ASSERT_EQ(0, coarse->search(VectorData{DenseVector{vectors[0].data()}}, + valid_param, &result)); + EXPECT_EQ(kTopk, result.doc_list_.size()); + + ASSERT_EQ(0, coarse->close()); + ASSERT_EQ(0, fine->close()); + zvec::test_util::RemoveTestFiles(coarse_path); + zvec::test_util::RemoveTestFiles(fine_path); +} + TEST(IndexInterface, BuilderChainingReturnsCorrectType) { HNSWIndexParamBuilder builder; auto &ref = builder.with_version(1) diff --git a/tests/core/interface/ivf_turbo_index_test.cc b/tests/core/interface/ivf_turbo_index_test.cc index e792e335ba..a76b206f30 100644 --- a/tests/core/interface/ivf_turbo_index_test.cc +++ b/tests/core/interface/ivf_turbo_index_test.cc @@ -15,17 +15,21 @@ #include #include #include +#include +#include #include #include #include #include #include +#include #include #include #include #include #include #include "algorithm/ivf/ivf_builder.h" +#include "algorithm/ivf/ivf_searcher.h" #include "tests/test_util.h" namespace zvec::core_interface { @@ -434,5 +438,418 @@ TEST(IVFTurboCompatibility, ReopensLegacyInt8AndMergesWithTurbo) { ASSERT_EQ(0, legacy->close()); } +constexpr uint32_t kPqDimension = 12; +constexpr uint32_t kPqCount = 321; +constexpr uint32_t kPqTopK = 7; +using PqCase = std::tuple; + +class IVFPqTest : public testing::TestWithParam { + protected: + void SetUp() override { + path_ = std::string("ivf_pq_") + + testing::UnitTest::GetInstance()->current_test_info()->name(); + std::replace(path_.begin(), path_.end(), '/', '_'); + std::mt19937 random(963); + std::normal_distribution value(0.0f, 2.0f); + for (uint32_t i = 0; i < kPqCount; ++i) { + vectors_.emplace_back(kPqDimension); + for (auto &v : vectors_.back()) v = value(random); + } + } + void TearDown() override { + for (const auto &suffix : {"", ".fine", ".merge"}) + test_util::RemoveTestFiles(path_ + suffix); + } + IVFIndexParam::Pointer param(bool pq = true) const { + auto builder = IVFIndexParamBuilder() + .with_dimension(kPqDimension) + .with_data_type(DataType::DT_FP32) + .with_metric_type(std::get<2>(GetParam())) + .with_n_list(3) + .with_n_iters(3); + if (pq) { + int mode = std::get<0>(GetParam()); + PqQuantizerParam quantizer(3, mode == 8 ? 8 : 4, std::get<1>(GetParam())); + quantizer.fast_scan = mode == 0; + quantizer.opq_iter = 2; + quantizer.opq_pq_iter = 2; + builder.with_quantizer_param(quantizer); + } + return builder.build(); + } + StorageOptions::StorageType storage() const { + return std::get<3>(GetParam()) ? StorageOptions::StorageType::kBufferPool + : StorageOptions::StorageType::kMMAP; + } + void build(const Index::Pointer &index) const { + for (uint32_t i = 0; i < kPqCount; ++i) { + ASSERT_EQ(0, index->add(VectorData{DenseVector{vectors_[i].data()}}, i)); + } + ASSERT_EQ(0, index->train()); + } + void same(const SearchResult &expected, const SearchResult &actual) const { + ASSERT_EQ(expected.doc_list_.size(), actual.doc_list_.size()); + std::map expected_by_key, actual_by_key; + for (size_t i = 0; i < expected.doc_list_.size(); ++i) { + // Quantized ties may be emitted in either order, but ranks and the + // complete set of (ID, score) pairs must agree. + EXPECT_FLOAT_EQ(expected.doc_list_[i].score(), + actual.doc_list_[i].score()); + expected_by_key.emplace(expected.doc_list_[i].key(), + expected.doc_list_[i].score()); + actual_by_key.emplace(actual.doc_list_[i].key(), + actual.doc_list_[i].score()); + } + ASSERT_EQ(expected_by_key.size(), actual_by_key.size()); + for (const auto &entry : expected_by_key) { + auto found = actual_by_key.find(entry.first); + ASSERT_NE(actual_by_key.end(), found); + EXPECT_FLOAT_EQ(entry.second, found->second); + } + } + std::string path_; + std::vector> vectors_; +}; + +TEST_P(IVFPqTest, ScoresPersistenceCandidatesAndRefinement) { + auto config = param(); + IVFIndexParam roundtrip; + ASSERT_TRUE(roundtrip.deserialize_from_json(config->serialize_to_json())); + const auto *pq = + dynamic_cast(roundtrip.quantizer_param.get()); + ASSERT_NE(nullptr, pq); + EXPECT_EQ(std::get<0>(GetParam()) == 0, pq->fast_scan); + EXPECT_EQ(std::get<1>(GetParam()), pq->enable_rotate); + EXPECT_EQ(2u, pq->opq_iter); + EXPECT_EQ(2u, pq->opq_pq_iter); + auto index = IndexFactory::CreateAndInitIndex(roundtrip); + ASSERT_NE(nullptr, index); + ASSERT_EQ(0, index->open(path_, {storage(), true})); + build(index); + const VectorData query{DenseVector{vectors_[7].data()}}; + auto query_param = + IVFQueryParamBuilder().with_nprobe(3).with_topk(kPqCount).build(); + SearchResult all; + ASSERT_EQ(0, index->search(query, query_param, &all)); + ASSERT_EQ(kPqCount, all.doc_list_.size()); + + // The merged internal-ID API must preserve PQ/OPQ/FastScan score semantics. + std::vector internal_ids(kPqCount); + std::vector internal_scores(kPqCount); + ASSERT_EQ(0, + index->search_internal_ids(query, query_param, internal_ids.data(), + internal_scores.data())); + for (size_t i = 0; i < all.doc_list_.size(); ++i) { + EXPECT_EQ(all.doc_list_[i].key(), static_cast(internal_ids[i])); + EXPECT_FLOAT_EQ(all.doc_list_[i].score(), internal_scores[i]); + } + + // Independently encode and score each row with the persisted Turbo state. + auto disk = core::IndexFactory::CreateStorage("MMapFileReadStorage"); + ASSERT_EQ(0, disk->init(ailego::Params{})); + ASSERT_EQ(0, disk->open(path_, false)); + core::IVFSearcher searcher; + ASSERT_EQ(0, searcher.init(ailego::Params{})); + ASSERT_EQ(0, searcher.load(disk, nullptr)); + auto quantizer = searcher.quantizer(); + ASSERT_NE(nullptr, quantizer); + std::string encoded_query(quantizer->quantized_query_vector_length(), '\0'); + quantizer->quantize_query(vectors_[7].data(), encoded_query.data()); + const size_t code_size = quantizer->quantized_datapoint_vector_length(); + std::string code(code_size, '\0'); + std::string packed(code_size * 32, '\0'); + auto *packer = dynamic_cast(quantizer.get()); + for (const auto &doc : all.doc_list_) { + ASSERT_LT(doc.key(), kPqCount); + quantizer->quantize_data(vectors_[doc.key()].data(), code.data()); + float expected; + if (packer) { + ASSERT_EQ(0, + packer->pack_codes(code.data(), 1, code_size, packed.data())); + packer->calc_distance_packed_block(packed.data(), 1, encoded_query.data(), + &expected); + } else { + expected = + quantizer->calc_distance_dp_query(code.data(), encoded_query.data()); + } + if (std::get<2>(GetParam()) == MetricType::kInnerProduct) + expected = -expected; + EXPECT_NEAR(expected, doc.score(), + 1e-4f * std::max(1.0f, std::abs(expected))); + } + ASSERT_EQ(0, searcher.unload()); + ASSERT_EQ(0, disk->close()); + + auto filtered_param = IVFQueryParamBuilder() + .with_nprobe(3) + .with_topk(kPqTopK) + .with_fetch_vector(true) + .build(); + auto set_filter = [](const IVFQueryParam::Pointer &p) { + p->filter = std::make_shared(); + p->filter->set([](uint64_t id) { return id % 3 == 0; }); + }; + set_filter(filtered_param); + SearchResult filtered; + ASSERT_EQ(0, index->search(query, filtered_param, &filtered)); + ASSERT_EQ(kPqTopK, filtered.doc_list_.size()); + for (size_t i = 0; i < filtered.doc_list_.size(); ++i) { + auto id = filtered.doc_list_[i].key(); + EXPECT_NE(0u, id % 3); + VectorDataBuffer fetched; + ASSERT_EQ(0, index->fetch(id, &fetched)); + auto decoded = std::get(fetched.vector_buffer).data; + EXPECT_EQ(decoded, filtered.reverted_vector_list_[i]); + if (packer) + EXPECT_EQ(0, std::memcmp(decoded.data(), vectors_[id].data(), + kPqDimension * sizeof(float))); + } + auto excluded = + IVFQueryParamBuilder().with_nprobe(3).with_topk(kPqTopK).build(); + excluded->filter = std::make_shared(); + excluded->filter->set([](uint64_t) { return true; }); + SearchResult empty; + ASSERT_EQ(0, index->search(query, excluded, &empty)); + EXPECT_TRUE(empty.doc_list_.empty()); + excluded->filter.reset(); + excluded->bf_pks = std::make_shared>(); + ASSERT_EQ(0, index->search(query, excluded, &empty)); + EXPECT_TRUE(empty.doc_list_.empty()); + + // Candidate-ID search must retain this index's own PQ scores. + auto selected = IVFQueryParamBuilder().with_topk(kPqTopK).build(); + selected->bf_pks = std::make_shared>(); + for (const auto &doc : filtered.doc_list_) + selected->bf_pks->push_back(doc.key()); + selected->bf_pks->push_back(kPqCount + 100); + SearchResult selected_result; + ASSERT_EQ(0, index->search(query, selected, &selected_result)); + same(filtered, selected_result); + + auto fine = IndexFactory::CreateAndInitIndex(*param(false)); + ASSERT_NE(nullptr, fine); + ASSERT_EQ(0, fine->open(path_ + ".fine", {storage(), true})); + build(fine); + auto refiner = std::make_shared(); + refiner->reference_index = fine; + for (float scale : {0.0f, 0.5f, 1.0f, 2.5f, 4.0f}) { + SCOPED_TRACE(scale); + auto candidates_param = + IVFQueryParamBuilder() + .with_topk(uint32_t(kPqTopK * std::max(1.0f, scale))) + .with_nprobe(3) + .build(); + set_filter(candidates_param); + SearchResult candidates; + ASSERT_EQ(0, index->search(query, candidates_param, &candidates)); + auto exact_param = IVFQueryParamBuilder() + .with_topk(kPqTopK) + .with_fetch_vector(true) + .build(); + exact_param->bf_pks = std::make_shared>(); + for (const auto &doc : candidates.doc_list_) + exact_param->bf_pks->push_back(doc.key()); + SearchResult expected, actual; + ASSERT_EQ(0, fine->search(query, exact_param, &expected)); + auto refined_param = IVFQueryParamBuilder() + .with_topk(kPqTopK) + .with_nprobe(3) + .with_fetch_vector(true) + .with_refiner_param(refiner) + .build(); + set_filter(refined_param); + refiner->scale_factor_ = scale; + ASSERT_EQ(0, index->search(query, refined_param, &actual)); + same(expected, actual); + EXPECT_EQ(expected.reverted_vector_list_, actual.reverted_vector_list_); + } + for (float invalid : {-1.0f, std::numeric_limits::infinity(), + std::numeric_limits::quiet_NaN(), + std::numeric_limits::max()}) { + refiner->scale_factor_ = invalid; + auto p = IVFQueryParamBuilder() + .with_topk(kPqTopK) + .with_refiner_param(refiner) + .build(); + SearchResult result; + EXPECT_NE(0, index->search(query, p, &result)); + } + ASSERT_EQ(0, fine->close()); + ASSERT_EQ(0, index->close()); + auto reopened = IndexFactory::CreateAndInitIndex(*param(false)); + ASSERT_EQ(0, reopened->open(path_, {storage(), false})); + SearchResult after; + ASSERT_EQ(0, reopened->search(query, query_param, &after)); + same(all, after); + // Merge must consume decoded/original input, never packed posting bytes. + std::vector decoded(kPqCount); + for (uint32_t id = 0; id < kPqCount; ++id) { + VectorDataBuffer fetched; + ASSERT_EQ(0, reopened->fetch(id, &fetched)); + decoded[id] = std::get(fetched.vector_buffer).data; + } + auto merged = IndexFactory::CreateAndInitIndex(*param(false)); + ASSERT_EQ(0, merged->open(path_ + ".merge", {storage(), true})); + IndexFilter merge_filter; + merge_filter.set([](uint64_t id) { return id == 7; }); + ASSERT_EQ(0, merged->merge({reopened}, merge_filter)); + EXPECT_EQ(kPqCount - 1, merged->get_doc_count()); + for (uint32_t id = 0; id < kPqCount; ++id) { + if (id == 7) continue; + VectorDataBuffer fetched; + // Merge's existing contract compacts surviving IDs in source order. + ASSERT_EQ(0, merged->fetch(id < 7 ? id : id - 1, &fetched)); + const auto &actual = + std::get(fetched.vector_buffer).data; + ASSERT_EQ(decoded[id].size(), actual.size()); + for (uint32_t d = 0; d < kPqDimension; ++d) { + float expected_value, actual_value; + std::memcpy(&expected_value, decoded[id].data() + d * sizeof(float), + sizeof(float)); + std::memcpy(&actual_value, actual.data() + d * sizeof(float), + sizeof(float)); + // Cosine FP32 storage normalizes and reconstructs its components. + EXPECT_NEAR(expected_value, actual_value, + 1e-5f * std::max(1.0f, std::abs(expected_value))); + } + } + ASSERT_EQ(0, merged->close()); + ASSERT_EQ(0, reopened->close()); +} + +INSTANTIATE_TEST_SUITE_P( + Turbo, IVFPqTest, + testing::Combine(testing::Values(8, 4, 0), testing::Bool(), + testing::Values(MetricType::kL2sq, + MetricType::kInnerProduct, + MetricType::kCosine), + testing::Bool())); + +TEST(IVFPqStorage, PackedBlocksAcrossPages) { + const std::string path = "ivf_pq_cross_page.index"; + struct Cleanup { + std::string path; + ~Cleanup() { + test_util::RemoveTestFiles(path); + test_util::RemoveTestFiles(path + ".merged"); + } + } cleanup{path}; + // Five chunks produce 96-byte blocks. The posting body exceeds a 16 KiB + // page and the block stride does not divide either 4 KiB or 16 KiB pages. + PqQuantizerParam pq(5, 4); + pq.fast_scan = true; + auto config = IVFIndexParamBuilder() + .with_dimension(30) + .with_data_type(DataType::DT_FP32) + .with_metric_type(MetricType::kL2sq) + .with_n_list(7) + .with_n_iters(3) + .with_quantizer_param(pq) + .build(); + auto mmap = IndexFactory::CreateAndInitIndex(*config); + ASSERT_NE(nullptr, mmap); + ASSERT_EQ(0, mmap->open(path, {StorageOptions::StorageType::kMMAP, true})); + std::mt19937 random(905); + std::normal_distribution value(0.0f, 1.0f); + std::vector vector(30); + for (uint32_t id = 0; id < 8193; ++id) { + for (auto &v : vector) v = value(random); + ASSERT_EQ(0, mmap->add(VectorData{DenseVector{vector.data()}}, id)); + } + ASSERT_EQ(0, mmap->train()); + auto pooled = IndexFactory::CreateAndInitIndex(*config); + ASSERT_EQ( + 0, pooled->open(path, {StorageOptions::StorageType::kBufferPool, false})); + auto run = [&](const Index::Pointer &index, uint32_t nprobe, bool filtered, + float radius = 0.0f) { + auto param = IVFQueryParamBuilder() + .with_nprobe(nprobe) + .with_topk(8193) + .with_radius(radius) + .build(); + if (filtered) { + param->filter = std::make_shared(); + param->filter->set([](uint64_t id) { return id % 7 == 0; }); + } + SearchResult result; + EXPECT_EQ(0, index->search(VectorData{DenseVector{vector.data()}}, param, + &result)); + std::map scores; + for (const auto &doc : result.doc_list_) + scores.emplace(doc.key(), doc.score()); + return scores; + }; + for (uint32_t nprobe : {1u, 7u, 0u}) { + SCOPED_TRACE(nprobe); + for (bool filtered : {false, true}) { + SCOPED_TRACE(filtered); + auto expected = run(mmap, nprobe, filtered); + ASSERT_FALSE(expected.empty()); + if (nprobe == 7 && !filtered) EXPECT_EQ(8193u, expected.size()); + if (nprobe == 1 || nprobe == 0) EXPECT_LT(expected.size(), 8193u); + for (int pass = 0; pass < 3; ++pass) { + // Cold contiguous reads and warm resident page spans must agree. + EXPECT_EQ(expected, run(pooled, nprobe, filtered)); + } + } + } + auto all_scores = run(mmap, 7, false); + // Restricted-key scoring must retain complete packed blocks across pages. + auto candidates = IVFQueryParamBuilder().with_topk(8193).build(); + candidates->bf_pks = std::make_shared>(); + for (const auto &entry : all_scores) { + if (entry.first % 3 == 0) candidates->bf_pks->push_back(entry.first); + } + candidates->bf_pks->push_back(10000); // Missing keys are ignored. + for (const auto &index : {mmap, pooled}) { + SearchResult selected; + ASSERT_EQ(0, index->search(VectorData{DenseVector{vector.data()}}, + candidates, &selected)); + ASSERT_EQ(candidates->bf_pks->size() - 1, selected.doc_list_.size()); + for (const auto &doc : selected.doc_list_) { + EXPECT_EQ(0u, doc.key() % 3); + EXPECT_FLOAT_EQ(all_scores.at(doc.key()), doc.score()); + } + } + run(pooled, 7, false, 1.0f); + EXPECT_EQ(all_scores, run(pooled, 7, false)); + auto merged = IndexFactory::CreateAndInitIndex(*config); + ASSERT_EQ(0, merged->open(path + ".merged", + {StorageOptions::StorageType::kBufferPool, true})); + ASSERT_EQ(0, merged->merge({pooled}, IndexFilter{})); + EXPECT_EQ(8193u, merged->get_doc_count()); + VectorDataBuffer fetched; + ASSERT_EQ(0, merged->fetch(8192, &fetched)); + EXPECT_EQ(std::string(reinterpret_cast(vector.data()), + vector.size() * sizeof(float)), + std::get(fetched.vector_buffer).data); + EXPECT_EQ(8193u, run(merged, 7, false).size()); + ASSERT_EQ(0, merged->close()); + ASSERT_EQ(0, pooled->close()); + ASSERT_EQ(0, mmap->close()); +} + +TEST(IVFPqValidation, InvalidParameters) { + for (const auto &pq : {PqQuantizerParam(0, 4), PqQuantizerParam(3, 6), + PqQuantizerParam(5, 4), PqQuantizerParam(13, 8)}) { + auto param = IVFIndexParamBuilder() + .with_dimension(kPqDimension) + .with_data_type(DataType::DT_FP32) + .with_metric_type(MetricType::kL2sq) + .with_quantizer_param(pq) + .build(); + EXPECT_EQ(nullptr, IndexFactory::CreateAndInitIndex(*param)); + } + PqQuantizerParam pq(3, 8); + pq.fast_scan = true; + auto param = IVFIndexParamBuilder() + .with_dimension(kPqDimension) + .with_data_type(DataType::DT_FP32) + .with_metric_type(MetricType::kL2sq) + .with_quantizer_param(pq) + .build(); + EXPECT_EQ(nullptr, IndexFactory::CreateAndInitIndex(*param)); +} + } // namespace } // namespace zvec::core_interface diff --git a/tests/core/metric/uniform_uint8_metric_test.cc b/tests/core/metric/uniform_uint8_metric_test.cc index b537b9e4ba..cbef92908f 100644 --- a/tests/core/metric/uniform_uint8_metric_test.cc +++ b/tests/core/metric/uniform_uint8_metric_test.cc @@ -199,7 +199,8 @@ class UniformUint8SeparateExtraValuesTest TEST_P(UniformUint8SeparateExtraValuesTest, QueryBatchConsumesSeparateExtraValues) { - // Exercise multiple batches, look-ahead prefetches, and the single-vector tail. + // Exercise multiple batches, look-ahead prefetches, and the single-vector + // tail. constexpr size_t kVectorCount = 11; const size_t kDimension = GetParam(); std::vector query_codes(kDimension); @@ -251,11 +252,12 @@ TEST_P(UniformUint8SeparateExtraValuesTest, } INSTANTIATE_TEST_SUITE_P(KernelLayouts, UniformUint8SeparateExtraValuesTest, - testing::Values(size_t{128}, size_t{129}, size_t{960})); + testing::Values(size_t{128}, size_t{129}, + size_t{960})); TEST(UniformUint8Metric, QueryPreprocessConvertsCanonicalLayout) { - const size_t dimensions[] = {1, 63, 64, 65, 127, 128, 129, 960, - MAX_DIMENSION}; + const size_t dimensions[] = {1, 63, 64, 65, 127, + 128, 129, 960, MAX_DIMENSION}; for (size_t dimension : dimensions) { auto metric = CreateMetric(dimension); ASSERT_TRUE(metric); @@ -263,8 +265,8 @@ TEST(UniformUint8Metric, QueryPreprocessConvertsCanonicalLayout) { ASSERT_TRUE(preprocess); for (size_t pattern = 0; pattern < 4; ++pattern) { - SCOPED_TRACE(testing::Message() << "dimension=" << dimension - << ", pattern=" << pattern); + SCOPED_TRACE(testing::Message() + << "dimension=" << dimension << ", pattern=" << pattern); const uint8_t constant_codes[] = {0, 128, 255}; std::vector query_codes(dimension); int64_t expected_correction = 0; @@ -279,7 +281,8 @@ TEST(UniformUint8Metric, QueryPreprocessConvertsCanonicalLayout) { auto query = EncodeQuery(query_codes); preprocess(query.data(), query.size()); const auto *raw_query = reinterpret_cast(query.data()); - EXPECT_TRUE(std::equal(query_codes.begin(), query_codes.end(), raw_query)); + EXPECT_TRUE( + std::equal(query_codes.begin(), query_codes.end(), raw_query)); EXPECT_EQ(expected_correction, ReadQueryCorrection(query, dimension)); } } @@ -296,16 +299,17 @@ TEST(UniformUint8Metric, QueryPreprocessPreservesOversizedFallback) { for (const uint8_t code : {uint8_t{128}, uint8_t{255}}) { SCOPED_TRACE(testing::Message() << "dimension=" << dimension - << ", code=" << static_cast(code)); + << ", code=" << static_cast(code)); const std::vector query_codes(dimension, code); auto query = EncodeQuery(query_codes); const uint32_t original_tail = ReadTail(query, dimension); - const int64_t correction = - static_cast(dimension) * code * (static_cast(code) - 256); + const int64_t correction = static_cast(dimension) * code * + (static_cast(code) - 256); preprocess(query.data(), query.size()); const auto *raw_query = reinterpret_cast(query.data()); - EXPECT_TRUE(std::equal(query_codes.begin(), query_codes.end(), raw_query)); + EXPECT_TRUE( + std::equal(query_codes.begin(), query_codes.end(), raw_query)); if (correction < (std::numeric_limits::min)() || correction > (std::numeric_limits::max)()) { EXPECT_EQ(original_tail, ReadTail(query, dimension)); diff --git a/tests/db/collection_test.cc b/tests/db/collection_test.cc index bac8702032..b6a0e5e47a 100644 --- a/tests/db/collection_test.cc +++ b/tests/db/collection_test.cc @@ -81,7 +81,7 @@ class CollectionTest : public ::testing::Test { } CollectionSchema make_schema(const std::string &name = "x", - const std::string &field = "value") { + const std::string &field = "value") { CollectionSchema schema(name); EXPECT_TRUE(schema .add_field(std::make_shared( @@ -104,7 +104,7 @@ class CollectionTest : public ::testing::Test { } Doc make_doc(const std::string &id, int32_t value, - const std::string &field = "value") { + const std::string &field = "value") { Doc doc; doc.set_pk(id); EXPECT_TRUE(doc.set(field, value)); @@ -120,7 +120,7 @@ class CollectionTest : public ::testing::Test { } void expect_value(const std::string &id, int32_t expected, - const std::string &field = "value") { + const std::string &field = "value") { auto result = collection_->fetch({id}); ASSERT_TRUE(result.has_value()) << result.error().message(); ASSERT_EQ(result.value().size(), 1u); diff --git a/tests/db/crash_recovery/utility.h b/tests/db/crash_recovery/utility.h index 826111ca24..ceb4b73645 100644 --- a/tests/db/crash_recovery/utility.h +++ b/tests/db/crash_recovery/utility.h @@ -22,7 +22,6 @@ #include #include #include - #include #include @@ -156,7 +155,6 @@ inline Doc CreateTestDoc(uint64_t doc_id, int version) { } - /** * @brief Locate a binary by name, searching common paths and TEST_BINARY_DIR. * diff --git a/tests/db/index/column/inverted_column/inverted_column_indexer_bool_test.cc b/tests/db/index/column/inverted_column/inverted_column_indexer_bool_test.cc index b26fe7123e..f22feab293 100644 --- a/tests/db/index/column/inverted_column/inverted_column_indexer_bool_test.cc +++ b/tests/db/index/column/inverted_column/inverted_column_indexer_bool_test.cc @@ -337,7 +337,8 @@ TEST_F(InvertedIndexTest, SEALED) { TEST_F(InvertedIndexTest, SNAPSHOT) { #ifdef __ANDROID__ - GTEST_SKIP() << "Skipped on Android: emulator filesystem lacks hardlink support (needed by RocksDB checkpoint)"; + GTEST_SKIP() << "Skipped on Android: emulator filesystem lacks hardlink " + "support (needed by RocksDB checkpoint)"; #endif ASSERT_TRUE(indexer_); diff --git a/tests/db/index/column/inverted_column/inverted_column_indexer_sequential_numbers_test.cc b/tests/db/index/column/inverted_column/inverted_column_indexer_sequential_numbers_test.cc index 0d2bda43ab..74dc0a88f9 100644 --- a/tests/db/index/column/inverted_column/inverted_column_indexer_sequential_numbers_test.cc +++ b/tests/db/index/column/inverted_column/inverted_column_indexer_sequential_numbers_test.cc @@ -698,7 +698,8 @@ TEST_F(InvertedIndexTest, SEALED) { TEST_F(InvertedIndexTest, CREATE_SNAPSHOT) { #ifdef __ANDROID__ - GTEST_SKIP() << "Skipped on Android: emulator filesystem lacks hardlink support (needed by RocksDB checkpoint)"; + GTEST_SKIP() << "Skipped on Android: emulator filesystem lacks hardlink " + "support (needed by RocksDB checkpoint)"; #endif ASSERT_TRUE(indexer_); diff --git a/tests/db/index/common/manifest_codec_golden_test.cc b/tests/db/index/common/manifest_codec_golden_test.cc index d8485b91ba..3987e51279 100644 --- a/tests/db/index/common/manifest_codec_golden_test.cc +++ b/tests/db/index/common/manifest_codec_golden_test.cc @@ -1145,3 +1145,24 @@ TEST(ManifestCodecGolden, WireReaderRejectsGroups) { EXPECT_FALSE(r.ok()); } } + +TEST(ManifestCodecTest, IvfPqParametersSurvivePersistence) { + for (bool rotate : {false, true}) { + for (int mode = 0; mode < 3; ++mode) { + const QuantizerParam qp(rotate, 4, mode == 0 ? 8 : 4, mode == 2, + rotate ? 3 : 0, rotate ? 2 : 0); + IVFIndexParams params(MetricType::COSINE, 16, 7, false, QuantizeType::PQ, + qp); + std::string encoded; + ManifestCodec::EncodeIndexParams(¶ms, &encoded); + const auto restored = std::dynamic_pointer_cast( + ManifestCodec::DecodeIndexParams(encoded)); + ASSERT_NE(nullptr, restored); + EXPECT_EQ(params, *restored); + EXPECT_EQ(qp, restored->quantizer_param()); + std::string reencoded; + ManifestCodec::EncodeIndexParams(restored.get(), &reencoded); + EXPECT_EQ(encoded, reencoded); + } + } +} diff --git a/tests/db/index/common/schema_test.cc b/tests/db/index/common/schema_test.cc index f9877d57a6..e20c7aa687 100644 --- a/tests/db/index/common/schema_test.cc +++ b/tests/db/index/common/schema_test.cc @@ -1356,3 +1356,58 @@ TEST(FieldSchemaTest, HnswRabitqIndexValidation_UnsupportedDataTypes) { << status.message(); } } + +TEST(FieldSchemaTest, IvfPqValidatesLayoutAndSupportedIndex) { + auto check = [](QuantizerParam qp, int dim = 16, + DataType type = DataType::VECTOR_FP32, + MetricType metric = MetricType::L2) { + return FieldSchema("vector", type, dim, false, + std::make_shared(metric, 4, 3, false, + QuantizeType::PQ, qp)) + .validate(); + }; + EXPECT_TRUE(check(QuantizerParam(true, 4, 4, true, 2, 2)).ok()); + EXPECT_TRUE(check(QuantizerParam(false, 3, 8), 16).ok()); + EXPECT_TRUE(check(QuantizerParam(false, 4, 4), 16, DataType::VECTOR_FP32, + MetricType::IP) + .ok()); + EXPECT_TRUE(check(QuantizerParam(true, 4, 8), 16, DataType::VECTOR_FP32, + MetricType::COSINE) + .ok()); + for (const auto &qp : + {QuantizerParam(false, 0), QuantizerParam(false, -1), + QuantizerParam(false, 17), QuantizerParam(false, 4, 3), + QuantizerParam(false, 3, 4), QuantizerParam(false, 4, 8, true), + QuantizerParam(true, 4, 4, false, 0, 2), + QuantizerParam(true, 4, 4, false, 2, 0)}) { + EXPECT_FALSE(check(qp).ok()); + } + EXPECT_FALSE(check(QuantizerParam(), 16, DataType::VECTOR_FP16).ok()); + EXPECT_FALSE( + check(QuantizerParam(), 16, DataType::VECTOR_FP32, MetricType::MIPSL2) + .ok()); + FieldSchema flat( + "vector", DataType::VECTOR_FP32, 16, false, + std::make_shared(MetricType::L2, QuantizeType::PQ)); + EXPECT_FALSE(flat.validate().ok()); +} + +TEST(FieldSchemaTest, HnswPqSupportsOpqButRejectsPackedCodes) { + for (auto metric : {MetricType::L2, MetricType::IP, MetricType::COSINE}) { + for (int bits : {4, 8}) { + auto params = std::make_shared( + metric, 16, 100, QuantizeType::PQ, false, + QuantizerParam(true, 4, bits, false, 2, 2)); + EXPECT_TRUE( + FieldSchema("vector", DataType::VECTOR_FP32, 16, false, params) + .validate() + .ok()); + } + } + auto packed = std::make_shared( + MetricType::L2, 16, 100, QuantizeType::PQ, false, + QuantizerParam(false, 4, 4, true)); + EXPECT_FALSE(FieldSchema("vector", DataType::VECTOR_FP32, 16, false, packed) + .validate() + .ok()); +} diff --git a/tests/db/reranker_test.cc b/tests/db/reranker_test.cc index 7d87962945..95c75a973d 100644 --- a/tests/db/reranker_test.cc +++ b/tests/db/reranker_test.cc @@ -126,7 +126,7 @@ TEST(RerankRrfTest, NegativeRankConstantError) { reranker::rerank(reranker::RrfParams{/*rank_constant=*/-1}, results, /*fields=*/{}, /*topn=*/10); ASSERT_FALSE(result.has_value()); - + result = reranker::rerank(reranker::RrfParams{/*rank_constant=*/0}, results, /*fields=*/{}, /*topn=*/10); ASSERT_TRUE(result.has_value()); diff --git a/tests/ios_test_sandbox.cc b/tests/ios_test_sandbox.cc index ca8bed2c9d..7356438f5a 100644 --- a/tests/ios_test_sandbox.cc +++ b/tests/ios_test_sandbox.cc @@ -7,39 +7,38 @@ #include #if TARGET_OS_IOS || TARGET_OS_SIMULATOR +#include #include -#include #include -#include +#include -__attribute__((constructor)) -static void ios_test_sandbox_setup() { - // 1. Raise file descriptor limit (iOS default is very low) - struct rlimit rl; - if (getrlimit(RLIMIT_NOFILE, &rl) == 0) { - rlim_t target = 65536; - if (rl.rlim_cur < target) { - rl.rlim_cur = (rl.rlim_max >= target) ? target : rl.rlim_max; - if (setrlimit(RLIMIT_NOFILE, &rl) == 0) { - fprintf(stderr, "[iOS] File descriptor limit raised to: %llu\n", - (unsigned long long)rl.rlim_cur); - } - } +__attribute__((constructor)) static void ios_test_sandbox_setup() { + // 1. Raise file descriptor limit (iOS default is very low) + struct rlimit rl; + if (getrlimit(RLIMIT_NOFILE, &rl) == 0) { + rlim_t target = 65536; + if (rl.rlim_cur < target) { + rl.rlim_cur = (rl.rlim_max >= target) ? target : rl.rlim_max; + if (setrlimit(RLIMIT_NOFILE, &rl) == 0) { + fprintf(stderr, "[iOS] File descriptor limit raised to: %llu\n", + (unsigned long long)rl.rlim_cur); + } } + } - // 2. Change CWD to writable sandbox directory - // TMPDIR is set by iOS to the app's writable sandbox tmp directory - const char* tmpdir = getenv("TMPDIR"); - if (tmpdir && chdir(tmpdir) == 0) { - fprintf(stderr, "[iOS] Working directory set to: %s\n", tmpdir); - } else { - // Fallback: try HOME directory - const char* home = getenv("HOME"); - if (home && chdir(home) == 0) { - fprintf(stderr, "[iOS] Working directory set to: %s\n", home); - } + // 2. Change CWD to writable sandbox directory + // TMPDIR is set by iOS to the app's writable sandbox tmp directory + const char *tmpdir = getenv("TMPDIR"); + if (tmpdir && chdir(tmpdir) == 0) { + fprintf(stderr, "[iOS] Working directory set to: %s\n", tmpdir); + } else { + // Fallback: try HOME directory + const char *home = getenv("HOME"); + if (home && chdir(home) == 0) { + fprintf(stderr, "[iOS] Working directory set to: %s\n", home); } + } } -#endif // TARGET_OS_IOS || TARGET_OS_SIMULATOR -#endif // __APPLE__ +#endif // TARGET_OS_IOS || TARGET_OS_SIMULATOR +#endif // __APPLE__ diff --git a/tests/turbo/turbo_pq_int4_quantizer_test.cc b/tests/turbo/turbo_pq_int4_quantizer_test.cc index 297bd7aee9..141f1fdab4 100644 --- a/tests/turbo/turbo_pq_int4_quantizer_test.cc +++ b/tests/turbo/turbo_pq_int4_quantizer_test.cc @@ -341,9 +341,15 @@ TEST(PqInt4Quantizer, SerializeDeserialize) { float adc2 = q2->calc_distance_dp_query(code2.data(), lut2.data()); EXPECT_NEAR(adc1, adc2, 1e-6f); - // Note: SDC (calc_distance_dp_dp) is intentionally NOT tested after - // deserialization because dist_table_ is a build-phase-only structure - // and is not persisted. + // Reopening for writes restores SDC separately; ADC above needs no prepare. + iter->next(); + ASSERT_TRUE(iter->is_valid()); + quantizer->quantize_data(iter->data(), code2.data()); + const float sdc = quantizer->calc_distance_dp_dp(code1.data(), code2.data()); + ASSERT_EQ(0, q2->prepare_build()); + EXPECT_NEAR(sdc, q2->calc_distance_dp_dp(code1.data(), code2.data()), 1e-6f); + ASSERT_EQ(0, q2->prepare_build()); // Safe to prepare again. + EXPECT_NEAR(sdc, q2->calc_distance_dp_dp(code1.data(), code2.data()), 1e-6f); } // --------------------------------------------------------------------------- diff --git a/tests/turbo/turbo_pq_int8_quantizer_test.cc b/tests/turbo/turbo_pq_int8_quantizer_test.cc index 4368b69abb..e9cea0a232 100644 --- a/tests/turbo/turbo_pq_int8_quantizer_test.cc +++ b/tests/turbo/turbo_pq_int8_quantizer_test.cc @@ -291,9 +291,15 @@ TEST(PqInt8Quantizer, SerializeDeserialize) { float adc2 = q2->calc_distance_dp_query(code2.data(), lut2.data()); EXPECT_NEAR(adc1, adc2, 1e-6f); - // Note: SDC (calc_distance_dp_dp) is intentionally NOT tested after - // deserialization because dist_table_ is a build-phase-only structure - // and is not persisted. + // Reopening for writes restores SDC separately; ADC above needs no prepare. + iter->next(); + ASSERT_TRUE(iter->is_valid()); + quantizer->quantize_data(iter->data(), code2.data()); + const float sdc = quantizer->calc_distance_dp_dp(code1.data(), code2.data()); + ASSERT_EQ(0, q2->prepare_build()); + EXPECT_NEAR(sdc, q2->calc_distance_dp_dp(code1.data(), code2.data()), 1e-6f); + ASSERT_EQ(0, q2->prepare_build()); // Safe to prepare again. + EXPECT_NEAR(sdc, q2->calc_distance_dp_dp(code1.data(), code2.data()), 1e-6f); } TEST(PqInt8Quantizer, OpqSerializeDeserializeUniformAndNonUniformChunks) { diff --git a/tools/core/local_builder.cc b/tools/core/local_builder.cc index 03aea9d983..f9f807fa10 100644 --- a/tools/core/local_builder.cc +++ b/tools/core/local_builder.cc @@ -523,9 +523,12 @@ int build_sparse_by_streamer(IndexStreamer::Pointer &streamer, return 0; } -int do_build_by_streamer(IndexStreamer::Pointer &streamer, - uint32_t thread_count, RetrievalMode retrieval_mode, - const IndexStorage::Pointer &storage = nullptr) { +int do_build_by_streamer( + IndexStreamer::Pointer &streamer, uint32_t thread_count, + RetrievalMode retrieval_mode, + const IndexStorage::Pointer &storage = nullptr, + const std::shared_ptr &quantizer = nullptr, + const IndexQueryMeta &quantized_qmeta = IndexQueryMeta()) { int ret; std::atomic finished{0}; std::atomic errcode{0}; @@ -597,6 +600,10 @@ int do_build_by_streamer(IndexStreamer::Pointer &streamer, } std::string ovec; IndexQueryMeta ometa; + std::string code_buffer; + if (quantizer) { + code_buffer.resize(quantizer->quantized_datapoint_vector_length()); + } for (uint32_t id = idx; id < end_cursor && !stop_now && errcode.load(std::memory_order_acquire) == 0; id += thread_count) { @@ -615,6 +622,13 @@ int do_build_by_streamer(IndexStreamer::Pointer &streamer, return; } task_ret = add_to_streamer(key, ovec.data(), ometa, ctx); + } else if (quantizer) { + // The streamer stores the quantized codes; quantize each original + // vector on the fly so the graph is built with quantizer distances. + quantizer->quantize_data(holder->get_vector_by_index(id), + code_buffer.data()); + task_ret = + add_to_streamer(key, code_buffer.data(), quantized_qmeta, ctx); } else { task_ret = add_to_streamer(key, holder->get_vector_by_index(id), qmeta, ctx); @@ -673,9 +687,10 @@ int do_build_by_streamer(IndexStreamer::Pointer &streamer, return 0; } -int build_by_streamer(IndexStreamer::Pointer &streamer, - YAML::Node &config_common, - const IndexConverter::Pointer &converter) { +int build_by_streamer( + IndexStreamer::Pointer &streamer, YAML::Node &config_common, + const IndexConverter::Pointer &converter, + const std::shared_ptr &quantizer = nullptr) { if (!config_common["IndexPath"]) { LOG_ERROR("Miss params IndexPath for Streamer"); return IndexError_InvalidArgument; @@ -728,9 +743,17 @@ int build_by_streamer(IndexStreamer::Pointer &streamer, retrieval_mode = RM_SPARSE; } + IndexQueryMeta quantized_qmeta; + if (quantizer) { + quantized_qmeta.set_meta(meta.data_type(), meta.dimension(), + static_cast(quantizer->type()), + meta.extra_meta_size()); + } + LOG_DEBUG("thread count: %zu, retrieval mode: %s", thread_count, retrieval_mode == 1 ? "Dense" : "Sparse"); - return do_build_by_streamer(streamer, thread_count, retrieval_mode, storage); + return do_build_by_streamer(streamer, thread_count, retrieval_mode, storage, + quantizer, quantized_qmeta); } IndexSparseHolder::Pointer convert_sparse_holder( @@ -863,7 +886,8 @@ IndexHolder::Pointer fill_quantized_holder( IndexHolder::Pointer quantize_holder( const std::string &name, const ailego::Params ¶ms, VecsIndexHolder::Pointer &in_holder, IndexMeta &index_meta, - std::shared_ptr *out_quantizer) { + std::shared_ptr *out_quantizer, + bool materialize = true) { IndexHolder::Pointer cast_holder = std::dynamic_pointer_cast(in_holder); if (name.empty()) { @@ -905,7 +929,7 @@ IndexHolder::Pointer quantize_holder( } uint32_t alloc_dim = static_cast(code_bytes / unit); - if (!quantizer->require_train()) { + if (!quantizer->require_train() || !materialize) { out_meta.set_quantizer(name, 0, params); } else { LOG_WARN( @@ -914,6 +938,12 @@ IndexHolder::Pointer quantize_holder( name.c_str()); } + if (out_quantizer) { + *out_quantizer = quantizer; + } + index_meta = out_meta; + if (!materialize) return cast_holder; + IndexHolder::Pointer result; switch (out_meta.data_type()) { case IndexMeta::DataType::DT_FP32: @@ -938,10 +968,6 @@ IndexHolder::Pointer quantize_holder( return IndexHolder::Pointer(); } - if (out_quantizer) { - *out_quantizer = quantizer; - } - index_meta = out_meta; return result; } @@ -1236,16 +1262,9 @@ int do_build(YAML::Node &config_root, YAML::Node &config_common) { std::shared_ptr build_quantizer; IndexHolder::Pointer cv_build_holder; if (!quantizer_name.empty()) { - // Quantizer path: only supported with IndexBuilder classes. The - // streamer path feeds raw vectors from the global holder and relies on - // reformer info in the meta, which a quantizer does not provide. - if (!builder) { - LOG_ERROR("QuantizerName is not supported with streamer class %s", - builder_class.c_str()); - return -1; - } - cv_build_holder = quantize_holder(quantizer_name, quantizer_params, - build_holder, meta, &build_quantizer); + cv_build_holder = + quantize_holder(quantizer_name, quantizer_params, build_holder, meta, + &build_quantizer, builder != nullptr); if (!cv_build_holder) { LOG_ERROR("Quantize holder failed."); return -1; @@ -1293,7 +1312,11 @@ int do_build(YAML::Node &config_root, YAML::Node &config_common) { ret = builder->init(meta, params); } } else { - ret = streamer->init(meta, params); + if (build_quantizer) { + ret = streamer->init(meta, params, build_quantizer); + } else { + ret = streamer->init(meta, params); + } } if (ret < 0) { LOG_ERROR("Failed to init builder, ret=%d", ret); @@ -1478,7 +1501,8 @@ int do_build(YAML::Node &config_root, YAML::Node &config_common) { retrieval_mode = "dense"; } - ret = build_by_streamer(streamer, config_common, build_converter); + ret = build_by_streamer(streamer, config_common, build_converter, + build_quantizer); } size_t build_time = timer.milli_seconds(); if (ret < 0) { diff --git a/tools/core/local_builder_original.cc b/tools/core/local_builder_original.cc index 7bd1032ccf..4225fd61a1 100644 --- a/tools/core/local_builder_original.cc +++ b/tools/core/local_builder_original.cc @@ -766,14 +766,10 @@ IndexHolder::Pointer quantize_holder( return IndexHolder::Pointer(); } - if (!quantizer->require_train()) { - out_meta.set_quantizer(name, 0, params); - } else { - cerr << "Quantizer " << name - << " requires training, query-side quantizer info is not recorded " - "in the index meta" - << endl; - } + // The descriptor must be recorded for trained quantizers too: the codebook + // is persisted in the index and restored at load time from this descriptor + // (IVFEntity::load_quantizer). Without it turbo IVF rejects the builder init. + out_meta.set_quantizer(name, 0, params); IndexHolder::Pointer result = std::make_shared(cast_holder,