diff --git a/python/tests/test_fast_query.py b/python/tests/test_fast_query.py new file mode 100644 index 000000000..12f26d230 --- /dev/null +++ b/python/tests/test_fast_query.py @@ -0,0 +1,606 @@ +"""Advanced dense search: collection queries, scores, fallback and lifetime.""" + +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier + +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.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"], +) +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]) + # 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] + ) + 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) + + +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) + 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() + + +@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/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..f6f4bded6 100644 --- a/python/zvec/model/collection.py +++ b/python/zvec/model/collection.py @@ -509,6 +509,42 @@ 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. + 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 + ``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..355708e0d 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 { +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"); + } + 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; +} + // 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,42 @@ 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 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. + 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, data_type, dimension); + } + 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. +Parameters and execution state belong to each call, as with Collection.query. +)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..46e975398 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 || + (!quantizer_ && !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..75aaebfba 100644 --- a/src/core/algorithm/flat/flat_streamer_entity.cc +++ b/src/core/algorithm/flat/flat_streamer_entity.cc @@ -13,7 +13,9 @@ // limitations under the License. #include "flat_streamer_entity.h" +#include #include +#include #include #include "flat_utility.h" @@ -596,6 +598,90 @@ 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; + 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); + 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 (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 { + for (size_t i = 0; i < count; ++i) { + 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.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); + 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].first); + if (output_scores) output_scores[i] = documents[i].second; + } + 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..1484e4da3 100644 --- a/src/core/algorithm/flat/flat_streamer_entity.h +++ b/src/core/algorithm/flat/flat_streamer_entity.h @@ -83,6 +83,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 +526,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..2bd105d33 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" @@ -186,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() { @@ -584,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; } @@ -801,6 +799,213 @@ 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_ && 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)) { + 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 +1214,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 +1232,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 +1243,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..48526c161 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 @@ -21,6 +22,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,13 @@ 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, + DataType query_data_type, + uint32_t query_dimension) const override; + Result group_by_query( const GroupByVectorQuery &query) const override; @@ -174,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(); @@ -303,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}; @@ -400,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); @@ -457,6 +486,9 @@ Status CollectionImpl::close_unsafe() { } } + fast_query_fields_.clear(); + read_only_segments_.clear(); + // always release resources regardless of flush outcome writing_segment_.reset(); segment_manager_.reset(); @@ -1816,6 +1848,136 @@ Result CollectionImpl::query(const SearchQuery &query) const { return query_unsafe(query); } +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)); + } +} + +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")); + } + + 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)); + } + 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 ((query_data_type != DataType::UNDEFINED || query_dimension != 0) && + (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 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}; + + 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; + 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(), [metric](const auto &a, const auto &b) { + return metric == MetricType::IP ? a.first > b.first + : a.first < b.first; + }); + for (size_t i = 0; i < keep; ++i) { + 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 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.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 c407b9bee..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,10 +38,21 @@ 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; + //! True when at least one backing vector indexer is available for search. + bool has_searchable_indexers() const { + return !indexers_.empty(); + } + 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 b82ad817b..35bcb8e10 100644 --- a/src/db/index/column/vector_column/engine_helper.hpp +++ b/src/db/index/column/vector_column/engine_helper.hpp @@ -95,20 +95,41 @@ 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) { + // 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)); + } + return true; + } + 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) { - 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,7 +156,90 @@ class ProximaEngineHelper { } public: - static Result> + // 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, const vector_column_params::QueryParams &query_params) { @@ -168,15 +272,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 +285,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 +298,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 +311,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 +324,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 +337,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/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 efc27f758..e993e4a0c 100644 --- a/src/db/index/segment/segment.cc +++ b/src/db/index/segment/segment.cc @@ -255,6 +255,13 @@ class SegmentImpl : public Segment, ExecBatchPtr fetch(const std::vector &columns, int segment_doc_id) const override; + bool has_identity_doc_ids() const override { + return has_identity_doc_ids_; + } + + // 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; @@ -410,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> @@ -543,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(); } @@ -4534,6 +4556,20 @@ BlockID SegmentImpl::allocate_block_id() { return block_id_allocator_.fetch_add(1); } +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 (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); + } + doc_id = static_cast(doc_ids_[doc_id]); + } + 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..1951579b5 100644 --- a/src/db/index/segment/segment.h +++ b/src/db/index/segment/segment.h @@ -171,6 +171,13 @@ class Segment { virtual ExecBatchPtr fetch(const std::vector &columns, int segment_doc_id) 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( 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..42ced30e3 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, @@ -229,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: @@ -251,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 0426658c4..00c021bb5 100644 --- a/src/include/zvec/db/collection.h +++ b/src/include/zvec/db/collection.h @@ -105,6 +105,21 @@ 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. 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 + // 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, + 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 c983ac4e1..c59a48581 100644 --- a/src/include/zvec/db/query.h +++ b/src/include/zvec/db/query.h @@ -132,6 +132,12 @@ 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; +}; + 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..146e69791 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,8 +3968,8 @@ TEST(IndexInterface, HNSWRabitqGeneral) { // mode. TEST(IndexInterface, ContiguousMemoryEndToEnd) { constexpr uint32_t kDimension = 32; - constexpr uint32_t kNumDocs = 500; - constexpr int kTopk = 10; + constexpr uint32_t kNumDocs = 1200; + 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 @@ -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..a754cb2ca --- /dev/null +++ b/tests/db/fast_query_test.cc @@ -0,0 +1,326 @@ +// 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 +#include +#include +#include +#include "db/common/file_helper.h" + +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); + 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(); + 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())); + EXPECT_FLOAT_EQ(actual->scores[i], expected.value()[i]->score()); + } + } + } + // 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/column/vector_column_indexer_test.cc b/tests/db/index/column/vector_column_indexer_test.cc index b0bfd9c81..b41c862bf 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,95 @@ 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()); +} + + +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; 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 c5afacb55..4b2adfce8 100644 --- a/tests/db/sqlengine/mock_segment.h +++ b/tests/db/sqlengine/mock_segment.h @@ -312,6 +312,14 @@ class MockSegment : public Segment { return 0; } + bool has_identity_doc_ids() const override { + return false; + } + + Status get_global_doc_ids(std::vector &) const override { + return Status::OK(); + } + TablePtr fetch(const std::vector &columns, const std::vector &indices) const override { std::string s = "";