From d68ee69cf5dc1a6d3c6fad296f4e4ca1fab8f885 Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Wed, 16 Sep 2026 15:12:09 +0800 Subject: [PATCH 1/8] feat(search): add fast_query for dense read-only collections Return owned internal-ID arrays and optional scores through C++ search_fast and Collection.fast_query with a thin NumPy binding. Reuse query parameters and graph/Flat fast paths, retaining filtered and multi-segment fallbacks. Preserve the existing scale_factor defaults and Uniform raw-Flat fallback. Cover native reference types, default and explicit refine budgets, all dense index families, parameter reuse, result ownership, and untrained/mixed Uniform segments. Validation: 634 related Python tests passed; 70 C++ tests passed with one existing skip. Ruff, changed-line clang-format, and diff checks passed. Adapted from alignment commit 9ddbd5aa082615226d91addb41cd60496143d15a onto upstream/main 43f72b186a0524b7f6b53c645bf90ac44ac05db3. --- python/tests/test_fast_query.py | 446 ++++++++++++++++++ python/tests/test_fast_query_refine.py | 165 +++++++ .../test_fast_query_uniform_lifecycle.py | 124 +++++ python/zvec/model/collection.py | 37 ++ src/binding/python/model/python_collection.cc | 84 ++++ src/core/algorithm/flat/flat_streamer.cc | 25 + src/core/algorithm/flat/flat_streamer.h | 8 + .../algorithm/flat/flat_streamer_entity.cc | 83 ++++ .../algorithm/flat/flat_streamer_entity.h | 20 + src/core/algorithm/hnsw/hnsw_context.h | 16 +- src/core/algorithm/hnsw/hnsw_streamer.cc | 8 +- src/core/algorithm/hnsw/hnsw_streamer.h | 1 + src/core/algorithm/vamana/vamana_context.cc | 16 +- src/core/algorithm/vamana/vamana_context.h | 3 +- src/core/algorithm/vamana/vamana_streamer.cc | 8 +- src/core/algorithm/vamana/vamana_streamer.h | 1 + src/core/interface/index.cc | 212 ++++++++- src/db/collection.cc | 368 +++++++++++++++ .../combined_vector_column_indexer.h | 20 + src/db/index/segment/segment.cc | 30 ++ src/db/index/segment/segment.h | 9 + .../zvec/core/framework/index_runner.h | 18 +- src/include/zvec/core/interface/index.h | 21 +- src/include/zvec/db/collection.h | 11 + src/include/zvec/db/query.h | 12 + .../core/algorithm/flat/flat_streamer_test.cc | 15 + tests/core/interface/index_interface_test.cc | 267 ++++++++++- tests/db/fast_query_test.cc | 82 ++++ tests/db/sqlengine/mock_segment.h | 6 + 29 files changed, 2096 insertions(+), 20 deletions(-) create mode 100644 python/tests/test_fast_query.py create mode 100644 python/tests/test_fast_query_refine.py create mode 100644 python/tests/test_fast_query_uniform_lifecycle.py create mode 100644 tests/db/fast_query_test.cc diff --git a/python/tests/test_fast_query.py b/python/tests/test_fast_query.py new file mode 100644 index 000000000..49069caf5 --- /dev/null +++ b/python/tests/test_fast_query.py @@ -0,0 +1,446 @@ +"""Advanced dense search: collection queries, scores, fallback and lifetime.""" + +import numpy as np +import pytest + +import zvec +from zvec import ( + CollectionOption, + CollectionSchema, + Doc, + Query, + HnswIndexParam, + HnswQueryParam, + VamanaIndexParam, + VamanaQueryParam, + VectorSchema, +) +from zvec.typing import DataType, MetricType, QuantizeType + + +@pytest.fixture( + params=[(128, False), (1200, False), (1200, True)], + ids=["brute_fallback", "vamana_graph", "hnsw_graph"], +) +def collection(tmp_path, request): + """Cross the Vamana brute-force threshold to exercise both result paths.""" + rng = np.random.default_rng(721) + vectors = rng.normal(size=(request.param[0], 32)).astype(np.float32) + schema = CollectionSchema( + name="fast_query", + vectors=[ + VectorSchema( + "vector", + DataType.VECTOR_FP32, + dimension=32, + index_param=HnswIndexParam( + metric_type=MetricType.L2, + m=16, + ef_construction=64, + use_contiguous_memory=True, + ) + if request.param[1] + else VamanaIndexParam( + metric_type=MetricType.L2, + max_degree=16, + search_list_size=64, + use_contiguous_memory=True, + two_pass_build=True, + quantize_type=QuantizeType.UNDEFINED, + ), + ) + ], + ) + path = str(tmp_path / "collection") + writer = zvec.create_and_open(path, schema) + for start in range(0, len(vectors), 512): + statuses = writer.insert( + [ + Doc(id=f"row-{i}", vectors={"vector": vectors[i].tolist()}) + for i in range(start, min(start + 512, len(vectors))) + ] + ) + assert all(status.ok() for status in statuses) + with pytest.raises(ValueError, match="read-only"): + writer.fast_query("vector", vectors[0]) + writer.optimize() + writer.close() + reader = zvec.open(path, CollectionOption(read_only=True, enable_mmap=True)) + try: + yield reader, vectors, HnswQueryParam if request.param[1] else VamanaQueryParam + finally: + reader.close() + + +def test_fast_query_matches_query_and_owns_results(collection): + coll, vectors, param_type = collection + query = np.ascontiguousarray(vectors[17] + 0.013, dtype=np.float32) + for ef in (24, 80): + param = param_type( + **({"ef": ef} if param_type is HnswQueryParam else {"ef_search": ef}), + is_using_refiner=False, + ) + expected_docs = coll.query( + Query(field_name="vector", vector=query, param=param), + topk=10, + output_fields=[], + ) + expected = np.array([int(doc.id[4:]) for doc in expected_docs]) + ids = coll.fast_query("vector", query, param, topk=10) + np.testing.assert_array_equal(ids, expected) + assert ids.dtype == np.int64 + scored_ids, scores = coll.fast_query( + "vector", query, param, topk=10, return_scores=True + ) + np.testing.assert_array_equal(scored_ids, expected) + np.testing.assert_allclose( + scores, [doc.score for doc in expected_docs], rtol=1e-5, atol=1e-5 + ) + # Another search must not overwrite a capsule-owned result array. + coll.fast_query("vector", vectors[25], param, topk=10) + np.testing.assert_array_equal(ids, expected) + + +def test_preconditions_and_invalid_vectors(collection): + coll, vectors, param_type = collection + with pytest.raises(ValueError, match="dense vector field"): + coll.fast_query("missing", vectors[0]) + wrong_param = ( + VamanaQueryParam() if param_type is HnswQueryParam else HnswQueryParam() + ) + with pytest.raises(ValueError, match="parameter type"): + coll.fast_query("vector", vectors[0], wrong_param) + param = param_type(is_using_refiner=False) + for invalid_query in ( + vectors[0, :-1], + vectors[0].astype(np.float64), + vectors[0].astype(np.int32), + vectors[0].astype(np.complex64), + np.zeros(32, dtype=[("x", np.float32)]), + vectors[0].astype(">f4"), + vectors[0, ::2], + vectors[0].reshape(1, -1), + ): + with pytest.raises(ValueError, match="dtype|dimension|1D"): + coll.fast_query("vector", invalid_query, param) + + +def test_reused_inline_and_default_params_and_close(collection): + coll, vectors, param_type = collection + query = np.ascontiguousarray(vectors[17] + 0.013) + for ef in (80, 16, 64): + settings = {"ef": ef} if param_type is HnswQueryParam else {"ef_search": ef} + param = param_type(**settings) + for topk in (1, 21, 3, 10): + ids, scores = coll.fast_query( + "vector", query, param, topk=topk, return_scores=True + ) + assert len(ids) == len(scores) == topk + assert scores.dtype == np.float32 + np.testing.assert_array_equal( + ids, coll.fast_query("vector", query, param_type(**settings), topk=topk) + ) + docs = coll.query( + Query(field_name="vector", vector=query, param=param), topk=topk + ) + np.testing.assert_array_equal(ids, [int(doc.id[4:]) for doc in docs]) + # None must restore defaults after a customized query. + docs = coll.query(Query(field_name="vector", vector=query), topk=10) + np.testing.assert_array_equal( + coll.fast_query("vector", query), [int(doc.id[4:]) for doc in docs] + ) + raw = coll._obj + saved_ids, saved_scores = ids.copy(), scores.copy() + for topk in (0, -1): + empty_ids, empty_scores = coll.fast_query( + "vector", query, param, topk=topk, return_scores=True + ) + assert empty_ids.shape == empty_scores.shape == (0,) + assert empty_ids.dtype == np.int64 + assert empty_scores.dtype == np.float32 + coll.close() + np.testing.assert_array_equal(ids, saved_ids) + np.testing.assert_array_equal(scores, saved_scores) + for obj in (coll, raw): + with pytest.raises(ValueError, match="closed"): + obj.fast_query("vector", query, param) + + +@pytest.mark.parametrize("compact", [False, True], ids=["delete_filter", "compacted"]) +def test_internal_ids_survive_deletion_and_compaction(tmp_path, compact): + vectors = np.random.default_rng(823).normal(size=(64, 32)).astype(np.float32) + schema = CollectionSchema( + name="deleted_ordinals", + vectors=[ + VectorSchema( + "vector", + DataType.VECTOR_FP32, + dimension=32, + index_param=VamanaIndexParam( + metric_type=MetricType.L2, + max_degree=16, + search_list_size=32, + use_contiguous_memory=True, + ), + ) + ], + ) + path = str(tmp_path / "deleted") + writer = zvec.create_and_open(path, schema) + assert all( + result.ok() + for result in writer.insert( + [ + Doc(id=f"row-{i}", vectors={"vector": vector.tolist()}) + for i, vector in enumerate(vectors) + ] + ) + ) + writer.optimize() + assert all(result.ok() for result in writer.delete([f"row-{i}" for i in range(8)])) + if compact: + writer.optimize() + writer.close() + coll = zvec.open(path, CollectionOption(read_only=True, enable_mmap=True)) + try: + param = VamanaQueryParam(ef_search=64, is_using_refiner=False) + query = vectors[12] + expected = [ + int(doc.id[4:]) + for doc in coll.query( + Query(field_name="vector", vector=query, param=param), topk=10 + ) + ] + assert all(row >= 8 for row in expected) + np.testing.assert_array_equal(coll.fast_query("vector", query, param), expected) + docs = coll.query( + Query(field_name="vector", vector=query, param=param), topk=10 + ) + ids, scores = coll.fast_query("vector", query, param, return_scores=True) + np.testing.assert_array_equal(ids, expected) + np.testing.assert_allclose(scores, [doc.score for doc in docs], rtol=1e-5) + # Padding and result ownership on the filtered / compacted routes. + padded_ids, padded_scores = coll.fast_query( + "vector", query, param, topk=70, return_scores=True + ) + assert len(padded_ids) == len(padded_scores) == 70 + assert np.all(padded_ids[56:] == -1) + assert np.all(np.isnan(padded_scores[56:])) + finally: + coll.close() + + +def test_ids_are_merged_across_segments(tmp_path): + vectors = np.random.default_rng(984).normal(size=(1100, 32)).astype(np.float32) + schema = CollectionSchema( + name="multiple_segments", + vectors=[ + VectorSchema( + "vector", + DataType.VECTOR_FP32, + dimension=32, + index_param=VamanaIndexParam( + metric_type=MetricType.L2, + max_degree=16, + search_list_size=32, + use_contiguous_memory=True, + ), + ) + ], + ) + # The native schema persists this limit, but the Python constructor does + # not expose it. Restore a schema with the smallest allowed segment size. + native_schema = schema._get_object() + name, fields, _ = native_schema.__getstate__() + restored = type(native_schema).__new__(type(native_schema)) + restored.__setstate__((name, fields, 1000)) + schema = CollectionSchema._from_core(restored) + path = str(tmp_path / "segments") + writer = zvec.create_and_open(path, schema) + for start in range(0, len(vectors), 500): + assert all( + result.ok() + for result in writer.insert( + [ + Doc(id=f"row-{i}", vectors={"vector": vectors[i].tolist()}) + for i in range(start, min(start + 500, len(vectors))) + ] + ) + ) + writer.close() + coll = zvec.open(path, CollectionOption(read_only=True, enable_mmap=True)) + try: + param = VamanaQueryParam(ef_search=64, is_using_refiner=False) + # Both the persisted segment and the final writing segment must + # participate, with the latter's segment-local IDs mapped globally. + for query_id in (12, 1050): + query = vectors[query_id] + expected = [ + int(doc.id[4:]) + for doc in coll.query( + Query(field_name="vector", vector=query, param=param), topk=10 + ) + ] + assert expected[0] == query_id + np.testing.assert_array_equal( + coll.fast_query("vector", query, param), expected + ) + docs = coll.query( + Query(field_name="vector", vector=query, param=param), topk=10 + ) + ids, scores = coll.fast_query("vector", query, param, return_scores=True) + np.testing.assert_array_equal(ids, expected) + np.testing.assert_allclose(scores, [doc.score for doc in docs], rtol=1e-5) + finally: + coll.close() + + +def test_field_and_collection_caches_are_independent(tmp_path): + rng = np.random.default_rng(726) + vectors = { + "l2": rng.normal(size=(64, 32)).astype(np.float32), + "ip": rng.normal(size=(64, 16)).astype(np.float32), + } + schema = CollectionSchema( + name="two_fields", + vectors=[ + VectorSchema( + name, + DataType.VECTOR_FP32, + dimension=values.shape[1], + index_param=zvec.FlatIndexParam( + metric_type=MetricType.L2 if name == "l2" else MetricType.IP + ), + ) + for name, values in vectors.items() + ], + ) + readers = [] + try: + for number in range(2): + path = str(tmp_path / f"collection-{number}") + writer = zvec.create_and_open(path, schema) + assert all( + s.ok() + for s in writer.insert( + [ + Doc( + id=str(i), + vectors={ + name: values[i if number == 0 else 63 - i].tolist() + for name, values in vectors.items() + }, + ) + for i in range(64) + ] + ) + ) + writer.optimize() + writer.close() + readers.append(zvec.open(path, CollectionOption(read_only=True))) + for name in ("l2", "ip", "l2", "ip"): + query = vectors[name][17] + for reader in readers: + docs = reader.query(Query(field_name=name, vector=query), topk=10) + ids, scores = reader.fast_query( + name, query, topk=10, return_scores=True + ) + np.testing.assert_array_equal(ids, [int(doc.id) for doc in docs]) + np.testing.assert_allclose( + scores, [doc.score for doc in docs], rtol=1e-5 + ) + finally: + for reader in readers: + reader.close() + + +@pytest.mark.parametrize( + "index_kind", + ["flat", "hnsw", "vamana", "ivf", "hnsw_rabitq", "ivf_rabitq", "diskann"], +) +@pytest.mark.parametrize("metric", [MetricType.L2, MetricType.IP, MetricType.COSINE]) +def test_fast_query_index_and_metric_dispatch(tmp_path, index_kind, metric): + """Exercise every dense index dispatch, including score normalization.""" + record = dict(metric_type=metric, quantize_type=QuantizeType.INT8) + factories = { + "flat": lambda: (zvec.FlatIndexParam(**record), None), + "hnsw": lambda: ( + HnswIndexParam(m=16, ef_construction=64, **record), + HnswQueryParam(ef=64), + ), + "vamana": lambda: ( + VamanaIndexParam(max_degree=16, search_list_size=64, **record), + VamanaQueryParam(ef_search=64), + ), + "ivf": lambda: ( + zvec.IVFIndexParam(n_list=4, n_iters=2, **record), + zvec.IVFQueryParam(nprobe=4), + ), + "hnsw_rabitq": lambda: ( + zvec.HnswRabitqIndexParam( + metric_type=metric, + m=16, + ef_construction=64, + total_bits=4, + num_clusters=4, + sample_count=256, + ), + zvec.HnswRabitqQueryParam(ef=64), + ), + "ivf_rabitq": lambda: ( + zvec.IvfRabitqIndexParam(metric_type=metric, nlist=4, total_bits=4), + zvec.IvfRabitqQueryParam(nprobe=4), + ), + "diskann": lambda: ( + zvec.DiskAnnIndexParam( + metric_type=metric, max_degree=16, list_size=64, pq_chunk_num=4 + ), + zvec.DiskAnnQueryParam(list_size=64), + ), + } + index, query_param = factories[index_kind]() + vectors = np.random.default_rng(754).normal(size=(512, 128)).astype(np.float32) + schema = CollectionSchema( + name="fast_index_dispatch", + vectors=[VectorSchema("vector", DataType.VECTOR_FP32, 128, index_param=index)], + ) + path = str(tmp_path / "collection") + try: + writer = zvec.create_and_open(path, schema) + except RuntimeError as exc: + if "not supported on this platform" in str(exc) or "RabitQ requires AVX" in str( + exc + ): + pytest.skip(str(exc)) + raise + try: + assert all( + s.ok() + for s in writer.insert( + [ + Doc(id=str(i), vectors={"vector": v.tolist()}) + for i, v in enumerate(vectors) + ] + ) + ) + writer.optimize() + finally: + writer.close() + reader = zvec.open(path, CollectionOption(read_only=True)) + try: + for param in (query_param, None, query_param): + for row in (12, 41): + query = np.ascontiguousarray(vectors[row] + np.float32(0.021)) + docs = reader.query(Query("vector", vector=query, param=param), topk=10) + ids, scores = reader.fast_query( + "vector", query, param, return_scores=True + ) + np.testing.assert_array_equal(ids, [int(d.id) for d in docs]) + np.testing.assert_allclose( + scores, [d.score for d in docs], rtol=2e-5, atol=2e-5 + ) + np.testing.assert_array_equal( + ids, reader.fast_query("vector", query, param) + ) + finally: + reader.close() diff --git a/python/tests/test_fast_query_refine.py b/python/tests/test_fast_query_refine.py new file mode 100644 index 000000000..10262953a --- /dev/null +++ b/python/tests/test_fast_query_refine.py @@ -0,0 +1,165 @@ +"""Fast search with native reference storage and optional refined scores.""" + +from functools import partial + +import numpy as np +import pytest + +import zvec +from zvec import CollectionOption, CollectionSchema, Doc, Query +from zvec import HnswIndexParam, HnswQueryParam +from zvec import VamanaIndexParam, VamanaQueryParam, VectorSchema +from zvec.typing import DataType, MetricType, QuantizeType + + +@pytest.mark.parametrize( + "quantizer,flat_type", + [ + (QuantizeType.INT8, DataType.VECTOR_FP16), + (QuantizeType.FP16, DataType.VECTOR_FP16), + (QuantizeType.INT4, DataType.VECTOR_FP16), + (QuantizeType.UNIFORM_UINT4, DataType.VECTOR_UINT8), + (QuantizeType.UNIFORM_UINT7, DataType.VECTOR_UINT8), + (QuantizeType.UNIFORM_UINT8, DataType.VECTOR_UINT8), + ], + ids=[ + "record_int8", + "record_fp16", + "record_int4", + "uniform4", + "uniform7", + "uniform8", + ], +) +@pytest.mark.parametrize("deleted", [False, True], ids=["direct", "delete_fallback"]) +@pytest.mark.parametrize("index_kind", ["vamana", "hnsw"]) +def test_refine_candidate_counts_and_parameter_switches( + tmp_path, quantizer, flat_type, deleted, index_kind +): + rng = np.random.default_rng(20260911) + vectors = rng.integers(0, 128, size=(1400, 128)).astype(np.float32) + if flat_type == DataType.VECTOR_FP16: + vectors = vectors / 129.3 + index_params = dict( + metric_type=MetricType.L2, + quantize_type=quantizer, + use_contiguous_memory=True, + use_flat_contiguous_memory=True, + flat_data_type=flat_type, + ) + if index_kind == "vamana": + index = VamanaIndexParam( + max_degree=32, search_list_size=100, two_pass_build=True, **index_params + ) + query_param = partial(VamanaQueryParam, ef_search=64) + else: + index = HnswIndexParam(m=32, ef_construction=100, **index_params) + query_param = partial(HnswQueryParam, ef=64) + schema = CollectionSchema( + name="fast_query_refine", + vectors=[ + VectorSchema( + "vector", + DataType.VECTOR_FP32, + dimension=128, + index_param=index, + ) + ], + ) + path = str(tmp_path / "index") + writer = zvec.create_and_open(path, schema) + for start in range(0, len(vectors), 200): + assert all( + s.ok() + for s in writer.insert( + [ + Doc(id=str(i), vectors={"vector": vectors[i].tolist()}) + for i in range(start, min(start + 200, len(vectors))) + ] + ) + ) + writer.optimize() + if deleted: + assert all(s.ok() for s in writer.delete(["0"])) + writer.close() + coll = zvec.open(path, CollectionOption(read_only=True, enable_mmap=True)) + try: + for refine, scale, candidates in [ + (True, None, 64), + (True, 0.0, 64), + (True, 0.5, 10), + (True, 1.0, 10), + (True, 1.1, 11), + (True, 1.9, 19), + (False, 0.0, 10), + (True, 1.4, 14), + (True, 2.6, 26), + (False, 2.6, 10), + ]: + # Default/zero preserves main's max(topk, ef) candidate budget; + # explicit positive factors use max(topk, floor(topk * factor)). + options = {} if scale is None else {"scale_factor": scale} + param = query_param(is_using_refiner=refine, **options) + for row in [11, 89, 531]: + query = np.ascontiguousarray(vectors[row] + np.float32(0.021)) + coarse_param = query_param() + coarse_docs = coll.query( + Query(field_name="vector", vector=query, param=coarse_param), + topk=candidates, + output_fields=[], + ) + coarse_ids = np.asarray( + [int(doc.id) for doc in coarse_docs], dtype=np.int64 + ) + output = coll.fast_query( + "vector", + query, + param, + topk=10, + ) + scored_ids, scores = coll.fast_query( + "vector", + query, + param, + topk=10, + return_scores=True, + ) + np.testing.assert_array_equal(scored_ids, output) + # Ordinary query and fast_query consume exactly the same params. + docs = coll.query( + Query(field_name="vector", vector=query, param=param), topk=10 + ) + np.testing.assert_array_equal(output, [int(doc.id) for doc in docs]) + np.testing.assert_allclose( + scores, [doc.score for doc in docs], rtol=2e-5, atol=2e-5 + ) + if refine: + native_dtype = ( + np.float16 if flat_type == DataType.VECTOR_FP16 else np.uint8 + ) + native_query = query.astype(native_dtype).astype(np.float32) + native_rows = ( + vectors[coarse_ids].astype(native_dtype).astype(np.float32) + ) + distances = np.square(native_rows - native_query).sum(axis=1) + order = np.lexsort((coarse_ids, distances))[: len(output)] + np.testing.assert_array_equal(output, coarse_ids[order]) + np.testing.assert_allclose( + scores, distances[order], rtol=2e-5, atol=2e-5 + ) + else: + np.testing.assert_array_equal(output, coarse_ids) + np.testing.assert_allclose( + scores, [doc.score for doc in coarse_docs], rtol=2e-5, atol=2e-5 + ) + for scale in (-1.0, float("nan"), float("inf")): + param = query_param(is_using_refiner=True, scale_factor=scale) + with pytest.raises((ValueError, RuntimeError)): + coll.query( + Query(field_name="vector", vector=vectors[11], param=param), topk=10 + ) + with pytest.raises((ValueError, RuntimeError)): + coll.fast_query("vector", vectors[11], param) + + finally: + coll.close() diff --git a/python/tests/test_fast_query_uniform_lifecycle.py b/python/tests/test_fast_query_uniform_lifecycle.py new file mode 100644 index 000000000..d04f516cf --- /dev/null +++ b/python/tests/test_fast_query_uniform_lifecycle.py @@ -0,0 +1,124 @@ +"""Fast queries must include Uniform segments that have not been trained yet.""" + +import numpy as np +import pytest + +import zvec + + +@pytest.mark.parametrize("index_kind", ["hnsw", "vamana"]) +@pytest.mark.parametrize( + "quantizer", + [ + zvec.QuantizeType.UNIFORM_UINT7, + zvec.QuantizeType.UNIFORM_UINT8, + zvec.QuantizeType.UNIFORM_UINT4, + ], +) +@pytest.mark.parametrize( + "flat_type", + [zvec.DataType.VECTOR_FP32, zvec.DataType.VECTOR_FP16, zvec.DataType.VECTOR_UINT8], +) +@pytest.mark.parametrize("enable_mmap", [False, True]) +def test_untrained_and_mixed_uniform_segments( + tmp_path, index_kind, quantizer, flat_type, enable_mmap +): + zvec.init() + vectors = np.random.default_rng(754).uniform(1, 240, (65, 17)).astype(np.float32) + index_type, query_type = ( + (zvec.HnswIndexParam, zvec.HnswQueryParam) + if index_kind == "hnsw" + else (zvec.VamanaIndexParam, zvec.VamanaQueryParam) + ) + schema = zvec.CollectionSchema( + name="fast_uniform_lifecycle", + vectors=[ + zvec.VectorSchema( + "vector", + zvec.DataType.VECTOR_FP32, + 17, + index_param=index_type( + metric_type=zvec.MetricType.L2, + quantize_type=quantizer, + flat_data_type=flat_type, + use_flat_contiguous_memory=True, + ), + ) + ], + ) + path = str(tmp_path / "collection") + writer_option = zvec.CollectionOption(enable_mmap=enable_mmap) + writer = zvec.create_and_open(path, schema, writer_option) + + def doc(i): + return zvec.Doc(id=f"row-{i}", vectors={"vector": vectors[i].tolist()}) + + def check_reader(count): + reader = zvec.open( + path, zvec.CollectionOption(read_only=True, enable_mmap=enable_mmap) + ) + try: + for param in ( + None, + query_type(), + query_type(is_using_refiner=True), + query_type(is_using_refiner=True, scale_factor=2.0), + query_type(is_linear=True), + query_type(is_linear=True, is_using_refiner=True), + ): + for row in (12, count - 1): + query = vectors[row] + for topk in (10, count + 3): + docs = reader.query( + zvec.Query("vector", vector=query, param=param), topk=topk + ) + expected = np.array([int(d.id[4:]) for d in docs]) + ids, scores = reader.fast_query( + "vector", query, param, topk=topk, return_scores=True + ) + assert docs[0].id == f"row-{row}" + actual = ids[: len(docs)] + assert set(actual) == set(expected) + # SQL and fast_query merge segments differently. Both + # order by score, without a shared tie-break rule. Only + # exactly equal scores may exchange positions. + expected_scores = {int(d.id[4:]): d.score for d in docs} + np.testing.assert_array_equal( + [expected_scores[i] for i in actual], + [d.score for d in docs], + ) + np.testing.assert_allclose( + scores[: len(docs)], + [expected_scores[i] for i in actual], + rtol=2e-5, + atol=2e-5, + ) + np.testing.assert_array_equal( + ids, reader.fast_query("vector", query, param, topk=topk) + ) + assert np.all(ids[len(docs) :] == -1) + assert np.all(np.isnan(scores[len(docs) :])) + if topk > count: + assert set(expected) == set(range(count)) + finally: + reader.close() + + try: + assert all(s.ok() for s in writer.insert([doc(i) for i in range(64)])) + writer.close() + check_reader(64) + writer = zvec.open(path, writer_option) + writer.optimize() + writer.close() + check_reader(64) + writer = zvec.open(path, writer_option) + assert writer.insert(doc(64)).ok() + writer.flush() + writer.close() + check_reader(65) + writer = zvec.open(path, writer_option) + writer.optimize() + writer.close() + check_reader(65) + finally: + writer.close() diff --git a/python/zvec/model/collection.py b/python/zvec/model/collection.py index 33620dd3f..e382033d8 100644 --- a/python/zvec/model/collection.py +++ b/python/zvec/model/collection.py @@ -509,6 +509,43 @@ def iter_docs( # ========== Collection DQL-Query Methods ========== + def fast_query( + self, + field_name: str, + vector, + param=None, + topk: int = 10, + return_scores: bool = False, + ): + """Query a dense field directly, returning internal numeric IDs. + + This advanced API requires a read-only collection. No preparation is + required: parameters may be constructed inline or reused across calls. + Index references are cached internally; parameters are read each time. + Calls to ``fast_query`` and collection close must be serial. + + ``vector`` must be a contiguous 1D NumPy array matching the field's input + dtype and dimension. The result is an owning int64 array. With + ``return_scores=True``, returns ``(ids, scores)`` with float32 scores, + including refinement when enabled. Missing results are padded with + ID -1 / score NaN. Refinement uses ``param.scale_factor`` with the same + candidate-count semantics as :meth:`query`. + + Use :meth:`query` for user string IDs, scalar filters, sparse queries, + group-by or fetching fields and vectors. Internal IDs must not be + stored across collection mutations or compaction. + + Examples: + >>> ids = collection.fast_query("vector", vector, param, topk=10) + >>> ids, scores = collection.fast_query( + ... "vector", vector, param, topk=10, return_scores=True + ... ) + """ + if self._obj is None: + msg = "fast query collection is closed" + raise ValueError(msg) + return self._obj.fast_query(field_name, vector, param, topk, return_scores) + def query( self, queries: Optional[Union[Query, list[Query]]] = None, diff --git a/src/binding/python/model/python_collection.cc b/src/binding/python/model/python_collection.cc index 2fe31ef7c..71a12a2b9 100644 --- a/src/binding/python/model/python_collection.cc +++ b/src/binding/python/model/python_collection.cc @@ -13,6 +13,8 @@ // limitations under the License. #include "python_collection.h" +#include +#include #include #include #include @@ -23,6 +25,41 @@ namespace zvec { namespace { +DenseQueryShape dense_query_shape(const py::array &vector) { + if (vector.ndim() != 1 || !(vector.flags() & py::array::c_style)) { + throw py::value_error("query vector must be a contiguous 1D array"); + } + if (vector.shape(0) > std::numeric_limits::max()) { + throw py::value_error("query vector dimension is too large"); + } + // Read NumPy metadata directly instead of allocating a temporary buffer + // descriptor, shape/stride vectors and a PEP 3118 format string per query. + const auto dtype = vector.dtype(); + const uint16_t endian_probe = 1; + const char native_order = + *reinterpret_cast(&endian_probe) ? '<' : '>'; + const char order = dtype.byteorder(); + if (order != '=' && order != '|' && order != native_order) { + throw py::value_error("query vector requires a native numeric dtype"); + } + const auto size = dtype.itemsize(); + DataType type; + if (dtype.kind() == 'f' && size == 4) { + type = DataType::VECTOR_FP32; + } else if (dtype.kind() == 'f' && size == 8) { + type = DataType::VECTOR_FP64; + } else if (dtype.kind() == 'f' && size == 2) { + type = DataType::VECTOR_FP16; + } else if (dtype.kind() == 'i' && size == 1) { + type = DataType::VECTOR_INT8; + } else if (dtype.kind() == 'u' && size == 1) { + type = DataType::VECTOR_UINT8; + } else { + throw py::value_error("unsupported query vector dtype"); + } + return {type, static_cast(vector.shape(0))}; +} + // Batch-materialize a DocPtrList into a list of (id, score, fields, vectors) // tuples in a single GIL-held section, avoiding per-doc _Doc wrappers and // per-doc Python->C++ crossings on the hot query path. The forward/vector field @@ -98,6 +135,19 @@ py::list execute_for_python(const Collection &collection, const Query &query) { return docs_to_tuples(snapshot.docs, *snapshot.schema); } +// Transfer ownership of the native buffer without a per-result Python loop. +template +py::array_t owned_array(std::vector values) { + auto buffer = std::make_unique>(std::move(values)); + const auto size = static_cast(buffer->size()); + auto *ptr = buffer.get(); + py::capsule owner(ptr, + [](void *p) { delete static_cast *>(p); }); + buffer.release(); + return py::array_t({size}, {static_cast(sizeof(T))}, + ptr->data(), owner); +} + void ZVecPyCollection::Initialize(pybind11::module_ &m) { py::class_(m, "_GroupResult") .def_readonly("group_by_value", &GroupResult::group_by_value_) @@ -340,6 +390,40 @@ void ZVecPyCollection::bind_dql_methods( "Execute a multi query with re-ranking and return results as a " "list of (id, score, fields, vectors) tuples materialized in one " "batch.") + .def( + "fast_query", + [](const Collection &self, const std::string &field_name, + const py::array &vector, QueryParams *params, int topk, + bool return_scores) -> py::object { + const auto shape = dense_query_shape(vector); + // Python keeps the argument alive for this call. The DB reads it + // synchronously and never retains it, so no shared ownership + // conversion or reference-count traffic is needed here. + const QueryParams::Ptr borrowed(QueryParams::Ptr{}, params); + Result result; + { + py::gil_scoped_release release; + result = self.fast_query(field_name, vector.data(), borrowed, + topk, return_scores, &shape); + } + auto output = unwrap_expected(std::move(result)); + auto ids = owned_array(std::move(output.ids)); + if (!return_scores) return ids; + return py::make_tuple(std::move(ids), + owned_array(std::move(output.scores))); + }, + py::arg("field_name"), py::arg("vector").noconvert(), + py::arg("param") = static_cast(nullptr), + py::arg("topk") = 10, py::arg("return_scores") = false, + R"doc(Advanced dense query on a read-only collection, with no preparation step. + +Pass a contiguous 1D NumPy vector and query parameters on each call. Parameters +may be constructed inline or reused. Returns an owning int64 internal ID array, +or (ids, float32 scores) with return_scores=True. Scores include refinement. +Missing results are padded with ID -1 / score NaN. Refinement uses param.scale_factor +with the same semantics as Collection.query. +Calls to fast_query and collection.close must be serial. +)doc") .def("GroupByQuery", [](const Collection &self, const GroupByVectorQuery &query) { Result result; diff --git a/src/core/algorithm/flat/flat_streamer.cc b/src/core/algorithm/flat/flat_streamer.cc index d233f425e..dc47710a9 100644 --- a/src/core/algorithm/flat/flat_streamer.cc +++ b/src/core/algorithm/flat/flat_streamer.cc @@ -449,6 +449,23 @@ int FlatStreamer::search_bf_by_p_keys_impl( return 0; } +template +int FlatStreamer::search_by_p_keys_fast( + const void *query, const std::vector &keys, int64_t *output_ids, + float *output_scores, size_t topk, const IndexQueryMeta &qmeta, + Context::UPointer &context) const { + if (!query || !output_ids || topk == 0 || !context || + !metric_->is_matched(meta_, qmeta)) { + return IndexError_InvalidArgument; + } + auto *flat_context = + dynamic_cast *>(context.get()); + if (!flat_context) return IndexError_InvalidArgument; + if (flat_context->magic() != magic_) flat_context->reset(this); + return entity_->search_by_p_keys_fast(query, keys, output_ids, output_scores, + topk, flat_context->search_scratch()); +} + template int FlatStreamer::group_by_search_impl( const void *query, const IndexQueryMeta &qmeta, uint32_t count, @@ -545,6 +562,14 @@ int FlatStreamer::group_by_search_p_keys_impl( return 0; } +// The direct operation is not virtual on the streamer. +template int FlatStreamer<16>::search_by_p_keys_fast( + const void *, const std::vector &, int64_t *, float *, size_t, + const IndexQueryMeta &, Context::UPointer &) const; +template int FlatStreamer<32>::search_by_p_keys_fast( + const void *, const std::vector &, int64_t *, float *, size_t, + const IndexQueryMeta &, Context::UPointer &) const; + INDEX_FACTORY_REGISTER_STREAMER_ALIAS(LinearStreamer, FlatStreamer<32>); INDEX_FACTORY_REGISTER_STREAMER_ALIAS(FlatStreamer, FlatStreamer<32>); INDEX_FACTORY_REGISTER_STREAMER_ALIAS(FlatStreamer16, FlatStreamer<16>); diff --git a/src/core/algorithm/flat/flat_streamer.h b/src/core/algorithm/flat/flat_streamer.h index 26823172d..056584068 100644 --- a/src/core/algorithm/flat/flat_streamer.h +++ b/src/core/algorithm/flat/flat_streamer.h @@ -97,6 +97,14 @@ class FlatStreamer : public IndexStreamer { const IndexQueryMeta &qmeta, uint32_t count, Context::UPointer &context) const override; + // Candidate-only output for a contiguous reference. Returns NotImplemented + // without writing output if its storage cannot serve every requested key. + int search_by_p_keys_fast(const void *query, + const std::vector &keys, + int64_t *output_ids, float *output_scores, + size_t topk, const IndexQueryMeta &qmeta, + Context::UPointer &context) const; + int group_by_search_impl(const void *query, const IndexQueryMeta &qmeta, uint32_t count, Context::UPointer &context) const; diff --git a/src/core/algorithm/flat/flat_streamer_entity.cc b/src/core/algorithm/flat/flat_streamer_entity.cc index 6f0e5ac86..aeaf3c360 100644 --- a/src/core/algorithm/flat/flat_streamer_entity.cc +++ b/src/core/algorithm/flat/flat_streamer_entity.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "flat_streamer_entity.h" +#include #include #include #include "flat_utility.h" @@ -596,6 +597,88 @@ int FlatContiguousStreamerEntity::search_by_p_keys( scratch, nullptr, heap); } +int FlatStreamerEntity::search_by_p_keys_fast( + const void * /*query*/, const std::vector & /*keys*/, + int64_t * /*output_ids*/, float * /*output_scores*/, size_t /*topk*/, + FlatSearchScratch * /*scratch*/) const { + return IndexError_NotImplemented; +} + +int FlatContiguousStreamerEntity::search_by_p_keys_fast( + const void *query, const std::vector &keys, int64_t *output_ids, + float *output_scores, size_t topk, FlatSearchScratch *scratch) const { + if (!query || !output_ids || topk == 0 || !scratch) { + return IndexError_InvalidArgument; + } + // Hold the same immutable generation for pointer lookup and evaluation. + // A concurrent degrade/add cannot free its rows while this call is active. + auto storage = load_contiguous_storage(); + if (!storage || !storage->vector_memory || keys.size() < topk) { + return IndexError_NotImplemented; + } + auto &ptrs = scratch->vector_ptrs; + auto &extras = scratch->extra_values; + auto &distances = scratch->distances; + auto &documents = scratch->candidate_documents; + const size_t count = keys.size(); + const size_t extra_size = extra_values_size(); + const bool has_extras = extra_size != 0; + const size_t vector_data_size = meta().element_size() - extra_size; + ptrs.resize(count); + extras.resize(has_extras ? count : 0); + distances.resize(count); + documents.resize(count); + for (size_t i = 0; i < count; ++i) { + ptrs[i] = get_vector_ptr_by_key(*storage, keys[i]); + if (!ptrs[i]) return IndexError_NotImplemented; + if (has_extras) { + extras[i] = static_cast(ptrs[i]) + vector_data_size; + } + } + // Use the native reference's existing metric dispatch, including any + // batch-query preprocessing and per-record extra values. + const void *batch_query = query; + if (const auto &preprocess = batch_query_preprocess(); + preprocess != nullptr) { + auto &buffer = scratch->query_buffer; + buffer.resize(meta().element_size()); + std::memcpy(buffer.data(), query, buffer.size()); + preprocess(buffer.data(), meta().dimension()); + batch_query = buffer.data(); + } + if (const auto &batch = batch_distance(); batch) { + batch(ptrs.data(), batch_query, count, meta().dimension(), distances.data(), + has_extras ? extras.data() : nullptr); + } else { + for (size_t i = 0; i < count; ++i) { + distance()(query, ptrs[i], meta().dimension(), distances.data() + i); + } + } + for (size_t i = 0; i < count; ++i) { + documents[i] = {keys[i], distances[i]}; + } + const auto better = [](const auto &lhs, const auto &rhs) { + return lhs.distance < rhs.distance || + (lhs.distance == rhs.distance && lhs.key < rhs.key); + }; + auto selected_end = documents.begin() + topk; + std::make_heap(documents.begin(), selected_end, better); + for (auto candidate = selected_end; candidate != documents.end(); + ++candidate) { + if (better(*candidate, documents.front())) { + std::pop_heap(documents.begin(), selected_end, better); + *(selected_end - 1) = *candidate; + std::push_heap(documents.begin(), selected_end, better); + } + } + std::sort_heap(documents.begin(), selected_end, better); + for (size_t i = 0; i < topk; ++i) { + output_ids[i] = static_cast(documents[i].key); + if (output_scores) output_scores[i] = documents[i].distance; + } + return 0; +} + int FlatContiguousStreamerEntity::build_contiguous_memory() { degrade_to_mmap(); diff --git a/src/core/algorithm/flat/flat_streamer_entity.h b/src/core/algorithm/flat/flat_streamer_entity.h index 80f4e01ac..4dbcfa653 100644 --- a/src/core/algorithm/flat/flat_streamer_entity.h +++ b/src/core/algorithm/flat/flat_streamer_entity.h @@ -34,6 +34,11 @@ namespace core { //! Reusable request-local buffers for storage-specific Flat search paths. struct FlatSearchScratch { + struct CandidateDocument { + uint64_t key; + float distance; + }; + std::vector candidate_documents{}; std::vector vector_ptrs{}; std::vector extra_values{}; std::vector vector_keys{}; @@ -83,6 +88,15 @@ class FlatStreamerEntity { FlatSearchScratch *scratch = nullptr, size_t batch_size = 0) const; + // A storage-specific candidate-only operation. Generic layouts request + // the streamer's existing full-result fallback without touching output. + // Scores use the metric's internal representation and are optional. + virtual int search_by_p_keys_fast(const void *query, + const std::vector &keys, + int64_t *output_ids, float *output_scores, + size_t topk, + FlatSearchScratch *scratch) const; + //! Search in a block void search_block(const void *query, const BlockLocation &bl, const BlockHeader *hd, float norm_val, @@ -517,6 +531,12 @@ class FlatContiguousStreamerEntity : public FlatStreamerEntity { FlatSearchScratch *scratch, size_t batch_size) const override; + int search_by_p_keys_fast(const void *query, + const std::vector &keys, + int64_t *output_ids, float *output_scores, + size_t topk, + FlatSearchScratch *scratch) const override; + bool is_contiguous() const { return !!load_contiguous_storage(); } diff --git a/src/core/algorithm/hnsw/hnsw_context.h b/src/core/algorithm/hnsw/hnsw_context.h index 08eaf8a74..ba5f6ecb2 100644 --- a/src/core/algorithm/hnsw/hnsw_context.h +++ b/src/core/algorithm/hnsw/hnsw_context.h @@ -169,11 +169,21 @@ class HnswContext : public IndexContext { } } - inline void topk_to_keys(std::vector &keys) { + inline void topk_to_keys(std::vector &keys, + std::vector *scores = nullptr) { keys.clear(); keys.reserve((std::min)(static_cast(topk_), search_heap_.size())); - collect_search_result( - [&](node_id_t id, dist_t) { keys.push_back(entity_->get_key(id)); }); + if (scores) { + scores->clear(); + scores->reserve(keys.capacity()); + collect_search_result([&](node_id_t id, dist_t score) { + keys.push_back(entity_->get_key(id)); + scores->push_back(score); + }); + } else { + collect_search_result( + [&](node_id_t id, dist_t) { keys.push_back(entity_->get_key(id)); }); + } } inline void recal_topk_dist() { diff --git a/src/core/algorithm/hnsw/hnsw_streamer.cc b/src/core/algorithm/hnsw/hnsw_streamer.cc index c9b11c450..1542447fe 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer.cc +++ b/src/core/algorithm/hnsw/hnsw_streamer.cc @@ -845,8 +845,10 @@ int HnswStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, int HnswStreamer::search_candidates_impl( const void *query, const IndexQueryMeta &qmeta, std::vector &keys, + std::vector *scores, IndexStreamer::Context::Pointer &context) const { keys.clear(); + if (scores) scores->clear(); int ret = check_params(query, qmeta); if (ailego_unlikely(ret != 0)) { return ret; @@ -861,7 +863,8 @@ int HnswStreamer::search_candidates_impl( } if (entity_->doc_cnt() <= ctx->get_bruteforce_threshold()) { - return IndexRunner::search_candidates_impl(query, qmeta, keys, context); + return IndexRunner::search_candidates_impl(query, qmeta, keys, scores, + context); } if (ctx->magic() != magic_) { @@ -880,10 +883,11 @@ int HnswStreamer::search_candidates_impl( LOG_ERROR("Hnsw searcher fast search failed"); return ret; } - ctx->topk_to_keys(keys); + ctx->topk_to_keys(keys, scores); if (ailego_unlikely(ctx->error())) { keys.clear(); + if (scores) scores->clear(); return IndexError_Runtime; } return 0; diff --git a/src/core/algorithm/hnsw/hnsw_streamer.h b/src/core/algorithm/hnsw/hnsw_streamer.h index 91ee055fd..2aa19e4e6 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer.h +++ b/src/core/algorithm/hnsw/hnsw_streamer.h @@ -122,6 +122,7 @@ class HnswStreamer : public IndexStreamer { int search_candidates_impl(const void *query, const IndexQueryMeta &qmeta, std::vector &keys, + std::vector *scores, Context::Pointer &context) const override; //! Similarity brute force search diff --git a/src/core/algorithm/vamana/vamana_context.cc b/src/core/algorithm/vamana/vamana_context.cc index a5ada5a7d..5c7693a26 100644 --- a/src/core/algorithm/vamana/vamana_context.cc +++ b/src/core/algorithm/vamana/vamana_context.cc @@ -209,11 +209,21 @@ void VamanaContext::topk_to_result(uint32_t idx) { }); } -void VamanaContext::topk_to_keys(std::vector &keys) { +void VamanaContext::topk_to_keys(std::vector &keys, + std::vector *scores) { keys.clear(); keys.reserve((std::min)(static_cast(topk_), search_heap_.size())); - collect_search_result( - [&](node_id_t id, dist_t) { keys.push_back(entity_->get_key(id)); }); + if (scores) { + scores->clear(); + scores->reserve(keys.capacity()); + collect_search_result([&](node_id_t id, dist_t score) { + keys.push_back(entity_->get_key(id)); + scores->push_back(score); + }); + } else { + collect_search_result( + [&](node_id_t id, dist_t) { keys.push_back(entity_->get_key(id)); }); + } } void VamanaContext::fill_random_to_topk_full() { diff --git a/src/core/algorithm/vamana/vamana_context.h b/src/core/algorithm/vamana/vamana_context.h index 0546257db..540d55f45 100644 --- a/src/core/algorithm/vamana/vamana_context.h +++ b/src/core/algorithm/vamana/vamana_context.h @@ -118,7 +118,8 @@ class VamanaContext : public IndexContext { void topk_to_result(uint32_t idx); - void topk_to_keys(std::vector &keys); + void topk_to_keys(std::vector &keys, + std::vector *scores = nullptr); inline void reset_query(const void *query) { if (auto query_preprocess_func = index_metric_->get_query_preprocess_func(); diff --git a/src/core/algorithm/vamana/vamana_streamer.cc b/src/core/algorithm/vamana/vamana_streamer.cc index 7a2fbdae6..8cf3c8af9 100644 --- a/src/core/algorithm/vamana/vamana_streamer.cc +++ b/src/core/algorithm/vamana/vamana_streamer.cc @@ -754,8 +754,10 @@ int VamanaStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, int VamanaStreamer::search_candidates_impl(const void *query, const IndexQueryMeta &qmeta, std::vector &keys, + std::vector *scores, Context::Pointer &context) const { keys.clear(); + if (scores) scores->clear(); int ret = check_params(query, qmeta); if (ailego_unlikely(ret != 0)) return ret; @@ -769,7 +771,8 @@ int VamanaStreamer::search_candidates_impl(const void *query, } if (entity_->doc_cnt() <= ctx->get_bruteforce_threshold()) { - return IndexRunner::search_candidates_impl(query, qmeta, keys, context); + return IndexRunner::search_candidates_impl(query, qmeta, keys, scores, + context); } if (ctx->magic() != magic_) { @@ -788,10 +791,11 @@ int VamanaStreamer::search_candidates_impl(const void *query, LOG_ERROR("Vamana search failed"); return ret; } - ctx->topk_to_keys(keys); + ctx->topk_to_keys(keys, scores); if (ailego_unlikely(ctx->error())) { keys.clear(); + if (scores) scores->clear(); return IndexError_Runtime; } return 0; diff --git a/src/core/algorithm/vamana/vamana_streamer.h b/src/core/algorithm/vamana/vamana_streamer.h index 4e6c55e09..165a80669 100644 --- a/src/core/algorithm/vamana/vamana_streamer.h +++ b/src/core/algorithm/vamana/vamana_streamer.h @@ -61,6 +61,7 @@ class VamanaStreamer : public IndexStreamer { int search_candidates_impl(const void *query, const IndexQueryMeta &qmeta, std::vector &keys, + std::vector *scores, Context::Pointer &context) const override; int search_bf_impl(const void *query, const IndexQueryMeta &qmeta, diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc index b142a62d8..09f0b157e 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -22,6 +22,7 @@ #include #include #include +#include "algorithm/flat/flat_streamer.h" #include "mixed_reducer/mixed_reducer_params.h" #include "utility/utility_params.h" @@ -801,6 +802,208 @@ int Index::search(const VectorData &vector_data, } +int Index::search_fast(const VectorData &vector_data, + const BaseIndexQueryParam::Pointer &search_param, + int64_t *output_ids, float *output_scores) { + if (!output_ids || !search_param || search_param->topk == 0) { + return core::IndexError_InvalidArgument; + } + const size_t topk = search_param->topk; + if (!is_open_) { + LOG_ERROR("Index is not open"); + return core::IndexError_Runtime; + } + if (is_sparse_ || has_group_by_search(search_param)) { + return core::IndexError_Unsupported; + } + + if (search_param->refiner_param) { + const int direct_ret = _search_refine_fast(vector_data, search_param, + output_ids, output_scores); + if (direct_ret != core::IndexError_NotImplemented) return direct_ret; + // Keep the public path for filtered/non-contiguous references and modes + // that cannot use the candidate-only operation. + thread_local SearchResult result; + const int ret = search(vector_data, search_param, &result); + if (ret != 0) return ret; + const size_t count = + std::min(static_cast(topk), result.doc_list_.size()); + for (size_t i = 0; i < count; ++i) { + output_ids[i] = static_cast(result.doc_list_[i].key()); + if (output_scores) output_scores[i] = result.doc_list_[i].score(); + } + std::fill(output_ids + count, output_ids + topk, int64_t{-1}); + if (output_scores) { + std::fill(output_scores + count, output_scores + topk, + std::numeric_limits::quiet_NaN()); + } + return 0; + } + + if (!is_trained_ && train() != 0) { + LOG_ERROR("Failed to train index"); + return core::IndexError_Runtime; + } + auto &context = acquire_context(); + if (!context) return core::IndexError_Runtime; + + int ret = _prepare_for_search(vector_data, search_param, context); + std::string transformed_query; + const void *query = nullptr; + core::IndexQueryMeta query_meta; + if (ret == 0) { + ret = _prepare_dense_query(vector_data, &transformed_query, &query, + &query_meta); + } + thread_local std::vector keys; + std::vector *scores = nullptr; + if (output_scores) { + thread_local std::vector score_buffer; + scores = &score_buffer; + } + if (ret == 0) { + ret = _execute_dense_search(query, query_meta, search_param, context, &keys, + scores); + } + if (ret == 0) { + const size_t count = std::min(static_cast(topk), keys.size()); + if (output_scores && scores->size() != keys.size()) { + ret = core::IndexError_Runtime; + } else { + for (size_t i = 0; i < count; ++i) { + output_ids[i] = static_cast(keys[i]); + if (output_scores) output_scores[i] = (*scores)[i]; + } + if (output_scores) { + ret = _normalize_buffer_scores(vector_data, output_ids, output_scores, + count); + } + } + if (ret == 0) { + std::fill(output_ids + count, output_ids + topk, int64_t{-1}); + if (output_scores) { + std::fill(output_scores + count, output_scores + topk, + std::numeric_limits::quiet_NaN()); + } + } + } + if (context) context->reset(); + return ret; +} + + +int Index::_normalize_buffer_scores(const VectorData &vector_data, + const int64_t *output_ids, + float *output_scores, size_t count) { + if (!output_scores || count == 0) return 0; + if (metric_->support_normalize()) { + for (size_t i = 0; i < count; ++i) { + metric_->normalize(output_scores + i); + } + } + if (!reformer_) return 0; + if (!std::holds_alternative(vector_data.vector)) { + return core::IndexError_Runtime; + } + core::IndexDocumentList documents; + documents.reserve(count); + for (size_t i = 0; i < count; ++i) { + documents.emplace_back(static_cast(output_ids[i]), + output_scores[i]); + } + const auto &dense_vector = std::get(vector_data.vector); + const int ret = + reformer_->normalize(dense_vector.data, input_vector_meta_, documents); + if (ret != 0) return ret; + for (size_t i = 0; i < count; ++i) { + output_scores[i] = documents[i].score(); + } + return 0; +} + + +int Index::_search_refine_fast(const VectorData &vector_data, + const BaseIndexQueryParam::Pointer &search_param, + int64_t *output_ids, float *output_scores) { + auto &reference = search_param->refiner_param->reference_index; + // Public Search remains responsible for these modes and their diagnostics. + // In particular, do not consume/move a user filter before taking fallback. + if (!reference || reference.get() == this || !reference->is_open_ || + reference->is_sparse_ || + reference->param_.index_type != IndexType::kFlat || + search_param->fetch_vector || + (search_param->filter && search_param->filter->is_valid())) { + return core::IndexError_NotImplemented; + } + const size_t topk = search_param->topk; + // Score-only normalization is known to preserve membership for raw L2/IP + // references. Other metrics/quantizers retain the public normalization path. + const auto &reference_quantizer = reference->param_.quantizer_param; + if ((reference->param_.metric_type != MetricType::kL2sq && + reference->param_.metric_type != MetricType::kInnerProduct) || + (reference_quantizer && + reference_quantizer->type != QuantizerType::kNone && + reference_quantizer->type != QuantizerType::kFP16)) { + return core::IndexError_NotImplemented; + } + const auto *flat = + dynamic_cast *>(reference->streamer_.get()); + if (!flat) return core::IndexError_NotImplemented; + if ((!is_trained_ && train() != 0) || + (!reference->is_trained_ && reference->train() != 0)) { + return core::IndexError_Runtime; + } + auto &context = acquire_context(); + if (!context) return core::IndexError_Runtime; + int ret = _prepare_for_search(vector_data, search_param, context); + const int coarse_topk = _get_coarse_search_topk(search_param); + if (ret != 0 || coarse_topk < 0) { + context->reset(); + return ret != 0 ? ret : coarse_topk; + } + context->set_topk(coarse_topk); + context->set_fetch_vector(false); + // As with the existing candidate vector, this scratch assumes a serial, + // non-reentrant query on each thread; this route invokes no user filters or + // group-by callbacks. + thread_local std::string transformed_query; + const void *query = nullptr; + core::IndexQueryMeta query_meta; + ret = _prepare_dense_query(vector_data, &transformed_query, &query, + &query_meta); + // The thread owns these buffers, so queries on different Index instances + // may retain capacity without retaining pointers into a previous index. + thread_local std::vector keys; + if (ret == 0) { + ret = + _execute_dense_search(query, query_meta, search_param, context, &keys); + } + if (ret == 0) { + auto &reference_context = reference->acquire_context(); + if (!reference_context) { + ret = core::IndexError_Runtime; + } else { + thread_local std::string reference_query_storage; + ret = reference->_prepare_dense_query( + vector_data, &reference_query_storage, &query, &query_meta); + if (ret == 0) { + ret = flat->search_by_p_keys_fast( + query, keys, output_ids, output_scores, static_cast(topk), + query_meta, reference_context); + if (ret == 0 && output_scores) { + ret = reference->_normalize_buffer_scores(vector_data, output_ids, + output_scores, + static_cast(topk)); + } + } + reference_context->reset(); + } + } + context->reset(); + return ret; +} + + int Index::_dense_fetch(const uint32_t doc_id, VectorDataBuffer *vector_data_buffer) { core::IndexStorage::MemoryBlock vector_block; @@ -1009,9 +1212,10 @@ int Index::_prepare_dense_query(const VectorData &vector_data, int Index::_execute_dense_search( const void *vector, const core::IndexQueryMeta &new_meta, const BaseIndexQueryParam::Pointer &search_param, - core::IndexContext::Pointer &context, - std::vector *candidate_keys) { + core::IndexContext::Pointer &context, std::vector *candidate_keys, + std::vector *candidate_scores) { if (candidate_keys) candidate_keys->clear(); + if (candidate_scores) candidate_scores->clear(); if (search_param->bf_pks != nullptr) { if (streamer_->search_bf_by_p_keys_impl( vector, std::vector>{*search_param->bf_pks}, @@ -1026,7 +1230,7 @@ int Index::_execute_dense_search( } } else if (candidate_keys) { return streamer_->search_candidates_impl(vector, new_meta, *candidate_keys, - context); + candidate_scores, context); } else { if (streamer_->search_impl(vector, new_meta, 1, context) != 0) { LOG_ERROR("Failed to search vector"); @@ -1037,8 +1241,10 @@ int Index::_execute_dense_search( if (candidate_keys) { const auto &documents = context->result(); candidate_keys->reserve(documents.size()); + if (candidate_scores) candidate_scores->reserve(documents.size()); for (const auto &document : documents) { candidate_keys->push_back(document.key()); + if (candidate_scores) candidate_scores->push_back(document.score()); } } return 0; diff --git a/src/db/collection.cc b/src/db/collection.cc index a96964494..004462748 100644 --- a/src/db/collection.cc +++ b/src/db/collection.cc @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include #include #include @@ -20,6 +21,7 @@ #include #include #include +#include #include #include #include @@ -44,6 +46,7 @@ #include "db/common/typedef.h" #include "db/common/utils.h" #include "db/doc_iterator_internal.h" +#include "db/index/column/vector_column/engine_helper.hpp" #include "db/index/common/delete_store.h" #include "db/index/common/id_map.h" #include "db/index/common/identifier_validation.h" @@ -132,6 +135,11 @@ class CollectionImpl : public Collection { Result query(const MultiQuery &query) const override; + Result fast_query( + const std::string &field_name, const void *query_vector, + const QueryParams::Ptr &query_params, int topk, bool return_scores, + const DenseQueryShape *query_shape) const override; + Result group_by_query( const GroupByVectorQuery &query) const override; @@ -291,6 +299,31 @@ class CollectionImpl : public Collection { ColumnOp op); private: + struct FastQuerySegment { + Segment::Ptr segment; + CombinedVectorColumnIndexer::Ptr indexer; + IndexFilter::Ptr filter; + }; + + struct FastQueryState { + DataType data_type{DataType::UNDEFINED}; + uint32_t dimension{0}; + IndexType index_type{IndexType::UNDEFINED}; + MetricType metric{MetricType::L2}; + std::vector segments; + core_interface::Index::Pointer raw_index; + core_interface::Index::Pointer raw_reference_index; + core_interface::BaseIndexQueryParam::Pointer engine_query_param; + core_interface::BaseIndexQueryParam::Pointer default_engine_query_param; + core_interface::RefinerParam::Pointer refiner_param; + }; + + Result resolve_fast_query( + const std::string &field_name) const; + Status update_fast_query_params(FastQueryState &state, + const QueryParams::Ptr ¶ms) const; + // Only fast_query uses these caches. Its calls and close must be serial. + mutable std::unordered_map fast_query_states_; std::string path_; bool destroyed_{false}; @@ -457,6 +490,8 @@ Status CollectionImpl::close_unsafe() { } } + fast_query_states_.clear(); + // always release resources regardless of flush outcome writing_segment_.reset(); segment_manager_.reset(); @@ -1816,6 +1851,339 @@ Result CollectionImpl::query(const SearchQuery &query) const { return query_unsafe(query); } +Result CollectionImpl::resolve_fast_query( + const std::string &field_name) const { + std::shared_lock schema_lock(schema_handle_mtx_); + const FieldSchema *field_schema = + field_name.empty() ? nullptr : schema_->get_field(field_name); + if (field_schema == nullptr || !field_schema->is_dense_vector()) { + return tl::make_unexpected(Status::InvalidArgument( + "fast search requires a dense vector field: ", field_name)); + } + + bool quantized = false; + MetricType metric = MetricType::L2; + if (const auto *vp = dynamic_cast( + field_schema->index_params().get())) { + metric = vp->metric_type(); + quantized = vp->quantize_type() != QuantizeType::UNDEFINED; + } + + FastQueryState cache; + cache.data_type = field_schema->data_type(); + cache.dimension = field_schema->dimension(); + cache.index_type = field_schema->index_params() + ? field_schema->index_params()->type() + : IndexType::UNDEFINED; + cache.metric = metric; + const auto segments = get_all_segments(); + cache.segments.reserve(segments.size()); + for (const auto &seg : segments) { + CombinedVectorColumnIndexer::Ptr indexer; + if (!quantized) { + indexer = seg->get_combined_vector_indexer(field_name); + } else { + indexer = seg->get_quant_combined_vector_indexer(field_name); + if (indexer != nullptr && !indexer->has_searchable_indexers()) { + indexer = seg->get_combined_vector_indexer(field_name); + } + } + if (!indexer || !indexer->has_searchable_indexers()) { + continue; + } + FastQuerySegment entry; + entry.segment = seg; + entry.indexer = std::move(indexer); + entry.filter = seg->get_filter(); + cache.segments.push_back(std::move(entry)); + } + + if (cache.segments.empty()) { + return tl::make_unexpected(Status::InvalidArgument( + "fast query: no searchable vector index for field ", field_name)); + } + + if (cache.segments.size() == 1) { + const auto &entry = cache.segments[0]; + if (entry.filter == nullptr && entry.indexer->is_single_block() && + entry.segment->has_identity_doc_ids()) { + auto primary = entry.indexer->primary_indexer(); + if (primary) { + cache.raw_index = primary->debug_get_index(); + auto reference = entry.indexer->reference_indexer(); + if (reference) { + cache.raw_reference_index = reference->debug_get_index(); + } + } + } + } + + if (!cache.raw_index) return cache; + auto primary = cache.segments[0].indexer->primary_indexer(); + vector_column_params::QueryParams qp; + qp.data_type = cache.data_type; + qp.dimension = cache.dimension; + auto engine_qp = ProximaEngineHelper::convert_to_engine_query_param( + primary->field_schema(), qp); + if (!engine_qp) return tl::make_unexpected(engine_qp.error()); + cache.engine_query_param = std::move(engine_qp.value()); + cache.default_engine_query_param = cache.engine_query_param->clone(); + if (cache.raw_reference_index) { + cache.refiner_param = std::make_shared(); + cache.refiner_param->reference_index = cache.raw_reference_index; + } + return cache; +} + +namespace { +template +bool UpdateFastQueryParam(const QueryParams::Ptr ¶ms, + core_interface::BaseIndexQueryParam *engine, + const core_interface::BaseIndexQueryParam *defaults, + Update update) { + const auto *p = dynamic_cast(params.get()); + if (params && !p) return false; + if (auto *e = dynamic_cast(engine)) { + update(p, e, static_cast(defaults)); + } + return true; +} +} // namespace + +Status CollectionImpl::update_fast_query_params( + FastQueryState &state, const QueryParams::Ptr ¶ms) const { + if (params && params->type() != state.index_type) { + return Status::InvalidArgument( + "fast query: query parameter type does not match the field index"); + } + auto *engine = state.engine_query_param.get(); + const auto *defaults = state.default_engine_query_param.get(); + bool valid = false; + switch (state.index_type) { + case IndexType::VAMANA: + valid = UpdateFastQueryParam( + params, engine, defaults, [](const auto *p, auto *e, const auto *d) { + e->ef_search = p ? p->ef_search() : d->ef_search; + e->prefetch_offset = p ? p->prefetch_offset() : d->prefetch_offset; + e->prefetch_lines = p ? p->prefetch_lines() : d->prefetch_lines; + }); + break; + case IndexType::HNSW: + valid = + UpdateFastQueryParam( + params, engine, defaults, + [](const auto *p, auto *e, const auto *d) { + e->ef_search = p ? p->ef() : d->ef_search; + e->prefetch_offset = + p ? p->prefetch_offset() : d->prefetch_offset; + e->prefetch_lines = p ? p->prefetch_lines() : d->prefetch_lines; + }); + break; + case IndexType::HNSW_RABITQ: + valid = UpdateFastQueryParam( + params, engine, defaults, [](const auto *p, auto *e, const auto *d) { + e->ef_search = p ? p->ef() : d->ef_search; + }); + break; + case IndexType::IVF: + valid = + UpdateFastQueryParam( + params, engine, defaults, + [](const auto *p, auto *e, const auto *d) { + e->nprobe = p ? p->nprobe() : d->nprobe; + }); + break; + case IndexType::IVF_RABITQ: + valid = UpdateFastQueryParam( + params, engine, defaults, [](const auto *p, auto *e, const auto *d) { + e->nprobe = p ? p->nprobe() : d->nprobe; + }); + break; + case IndexType::DISKANN: + valid = UpdateFastQueryParam( + params, engine, defaults, [](const auto *p, auto *e, const auto *d) { + e->list_size = p ? p->list_size() : d->list_size; + }); + break; + case IndexType::FLAT: + valid = + UpdateFastQueryParam( + params, engine, defaults, + [](const auto *, auto *, const auto *) {}); + break; + default: + break; + } + if (!valid) + return Status::InvalidArgument( + "fast query: unsupported query parameter type"); + if (engine) { + engine->radius = params ? params->radius() : defaults->radius; + engine->is_linear = params ? params->is_linear() : defaults->is_linear; + engine->refiner_param = + params && params->is_using_refiner() ? state.refiner_param : nullptr; + if (engine->refiner_param) { + engine->refiner_param->scale_factor_ = params->scale_factor(); + } + } + return Status::OK(); +} + +Result CollectionImpl::fast_query( + const std::string &field_name, const void *query_vector, + const QueryParams::Ptr &query_params, int topk, bool return_scores, + const DenseQueryShape *query_shape) const { + CHECK_DESTROY_RETURN_STATUS_EXPECTED(destroyed_, false); + CHECK_CLOSED_RETURN_STATUS_EXPECTED(closed_, false); + if (!options_.read_only_) { + return tl::make_unexpected( + Status::InvalidArgument("fast query requires a read-only collection")); + } + auto cached = fast_query_states_.find(field_name); + if (cached == fast_query_states_.end()) { + auto resolved = resolve_fast_query(field_name); + if (!resolved) return tl::make_unexpected(resolved.error()); + cached = fast_query_states_.emplace(field_name, std::move(resolved.value())) + .first; + } + auto &state = cached->second; + const auto param_status = update_fast_query_params(state, query_params); + CHECK_RETURN_STATUS_EXPECTED(param_status); + if (query_vector == nullptr) { + return tl::make_unexpected( + Status::InvalidArgument("fast query: query_vector is null")); + } + if (topk <= 0) { + return FastQueryResult{}; + } + if (query_shape && (query_shape->data_type != state.data_type || + query_shape->dimension != state.dimension)) { + return tl::make_unexpected(Status::InvalidArgument( + "query vector dtype or dimension does not match the field")); + } + const bool refine = query_params && query_params->is_using_refiner(); + if (state.raw_index && state.engine_query_param && + (!refine || state.raw_reference_index)) { + state.engine_query_param->topk = static_cast(topk); + core_interface::DenseVector dense_query{query_vector}; + core_interface::VectorData query_data{dense_query}; + FastQueryResult out; + out.ids.resize(static_cast(topk), int64_t{-1}); + if (return_scores) out.scores.resize(static_cast(topk)); + const int ret = state.raw_index->search_fast( + query_data, state.engine_query_param, out.ids.data(), + return_scores ? out.scores.data() : nullptr); + if (ret != 0) { + return tl::make_unexpected( + Status::InternalError("fast query: index search failed")); + } + return out; + } + + const auto &cache = state; + + const int search_topk = topk; + + const MetricType metric = cache.metric; + auto better = [metric](float a, float b) { + return metric == MetricType::IP ? a > b : a < b; + }; + + vector_column_params::QueryParams qp; + qp.topk = static_cast(search_topk); + qp.data_type = cache.data_type; + qp.dimension = cache.dimension; + qp.query_params = query_params; + + vector_column_params::VectorData vector_data; + vector_data.vector = vector_column_params::DenseVector{query_vector}; + + auto search_one = [&](const FastQuerySegment &entry, + FastQueryResult *out) -> Status { + qp.filter = entry.filter ? entry.filter.get() : nullptr; + + auto res = entry.indexer->Search(vector_data, qp); + if (!res) { + return res.error(); + } + IndexResults::Ptr results = std::move(res.value()); + + std::vector indices; + indices.reserve(search_topk); + const size_t score_begin = out->scores.size(); + for (auto it = results->create_iterator(); it->valid(); it->next()) { + indices.push_back(static_cast(it->doc_id())); + out->scores.push_back(it->score()); + } + if (indices.empty()) { + return Status::OK(); + } + + std::vector ids; + auto s = entry.segment->get_global_doc_ids(indices, ids); + if (!s.ok()) { + out->scores.resize(score_begin); + return s; + } + out->ids.insert(out->ids.end(), ids.begin(), ids.end()); + return Status::OK(); + }; + + auto finish = [topk, return_scores](FastQueryResult out) { + out.ids.resize(static_cast(topk), int64_t{-1}); + if (return_scores) { + out.scores.resize(static_cast(topk), + std::numeric_limits::quiet_NaN()); + } else { + out.scores.clear(); + } + return out; + }; + + if (cache.segments.size() == 1) { + FastQueryResult out; + auto s = search_one(cache.segments[0], &out); + CHECK_RETURN_STATUS_EXPECTED(s); + if (static_cast(out.ids.size()) > topk) { + out.ids.resize(topk); + out.scores.resize(topk); + } + return finish(std::move(out)); + } + + std::vector> candidates; + candidates.reserve(static_cast(search_topk) * cache.segments.size()); + for (const auto &entry : cache.segments) { + FastQueryResult seg_out; + auto s = search_one(entry, &seg_out); + CHECK_RETURN_STATUS_EXPECTED(s); + for (size_t i = 0; i < seg_out.ids.size(); ++i) { + candidates.emplace_back(seg_out.scores[i], seg_out.ids[i]); + } + } + + const size_t keep = std::min(static_cast(topk), candidates.size()); + std::partial_sort(candidates.begin(), candidates.begin() + keep, + candidates.end(), + [&better](const std::pair &a, + const std::pair &b) { + return better(a.first, b.first); + }); + + FastQueryResult out; + out.ids.reserve(keep); + out.scores.reserve(keep); + for (size_t i = 0; i < keep; ++i) { + out.scores.push_back(candidates[i].first); + out.ids.push_back(candidates[i].second); + } + return finish(std::move(out)); +} + Result CollectionImpl::query(const MultiQuery &query) const { std::shared_lock lock(schema_handle_mtx_); diff --git a/src/db/index/column/vector_column/combined_vector_column_indexer.h b/src/db/index/column/vector_column/combined_vector_column_indexer.h index c407b9bee..26e411073 100644 --- a/src/db/index/column/vector_column/combined_vector_column_indexer.h +++ b/src/db/index/column/vector_column/combined_vector_column_indexer.h @@ -42,6 +42,26 @@ class CombinedVectorColumnIndexer { uint32_t segment_doc_id) const; + //! True when at least one backing vector indexer is available for search. + bool has_searchable_indexers() const { + return !indexers_.empty(); + } + + //! True when one block starts at segment row zero. + bool is_single_block() const { + return indexers_.size() == 1 && block_offsets_[0] == 0; + } + + //! Primary block indexer (valid when ``is_single_block()``). + VectorColumnIndexer::Ptr primary_indexer() const { + return indexers_.empty() ? nullptr : indexers_[0]; + } + + //! Raw-vector reference block used by the primary index's refiner. + VectorColumnIndexer::Ptr reference_indexer() const { + return normal_indexers_.empty() ? nullptr : normal_indexers_[0]; + } + protected: /** * A filter wrapper that applies an offset to document IDs before diff --git a/src/db/index/segment/segment.cc b/src/db/index/segment/segment.cc index efc27f758..2d72020a8 100644 --- a/src/db/index/segment/segment.cc +++ b/src/db/index/segment/segment.cc @@ -255,6 +255,12 @@ class SegmentImpl : public Segment, ExecBatchPtr fetch(const std::vector &columns, int segment_doc_id) const override; + // Gather stable insertion ordinals without Arrow/user-ID materialization. + Status get_global_doc_ids(const std::vector &segment_doc_ids, + std::vector &out) const override; + + bool has_identity_doc_ids() const override; + RecordBatchReaderPtr scan( const std::vector &columns) const override; @@ -4534,6 +4540,30 @@ BlockID SegmentImpl::allocate_block_id() { return block_id_allocator_.fetch_add(1); } +bool SegmentImpl::has_identity_doc_ids() const { + std::lock_guard lock(seg_mtx_); + for (size_t i = 0; i < doc_ids_.size(); ++i) { + if (doc_ids_[i] != i) return false; + } + return true; +} + +Status SegmentImpl::get_global_doc_ids(const std::vector &segment_doc_ids, + std::vector &out) const { + out.resize(segment_doc_ids.size()); + std::lock_guard lock(seg_mtx_); + const size_t n = doc_ids_.size(); + for (size_t i = 0; i < segment_doc_ids.size(); ++i) { + const int sid = segment_doc_ids[i]; + if (sid < 0 || static_cast(sid) >= n) { + return Status::InvalidArgument("segment_doc_id out of range: ", sid); + } + out[i] = static_cast(doc_ids_[sid]); + } + return Status::OK(); +} + + Result SegmentImpl::get_global_doc_id(uint32_t segment_doc_id) const { // Read-only lookup into doc_ids_. std::shared_lock lock(seg_mtx_); diff --git a/src/db/index/segment/segment.h b/src/db/index/segment/segment.h index 361fc3a7d..38b022453 100644 --- a/src/db/index/segment/segment.h +++ b/src/db/index/segment/segment.h @@ -76,6 +76,11 @@ class Segment { // Count documents visible to an optional global-doc-ID filter. virtual uint64_t doc_count(const IndexFilter::Ptr filter = nullptr) = 0; + // Validates whether block keys can be returned as global document IDs. + virtual bool has_identity_doc_ids() const { + return false; + } + virtual bool has_record() = 0; // ---- Schema and index mutation ----------------------------------------- @@ -171,6 +176,10 @@ class Segment { virtual ExecBatchPtr fetch(const std::vector &columns, int segment_doc_id) const = 0; + // Gather stable insertion ordinals without Arrow/user-ID materialization. + virtual Status get_global_doc_ids(const std::vector &segment_doc_ids, + std::vector &out) const = 0; + // Keep Segment alive while consuming the returned reader. virtual RecordBatchReaderPtr scan( const std::vector &columns) const = 0; diff --git a/src/include/zvec/core/framework/index_runner.h b/src/include/zvec/core/framework/index_runner.h index 825232c65..58e296d70 100644 --- a/src/include/zvec/core/framework/index_runner.h +++ b/src/include/zvec/core/framework/index_runner.h @@ -501,19 +501,29 @@ class IndexRunner : public IndexModule { //! Search one ungrouped dense query, preserving search_impl's primary keys //! and result order (including topk, filters and threshold). The caller owns - //! the reusable output buffer, which is cleared on entry. Only keys are - //! required; context document/vector results are unspecified. Algorithms may - //! override this to export their retained pool without materializing scores. + //! the reusable output buffers, which are cleared on entry. Scores are + //! optional; context document/vector results are unspecified. Algorithms may + //! override this to export their retained pool without materializing docs. virtual int search_candidates_impl(const void *query, const IndexQueryMeta &qmeta, std::vector &keys, + std::vector *scores, Context::Pointer &context) const { keys.clear(); + if (scores) scores->clear(); const int ret = search_impl(query, qmeta, 1, context); if (ret != 0) return ret; const auto &result = context->result(); keys.reserve(result.size()); - for (const auto &document : result) keys.push_back(document.key()); + if (scores) { + scores->reserve(result.size()); + for (const auto &document : result) { + keys.push_back(document.key()); + scores->push_back(document.score()); + } + } else { + for (const auto &document : result) keys.push_back(document.key()); + } return 0; } //! Similarity search diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index 8ad3879ab..5354af658 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -141,6 +141,18 @@ class ZVEC_CORE_API Index { const BaseIndexQueryParam::Pointer &search_param, SearchResult *result); + // Advanced dense ANN API for latency-sensitive callers. Writes sorted keys + // and optional public scores into caller-owned storage. Direct paths avoid + // SearchResult materialization; other modes fall back to search(). Sparse + // and group-by searches are unsupported. Passing + // nullptr for output_scores keeps the ID-only path. Output buffers must hold + // at least search_param->topk elements. Refinement uses the same candidate + // selection and scoring as search(). The direct contiguous Flat path breaks + // equal-score ties by key. + int search_fast(const VectorData &query, + const BaseIndexQueryParam::Pointer &search_param, + int64_t *output_ids, float *output_scores); + virtual int add_with_source(const VectorData &vector, uint32_t doc_id, const core::VectorSource &src); virtual int search_with_source( @@ -190,12 +202,19 @@ class ZVEC_CORE_API Index { const core::IndexQueryMeta &query_meta, const BaseIndexQueryParam::Pointer &search_param, core::IndexContext::Pointer &context, - std::vector *candidate_keys = nullptr); + std::vector *candidate_keys = nullptr, + std::vector *candidate_scores = nullptr); + int _normalize_buffer_scores(const VectorData &query, + const int64_t *output_ids, float *output_scores, + size_t count); int _collect_dense_result(const VectorData &query, const core::IndexQueryMeta &query_meta, const BaseIndexQueryParam::Pointer &search_param, SearchResult *result, core::IndexContext::Pointer &context); + int _search_refine_fast(const VectorData &query, + const BaseIndexQueryParam::Pointer &search_param, + int64_t *output_ids, float *output_scores); int _refine_dense_candidates(const VectorData &query, const BaseIndexQueryParam::Pointer &search_param, const std::vector> &keys, diff --git a/src/include/zvec/db/collection.h b/src/include/zvec/db/collection.h index 0426658c4..946506ece 100644 --- a/src/include/zvec/db/collection.h +++ b/src/include/zvec/db/collection.h @@ -105,6 +105,17 @@ class ZVEC_API Collection { virtual Result query(const MultiQuery &query) const = 0; + // Advanced dense search returning internal numeric IDs and optional scores. + // Requires a read-only collection; fast_query calls and close must be serial. + // Parameters are read on every call. Index references are cached internally. + // Missing neighbors are padded with ID -1 and score NaN. Refinement uses + // query_params->scale_factor(), with the same semantics as query(). + virtual Result fast_query( + const std::string &field_name, const void *query_vector, + const QueryParams::Ptr &query_params = nullptr, int topk = 10, + bool return_scores = false, + const DenseQueryShape *query_shape = nullptr) const = 0; + virtual Result group_by_query( const GroupByVectorQuery &query) const = 0; diff --git a/src/include/zvec/db/query.h b/src/include/zvec/db/query.h index c983ac4e1..84922363b 100644 --- a/src/include/zvec/db/query.h +++ b/src/include/zvec/db/query.h @@ -132,6 +132,18 @@ inline void QueryTarget::set_sparse_vector(std::string indices, vc.sparse_values_ = std::move(values); } +// Buffers returned by the advanced dense query API. +struct FastQueryResult { + std::vector ids; + std::vector scores; +}; + +// Optional buffer metadata supplied by bindings before passing raw pointers. +struct DenseQueryShape { + DataType data_type; + uint32_t dimension; +}; + struct ZVEC_API SearchQuery { QueryTarget target_; int topk_{0}; diff --git a/tests/core/algorithm/flat/flat_streamer_test.cc b/tests/core/algorithm/flat/flat_streamer_test.cc index 1345a25e3..d709cb97e 100644 --- a/tests/core/algorithm/flat/flat_streamer_test.cc +++ b/tests/core/algorithm/flat/flat_streamer_test.cc @@ -471,6 +471,21 @@ TEST_F(FlatStreamerTest, TestContiguousCandidateSearchAndInsertFallback) { ASSERT_EQ(2, context->result().size()); EXPECT_EQ(17, context->result()[0].key()); + std::array direct_ids{{-1, -1}}; + std::array direct_scores{{-1.0F, -1.0F}}; + ASSERT_EQ(0, flat->search_by_p_keys_fast( + query.data(), keys[0], direct_ids.data(), + direct_scores.data(), direct_ids.size(), qmeta, context)); + EXPECT_EQ((std::array{{17, 5}}), direct_ids); + EXPECT_FLOAT_EQ(0.0F, direct_scores[0]); + EXPECT_FLOAT_EQ(static_cast(dim * 12 * 12), direct_scores[1]); + + std::array ids_without_scores{{-1, -1}}; + ASSERT_EQ(0, flat->search_by_p_keys_fast( + query.data(), keys[0], ids_without_scores.data(), nullptr, + ids_without_scores.size(), qmeta, context)); + EXPECT_EQ(direct_ids, ids_without_scores); + IndexStorage::MemoryBlock contiguous_block; ASSERT_EQ(0, streamer->get_vector_by_id(17, contiguous_block)); ASSERT_NE(nullptr, contiguous_block.data()); diff --git a/tests/core/interface/index_interface_test.cc b/tests/core/interface/index_interface_test.cc index 5d7381346..80ef841d3 100644 --- a/tests/core/interface/index_interface_test.cc +++ b/tests/core/interface/index_interface_test.cc @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. #include +#include #include #include #include @@ -3967,7 +3968,7 @@ TEST(IndexInterface, HNSWRabitqGeneral) { // mode. TEST(IndexInterface, ContiguousMemoryEndToEnd) { constexpr uint32_t kDimension = 32; - constexpr uint32_t kNumDocs = 500; + constexpr uint32_t kNumDocs = 1200; constexpr int kTopk = 10; const std::string index_name{"test_contiguous.index"}; @@ -4017,6 +4018,23 @@ TEST(IndexInterface, ContiguousMemoryEndToEnd) { ASSERT_EQ(0, index->search(query, query_param, &result)); ASSERT_GT(result.doc_list_.size(), 0UL); ASSERT_EQ(i, result.doc_list_[0].key()); + + std::array ids; + std::array scores; + ASSERT_EQ(0, index->search_fast(query, query_param, ids.data(), + scores.data())); + ASSERT_EQ(kTopk, result.doc_list_.size()); + for (size_t rank = 0; rank < result.doc_list_.size(); ++rank) { + EXPECT_EQ(static_cast(result.doc_list_[rank].key()), + ids[rank]); + EXPECT_FLOAT_EQ(result.doc_list_[rank].score(), scores[rank]); + } + + std::array ids_without_scores; + ASSERT_EQ(0, + index->search_fast(query, query_param, + ids_without_scores.data(), nullptr)); + EXPECT_EQ(ids, ids_without_scores); } ASSERT_EQ(0, index->close()); } @@ -4586,3 +4604,250 @@ TEST(IndexInterface, BuilderChainingReturnsCorrectType) { #if defined(__GNUC__) || defined(__GNUG__) #pragma GCC diagnostic pop #endif + + +TEST(IndexInterface, RefineFastNativeTypesFallbackAndTieOrder) { + constexpr uint32_t kDimension = 128; + const std::string coarse_path = "docids_refine_coarse.index"; + const std::string fine_path = "docids_refine_fine.index"; + std::vector values(kDimension, 0.0f); + auto coarse_param = FlatIndexParamBuilder() + .with_metric_type(MetricType::kL2sq) + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .build(); + zvec::test_util::RemoveTestFiles(coarse_path); + auto coarse = IndexFactory::CreateAndInitIndex(*coarse_param); + ASSERT_TRUE(coarse); + ASSERT_EQ( + 0, coarse->open(coarse_path, {StorageOptions::StorageType::kMMAP, true})); + for (uint32_t key = 0; key < 8; ++key) { + values[0] = static_cast(key); + ASSERT_EQ(0, coarse->add(VectorData{DenseVector{values.data()}}, key)); + } + values[0] = 0.0f; + const VectorData query{DenseVector{values.data()}}; + for (auto type : {DataType::DT_FP32, DataType::DT_FP16, DataType::DT_UINT8}) { + for (bool contiguous : {false, true}) { + for (bool tied : {false, true}) { + SCOPED_TRACE(static_cast(type)); + SCOPED_TRACE(contiguous); + SCOPED_TRACE(tied); + zvec::test_util::RemoveTestFiles(fine_path); + auto fine_param = FlatIndexParamBuilder() + .with_metric_type(MetricType::kL2sq) + .with_data_type(DataType::DT_FP32) + .with_storage_data_type(type) + .with_dimension(kDimension) + .with_use_contiguous_memory(contiguous) + .build(); + auto fine = IndexFactory::CreateAndInitIndex(*fine_param); + ASSERT_TRUE(fine); + ASSERT_EQ(0, fine->open(fine_path, + {StorageOptions::StorageType::kMMAP, true})); + std::vector fine_values(kDimension, 0.0f); + for (uint32_t key = 0; key < 8; ++key) { + fine_values[0] = tied ? 1.0f : float(8 - key); + ASSERT_EQ( + 0, fine->add(VectorData{DenseVector{fine_values.data()}}, key)); + } + // The contiguous reference is built on reopen, not on add(). + ASSERT_EQ(0, fine->close()); + fine = IndexFactory::CreateAndInitIndex(*fine_param); + ASSERT_TRUE(fine); + ASSERT_EQ(0, fine->open(fine_path, + {StorageOptions::StorageType::kMMAP, false})); + auto refiner = std::make_shared(); + refiner->reference_index = fine; + refiner->scale_factor_ = 2.0f; + auto param = FlatQueryParamBuilder() + .with_topk(3) + .with_refiner_param(refiner) + .build(); + std::array ids{{-2, -2, -2}}; + for (int repeat = 0; repeat < 3; ++repeat) { + ASSERT_EQ(0, coarse->search_fast(query, param, ids.data(), nullptr)); + if (!tied) { + // Coarse selects keys 0..5; the reference reverses their order. + EXPECT_EQ((std::array{{5, 4, 3}}), ids); + } else if (contiguous) { + // Refined results use deterministic (distance, key) ordering. + EXPECT_EQ((std::array{{0, 1, 2}}), ids); + } else { + // The generic public fallback retains its existing tie rule. + for (int64_t key : ids) EXPECT_TRUE(key >= 0 && key < 6); + } + } + std::array scored_ids{{-2, -2, -2}}; + std::array scores; + ASSERT_EQ(0, coarse->search_fast(query, param, scored_ids.data(), + scores.data())); + EXPECT_EQ(ids, scored_ids); + for (size_t rank = 0; rank < scores.size(); ++rank) { + const float expected = + tied ? 1.0F : static_cast((3 + rank) * (3 + rank)); + EXPECT_FLOAT_EQ(expected, scores[rank]); + } + // Fewer coarse candidates than k takes the unchanged public fallback + // and must fill the remainder instead of exposing stale output. + param->radius = 0.5f; + ASSERT_EQ(0, coarse->search_fast(query, param, ids.data(), nullptr)); + EXPECT_EQ((std::array{{0, -1, -1}}), ids); + param->radius = 0.0f; + param->filter = std::make_shared(); + param->filter->set([](uint64_t key) { return key >= 2; }); + ASSERT_EQ(0, coarse->search_fast(query, param, ids.data(), nullptr)); + EXPECT_EQ(-1, ids[2]); + std::array filtered{{ids[0], ids[1]}}; + std::sort(filtered.begin(), filtered.end()); + EXPECT_EQ((std::array{{0, 1}}), filtered); + param->filter.reset(); + ASSERT_EQ(0, coarse->search_fast(query, param, ids.data(), nullptr)); + if (!tied) { + EXPECT_EQ((std::array{{5, 4, 3}}), ids); + } + ASSERT_EQ(0, fine->close()); + zvec::test_util::RemoveTestFiles(fine_path); + } + } + } + ASSERT_EQ(0, coarse->close()); + zvec::test_util::RemoveTestFiles(coarse_path); +} + +TEST(IndexInterface, RefineFastInnerProductUsesDescendingPublicScore) { + constexpr uint32_t kDimension = 128; + const std::string coarse_path = "docids_ip_coarse.index"; + const std::string fine_path = "docids_ip_fine.index"; + zvec::test_util::RemoveTestFiles(coarse_path); + auto coarse_param = FlatIndexParamBuilder() + .with_metric_type(MetricType::kInnerProduct) + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .build(); + auto coarse = IndexFactory::CreateAndInitIndex(*coarse_param); + ASSERT_TRUE(coarse); + ASSERT_EQ( + 0, coarse->open(coarse_path, {StorageOptions::StorageType::kMMAP, true})); + std::vector values(kDimension, 0.0f); + for (uint32_t key = 0; key < 8; ++key) { + values[0] = float(key); + ASSERT_EQ(0, coarse->add(VectorData{DenseVector{values.data()}}, key)); + } + values[0] = 1.0f; + const VectorData query{DenseVector{values.data()}}; + for (auto type : {DataType::DT_FP32, DataType::DT_FP16}) { + for (bool tied : {false, true}) { + SCOPED_TRACE(static_cast(type)); + SCOPED_TRACE(tied); + zvec::test_util::RemoveTestFiles(fine_path); + auto fine_param = FlatIndexParamBuilder() + .with_metric_type(MetricType::kInnerProduct) + .with_data_type(DataType::DT_FP32) + .with_storage_data_type(type) + .with_dimension(kDimension) + .with_use_contiguous_memory(true) + .build(); + auto fine = IndexFactory::CreateAndInitIndex(*fine_param); + ASSERT_TRUE(fine); + ASSERT_EQ( + 0, fine->open(fine_path, {StorageOptions::StorageType::kMMAP, true})); + std::vector fine_values(kDimension, 0.0f); + for (uint32_t key = 0; key < 8; ++key) { + fine_values[0] = tied ? 1.0f : float(8 - key); + ASSERT_EQ(0, + fine->add(VectorData{DenseVector{fine_values.data()}}, key)); + } + ASSERT_EQ(0, fine->close()); + fine = IndexFactory::CreateAndInitIndex(*fine_param); + ASSERT_TRUE(fine); + ASSERT_EQ(0, fine->open(fine_path, + {StorageOptions::StorageType::kMMAP, false})); + auto refiner = std::make_shared(); + refiner->reference_index = fine; + refiner->scale_factor_ = 2.0f; + auto param = FlatQueryParamBuilder() + .with_topk(3) + .with_refiner_param(refiner) + .build(); + std::array ids{{-2, -2, -2}}; + std::array scores; + ASSERT_EQ(0, + coarse->search_fast(query, param, ids.data(), scores.data())); + // Coarse picks keys 7..2 by descending IP; refine uses their inverse + // scores, selecting 2,3,4. Equal scores use ascending key as baseline. + EXPECT_EQ((std::array{{2, 3, 4}}), ids); + if (tied) { + EXPECT_EQ((std::array{{1.0F, 1.0F, 1.0F}}), scores); + } else { + EXPECT_EQ((std::array{{6.0F, 5.0F, 4.0F}}), scores); + } + if (!tied) { + SearchResult result; + ASSERT_EQ(0, coarse->search(query, param, &result)); + ASSERT_EQ(3U, result.doc_list_.size()); + for (size_t i = 0; i < 3; ++i) { + EXPECT_EQ(static_cast(ids[i]), result.doc_list_[i].key()); + EXPECT_FLOAT_EQ(float(6 - i), result.doc_list_[i].score()); + } + } + // search_param owns topk; the index writes exactly that many entries. + std::array larger{{-2, -2, -2, -2}}; + ASSERT_EQ(0, coarse->search_fast(query, param, larger.data(), nullptr)); + EXPECT_EQ(-2, larger[3]); + ASSERT_EQ(0, fine->close()); + zvec::test_util::RemoveTestFiles(fine_path); + } + } + ASSERT_EQ(0, coarse->close()); + zvec::test_util::RemoveTestFiles(coarse_path); +} + +TEST(IndexInterface, NativeFp16CandidateBatchAroundTwelveIsExact) { + const std::string path = "flat_fp16_twelve_batch.index"; + for (uint32_t dimension : {128U, 960U}) { + SCOPED_TRACE(dimension); + zvec::test_util::RemoveTestFiles(path); + auto param = FlatIndexParamBuilder() + .with_metric_type(MetricType::kL2sq) + .with_data_type(DataType::DT_FP32) + .with_storage_data_type(DataType::DT_FP16) + .with_dimension(dimension) + .with_use_contiguous_memory(true) + .build(); + auto index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_TRUE(index); + ASSERT_EQ(0, index->open(path, {StorageOptions::StorageType::kMMAP, true})); + std::vector values(dimension, 0.0f); + for (uint32_t key = 0; key < 14; ++key) { + values[0] = float(14 - key); + ASSERT_EQ(0, index->add(VectorData{DenseVector{values.data()}}, key)); + } + ASSERT_EQ(0, index->close()); + index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_TRUE(index); + ASSERT_EQ(0, + index->open(path, {StorageOptions::StorageType::kMMAP, false})); + values[0] = 0.0f; + for (uint32_t candidates : {11U, 12U, 13U, 12U}) { + SCOPED_TRACE(candidates); + auto query_param = FlatQueryParamBuilder().with_topk(10).build(); + query_param->bf_pks = std::make_shared>(); + for (uint32_t key = 0; key < candidates; ++key) { + query_param->bf_pks->push_back(key); + } + SearchResult result; + ASSERT_EQ(0, index->search(VectorData{DenseVector{values.data()}}, + query_param, &result)); + ASSERT_EQ(10U, result.doc_list_.size()); + for (uint32_t i = 0; i < 10; ++i) { + const uint32_t key = candidates - 1 - i; + EXPECT_EQ(key, result.doc_list_[i].key()); + EXPECT_FLOAT_EQ(float((14 - key) * (14 - key)), + result.doc_list_[i].score()); + } + } + ASSERT_EQ(0, index->close()); + zvec::test_util::RemoveTestFiles(path); + } +} diff --git a/tests/db/fast_query_test.cc b/tests/db/fast_query_test.cc new file mode 100644 index 000000000..e82b032e9 --- /dev/null +++ b/tests/db/fast_query_test.cc @@ -0,0 +1,82 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include "db/common/file_helper.h" + +using namespace zvec; + +TEST(FastQueryTest, ReadsRefineParametersOnEveryCall) { + const std::string path = "test_fast_search_refine_scale"; + FileHelper::RemoveDirectory(path); + ailego::MemoryLimitPool::get_instance().init(2 * 1024ll * 1024ll * 1024ll); + CollectionSchema schema("fast_search"); + schema.add_field(std::make_shared( + "vector", DataType::VECTOR_FP32, uint32_t{32}, false, + std::make_shared(MetricType::L2, QuantizeType::INT8))); + auto created = Collection::CreateAndOpen(path, schema, CollectionOptions{}); + ASSERT_TRUE(created) << created.error().message(); + auto writer = std::move(created.value()); + std::mt19937 rng(712); + std::normal_distribution normal; + std::vector docs; + for (int i = 0; i < 512; ++i) { + Doc doc; + doc.set_pk(std::to_string(i)); + std::vector vector(32); + for (auto &v : vector) v = normal(rng); + doc.set>("vector", vector); + docs.push_back(std::move(doc)); + } + auto inserted = writer->insert(docs); + ASSERT_TRUE(inserted) << inserted.error().message(); + for (const auto &status : inserted.value()) ASSERT_TRUE(status.ok()); + ASSERT_TRUE(writer->optimize(OptimizeOptions{1}).ok()); + ASSERT_TRUE(writer->close().ok()); + auto opened = Collection::Open(path, CollectionOptions{true, true}); + ASSERT_TRUE(opened) << opened.error().message(); + auto reader = std::move(opened.value()); + auto param = std::make_shared(true); + for (float scale : {1.0f, 3.0f, 7.0f, 1.0f}) { + param->set_scale_factor(scale); + for (int repeat = 0; repeat < 10; ++repeat) { + std::vector vector(32); + for (auto &v : vector) v = normal(rng); + SearchQuery query; + query.topk_ = 10; + query.target_.field_name_ = "vector"; + query.target_.set_vector( + std::string(reinterpret_cast(vector.data()), + vector.size() * sizeof(float))); + query.target_.query_params_ = + std::make_shared(true, scale); + auto expected = reader->query(query); + ASSERT_TRUE(expected) << expected.error().message(); + auto actual = + reader->fast_query("vector", vector.data(), param, 10, true); + ASSERT_TRUE(actual) << actual.error().message(); + ASSERT_EQ(actual->ids.size(), expected->size()); + for (size_t i = 0; i < expected->size(); ++i) { + EXPECT_EQ(actual->ids[i], std::stoll(expected.value()[i]->pk())); + EXPECT_FLOAT_EQ(actual->scores[i], expected.value()[i]->score()); + } + } + } + ASSERT_TRUE(reader->close().ok()); + reader.reset(); + FileHelper::RemoveDirectory(path); +} diff --git a/tests/db/sqlengine/mock_segment.h b/tests/db/sqlengine/mock_segment.h index c5afacb55..e64f4566c 100644 --- a/tests/db/sqlengine/mock_segment.h +++ b/tests/db/sqlengine/mock_segment.h @@ -312,6 +312,12 @@ class MockSegment : public Segment { return 0; } + Status get_global_doc_ids(const std::vector &segment_doc_ids, + std::vector &out) const override { + out.assign(segment_doc_ids.begin(), segment_doc_ids.end()); + return Status::OK(); + } + TablePtr fetch(const std::vector &columns, const std::vector &indices) const override { std::string s = ""; From a858c318326b721b7fb5eac86fd8456733722df6 Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Wed, 16 Sep 2026 19:43:16 +0800 Subject: [PATCH 2/8] fix(test): make lambda array extent constant on MSVC --- tests/core/interface/index_interface_test.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/core/interface/index_interface_test.cc b/tests/core/interface/index_interface_test.cc index 80ef841d3..146e69791 100644 --- a/tests/core/interface/index_interface_test.cc +++ b/tests/core/interface/index_interface_test.cc @@ -3969,7 +3969,7 @@ TEST(IndexInterface, HNSWRabitqGeneral) { TEST(IndexInterface, ContiguousMemoryEndToEnd) { constexpr uint32_t kDimension = 32; constexpr uint32_t kNumDocs = 1200; - constexpr int kTopk = 10; + static constexpr int kTopk = 10; const std::string index_name{"test_contiguous.index"}; // build_then_search builds an index from scratch (with use_contiguous_memory From 5171f77881d4bb19585224635829ec930e71995c Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Wed, 16 Sep 2026 20:25:28 +0800 Subject: [PATCH 3/8] refactor(flat): use a local pair vector for fast search candidates --- src/core/algorithm/flat/flat_streamer_entity.cc | 12 ++++++------ src/core/algorithm/flat/flat_streamer_entity.h | 5 ----- 2 files changed, 6 insertions(+), 11 deletions(-) diff --git a/src/core/algorithm/flat/flat_streamer_entity.cc b/src/core/algorithm/flat/flat_streamer_entity.cc index aeaf3c360..a7d623374 100644 --- a/src/core/algorithm/flat/flat_streamer_entity.cc +++ b/src/core/algorithm/flat/flat_streamer_entity.cc @@ -15,6 +15,7 @@ #include "flat_streamer_entity.h" #include #include +#include #include #include "flat_utility.h" @@ -619,7 +620,6 @@ int FlatContiguousStreamerEntity::search_by_p_keys_fast( auto &ptrs = scratch->vector_ptrs; auto &extras = scratch->extra_values; auto &distances = scratch->distances; - auto &documents = scratch->candidate_documents; const size_t count = keys.size(); const size_t extra_size = extra_values_size(); const bool has_extras = extra_size != 0; @@ -627,7 +627,6 @@ int FlatContiguousStreamerEntity::search_by_p_keys_fast( ptrs.resize(count); extras.resize(has_extras ? count : 0); distances.resize(count); - documents.resize(count); for (size_t i = 0; i < count; ++i) { ptrs[i] = get_vector_ptr_by_key(*storage, keys[i]); if (!ptrs[i]) return IndexError_NotImplemented; @@ -654,12 +653,13 @@ int FlatContiguousStreamerEntity::search_by_p_keys_fast( distance()(query, ptrs[i], meta().dimension(), distances.data() + i); } } + std::vector> documents(count); for (size_t i = 0; i < count; ++i) { documents[i] = {keys[i], distances[i]}; } const auto better = [](const auto &lhs, const auto &rhs) { - return lhs.distance < rhs.distance || - (lhs.distance == rhs.distance && lhs.key < rhs.key); + return lhs.second < rhs.second || + (lhs.second == rhs.second && lhs.first < rhs.first); }; auto selected_end = documents.begin() + topk; std::make_heap(documents.begin(), selected_end, better); @@ -673,8 +673,8 @@ int FlatContiguousStreamerEntity::search_by_p_keys_fast( } std::sort_heap(documents.begin(), selected_end, better); for (size_t i = 0; i < topk; ++i) { - output_ids[i] = static_cast(documents[i].key); - if (output_scores) output_scores[i] = documents[i].distance; + output_ids[i] = static_cast(documents[i].first); + if (output_scores) output_scores[i] = documents[i].second; } return 0; } diff --git a/src/core/algorithm/flat/flat_streamer_entity.h b/src/core/algorithm/flat/flat_streamer_entity.h index 4dbcfa653..1484e4da3 100644 --- a/src/core/algorithm/flat/flat_streamer_entity.h +++ b/src/core/algorithm/flat/flat_streamer_entity.h @@ -34,11 +34,6 @@ namespace core { //! Reusable request-local buffers for storage-specific Flat search paths. struct FlatSearchScratch { - struct CandidateDocument { - uint64_t key; - float distance; - }; - std::vector candidate_documents{}; std::vector vector_ptrs{}; std::vector extra_values{}; std::vector vector_keys{}; From 9084f6de28f83d5a0310181ce73bc40534d398dd Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Wed, 16 Sep 2026 20:36:46 +0800 Subject: [PATCH 4/8] refactor(query): pass dense query metadata as direct parameters --- src/binding/python/model/python_collection.cc | 12 +++++++----- src/db/collection.cc | 17 ++++++++++------- src/include/zvec/db/collection.h | 5 ++++- src/include/zvec/db/query.h | 6 ------ tests/db/fast_query_test.cc | 19 +++++++++++++++++++ 5 files changed, 40 insertions(+), 19 deletions(-) diff --git a/src/binding/python/model/python_collection.cc b/src/binding/python/model/python_collection.cc index 71a12a2b9..9f302cd84 100644 --- a/src/binding/python/model/python_collection.cc +++ b/src/binding/python/model/python_collection.cc @@ -25,7 +25,7 @@ namespace zvec { namespace { -DenseQueryShape dense_query_shape(const py::array &vector) { +DataType dense_query_data_type(const py::array &vector) { if (vector.ndim() != 1 || !(vector.flags() & py::array::c_style)) { throw py::value_error("query vector must be a contiguous 1D array"); } @@ -57,7 +57,7 @@ DenseQueryShape dense_query_shape(const py::array &vector) { } else { throw py::value_error("unsupported query vector dtype"); } - return {type, static_cast(vector.shape(0))}; + return type; } // Batch-materialize a DocPtrList into a list of (id, score, fields, vectors) @@ -395,7 +395,8 @@ void ZVecPyCollection::bind_dql_methods( [](const Collection &self, const std::string &field_name, const py::array &vector, QueryParams *params, int topk, bool return_scores) -> py::object { - const auto shape = dense_query_shape(vector); + const auto data_type = dense_query_data_type(vector); + const auto dimension = static_cast(vector.shape(0)); // Python keeps the argument alive for this call. The DB reads it // synchronously and never retains it, so no shared ownership // conversion or reference-count traffic is needed here. @@ -403,8 +404,9 @@ void ZVecPyCollection::bind_dql_methods( Result result; { py::gil_scoped_release release; - result = self.fast_query(field_name, vector.data(), borrowed, - topk, return_scores, &shape); + result = + self.fast_query(field_name, vector.data(), borrowed, topk, + return_scores, data_type, dimension); } auto output = unwrap_expected(std::move(result)); auto ids = owned_array(std::move(output.ids)); diff --git a/src/db/collection.cc b/src/db/collection.cc index 004462748..3056d15e5 100644 --- a/src/db/collection.cc +++ b/src/db/collection.cc @@ -135,10 +135,12 @@ class CollectionImpl : public Collection { Result query(const MultiQuery &query) const override; - Result fast_query( - const std::string &field_name, const void *query_vector, - const QueryParams::Ptr &query_params, int topk, bool return_scores, - const DenseQueryShape *query_shape) const override; + Result fast_query(const std::string &field_name, + const void *query_vector, + const QueryParams::Ptr &query_params, + int topk, bool return_scores, + DataType query_data_type, + uint32_t query_dimension) const override; Result group_by_query( const GroupByVectorQuery &query) const override; @@ -2036,7 +2038,7 @@ Status CollectionImpl::update_fast_query_params( Result CollectionImpl::fast_query( const std::string &field_name, const void *query_vector, const QueryParams::Ptr &query_params, int topk, bool return_scores, - const DenseQueryShape *query_shape) const { + DataType query_data_type, uint32_t query_dimension) const { CHECK_DESTROY_RETURN_STATUS_EXPECTED(destroyed_, false); CHECK_CLOSED_RETURN_STATUS_EXPECTED(closed_, false); if (!options_.read_only_) { @@ -2060,8 +2062,9 @@ Result CollectionImpl::fast_query( if (topk <= 0) { return FastQueryResult{}; } - if (query_shape && (query_shape->data_type != state.data_type || - query_shape->dimension != state.dimension)) { + if ((query_data_type != DataType::UNDEFINED || query_dimension != 0) && + (query_data_type != state.data_type || + query_dimension != state.dimension)) { return tl::make_unexpected(Status::InvalidArgument( "query vector dtype or dimension does not match the field")); } diff --git a/src/include/zvec/db/collection.h b/src/include/zvec/db/collection.h index 946506ece..b91641c09 100644 --- a/src/include/zvec/db/collection.h +++ b/src/include/zvec/db/collection.h @@ -110,11 +110,14 @@ class ZVEC_API Collection { // Parameters are read on every call. Index references are cached internally. // Missing neighbors are padded with ID -1 and score NaN. Refinement uses // query_params->scale_factor(), with the same semantics as query(). + // Supply both input dtype and dimension for validation; omitting both trusts + // the caller to provide a buffer matching the field. virtual Result fast_query( const std::string &field_name, const void *query_vector, const QueryParams::Ptr &query_params = nullptr, int topk = 10, bool return_scores = false, - const DenseQueryShape *query_shape = nullptr) const = 0; + DataType query_data_type = DataType::UNDEFINED, + uint32_t query_dimension = 0) const = 0; virtual Result group_by_query( const GroupByVectorQuery &query) const = 0; diff --git a/src/include/zvec/db/query.h b/src/include/zvec/db/query.h index 84922363b..c59a48581 100644 --- a/src/include/zvec/db/query.h +++ b/src/include/zvec/db/query.h @@ -138,12 +138,6 @@ struct FastQueryResult { std::vector scores; }; -// Optional buffer metadata supplied by bindings before passing raw pointers. -struct DenseQueryShape { - DataType data_type; - uint32_t dimension; -}; - struct ZVEC_API SearchQuery { QueryTarget target_; int topk_{0}; diff --git a/tests/db/fast_query_test.cc b/tests/db/fast_query_test.cc index e82b032e9..a03717910 100644 --- a/tests/db/fast_query_test.cc +++ b/tests/db/fast_query_test.cc @@ -13,6 +13,7 @@ // limitations under the License. #include +#include #include #include #include @@ -69,6 +70,24 @@ TEST(FastQueryTest, ReadsRefineParametersOnEveryCall) { auto actual = reader->fast_query("vector", vector.data(), param, 10, true); ASSERT_TRUE(actual) << actual.error().message(); + if (repeat == 0) { + auto checked = reader->fast_query("vector", vector.data(), param, 10, + true, DataType::VECTOR_FP32, 32); + ASSERT_TRUE(checked) << checked.error().message(); + EXPECT_EQ(actual->ids, checked->ids); + EXPECT_EQ(actual->scores, checked->scores); + // Wrong or partially supplied metadata must fail before vector reads. + for (const auto &[type, dimension] : + {std::pair{DataType::VECTOR_FP64, 32U}, + std::pair{DataType::VECTOR_FP32, 31U}, + std::pair{DataType::VECTOR_FP32, 0U}, + std::pair{DataType::UNDEFINED, 32U}}) { + auto invalid = reader->fast_query("vector", vector.data(), param, 10, + true, type, dimension); + ASSERT_FALSE(invalid); + EXPECT_EQ(StatusCode::INVALID_ARGUMENT, invalid.error().code()); + } + } ASSERT_EQ(actual->ids.size(), expected->size()); for (size_t i = 0; i < expected->size(); ++i) { EXPECT_EQ(actual->ids[i], std::stoll(expected.value()[i]->pk())); From dd5d1fa8e1ff5ff10f28b56d33a47fa99cda921b Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Wed, 16 Sep 2026 20:52:01 +0800 Subject: [PATCH 5/8] refactor(query): remove redundant raw index fallback --- src/db/collection.cc | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/db/collection.cc b/src/db/collection.cc index 3056d15e5..7120e01b6 100644 --- a/src/db/collection.cc +++ b/src/db/collection.cc @@ -1886,9 +1886,6 @@ Result CollectionImpl::resolve_fast_query( indexer = seg->get_combined_vector_indexer(field_name); } else { indexer = seg->get_quant_combined_vector_indexer(field_name); - if (indexer != nullptr && !indexer->has_searchable_indexers()) { - indexer = seg->get_combined_vector_indexer(field_name); - } } if (!indexer || !indexer->has_searchable_indexers()) { continue; From 6650a56847be25ad679c5b0aab8ea3b821f4aa0e Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Thu, 17 Sep 2026 10:16:19 +0800 Subject: [PATCH 6/8] refactor(query): share parameter mapping and simplify fast query state Reuse query parameter mapping across ordinary and fast queries, remove redundant cached references and result copies, and handle empty collections consistently. Add coverage for reusable parameters, empty collections, and Uniform quantizer fallback before and after optimize and subsequent inserts. Validation: 151 related tests passed; the complete 780-point SIFT/GIST comparison stayed within the established QPS and recall tolerances. --- python/tests/test_fast_query.py | 124 +++++++++++++++ src/db/collection.cc | 141 +++-------------- .../column/vector_column/engine_helper.hpp | 145 ++++++++++++------ .../column/vector_column_indexer_test.cc | 63 ++++++++ 4 files changed, 308 insertions(+), 165 deletions(-) diff --git a/python/tests/test_fast_query.py b/python/tests/test_fast_query.py index 49069caf5..c7957d0f9 100644 --- a/python/tests/test_fast_query.py +++ b/python/tests/test_fast_query.py @@ -18,6 +18,47 @@ from zvec.typing import DataType, MetricType, QuantizeType +@pytest.mark.parametrize( + "index_type", [zvec.FlatIndexParam, HnswIndexParam, VamanaIndexParam] +) +def test_empty_collection(tmp_path, index_type): + schema = CollectionSchema( + name="empty_fast_query", + vectors=[ + VectorSchema("vector", DataType.VECTOR_FP32, 32, index_param=index_type()) + ], + ) + path = str(tmp_path / "empty") + writer = zvec.create_and_open(path, schema) + writer.close() + reader = zvec.open(path, CollectionOption(read_only=True)) + vector = np.zeros(32, dtype=np.float32) + try: + assert reader.query(Query("vector", vector=vector), topk=3) == [] + for topk in (3, 0, 1): + ids, scores = reader.fast_query( + "vector", vector, topk=topk, return_scores=True + ) + np.testing.assert_array_equal(ids, np.full(topk, -1, dtype=np.int64)) + assert scores.shape == (topk,) + assert np.all(np.isnan(scores)) + np.testing.assert_array_equal( + ids, reader.fast_query("vector", vector, topk=topk) + ) + # Empty state must still validate the field, vector and parameter type. + with pytest.raises(ValueError, match="dense vector field"): + reader.fast_query("missing", vector) + with pytest.raises(ValueError, match="dtype|dimension"): + reader.fast_query("vector", vector[:-1]) + wrong_param = ( + HnswQueryParam() if index_type is not HnswIndexParam else VamanaQueryParam() + ) + with pytest.raises(ValueError, match="parameter type"): + reader.fast_query("vector", vector, wrong_param) + finally: + reader.close() + + @pytest.fixture( params=[(128, False), (1200, False), (1200, True)], ids=["brute_fallback", "vamana_graph", "hnsw_graph"], @@ -444,3 +485,86 @@ def test_fast_query_index_and_metric_dispatch(tmp_path, index_kind, metric): ) finally: reader.close() + + +@pytest.mark.parametrize("index_kind", ["vamana", "hnsw"]) +@pytest.mark.parametrize( + "quantizer", + [ + zvec.QuantizeType.UNIFORM_UINT4, + zvec.QuantizeType.UNIFORM_UINT7, + zvec.QuantizeType.UNIFORM_UINT8, + ], +) +def test_uniform_raw_fallback(tmp_path, index_kind, quantizer): + vectors = np.random.default_rng(945).normal(size=(130, 32)).astype(np.float32) + options = dict( + metric_type=zvec.MetricType.L2, + quantize_type=quantizer, + flat_data_type=zvec.DataType.VECTOR_FP16, + use_flat_contiguous_memory=True, + ) + if index_kind == "vamana": + index = zvec.VamanaIndexParam(max_degree=16, search_list_size=64, **options) + param_type = zvec.VamanaQueryParam + else: + index = zvec.HnswIndexParam(m=16, ef_construction=64, **options) + param_type = zvec.HnswQueryParam + schema = zvec.CollectionSchema( + name="uniform_fallback", + vectors=[ + zvec.VectorSchema( + "vector", zvec.DataType.VECTOR_FP32, 32, index_param=index + ) + ], + ) + path = str(tmp_path / "collection") + writer = zvec.create_and_open(path, schema) + assert all( + s.ok() + for s in writer.insert( + [ + zvec.Doc(id=str(i), vectors={"vector": vectors[i].tolist()}) + for i in range(128) + ] + ) + ) + writer.close() + for phase in ("untrained", "optimized", "new_writes"): + if phase != "untrained": + writer = zvec.open(path) + if phase == "optimized": + writer.optimize() + else: + assert all( + s.ok() + for s in writer.insert( + [ + zvec.Doc(id=str(i), vectors={"vector": vectors[i].tolist()}) + for i in (128, 129) + ] + ) + ) + writer.close() + reader = zvec.open(path, zvec.CollectionOption(read_only=True)) + try: + for refine in (False, True): + param = param_type(is_using_refiner=refine) + query = vectors[128] if phase == "new_writes" else vectors[17] + docs = reader.query( + zvec.Query("vector", vector=query, param=param), topk=10 + ) + ids, scores = reader.fast_query( + "vector", query, param, return_scores=True + ) + np.testing.assert_array_equal(ids, [int(d.id) for d in docs]) + np.testing.assert_allclose( + scores, [d.score for d in docs], rtol=2e-5, atol=2e-5 + ) + np.testing.assert_array_equal( + ids, reader.fast_query("vector", query, param) + ) + if phase == "new_writes": + assert ids[0] == 128 + finally: + reader.close() diff --git a/src/db/collection.cc b/src/db/collection.cc index 7120e01b6..af054624c 100644 --- a/src/db/collection.cc +++ b/src/db/collection.cc @@ -21,7 +21,6 @@ #include #include #include -#include #include #include #include @@ -314,7 +313,6 @@ class CollectionImpl : public Collection { MetricType metric{MetricType::L2}; std::vector segments; core_interface::Index::Pointer raw_index; - core_interface::Index::Pointer raw_reference_index; core_interface::BaseIndexQueryParam::Pointer engine_query_param; core_interface::BaseIndexQueryParam::Pointer default_engine_query_param; core_interface::RefinerParam::Pointer refiner_param; @@ -1879,6 +1877,7 @@ Result CollectionImpl::resolve_fast_query( : IndexType::UNDEFINED; cache.metric = metric; const auto segments = get_all_segments(); + if (segments.empty()) return cache; cache.segments.reserve(segments.size()); for (const auto &seg : segments) { CombinedVectorColumnIndexer::Ptr indexer; @@ -1910,8 +1909,12 @@ Result CollectionImpl::resolve_fast_query( if (primary) { cache.raw_index = primary->debug_get_index(); auto reference = entry.indexer->reference_indexer(); - if (reference) { - cache.raw_reference_index = reference->debug_get_index(); + auto reference_index = + reference ? reference->debug_get_index() : nullptr; + if (reference_index) { + cache.refiner_param = + std::make_shared(); + cache.refiner_param->reference_index = std::move(reference_index); } } } @@ -1927,102 +1930,16 @@ Result CollectionImpl::resolve_fast_query( if (!engine_qp) return tl::make_unexpected(engine_qp.error()); cache.engine_query_param = std::move(engine_qp.value()); cache.default_engine_query_param = cache.engine_query_param->clone(); - if (cache.raw_reference_index) { - cache.refiner_param = std::make_shared(); - cache.refiner_param->reference_index = cache.raw_reference_index; - } return cache; } -namespace { -template -bool UpdateFastQueryParam(const QueryParams::Ptr ¶ms, - core_interface::BaseIndexQueryParam *engine, - const core_interface::BaseIndexQueryParam *defaults, - Update update) { - const auto *p = dynamic_cast(params.get()); - if (params && !p) return false; - if (auto *e = dynamic_cast(engine)) { - update(p, e, static_cast(defaults)); - } - return true; -} -} // namespace - Status CollectionImpl::update_fast_query_params( FastQueryState &state, const QueryParams::Ptr ¶ms) const { - if (params && params->type() != state.index_type) { - return Status::InvalidArgument( - "fast query: query parameter type does not match the field index"); - } auto *engine = state.engine_query_param.get(); - const auto *defaults = state.default_engine_query_param.get(); - bool valid = false; - switch (state.index_type) { - case IndexType::VAMANA: - valid = UpdateFastQueryParam( - params, engine, defaults, [](const auto *p, auto *e, const auto *d) { - e->ef_search = p ? p->ef_search() : d->ef_search; - e->prefetch_offset = p ? p->prefetch_offset() : d->prefetch_offset; - e->prefetch_lines = p ? p->prefetch_lines() : d->prefetch_lines; - }); - break; - case IndexType::HNSW: - valid = - UpdateFastQueryParam( - params, engine, defaults, - [](const auto *p, auto *e, const auto *d) { - e->ef_search = p ? p->ef() : d->ef_search; - e->prefetch_offset = - p ? p->prefetch_offset() : d->prefetch_offset; - e->prefetch_lines = p ? p->prefetch_lines() : d->prefetch_lines; - }); - break; - case IndexType::HNSW_RABITQ: - valid = UpdateFastQueryParam( - params, engine, defaults, [](const auto *p, auto *e, const auto *d) { - e->ef_search = p ? p->ef() : d->ef_search; - }); - break; - case IndexType::IVF: - valid = - UpdateFastQueryParam( - params, engine, defaults, - [](const auto *p, auto *e, const auto *d) { - e->nprobe = p ? p->nprobe() : d->nprobe; - }); - break; - case IndexType::IVF_RABITQ: - valid = UpdateFastQueryParam( - params, engine, defaults, [](const auto *p, auto *e, const auto *d) { - e->nprobe = p ? p->nprobe() : d->nprobe; - }); - break; - case IndexType::DISKANN: - valid = UpdateFastQueryParam( - params, engine, defaults, [](const auto *p, auto *e, const auto *d) { - e->list_size = p ? p->list_size() : d->list_size; - }); - break; - case IndexType::FLAT: - valid = - UpdateFastQueryParam( - params, engine, defaults, - [](const auto *, auto *, const auto *) {}); - break; - default: - break; - } - if (!valid) - return Status::InvalidArgument( - "fast query: unsupported query parameter type"); + auto status = ProximaEngineHelper::update_engine_query_param( + state.index_type, params, engine, state.default_engine_query_param.get()); + CHECK_RETURN_STATUS(status); if (engine) { - engine->radius = params ? params->radius() : defaults->radius; - engine->is_linear = params ? params->is_linear() : defaults->is_linear; engine->refiner_param = params && params->is_using_refiner() ? state.refiner_param : nullptr; if (engine->refiner_param) { @@ -2067,7 +1984,7 @@ Result CollectionImpl::fast_query( } const bool refine = query_params && query_params->is_using_refiner(); if (state.raw_index && state.engine_query_param && - (!refine || state.raw_reference_index)) { + (!refine || state.refiner_param)) { state.engine_query_param->topk = static_cast(topk); core_interface::DenseVector dense_query{query_vector}; core_interface::VectorData query_data{dense_query}; @@ -2084,19 +2001,15 @@ Result CollectionImpl::fast_query( return out; } - const auto &cache = state; - - const int search_topk = topk; - - const MetricType metric = cache.metric; + const MetricType metric = state.metric; auto better = [metric](float a, float b) { return metric == MetricType::IP ? a > b : a < b; }; vector_column_params::QueryParams qp; - qp.topk = static_cast(search_topk); - qp.data_type = cache.data_type; - qp.dimension = cache.dimension; + qp.topk = static_cast(topk); + qp.data_type = state.data_type; + qp.dimension = state.dimension; qp.query_params = query_params; vector_column_params::VectorData vector_data; @@ -2113,8 +2026,7 @@ Result CollectionImpl::fast_query( IndexResults::Ptr results = std::move(res.value()); std::vector indices; - indices.reserve(search_topk); - const size_t score_begin = out->scores.size(); + indices.reserve(topk); for (auto it = results->create_iterator(); it->valid(); it->next()) { indices.push_back(static_cast(it->doc_id())); out->scores.push_back(it->score()); @@ -2123,14 +2035,7 @@ Result CollectionImpl::fast_query( return Status::OK(); } - std::vector ids; - auto s = entry.segment->get_global_doc_ids(indices, ids); - if (!s.ok()) { - out->scores.resize(score_begin); - return s; - } - out->ids.insert(out->ids.end(), ids.begin(), ids.end()); - return Status::OK(); + return entry.segment->get_global_doc_ids(indices, out->ids); }; auto finish = [topk, return_scores](FastQueryResult out) { @@ -2144,20 +2049,16 @@ Result CollectionImpl::fast_query( return out; }; - if (cache.segments.size() == 1) { + if (state.segments.size() == 1) { FastQueryResult out; - auto s = search_one(cache.segments[0], &out); + auto s = search_one(state.segments[0], &out); CHECK_RETURN_STATUS_EXPECTED(s); - if (static_cast(out.ids.size()) > topk) { - out.ids.resize(topk); - out.scores.resize(topk); - } return finish(std::move(out)); } std::vector> candidates; - candidates.reserve(static_cast(search_topk) * cache.segments.size()); - for (const auto &entry : cache.segments) { + candidates.reserve(static_cast(topk) * state.segments.size()); + for (const auto &entry : state.segments) { FastQueryResult seg_out; auto s = search_one(entry, &seg_out); CHECK_RETURN_STATUS_EXPECTED(s); diff --git a/src/db/index/column/vector_column/engine_helper.hpp b/src/db/index/column/vector_column/engine_helper.hpp index b82ad817b..a4232b526 100644 --- a/src/db/index/column/vector_column/engine_helper.hpp +++ b/src/db/index/column/vector_column/engine_helper.hpp @@ -95,6 +95,19 @@ class ProximaEngineHelper { } private: + template + static bool _update_query_param( + const QueryParams::Ptr ¶ms, + core_interface::BaseIndexQueryParam *engine, + const core_interface::BaseIndexQueryParam *defaults, Update update) { + const auto *p = dynamic_cast(params.get()); + if (params && !p) return false; + if (auto *e = dynamic_cast(engine)) { + update(p, e, static_cast(defaults)); + } + return true; + } + template static Result> _build_common_query_param( @@ -107,8 +120,10 @@ class ProximaEngineHelper { convert_to_engine_filter(db_query_params.filter); if (db_query_params.query_params) { - engine_query_param->radius = db_query_params.query_params->radius(); - engine_query_param->is_linear = db_query_params.query_params->is_linear(); + auto status = update_engine_query_param( + db_query_params.query_params->type(), db_query_params.query_params, + engine_query_param.get(), nullptr); + if (!status.ok()) return tl::make_unexpected(status); } if (db_query_params.refiner_param) { { @@ -135,6 +150,89 @@ class ProximaEngineHelper { } public: + // Update an existing engine parameter object without allocating. The DB + // type can differ from the engine type while an untrained graph uses Flat. + // A null engine validates only; defaults restores values when params is null. + static Status update_engine_query_param( + IndexType index_type, const zvec::QueryParams::Ptr ¶ms, + core_interface::BaseIndexQueryParam *engine, + const core_interface::BaseIndexQueryParam *defaults) { + if (params && params->type() != index_type) { + return Status::InvalidArgument( + "query parameter type does not match the field index"); + } + bool valid = false; + switch (index_type) { + case IndexType::VAMANA: + valid = _update_query_param( + params, engine, defaults, + [](const auto *p, auto *e, const auto *d) { + e->ef_search = p ? p->ef_search() : d->ef_search; + e->prefetch_offset = + p ? p->prefetch_offset() : d->prefetch_offset; + e->prefetch_lines = p ? p->prefetch_lines() : d->prefetch_lines; + }); + break; + case IndexType::HNSW: + valid = _update_query_param( + params, engine, defaults, + [](const auto *p, auto *e, const auto *d) { + e->ef_search = p ? p->ef() : d->ef_search; + e->prefetch_offset = + p ? p->prefetch_offset() : d->prefetch_offset; + e->prefetch_lines = p ? p->prefetch_lines() : d->prefetch_lines; + }); + break; + case IndexType::HNSW_RABITQ: + valid = _update_query_param( + params, engine, defaults, + [](const auto *p, auto *e, const auto *d) { + e->ef_search = p ? p->ef() : d->ef_search; + }); + break; + case IndexType::IVF: + valid = + _update_query_param( + params, engine, defaults, + [](const auto *p, auto *e, const auto *d) { + e->nprobe = p ? p->nprobe() : d->nprobe; + }); + break; + case IndexType::IVF_RABITQ: + valid = _update_query_param( + params, engine, defaults, + [](const auto *p, auto *e, const auto *d) { + e->nprobe = p ? p->nprobe() : d->nprobe; + }); + break; + case IndexType::DISKANN: + valid = _update_query_param( + params, engine, defaults, + [](const auto *p, auto *e, const auto *d) { + e->list_size = p ? p->list_size() : d->list_size; + }); + break; + case IndexType::FLAT: + // Flat has no index-specific fields; base QueryParams is sufficient. + valid = true; + break; + default: + break; + } + if (!valid) + return Status::InvalidArgument("unsupported query parameter type"); + if (engine) { + engine->radius = params ? params->radius() : defaults->radius; + engine->is_linear = params ? params->is_linear() : defaults->is_linear; + } + return Status::OK(); + } + static Result> convert_to_engine_query_param( const FieldSchema &field_schema, @@ -168,15 +266,6 @@ class ProximaEngineHelper { hnsw_query_param_result.error().message())); } auto &hnsw_query_param = hnsw_query_param_result.value(); - if (query_params.query_params) { - auto db_hnsw_query_params = dynamic_cast( - query_params.query_params.get()); - hnsw_query_param->ef_search = db_hnsw_query_params->ef(); - hnsw_query_param->prefetch_offset = - db_hnsw_query_params->prefetch_offset(); - hnsw_query_param->prefetch_lines = - db_hnsw_query_params->prefetch_lines(); - } return std::move(hnsw_query_param); } @@ -190,12 +279,6 @@ class ProximaEngineHelper { hnsw_query_param_result.error().message())); } auto &hnsw_query_param = hnsw_query_param_result.value(); - if (query_params.query_params) { - auto db_hnsw_rabitq_query_params = - dynamic_cast( - query_params.query_params.get()); - hnsw_query_param->ef_search = db_hnsw_rabitq_query_params->ef(); - } return std::move(hnsw_query_param); } @@ -209,11 +292,6 @@ class ProximaEngineHelper { ivf_query_param_result.error().message())); } auto &ivf_query_param = ivf_query_param_result.value(); - if (query_params.query_params) { - auto db_ivf_query_params = dynamic_cast( - query_params.query_params.get()); - ivf_query_param->nprobe = db_ivf_query_params->nprobe(); - } return std::move(ivf_query_param); } @@ -227,12 +305,6 @@ class ProximaEngineHelper { ivf_rabitq_query_param_result.error().message())); } auto &ivf_rabitq_query_param = ivf_rabitq_query_param_result.value(); - if (query_params.query_params) { - auto db_ivf_rabitq_query_params = - dynamic_cast( - query_params.query_params.get()); - ivf_rabitq_query_param->nprobe = db_ivf_rabitq_query_params->nprobe(); - } return std::move(ivf_rabitq_query_param); } @@ -246,13 +318,6 @@ class ProximaEngineHelper { diskann_query_param_result.error().message())); } auto &diskann_query_param = diskann_query_param_result.value(); - if (query_params.query_params) { - auto db_diskann_query_params = - dynamic_cast( - query_params.query_params.get()); - diskann_query_param->list_size = - static_cast(db_diskann_query_params->list_size()); - } return std::move(diskann_query_param); } @@ -266,16 +331,6 @@ class ProximaEngineHelper { vamana_query_param_result.error().message())); } auto &vamana_query_param = vamana_query_param_result.value(); - if (query_params.query_params) { - auto db_vamana_query_params = dynamic_cast( - query_params.query_params.get()); - vamana_query_param->ef_search = - static_cast(db_vamana_query_params->ef_search()); - vamana_query_param->prefetch_offset = - db_vamana_query_params->prefetch_offset(); - vamana_query_param->prefetch_lines = - db_vamana_query_params->prefetch_lines(); - } return std::move(vamana_query_param); } diff --git a/tests/db/index/column/vector_column_indexer_test.cc b/tests/db/index/column/vector_column_indexer_test.cc index b0bfd9c81..e5019a834 100644 --- a/tests/db/index/column/vector_column_indexer_test.cc +++ b/tests/db/index/column/vector_column_indexer_test.cc @@ -18,6 +18,7 @@ #include #include #include +#include "db/index/column/vector_column/engine_helper.hpp" #include "db/index/column/vector_column/vector_column_params.h" #include "tests/test_util.h" #include "zvec/ailego/utility/float_helper.h" @@ -32,6 +33,68 @@ using namespace zvec; +TEST(VectorColumnIndexerTest, ReusedQueryParametersAndFlatFallback) { + FieldSchema field("vector", DataType::VECTOR_FP32, 32, false, + std::make_shared(MetricType::L2)); + vector_column_params::QueryParams query; + auto result = + ProximaEngineHelper::convert_to_engine_query_param(field, query); + ASSERT_TRUE(result); + auto engine = std::move(result.value()); + auto defaults = engine->clone(); + auto *typed = dynamic_cast(engine.get()); + ASSERT_NE(typed, nullptr); + auto params = std::make_shared(); + for (int ef : {128, 32}) { + params->set_ef_search(ef); + params->set_prefetch_offset(3); + params->set_prefetch_lines(2); + params->set_radius(1.25f); + params->set_is_linear(true); + ASSERT_TRUE(ProximaEngineHelper::update_engine_query_param( + IndexType::VAMANA, params, engine.get(), defaults.get()) + .ok()); + EXPECT_EQ(typed->ef_search, ef); + EXPECT_EQ(typed->prefetch_offset, 3U); + EXPECT_EQ(typed->prefetch_lines, 2U); + EXPECT_FLOAT_EQ(typed->radius, 1.25f); + EXPECT_TRUE(typed->is_linear); + query.query_params = params; + auto fresh = + ProximaEngineHelper::convert_to_engine_query_param(field, query); + ASSERT_TRUE(fresh); + auto *fresh_typed = + dynamic_cast(fresh->get()); + ASSERT_NE(fresh_typed, nullptr); + EXPECT_EQ(fresh_typed->ef_search, typed->ef_search); + EXPECT_EQ(fresh_typed->prefetch_offset, typed->prefetch_offset); + EXPECT_EQ(fresh_typed->prefetch_lines, typed->prefetch_lines); + EXPECT_EQ(fresh_typed->radius, typed->radius); + EXPECT_EQ(fresh_typed->is_linear, typed->is_linear); + ASSERT_TRUE(ProximaEngineHelper::update_engine_query_param( + IndexType::VAMANA, nullptr, engine.get(), defaults.get()) + .ok()); + EXPECT_EQ(typed->ef_search, core_interface::kDefaultVamanaEfSearch); + EXPECT_EQ(typed->prefetch_offset, core_interface::kDefaultPrefetchOffset); + EXPECT_EQ(typed->prefetch_lines, core_interface::kDefaultPrefetchLines); + EXPECT_FLOAT_EQ(typed->radius, 0.0f); + EXPECT_FALSE(typed->is_linear); + } + // Untrained graph fields use Flat storage but retain graph query parameters. + field.set_index_params(std::make_shared(MetricType::L2)); + auto flat = ProximaEngineHelper::convert_to_engine_query_param(field, query); + ASSERT_TRUE(flat); + EXPECT_NE(dynamic_cast(flat->get()), + nullptr); + EXPECT_FLOAT_EQ(flat.value()->radius, 1.25f); + EXPECT_TRUE(flat.value()->is_linear); + auto invalid = std::make_shared(IndexType::VAMANA); + EXPECT_FALSE(ProximaEngineHelper::update_engine_query_param( + IndexType::VAMANA, invalid, engine.get(), defaults.get()) + .ok()); +} + + std::string print_dense_vector(const void *vector, size_t dim, DataType data_type) { std::stringstream ss; From 6e6aa966ffae06c9c1d8a964528a6a236ef65fb2 Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Thu, 17 Sep 2026 12:19:13 +0800 Subject: [PATCH 7/8] fix(search): restore fast query CI after Turbo Flat migration Use Turbo distance evaluation and score normalization in fast Flat searches. Preserve base query parameters when graph indexes fall back to Flat, and keep strict validation for concrete graph query paths. Stabilize the queue recovery test by leaving background eviction headroom. Validation: seven related C++ test executables passed; Python regressions passed 277 tests with seven platform skips; queue recovery passed 500 repetitions; clang-format and git diff checks passed. --- src/core/algorithm/flat/flat_streamer.cc | 2 +- .../algorithm/flat/flat_streamer_entity.cc | 5 +++- src/core/interface/index.cc | 7 ++++- .../column/vector_column/engine_helper.hpp | 6 ++++- .../column/vector_column_indexer_test.cc | 27 +++++++++++++++++++ 5 files changed, 43 insertions(+), 4 deletions(-) diff --git a/src/core/algorithm/flat/flat_streamer.cc b/src/core/algorithm/flat/flat_streamer.cc index dc47710a9..46e975398 100644 --- a/src/core/algorithm/flat/flat_streamer.cc +++ b/src/core/algorithm/flat/flat_streamer.cc @@ -455,7 +455,7 @@ int FlatStreamer::search_by_p_keys_fast( float *output_scores, size_t topk, const IndexQueryMeta &qmeta, Context::UPointer &context) const { if (!query || !output_ids || topk == 0 || !context || - !metric_->is_matched(meta_, qmeta)) { + (!quantizer_ && !metric_->is_matched(meta_, qmeta))) { return IndexError_InvalidArgument; } auto *flat_context = diff --git a/src/core/algorithm/flat/flat_streamer_entity.cc b/src/core/algorithm/flat/flat_streamer_entity.cc index a7d623374..75aaebfba 100644 --- a/src/core/algorithm/flat/flat_streamer_entity.cc +++ b/src/core/algorithm/flat/flat_streamer_entity.cc @@ -645,7 +645,10 @@ int FlatContiguousStreamerEntity::search_by_p_keys_fast( preprocess(buffer.data(), meta().dimension()); batch_query = buffer.data(); } - if (const auto &batch = batch_distance(); batch) { + if (quantizer()) { + quantizer()->calc_distance_dp_query_batch( + ptrs.data(), static_cast(count), query, distances.data()); + } else if (const auto &batch = batch_distance(); batch) { batch(ptrs.data(), batch_query, count, meta().dimension(), distances.data(), has_extras ? extras.data() : nullptr); } else { diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc index 09f0b157e..54020fd20 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -896,10 +896,15 @@ int Index::_normalize_buffer_scores(const VectorData &vector_data, const int64_t *output_ids, float *output_scores, size_t count) { if (!output_scores || count == 0) return 0; - if (metric_->support_normalize()) { + if (metric_ && metric_->support_normalize()) { for (size_t i = 0; i < count; ++i) { metric_->normalize(output_scores + i); } + } else if (turbo_quantizer_ && + turbo_quantizer_->support_score_normalization()) { + for (size_t i = 0; i < count; ++i) { + turbo_quantizer_->normalize_score(output_scores + i); + } } if (!reformer_) return 0; if (!std::holds_alternative(vector_data.vector)) { diff --git a/src/db/index/column/vector_column/engine_helper.hpp b/src/db/index/column/vector_column/engine_helper.hpp index a4232b526..3ee50533d 100644 --- a/src/db/index/column/vector_column/engine_helper.hpp +++ b/src/db/index/column/vector_column/engine_helper.hpp @@ -101,7 +101,11 @@ class ProximaEngineHelper { core_interface::BaseIndexQueryParam *engine, const core_interface::BaseIndexQueryParam *defaults, Update update) { const auto *p = dynamic_cast(params.get()); - if (params && !p) return false; + if (params && !p) { + // Flat fallback consumes only common fields, so base QueryParams is + // sufficient even when it carries the original graph index type. + return dynamic_cast(engine) != nullptr; + } if (auto *e = dynamic_cast(engine)) { update(p, e, static_cast(defaults)); } diff --git a/tests/db/index/column/vector_column_indexer_test.cc b/tests/db/index/column/vector_column_indexer_test.cc index e5019a834..b41c862bf 100644 --- a/tests/db/index/column/vector_column_indexer_test.cc +++ b/tests/db/index/column/vector_column_indexer_test.cc @@ -95,6 +95,33 @@ TEST(VectorColumnIndexerTest, ReusedQueryParametersAndFlatFallback) { } +TEST(VectorColumnIndexerTest, BaseQueryParametersForFlatFallback) { + FieldSchema field("vector", DataType::VECTOR_FP32, 32, false, + std::make_shared(MetricType::L2)); + for (auto type : {IndexType::HNSW, IndexType::HNSW_RABITQ, IndexType::VAMANA, + IndexType::IVF, IndexType::IVF_RABITQ, IndexType::DISKANN}) { + SCOPED_TRACE(static_cast(type)); + vector_column_params::QueryParams query; + query.topk = 7; + query.query_params = std::make_shared(type); + query.query_params->set_radius(1.25f); + query.query_params->set_is_linear(true); + auto result = + ProximaEngineHelper::convert_to_engine_query_param(field, query); + ASSERT_TRUE(result) << result.error().message(); + ASSERT_NE(nullptr, + dynamic_cast(result->get())); + EXPECT_EQ(7U, result.value()->topk); + EXPECT_FLOAT_EQ(1.25f, result.value()->radius); + EXPECT_TRUE(result.value()->is_linear); + // Validation without a Flat fallback still requires the concrete type. + EXPECT_FALSE(ProximaEngineHelper::update_engine_query_param( + type, query.query_params, nullptr, nullptr) + .ok()); + } +} + + std::string print_dense_vector(const void *vector, size_t dim, DataType data_type) { std::stringstream ss; From 9ddd8efd009cad59c47c06ff74667d024a35f613 Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Fri, 18 Sep 2026 11:18:11 +0800 Subject: [PATCH 8/8] refactor(query): prepare immutable fast query metadata at open --- python/tests/test_fast_query.py | 38 +- python/zvec/model/collection.py | 3 +- src/binding/python/model/python_collection.cc | 2 +- src/core/interface/index.cc | 31 +- src/db/collection.cc | 331 ++++++------------ .../combined_vector_column_indexer.cc | 44 +++ .../combined_vector_column_indexer.h | 21 +- .../column/vector_column/engine_helper.hpp | 12 +- .../vector_column/vector_column_indexer.cc | 21 ++ .../vector_column/vector_column_indexer.h | 6 + src/db/index/segment/segment.cc | 48 +-- src/db/index/segment/segment.h | 14 +- src/include/zvec/core/interface/index.h | 2 - src/include/zvec/db/collection.h | 5 +- tests/db/fast_query_test.cc | 225 ++++++++++++ tests/db/index/segment/segment_row_id_test.cc | 84 +++++ tests/db/sqlengine/mock_segment.h | 8 +- 17 files changed, 599 insertions(+), 296 deletions(-) diff --git a/python/tests/test_fast_query.py b/python/tests/test_fast_query.py index c7957d0f9..12f26d230 100644 --- a/python/tests/test_fast_query.py +++ b/python/tests/test_fast_query.py @@ -1,5 +1,8 @@ """Advanced dense search: collection queries, scores, fallback and lifetime.""" +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier + import numpy as np import pytest @@ -185,7 +188,7 @@ def test_reused_inline_and_default_params_and_close(collection): Query(field_name="vector", vector=query, param=param), topk=topk ) np.testing.assert_array_equal(ids, [int(doc.id[4:]) for doc in docs]) - # None must restore defaults after a customized query. + # A call with None uses defaults regardless of the preceding query. docs = coll.query(Query(field_name="vector", vector=query), topk=10) np.testing.assert_array_equal( coll.fast_query("vector", query), [int(doc.id[4:]) for doc in docs] @@ -207,6 +210,39 @@ def test_reused_inline_and_default_params_and_close(collection): obj.fast_query("vector", query, param) +def test_concurrent_calls_keep_query_parameters_local(collection): + coll, vectors, param_type = collection + queries = np.ascontiguousarray(vectors[[3, 17, 29, 41]] + 0.013) + params = [None] + for ef in (16, 64, 80): + settings = {"ef": ef} if param_type is HnswQueryParam else {"ef_search": ef} + params.append(param_type(**settings)) + topks = [1, 21, 3, 10] + expected = [] + for query, param, topk in zip(queries, params, topks): + docs = coll.query(Query("vector", vector=query, param=param), topk=topk) + expected.append( + ([int(doc.id[4:]) for doc in docs], [doc.score for doc in docs]) + ) + + ready = Barrier(4) + + def search(worker): + # Exercise the first calls concurrently, as well as repeated queries + # with different topk, defaults and graph search parameters. + ready.wait() + for repeat in range(40): + i = (worker + repeat) % len(queries) + ids, scores = coll.fast_query( + "vector", queries[i], params[i], topk=topks[i], return_scores=True + ) + np.testing.assert_array_equal(ids, expected[i][0]) + np.testing.assert_allclose(scores, expected[i][1], rtol=2e-5, atol=2e-5) + + with ThreadPoolExecutor(max_workers=4) as pool: + list(pool.map(search, range(4))) + + @pytest.mark.parametrize("compact", [False, True], ids=["delete_filter", "compacted"]) def test_internal_ids_survive_deletion_and_compaction(tmp_path, compact): vectors = np.random.default_rng(823).normal(size=(64, 32)).astype(np.float32) diff --git a/python/zvec/model/collection.py b/python/zvec/model/collection.py index e382033d8..f6f4bded6 100644 --- a/python/zvec/model/collection.py +++ b/python/zvec/model/collection.py @@ -521,8 +521,7 @@ def fast_query( This advanced API requires a read-only collection. No preparation is required: parameters may be constructed inline or reused across calls. - Index references are cached internally; parameters are read each time. - Calls to ``fast_query`` and collection close must be serial. + Parameters are read each time; execution state belongs to each call. ``vector`` must be a contiguous 1D NumPy array matching the field's input dtype and dimension. The result is an owning int64 array. With diff --git a/src/binding/python/model/python_collection.cc b/src/binding/python/model/python_collection.cc index 9f302cd84..355708e0d 100644 --- a/src/binding/python/model/python_collection.cc +++ b/src/binding/python/model/python_collection.cc @@ -424,7 +424,7 @@ may be constructed inline or reused. Returns an owning int64 internal ID array, or (ids, float32 scores) with return_scores=True. Scores include refinement. Missing results are padded with ID -1 / score NaN. Refinement uses param.scale_factor with the same semantics as Collection.query. -Calls to fast_query and collection.close must be serial. +Parameters and execution state belong to each call, as with Collection.query. )doc") .def("GroupByQuery", [](const Collection &self, const GroupByVectorQuery &query) { diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc index 54020fd20..2bd105d33 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -187,22 +187,16 @@ thread_local static std::array(is_sparse_); - if (_context_list[context_index_] == nullptr) { - if ((_context_list[context_index_] = streamer_->create_context()) == - nullptr) { - LOG_ERROR("Failed to create context"); - return false; - } - } - return true; -} - core::IndexContext::Pointer &Index::acquire_context() { - init_context(); - return _context_list[context_index_]; + const size_t context_index = + (magic_enum::enum_integer(param_.index_type) - 1) * 2 + + static_cast(is_sparse_); + auto &context = _context_list[context_index]; + if (!context) { + context = streamer_->create_context(); + if (!context) LOG_ERROR("Failed to create context"); + } + return context; } int Index::train() { @@ -585,14 +579,17 @@ 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 (!acquire_context()) { LOG_ERROR("Failed to init context"); return core::IndexError_Runtime; } is_open_ = true; is_read_only_ = storage_options.read_only; + // These streamer indexes use the no-op base train(). Finish that state + // transition before publication so concurrent first searches only read it. + // Builder-based indexes override open() and manage their own training. + is_trained_ = true; return 0; } diff --git a/src/db/collection.cc b/src/db/collection.cc index af054624c..48526c161 100644 --- a/src/db/collection.cc +++ b/src/db/collection.cc @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -183,6 +184,9 @@ class CollectionImpl : public Collection { Status recovery(); + // Called only during a read-only open, after all segments have recovered. + void prepare_fast_query(); + Status create_idmap_and_delete_store(); Status recover_idmap_and_delete_store(); @@ -300,30 +304,6 @@ class CollectionImpl : public Collection { ColumnOp op); private: - struct FastQuerySegment { - Segment::Ptr segment; - CombinedVectorColumnIndexer::Ptr indexer; - IndexFilter::Ptr filter; - }; - - struct FastQueryState { - DataType data_type{DataType::UNDEFINED}; - uint32_t dimension{0}; - IndexType index_type{IndexType::UNDEFINED}; - MetricType metric{MetricType::L2}; - std::vector segments; - core_interface::Index::Pointer raw_index; - core_interface::BaseIndexQueryParam::Pointer engine_query_param; - core_interface::BaseIndexQueryParam::Pointer default_engine_query_param; - core_interface::RefinerParam::Pointer refiner_param; - }; - - Result resolve_fast_query( - const std::string &field_name) const; - Status update_fast_query_params(FastQueryState &state, - const QueryParams::Ptr ¶ms) const; - // Only fast_query uses these caches. Its calls and close must be serial. - mutable std::unordered_map fast_query_states_; std::string path_; bool destroyed_{false}; @@ -336,6 +316,18 @@ class CollectionImpl : public Collection { CollectionOptions options_; + struct FastQueryField { + // Borrowed from schema_; the owning collection is read-only. + const FieldSchema *schema; + // Same order as read_only_segments_; null means no index in that segment. + std::vector indexers; + }; + + // Prepared during open and immutable until close. Queries hold + // schema_handle_mtx_ shared; close holds it exclusively before clearing. + std::vector read_only_segments_; + std::unordered_map fast_query_fields_; + mutable std::shared_mutex schema_handle_mtx_; // Number of open iterators, guarded by schema_handle_mtx_ (exclusive). int active_iterators_{0}; @@ -433,6 +425,10 @@ Status CollectionImpl::open(const CollectionOptions &options) { s = create(); } + if (s.ok() && options_.read_only_) { + prepare_fast_query(); + } + auto profiler = std::make_shared(); sql_engine_ = sqlengine::SQLEngine::create(profiler); @@ -490,7 +486,8 @@ Status CollectionImpl::close_unsafe() { } } - fast_query_states_.clear(); + fast_query_fields_.clear(); + read_only_segments_.clear(); // always release resources regardless of flush outcome writing_segment_.reset(); @@ -1851,238 +1848,134 @@ Result CollectionImpl::query(const SearchQuery &query) const { return query_unsafe(query); } -Result CollectionImpl::resolve_fast_query( - const std::string &field_name) const { - std::shared_lock schema_lock(schema_handle_mtx_); - const FieldSchema *field_schema = - field_name.empty() ? nullptr : schema_->get_field(field_name); - if (field_schema == nullptr || !field_schema->is_dense_vector()) { - return tl::make_unexpected(Status::InvalidArgument( - "fast search requires a dense vector field: ", field_name)); - } - - bool quantized = false; - MetricType metric = MetricType::L2; - if (const auto *vp = dynamic_cast( - field_schema->index_params().get())) { - metric = vp->metric_type(); - quantized = vp->quantize_type() != QuantizeType::UNDEFINED; - } - - FastQueryState cache; - cache.data_type = field_schema->data_type(); - cache.dimension = field_schema->dimension(); - cache.index_type = field_schema->index_params() - ? field_schema->index_params()->type() - : IndexType::UNDEFINED; - cache.metric = metric; - const auto segments = get_all_segments(); - if (segments.empty()) return cache; - cache.segments.reserve(segments.size()); - for (const auto &seg : segments) { - CombinedVectorColumnIndexer::Ptr indexer; - if (!quantized) { - indexer = seg->get_combined_vector_indexer(field_name); - } else { - indexer = seg->get_quant_combined_vector_indexer(field_name); - } - if (!indexer || !indexer->has_searchable_indexers()) { - continue; - } - FastQuerySegment entry; - entry.segment = seg; - entry.indexer = std::move(indexer); - entry.filter = seg->get_filter(); - cache.segments.push_back(std::move(entry)); - } - - if (cache.segments.empty()) { - return tl::make_unexpected(Status::InvalidArgument( - "fast query: no searchable vector index for field ", field_name)); - } - - if (cache.segments.size() == 1) { - const auto &entry = cache.segments[0]; - if (entry.filter == nullptr && entry.indexer->is_single_block() && - entry.segment->has_identity_doc_ids()) { - auto primary = entry.indexer->primary_indexer(); - if (primary) { - cache.raw_index = primary->debug_get_index(); - auto reference = entry.indexer->reference_indexer(); - auto reference_index = - reference ? reference->debug_get_index() : nullptr; - if (reference_index) { - cache.refiner_param = - std::make_shared(); - cache.refiner_param->reference_index = std::move(reference_index); - } +void CollectionImpl::prepare_fast_query() { + read_only_segments_ = get_all_segments(); + for (const auto &field : schema_->vector_fields()) { + if (!field->is_dense_vector()) continue; + const auto *index_params = + dynamic_cast(field->index_params().get()); + FastQueryField resolved{field.get(), {}}; + resolved.indexers.reserve(read_only_segments_.size()); + for (const auto &segment : read_only_segments_) { + CombinedVectorColumnIndexer::Ptr indexer; + if (index_params) { + indexer = + index_params->quantize_type() == QuantizeType::UNDEFINED + ? segment->get_combined_vector_indexer(field->name()) + : segment->get_quant_combined_vector_indexer(field->name()); } + resolved.indexers.push_back(std::move(indexer)); } + fast_query_fields_.emplace(field->name(), std::move(resolved)); } - - if (!cache.raw_index) return cache; - auto primary = cache.segments[0].indexer->primary_indexer(); - vector_column_params::QueryParams qp; - qp.data_type = cache.data_type; - qp.dimension = cache.dimension; - auto engine_qp = ProximaEngineHelper::convert_to_engine_query_param( - primary->field_schema(), qp); - if (!engine_qp) return tl::make_unexpected(engine_qp.error()); - cache.engine_query_param = std::move(engine_qp.value()); - cache.default_engine_query_param = cache.engine_query_param->clone(); - return cache; -} - -Status CollectionImpl::update_fast_query_params( - FastQueryState &state, const QueryParams::Ptr ¶ms) const { - auto *engine = state.engine_query_param.get(); - auto status = ProximaEngineHelper::update_engine_query_param( - state.index_type, params, engine, state.default_engine_query_param.get()); - CHECK_RETURN_STATUS(status); - if (engine) { - engine->refiner_param = - params && params->is_using_refiner() ? state.refiner_param : nullptr; - if (engine->refiner_param) { - engine->refiner_param->scale_factor_ = params->scale_factor(); - } - } - return Status::OK(); } Result CollectionImpl::fast_query( const std::string &field_name, const void *query_vector, const QueryParams::Ptr &query_params, int topk, bool return_scores, DataType query_data_type, uint32_t query_dimension) const { + std::shared_lock lock(schema_handle_mtx_); CHECK_DESTROY_RETURN_STATUS_EXPECTED(destroyed_, false); CHECK_CLOSED_RETURN_STATUS_EXPECTED(closed_, false); if (!options_.read_only_) { return tl::make_unexpected( Status::InvalidArgument("fast query requires a read-only collection")); } - auto cached = fast_query_states_.find(field_name); - if (cached == fast_query_states_.end()) { - auto resolved = resolve_fast_query(field_name); - if (!resolved) return tl::make_unexpected(resolved.error()); - cached = fast_query_states_.emplace(field_name, std::move(resolved.value())) - .first; + + const auto field = fast_query_fields_.find(field_name); + if (field == fast_query_fields_.end()) { + return tl::make_unexpected(Status::InvalidArgument( + "fast search requires a dense vector field: ", field_name)); } - auto &state = cached->second; - const auto param_status = update_fast_query_params(state, query_params); + const auto *field_schema = field->second.schema; + const auto *index_params = dynamic_cast( + field_schema->index_params().get()); + const IndexType index_type = field_schema->index_params() + ? field_schema->index_params()->type() + : IndexType::UNDEFINED; + const auto param_status = ProximaEngineHelper::update_engine_query_param( + index_type, query_params, nullptr, nullptr); CHECK_RETURN_STATUS_EXPECTED(param_status); if (query_vector == nullptr) { return tl::make_unexpected( Status::InvalidArgument("fast query: query_vector is null")); } - if (topk <= 0) { - return FastQueryResult{}; - } + if (topk <= 0) return FastQueryResult{}; if ((query_data_type != DataType::UNDEFINED || query_dimension != 0) && - (query_data_type != state.data_type || - query_dimension != state.dimension)) { + (query_data_type != field_schema->data_type() || + query_dimension != field_schema->dimension())) { return tl::make_unexpected(Status::InvalidArgument( "query vector dtype or dimension does not match the field")); } - const bool refine = query_params && query_params->is_using_refiner(); - if (state.raw_index && state.engine_query_param && - (!refine || state.refiner_param)) { - state.engine_query_param->topk = static_cast(topk); - core_interface::DenseVector dense_query{query_vector}; - core_interface::VectorData query_data{dense_query}; - FastQueryResult out; - out.ids.resize(static_cast(topk), int64_t{-1}); - if (return_scores) out.scores.resize(static_cast(topk)); - const int ret = state.raw_index->search_fast( - query_data, state.engine_query_param, out.ids.data(), - return_scores ? out.scores.data() : nullptr); - if (ret != 0) { - return tl::make_unexpected( - Status::InternalError("fast query: index search failed")); - } - return out; - } - - const MetricType metric = state.metric; - auto better = [metric](float a, float b) { - return metric == MetricType::IP ? a > b : a < b; - }; - - vector_column_params::QueryParams qp; - qp.topk = static_cast(topk); - qp.data_type = state.data_type; - qp.dimension = state.dimension; - qp.query_params = query_params; + const auto &segments = read_only_segments_; + const auto &indexers = field->second.indexers; + const MetricType metric = + index_params ? index_params->metric_type() : MetricType::L2; + vector_column_params::QueryParams params; + params.topk = static_cast(topk); + params.data_type = field_schema->data_type(); + params.dimension = field_schema->dimension(); + params.query_params = query_params; vector_column_params::VectorData vector_data; vector_data.vector = vector_column_params::DenseVector{query_vector}; - auto search_one = [&](const FastQuerySegment &entry, - FastQueryResult *out) -> Status { - qp.filter = entry.filter ? entry.filter.get() : nullptr; - - auto res = entry.indexer->Search(vector_data, qp); - if (!res) { - return res.error(); - } - IndexResults::Ptr results = std::move(res.value()); - - std::vector indices; - indices.reserve(topk); - for (auto it = results->create_iterator(); it->valid(); it->next()) { - indices.push_back(static_cast(it->doc_id())); - out->scores.push_back(it->score()); - } - if (indices.empty()) { - return Status::OK(); - } - - return entry.segment->get_global_doc_ids(indices, out->ids); - }; - - auto finish = [topk, return_scores](FastQueryResult out) { - out.ids.resize(static_cast(topk), int64_t{-1}); - if (return_scores) { - out.scores.resize(static_cast(topk), - std::numeric_limits::quiet_NaN()); - } else { - out.scores.clear(); - } - return out; - }; - - if (state.segments.size() == 1) { - FastQueryResult out; - auto s = search_one(state.segments[0], &out); - CHECK_RETURN_STATUS_EXPECTED(s); - return finish(std::move(out)); + FastQueryResult out; + out.ids.resize(static_cast(topk), int64_t{-1}); + // Segment merging needs scores even when the caller only requests IDs. + if (return_scores || segments.size() > 1) { + out.scores.resize(static_cast(topk), + std::numeric_limits::quiet_NaN()); } - std::vector> candidates; - candidates.reserve(static_cast(topk) * state.segments.size()); - for (const auto &entry : state.segments) { - FastQueryResult seg_out; - auto s = search_one(entry, &seg_out); - CHECK_RETURN_STATUS_EXPECTED(s); - for (size_t i = 0; i < seg_out.ids.size(); ++i) { - candidates.emplace_back(seg_out.scores[i], seg_out.ids[i]); - } + if (segments.size() > 1) { + candidates.reserve(static_cast(topk) * segments.size()); + } + bool searched = false; + for (size_t segment_id = 0; segment_id < segments.size(); ++segment_id) { + const auto &segment = segments[segment_id]; + const auto &indexer = indexers[segment_id]; + if (!indexer || !indexer->has_searchable_indexers()) continue; + searched = true; + auto filter = segment->get_filter(); + params.filter = filter.get(); + auto status = + indexer->SearchFast(vector_data, params, out.ids.data(), + out.scores.empty() ? nullptr : out.scores.data()); + CHECK_RETURN_STATUS_EXPECTED(status); + // The read-only collection contract makes this property immutable. + if (!segment->has_identity_doc_ids()) { + status = segment->get_global_doc_ids(out.ids); + CHECK_RETURN_STATUS_EXPECTED(status); + } + if (segments.size() == 1) return out; + for (size_t i = 0; i < out.ids.size() && out.ids[i] != -1; ++i) { + candidates.emplace_back(out.scores[i], out.ids[i]); + } + } + if (!searched && !segments.empty()) { + return tl::make_unexpected(Status::InvalidArgument( + "fast query: no searchable vector index for field ", field_name)); } const size_t keep = std::min(static_cast(topk), candidates.size()); std::partial_sort(candidates.begin(), candidates.begin() + keep, - candidates.end(), - [&better](const std::pair &a, - const std::pair &b) { - return better(a.first, b.first); + candidates.end(), [metric](const auto &a, const auto &b) { + return metric == MetricType::IP ? a.first > b.first + : a.first < b.first; }); - - FastQueryResult out; - out.ids.reserve(keep); - out.scores.reserve(keep); for (size_t i = 0; i < keep; ++i) { - out.scores.push_back(candidates[i].first); - out.ids.push_back(candidates[i].second); + out.ids[i] = candidates[i].second; + out.scores[i] = candidates[i].first; + } + std::fill(out.ids.begin() + keep, out.ids.end(), int64_t{-1}); + if (return_scores) { + std::fill(out.scores.begin() + keep, out.scores.end(), + std::numeric_limits::quiet_NaN()); + } else { + out.scores.clear(); } - return finish(std::move(out)); + return out; } Result CollectionImpl::query(const MultiQuery &query) const { diff --git a/src/db/index/column/vector_column/combined_vector_column_indexer.cc b/src/db/index/column/vector_column/combined_vector_column_indexer.cc index 31a38526f..186042a01 100644 --- a/src/db/index/column/vector_column/combined_vector_column_indexer.cc +++ b/src/db/index/column/vector_column/combined_vector_column_indexer.cc @@ -14,6 +14,7 @@ #include "combined_vector_column_indexer.h" #include #include +#include #include #include @@ -413,6 +414,49 @@ Result CombinedVectorColumnIndexer::Search( query_params.topk); } +Status CombinedVectorColumnIndexer::SearchFast( + const vector_column_params::VectorData &vector_data, + const vector_column_params::QueryParams &query_params, int64_t *output_ids, + float *output_scores) { + if (indexers_.size() == 1 && block_offsets_[0] == 0 && + query_params.filter == nullptr) { + vector_column_params::QueryParams params; + params.data_type = query_params.data_type; + params.dimension = query_params.dimension; + params.topk = query_params.topk; + params.query_params = query_params.query_params; + if (params.query_params && params.query_params->is_using_refiner()) { + if (normal_indexers_.size() != indexers_.size()) { + return Status::InvalidArgument( + "normal indexers size[", normal_indexers_.size(), + "] not match indexers size[", indexers_.size(), "]"); + } + params.refiner_param = + std::make_shared( + vector_column_params::RefinerParam{ + params.query_params->scale_factor(), normal_indexers_[0]}); + } + return indexers_[0]->SearchFast(vector_data, params, output_ids, + output_scores); + } + + // Reuse Search's offset/filter/refiner handling for composite segments. + auto result = Search(vector_data, query_params); + if (!result) return result.error(); + size_t count = 0; + for (auto it = result.value()->create_iterator(); + it->valid() && count < query_params.topk; it->next(), ++count) { + output_ids[count] = static_cast(it->doc_id()); + if (output_scores) output_scores[count] = it->score(); + } + std::fill(output_ids + count, output_ids + query_params.topk, int64_t{-1}); + if (output_scores) { + std::fill(output_scores + count, output_scores + query_params.topk, + std::numeric_limits::quiet_NaN()); + } + return Status::OK(); +} + Result CombinedVectorColumnIndexer::Fetch(uint32_t segment_doc_id) const { int32_t target_block_doc_id = -1; diff --git a/src/db/index/column/vector_column/combined_vector_column_indexer.h b/src/db/index/column/vector_column/combined_vector_column_indexer.h index 26e411073..d3a3a702b 100644 --- a/src/db/index/column/vector_column/combined_vector_column_indexer.h +++ b/src/db/index/column/vector_column/combined_vector_column_indexer.h @@ -38,6 +38,12 @@ class CombinedVectorColumnIndexer { const vector_column_params::VectorData &vector_data, const vector_column_params::QueryParams &query_params); + // Dense top-k search without group-by, brute-force keys or vector fetching. + // Uses the same block/refiner handling; output buffers hold topk elements. + Status SearchFast(const vector_column_params::VectorData &vector_data, + const vector_column_params::QueryParams &query_params, + int64_t *output_ids, float *output_scores); + virtual Result Fetch( uint32_t segment_doc_id) const; @@ -47,21 +53,6 @@ class CombinedVectorColumnIndexer { return !indexers_.empty(); } - //! True when one block starts at segment row zero. - bool is_single_block() const { - return indexers_.size() == 1 && block_offsets_[0] == 0; - } - - //! Primary block indexer (valid when ``is_single_block()``). - VectorColumnIndexer::Ptr primary_indexer() const { - return indexers_.empty() ? nullptr : indexers_[0]; - } - - //! Raw-vector reference block used by the primary index's refiner. - VectorColumnIndexer::Ptr reference_indexer() const { - return normal_indexers_.empty() ? nullptr : normal_indexers_[0]; - } - protected: /** * A filter wrapper that applies an offset to document IDs before diff --git a/src/db/index/column/vector_column/engine_helper.hpp b/src/db/index/column/vector_column/engine_helper.hpp index 3ee50533d..35bcb8e10 100644 --- a/src/db/index/column/vector_column/engine_helper.hpp +++ b/src/db/index/column/vector_column/engine_helper.hpp @@ -113,15 +113,17 @@ class ProximaEngineHelper { } template - static Result> + static Result> _build_common_query_param( const vector_column_params::QueryParams &db_query_params) { - auto engine_query_param = std::make_unique(); + auto engine_query_param = std::make_shared(); engine_query_param->topk = db_query_params.topk; engine_query_param->fetch_vector = db_query_params.fetch_vector; - engine_query_param->filter = - convert_to_engine_filter(db_query_params.filter); + if (db_query_params.filter) { + engine_query_param->filter = + convert_to_engine_filter(db_query_params.filter); + } if (db_query_params.query_params) { auto status = update_engine_query_param( @@ -237,7 +239,7 @@ class ProximaEngineHelper { return Status::OK(); } - static Result> + static Result convert_to_engine_query_param( const FieldSchema &field_schema, const vector_column_params::QueryParams &query_params) { diff --git a/src/db/index/column/vector_column/vector_column_indexer.cc b/src/db/index/column/vector_column/vector_column_indexer.cc index 1caa9964b..793e7c686 100644 --- a/src/db/index/column/vector_column/vector_column_indexer.cc +++ b/src/db/index/column/vector_column/vector_column_indexer.cc @@ -214,4 +214,25 @@ Result VectorColumnIndexer::Search( return result; } +Status VectorColumnIndexer::SearchFast( + const vector_column_params::VectorData &vector_data, + const vector_column_params::QueryParams &query_params, int64_t *output_ids, + float *output_scores) { + if (index == nullptr) { + return Status::InvalidArgument("Index not opened"); + } + auto engine_vector_data = + ProximaEngineHelper::convert_to_engine_vector(vector_data, is_sparse_); + if (!engine_vector_data) return engine_vector_data.error(); + auto engine_query_param = ProximaEngineHelper::convert_to_engine_query_param( + field_schema_, query_params); + if (!engine_query_param) return engine_query_param.error(); + if (index->search_fast(engine_vector_data.value(), + std::move(engine_query_param.value()), output_ids, + output_scores) != 0) { + return Status::InternalError("Failed to search vector"); + } + return Status::OK(); +} + } // namespace zvec diff --git a/src/db/index/column/vector_column/vector_column_indexer.h b/src/db/index/column/vector_column/vector_column_indexer.h index c37d37a37..3008c0871 100644 --- a/src/db/index/column/vector_column/vector_column_indexer.h +++ b/src/db/index/column/vector_column/vector_column_indexer.h @@ -81,6 +81,12 @@ class VectorColumnIndexer { virtual Result Search( const vector_column_params::VectorData &vector_data, const vector_column_params::QueryParams &query_params); + + // Dense top-k search without group-by, brute-force keys or vector fetching. + // Writes IDs and optional scores into topk-sized caller-owned buffers. + Status SearchFast(const vector_column_params::VectorData &vector_data, + const vector_column_params::QueryParams &query_params, + int64_t *output_ids, float *output_scores); // Result> BatchSearch( // const VectorDataset &vector_data, // const vector_column_params::QueryParams &query_params); diff --git a/src/db/index/segment/segment.cc b/src/db/index/segment/segment.cc index 2d72020a8..e993e4a0c 100644 --- a/src/db/index/segment/segment.cc +++ b/src/db/index/segment/segment.cc @@ -255,11 +255,12 @@ class SegmentImpl : public Segment, ExecBatchPtr fetch(const std::vector &columns, int segment_doc_id) const override; - // Gather stable insertion ordinals without Arrow/user-ID materialization. - Status get_global_doc_ids(const std::vector &segment_doc_ids, - std::vector &out) const override; + bool has_identity_doc_ids() const override { + return has_identity_doc_ids_; + } - bool has_identity_doc_ids() const override; + // Gather stable insertion ordinals without Arrow/user-ID materialization. + Status get_global_doc_ids(std::vector &doc_ids) const override; RecordBatchReaderPtr scan( const std::vector &columns) const override; @@ -416,6 +417,9 @@ class SegmentImpl : public Segment, // Maps segment-local doc ID (array index) to global doc ID (stored value) std::vector doc_ids_; + // Only valid for a read-only collection; queries never modify this flag. + bool has_identity_doc_ids_{false}; + std::array, std::unordered_map>>, static_cast(BlockType::VECTOR_INDEX_QUANTIZE) + 1> @@ -549,6 +553,18 @@ Status SegmentImpl::Open(const SegmentOptions &options) { fresh_persist_chunked_array(); + // WAL recovery can append doc IDs even for a read-only open. Initialize + // this property only after recovery has finished, before publication. + if (options_.read_only_) { + has_identity_doc_ids_ = true; + for (size_t i = 0; i < doc_ids_.size(); ++i) { + if (doc_ids_[i] != i) { + has_identity_doc_ids_ = false; + break; + } + } + } + return Status::OK(); } @@ -4540,25 +4556,15 @@ BlockID SegmentImpl::allocate_block_id() { return block_id_allocator_.fetch_add(1); } -bool SegmentImpl::has_identity_doc_ids() const { - std::lock_guard lock(seg_mtx_); - for (size_t i = 0; i < doc_ids_.size(); ++i) { - if (doc_ids_[i] != i) return false; - } - return true; -} - -Status SegmentImpl::get_global_doc_ids(const std::vector &segment_doc_ids, - std::vector &out) const { - out.resize(segment_doc_ids.size()); - std::lock_guard lock(seg_mtx_); +Status SegmentImpl::get_global_doc_ids(std::vector &doc_ids) const { + std::shared_lock lock(seg_mtx_); const size_t n = doc_ids_.size(); - for (size_t i = 0; i < segment_doc_ids.size(); ++i) { - const int sid = segment_doc_ids[i]; - if (sid < 0 || static_cast(sid) >= n) { - return Status::InvalidArgument("segment_doc_id out of range: ", sid); + for (auto &doc_id : doc_ids) { + if (doc_id == -1) continue; + if (doc_id < 0 || static_cast(doc_id) >= n) { + return Status::InvalidArgument("segment_doc_id out of range: ", doc_id); } - out[i] = static_cast(doc_ids_[sid]); + doc_id = static_cast(doc_ids_[doc_id]); } return Status::OK(); } diff --git a/src/db/index/segment/segment.h b/src/db/index/segment/segment.h index 38b022453..1951579b5 100644 --- a/src/db/index/segment/segment.h +++ b/src/db/index/segment/segment.h @@ -76,11 +76,6 @@ class Segment { // Count documents visible to an optional global-doc-ID filter. virtual uint64_t doc_count(const IndexFilter::Ptr filter = nullptr) = 0; - // Validates whether block keys can be returned as global document IDs. - virtual bool has_identity_doc_ids() const { - return false; - } - virtual bool has_record() = 0; // ---- Schema and index mutation ----------------------------------------- @@ -176,9 +171,12 @@ class Segment { virtual ExecBatchPtr fetch(const std::vector &columns, int segment_doc_id) const = 0; - // Gather stable insertion ordinals without Arrow/user-ID materialization. - virtual Status get_global_doc_ids(const std::vector &segment_doc_ids, - std::vector &out) const = 0; + // Valid only while the owning collection is open read-only. Computed by + // Open after recovery, before queries can run; not maintained for writes. + virtual bool has_identity_doc_ids() const = 0; + + // Map segment row IDs to insertion ordinals in place; preserve -1 padding. + virtual Status get_global_doc_ids(std::vector &doc_ids) const = 0; // Keep Segment alive while consuming the returned reader. virtual RecordBatchReaderPtr scan( diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index 5354af658..42ced30e3 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -248,7 +248,6 @@ class ZVEC_CORE_API Index { virtual int CreateAndInitStreamer(const BaseIndexParam ¶m) = 0; protected: - bool init_context(); core::IndexContext::Pointer &acquire_context(); protected: @@ -270,7 +269,6 @@ class ZVEC_CORE_API Index { // converter_/reformer_/metric_ stay null. std::shared_ptr turbo_quantizer_{}; - size_t context_index_; core::IndexStorage::Pointer storage_{}; bool is_open_{false}; diff --git a/src/include/zvec/db/collection.h b/src/include/zvec/db/collection.h index b91641c09..00c021bb5 100644 --- a/src/include/zvec/db/collection.h +++ b/src/include/zvec/db/collection.h @@ -106,8 +106,9 @@ class ZVEC_API Collection { virtual Result query(const MultiQuery &query) const = 0; // Advanced dense search returning internal numeric IDs and optional scores. - // Requires a read-only collection; fast_query calls and close must be serial. - // Parameters are read on every call. Index references are cached internally. + // Requires a read-only collection. Per-field index information is prepared + // at open and shared unchanged. Mutable parameters and execution state are + // local to each call, with the same lifetime locking as query(). // Missing neighbors are padded with ID -1 and score NaN. Refinement uses // query_params->scale_factor(), with the same semantics as query(). // Supply both input dtype and dimension for validation; omitting both trusts diff --git a/tests/db/fast_query_test.cc b/tests/db/fast_query_test.cc index a03717910..a754cb2ca 100644 --- a/tests/db/fast_query_test.cc +++ b/tests/db/fast_query_test.cc @@ -12,7 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include +#include #include +#include #include #include #include @@ -21,6 +24,194 @@ using namespace zvec; +TEST(FastQueryTest, ConcurrentFieldsFromFirstQueryThroughClose) { + const std::string path = "test_fast_query_concurrent_fields"; + FileHelper::RemoveDirectory(path); + ailego::MemoryLimitPool::get_instance().init(2 * 1024ll * 1024ll * 1024ll); + CollectionSchema schema("fast_query_fields"); + schema.add_field(std::make_shared( + "flat", DataType::VECTOR_FP32, uint32_t{32}, false, + std::make_shared(MetricType::L2))); + schema.add_field(std::make_shared( + "graph", DataType::VECTOR_FP32, uint32_t{64}, false, + std::make_shared(MetricType::L2, 16, 64))); + auto created = Collection::CreateAndOpen(path, schema, CollectionOptions{}); + ASSERT_TRUE(created) << created.error().message(); + auto writer = std::move(created.value()); + std::mt19937 rng(761); + std::normal_distribution normal; + auto make_vector = [&](size_t dimension) { + std::vector vector(dimension); + for (auto &value : vector) value = normal(rng); + return vector; + }; + std::vector docs; + for (int i = 0; i < 128; ++i) { + Doc doc; + doc.set_pk(std::to_string(i)); + doc.set>("flat", make_vector(32)); + doc.set>("graph", make_vector(64)); + docs.push_back(std::move(doc)); + } + auto inserted = writer->insert(docs); + ASSERT_TRUE(inserted) << inserted.error().message(); + for (const auto &status : inserted.value()) ASSERT_TRUE(status.ok()); + ASSERT_TRUE(writer->optimize(OptimizeOptions{1}).ok()); + + struct Request { + std::string field; + std::vector vector; + QueryParams::Ptr params; + int topk; + FastQueryResult expected; + }; + std::vector requests{ + {"flat", make_vector(32), nullptr, 1, {}}, + {"graph", make_vector(64), std::make_shared(32), 7, {}}, + {"graph", make_vector(64), std::make_shared(64), 13, {}}, + {"graph", make_vector(64), nullptr, 17, {}}}; + // Compute expectations on the writer so the reader below has never searched. + for (auto &request : requests) { + SearchQuery query; + query.topk_ = request.topk; + query.target_.field_name_ = request.field; + query.target_.query_params_ = request.params; + query.target_.set_vector( + std::string(reinterpret_cast(request.vector.data()), + request.vector.size() * sizeof(float))); + auto expected = writer->query(query); + ASSERT_TRUE(expected) << expected.error().message(); + ASSERT_EQ(expected->size(), request.topk); + for (const auto &doc : expected.value()) { + request.expected.ids.push_back(std::stoll(doc->pk())); + request.expected.scores.push_back(doc->score()); + } + } + ASSERT_TRUE(writer->close().ok()); + writer.reset(); + auto opened = Collection::Open(path, CollectionOptions{true, true}); + ASSERT_TRUE(opened) << opened.error().message(); + auto reader = std::move(opened.value()); + + std::atomic start{false}, closing{false}, stop{false}, failed{false}; + std::atomic completed{0}; + std::vector threads; + for (size_t worker = 0; worker < requests.size(); ++worker) { + threads.emplace_back([&, worker] { + while (!start.load()) std::this_thread::yield(); + for (size_t repeat = 0; !stop.load(); ++repeat) { + const auto &request = requests[(worker + repeat) % requests.size()]; + auto result = reader->fast_query( + request.field, request.vector.data(), request.params, request.topk, + true, DataType::VECTOR_FP32, request.vector.size()); + if (!result) { + EXPECT_TRUE(closing.load()) << result.error().message(); + EXPECT_EQ(result.error().code(), StatusCode::INVALID_ARGUMENT); + EXPECT_NE(result.error().message().find("closed"), std::string::npos); + failed.store(true); + break; + } + EXPECT_EQ(result->ids, request.expected.ids); + for (size_t i = 0; i < result->scores.size(); ++i) { + EXPECT_NEAR(result->scores[i], request.expected.scores[i], 1e-4f); + } + ++completed; + } + }); + } + start.store(true); + while (completed.load() < 128 && !failed.load()) std::this_thread::yield(); + closing.store(true); + const auto closed = reader->close(); + stop.store(true); + for (auto &thread : threads) thread.join(); + EXPECT_TRUE(closed.ok()) << closed.message(); + auto after_close = reader->fast_query("flat", requests[0].vector.data()); + ASSERT_FALSE(after_close); + EXPECT_NE(after_close.error().message().find("closed"), std::string::npos); + reader.reset(); + FileHelper::RemoveDirectory(path); +} + +TEST(FastQueryTest, PreservesOrdinalsAfterReopenAndCompaction) { + const std::string path = "test_fast_query_identity_doc_ids"; + FileHelper::RemoveDirectory(path); + ailego::MemoryLimitPool::get_instance().init(2 * 1024ll * 1024ll * 1024ll); + CollectionSchema schema("fast_query_identity"); + schema.add_field(std::make_shared( + "vector", DataType::VECTOR_FP32, uint32_t{32}, false, + std::make_shared(MetricType::L2))); + auto created = Collection::CreateAndOpen(path, schema, CollectionOptions{}); + ASSERT_TRUE(created) << created.error().message(); + auto writer = std::move(created.value()); + std::vector docs; + for (int i = 0; i < 16; ++i) { + Doc doc; + doc.set_pk(std::to_string(i)); + doc.set>("vector", std::vector(32, i)); + docs.push_back(std::move(doc)); + } + auto inserted = writer->insert(docs); + ASSERT_TRUE(inserted) << inserted.error().message(); + for (const auto &status : inserted.value()) ASSERT_TRUE(status.ok()); + ASSERT_TRUE(writer->optimize(OptimizeOptions{1}).ok()); + ASSERT_TRUE(writer->close().ok()); + writer.reset(); + + // Reopen with identity IDs, a delete filter, then compacted IDs with gaps. + for (int phase = 0; phase < 3; ++phase) { + SCOPED_TRACE(phase); + if (phase != 0) { + auto reopened = Collection::Open(path, CollectionOptions{}); + ASSERT_TRUE(reopened) << reopened.error().message(); + writer = std::move(reopened.value()); + if (phase == 1) { + // Keep ID 0 so the identity check must also detect internal gaps. + auto deleted = writer->delete_({"7", "9"}); + ASSERT_TRUE(deleted) << deleted.error().message(); + for (const auto &status : deleted.value()) ASSERT_TRUE(status.ok()); + } else { + ASSERT_TRUE(writer->optimize(OptimizeOptions{1}).ok()); + } + ASSERT_TRUE(writer->close().ok()); + writer.reset(); + } + + auto opened = Collection::Open(path, CollectionOptions{true, true}); + ASSERT_TRUE(opened) << opened.error().message(); + auto reader = std::move(opened.value()); + std::vector vector(32, 4.1f); + SearchQuery query; + query.topk_ = 20; + query.target_.field_name_ = "vector"; + query.target_.set_vector( + std::string(reinterpret_cast(vector.data()), + vector.size() * sizeof(float))); + auto expected = reader->query(query); + ASSERT_TRUE(expected) << expected.error().message(); + ASSERT_EQ(expected->size(), phase == 0 ? 16 : 14); + for (bool scores : {false, true}) { + auto actual = reader->fast_query("vector", vector.data(), nullptr, + query.topk_, scores); + ASSERT_TRUE(actual) << actual.error().message(); + ASSERT_EQ(actual->ids.size(), query.topk_); + for (size_t i = 0; i < actual->ids.size(); ++i) { + if (i < expected->size()) { + EXPECT_EQ(actual->ids[i], std::stoll(expected.value()[i]->pk())); + if (scores) { + EXPECT_FLOAT_EQ(actual->scores[i], expected.value()[i]->score()); + } + } else { + EXPECT_EQ(actual->ids[i], -1); + if (scores) EXPECT_TRUE(std::isnan(actual->scores[i])); + } + } + } + ASSERT_TRUE(reader->close().ok()); + } + FileHelper::RemoveDirectory(path); +} + TEST(FastQueryTest, ReadsRefineParametersOnEveryCall) { const std::string path = "test_fast_search_refine_scale"; FileHelper::RemoveDirectory(path); @@ -95,6 +286,40 @@ TEST(FastQueryTest, ReadsRefineParametersOnEveryCall) { } } } + // Simultaneous calls must not share mutable topk/refiner parameters, even + // when one caller requests defaults and another switches refinement off. + const std::vector topks{1, 17, 3, 10}; + const std::vector params{ + nullptr, std::make_shared(false), + std::make_shared(true, 3.0f), + std::make_shared(true, 7.0f)}; + std::vector> queries(4, std::vector(32)); + std::vector expected; + for (size_t i = 0; i < queries.size(); ++i) { + for (auto &value : queries[i]) value = normal(rng); + auto result = reader->fast_query("vector", queries[i].data(), params[i], + topks[i], true); + ASSERT_TRUE(result) << result.error().message(); + expected.push_back(std::move(result.value())); + } + std::atomic start{false}; + std::vector threads; + for (size_t worker = 0; worker < queries.size(); ++worker) { + threads.emplace_back([&, worker] { + while (!start.load()) std::this_thread::yield(); + for (size_t repeat = 0; repeat < 40; ++repeat) { + const size_t i = (worker + repeat) % queries.size(); + auto result = reader->fast_query("vector", queries[i].data(), params[i], + topks[i], true); + ASSERT_TRUE(result) << result.error().message(); + EXPECT_EQ(result->ids, expected[i].ids); + EXPECT_EQ(result->scores, expected[i].scores); + } + }); + } + start.store(true); + for (auto &thread : threads) thread.join(); + ASSERT_TRUE(reader->close().ok()); reader.reset(); FileHelper::RemoveDirectory(path); diff --git a/tests/db/index/segment/segment_row_id_test.cc b/tests/db/index/segment/segment_row_id_test.cc index c64ccd5bc..2987f1be0 100644 --- a/tests/db/index/segment/segment_row_id_test.cc +++ b/tests/db/index/segment/segment_row_id_test.cc @@ -21,6 +21,7 @@ #include #include #include "db/common/constants.h" +#include "db/index/storage/wal/wal_file.h" #include "segment_test_fixture.h" using namespace zvec; @@ -231,4 +232,87 @@ TEST_P(SegmentTest, DocCountDeleteFilterWithNonZeroGlobalDocID) { } +TEST_P(SegmentTest, MapGlobalDocIDsInPlacePreservesPadding) { + auto segment = test::TestHelper::CreateSegmentWithDoc( + col_path_, *schema_, 0, 100, id_map_, delete_store_, version_manager_, + options_, 0, 10); + ASSERT_NE(segment, nullptr); + + std::vector ids{9, 0, 3, 0, -1}; + ASSERT_TRUE(segment->get_global_doc_ids(ids).ok()); + EXPECT_EQ(ids, (std::vector{109, 100, 103, 100, -1})); + for (int64_t invalid : {-2, 10}) { + ids = {invalid}; + EXPECT_EQ(segment->get_global_doc_ids(ids).code(), + StatusCode::INVALID_ARGUMENT); + } + ids.clear(); + EXPECT_TRUE(segment->get_global_doc_ids(ids).ok()); +} + +TEST_P(SegmentTest, ReadOnlyIdentityDocIDsInitializedDuringOpen) { + struct Case { + uint64_t first_doc_id; + uint32_t count; + bool identity; + }; + const std::vector cases{{0, 0, true}, {0, 4, true}, {100, 4, false}}; + for (size_t i = 0; i < cases.size(); ++i) { + const auto &test_case = cases[i]; + SCOPED_TRACE(i); + auto writer = test::TestHelper::CreateSegmentWithDoc( + col_path_, *schema_, i, test_case.first_doc_id, id_map_, delete_store_, + version_manager_, options_, test_case.first_doc_id, test_case.count); + ASSERT_NE(writer, nullptr); + ASSERT_TRUE(writer->flush().ok()); + auto meta = writer->meta(); + writer.reset(); + + auto read_options = options_; + read_options.read_only_ = true; + auto opened = Segment::Open(col_path_, *schema_, *meta, id_map_, + delete_store_, version_manager_, read_options); + ASSERT_TRUE(opened) << opened.error().message(); + EXPECT_EQ(opened.value()->doc_count(), test_case.count); + EXPECT_EQ(opened.value()->has_identity_doc_ids(), test_case.identity); + } +} + +TEST_P(SegmentTest, ReadOnlyIdentityDocIDsIncludeRecoveredWal) { + CollectionSchema wal_schema(col_name_); + wal_schema.add_field( + std::make_shared("id", DataType::INT32, false)); + auto writer = test::TestHelper::CreateSegmentWithDoc( + col_path_, wal_schema, 0, 100, id_map_, delete_store_, version_manager_, + options_, 100, 0); + ASSERT_NE(writer, nullptr); + auto meta = writer->meta(); + writer.reset(); + + // Persisted doc IDs are empty (identity), but WAL recovery adds ID 100. + // Computing the flag before recovery would incorrectly keep it true. + const auto wal_path = FileHelper::MakeWalPath( + col_path_, meta->id(), meta->writing_forward_block()->id()); + WalFilePtr wal; + ASSERT_EQ(WalFile::CreateAndOpen(wal_path, WalOptions{0, true}, &wal), 0); + ASSERT_NE(wal, nullptr); + auto doc = test::TestHelper::CreateDoc(100, wal_schema); + doc.set_operator(Operator::INSERT); + auto serialized = doc.serialize(); + ASSERT_EQ(wal->append(std::string(serialized.begin(), serialized.end())), 0); + ASSERT_EQ(wal->close(), 0); + wal.reset(); + + auto read_options = options_; + read_options.read_only_ = true; + auto opened = Segment::Open(col_path_, wal_schema, *meta, id_map_, + delete_store_, version_manager_, read_options); + ASSERT_TRUE(opened) << opened.error().message(); + EXPECT_EQ(opened.value()->doc_count(), 1); + EXPECT_FALSE(opened.value()->has_identity_doc_ids()); + std::vector ids{0, -1}; + ASSERT_TRUE(opened.value()->get_global_doc_ids(ids).ok()); + EXPECT_EQ(ids, (std::vector{100, -1})); +} + INSTANTIATE_TEST_SUITE_P(MMapTest, SegmentTest, testing::Values(true, false)); diff --git a/tests/db/sqlengine/mock_segment.h b/tests/db/sqlengine/mock_segment.h index e64f4566c..4b2adfce8 100644 --- a/tests/db/sqlengine/mock_segment.h +++ b/tests/db/sqlengine/mock_segment.h @@ -312,9 +312,11 @@ class MockSegment : public Segment { return 0; } - Status get_global_doc_ids(const std::vector &segment_doc_ids, - std::vector &out) const override { - out.assign(segment_doc_ids.begin(), segment_doc_ids.end()); + bool has_identity_doc_ids() const override { + return false; + } + + Status get_global_doc_ids(std::vector &) const override { return Status::OK(); }