From a0838209ab84ffe0024d209ee160b05c3443b5e3 Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Fri, 16 Oct 2026 23:03:42 +0800 Subject: [PATCH] perf(search): output refined internal IDs directly Refine coarse candidates directly into ID and optional score buffers for supported contiguous Flat references. Reuse the existing legacy and Turbo distance kernels and preserve score ordering and normalization. Retain the generic refiner for unsupported requests and storage layouts. When direct refinement falls back after coarse search, reuse its candidates instead of repeating the coarse search. Add coverage for direct output, concurrent queries, storage modes and single-coarse-search behavior across late fallback cases. --- .../tests/test_query_internal_ids_refine.py | 41 ++++ src/core/algorithm/flat/flat_streamer.cc | 39 ++++ src/core/algorithm/flat/flat_streamer.h | 8 + .../algorithm/flat/flat_streamer_entity.cc | 80 +++++++ .../algorithm/flat/flat_streamer_entity.h | 15 ++ .../interface/index_search_internal_ids.cc | 107 +++++++++ src/include/zvec/core/interface/index.h | 3 + .../core/algorithm/flat/flat_streamer_test.cc | 104 +++++++++ tests/core/interface/index_interface_test.cc | 213 +++++++++++++++++- 9 files changed, 608 insertions(+), 2 deletions(-) diff --git a/python/tests/test_query_internal_ids_refine.py b/python/tests/test_query_internal_ids_refine.py index 6c276ef725..32b4e81042 100644 --- a/python/tests/test_query_internal_ids_refine.py +++ b/python/tests/test_query_internal_ids_refine.py @@ -1,6 +1,8 @@ """Fast search with native reference storage and optional refined scores.""" +from concurrent.futures import ThreadPoolExecutor from functools import partial +from threading import Barrier import numpy as np import pytest @@ -169,6 +171,45 @@ def test_refine_candidate_counts_and_parameter_switches( np.testing.assert_allclose( scores, [doc.score for doc in coarse_docs], rtol=2e-5, atol=2e-5 ) + concurrent_cases = [] + for row, refine, scale, topk in [ + (11, True, 1.1, 10), + (89, True, 2.6, 3), + (531, False, 2.6, 10), + (11, True, 1.4, 7), + ]: + query = np.ascontiguousarray(vectors[row] + np.float32(0.021)) + param = query_param(is_using_refiner=refine, scale_factor=scale) + docs = coll.query( + Query(field_name="vector", vector=query, param=param), topk=topk + ) + concurrent_cases.append( + (query, param, topk, [int(d.id) for d in docs], [d.score for d in docs]) + ) + ready = Barrier(4) + + def search(worker): + ready.wait() + for repeat in range(8): + query, param, topk, expected_ids, expected_scores = concurrent_cases[ + (worker + repeat) % len(concurrent_cases) + ] + return_scores = repeat % 2 == 0 + output = _query_internal_ids( + coll, "vector", query, param, topk, return_scores + ) + if return_scores: + ids, scores = output + np.testing.assert_allclose( + scores, expected_scores, rtol=2e-5, atol=2e-5 + ) + else: + ids = output + np.testing.assert_array_equal(ids, expected_ids) + + with ThreadPoolExecutor(max_workers=4) as pool: + list(pool.map(search, range(4))) + for scale in (-1.0, float("nan"), float("inf")): param = query_param(is_using_refiner=True, scale_factor=scale) with pytest.raises((ValueError, RuntimeError)): diff --git a/src/core/algorithm/flat/flat_streamer.cc b/src/core/algorithm/flat/flat_streamer.cc index d233f425ed..7ef06fe5ac 100644 --- a/src/core/algorithm/flat/flat_streamer.cc +++ b/src/core/algorithm/flat/flat_streamer.cc @@ -449,6 +449,37 @@ 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 (state_ != STATE_OPENED) return IndexError_NoReady; + if (!query || !output_ids || topk == 0 || !context) { + return IndexError_InvalidArgument; + } + auto *flat_context = + dynamic_cast *>(context.get()); + if (!flat_context) return IndexError_InvalidArgument; + if (flat_context->filter().is_valid() || flat_context->group_by_search() || + flat_context->threshold_set() || flat_context->fetch_vector()) { + return IndexError_NotImplemented; + } + if (quantizer_) { + // The caller has already prepared the query using this quantizer. + if (qmeta.dimension() != static_cast(quantizer_->dim()) || + qmeta.quantize_type() != static_cast(quantizer_->type()) || + qmeta.element_size() != quantizer_->quantized_query_vector_length()) { + return IndexError_InvalidArgument; + } + } else if (!metric_ || !metric_->is_matched(meta_, qmeta)) { + 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 +576,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 26823172d7..b5f28480ba 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 the request or storage needs full-result search. + 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 a4601f63eb..a54b1c44b5 100644 --- a/src/core/algorithm/flat/flat_streamer_entity.cc +++ b/src/core/algorithm/flat/flat_streamer_entity.cc @@ -596,6 +596,86 @@ 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 || + (quantizer() && + keys.size() > static_cast(std::numeric_limits::max()))) { + 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 same prepared-query dispatch as the regular Flat search. + if (quantizer()) { + quantizer()->calc_distance_dp_query_batch( + ptrs.data(), static_cast(count), query, distances.data()); + } else { + const void *batch_query = query; + if (const auto &preprocess = batch_query_preprocess(); + preprocess != nullptr) { + auto &buffer = scratch->query_buffer; + buffer.resize(meta().element_size()); + std::memcpy(buffer.data(), query, buffer.size()); + preprocess(buffer.data(), meta().dimension()); + batch_query = buffer.data(); + } + if (const auto &batch = batch_distance(); batch) { + batch(ptrs.data(), batch_query, count, meta().dimension(), + distances.data(), has_extras ? extras.data() : nullptr); + } else { + for (size_t i = 0; i < count; ++i) { + distance()(query, ptrs[i], meta().dimension(), distances.data() + i); + } + } + } + // Use the same heap operations and score-only ordering as IndexDocumentHeap, + // with lightweight key/score entries rather than full result documents. + ailego::KeyValueHeap documents(topk); + for (size_t i = 0; i < count; ++i) { + // Match the regular heap: NaN and positive overflow are not candidates. + // Let that path decide the result count without modifying caller output. + if (!(distances[i] <= std::numeric_limits::max())) { + return IndexError_NotImplemented; + } + documents.emplace(keys[i], distances[i]); + } + documents.sort(); + 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 fb4c8c666e..76e53e8d14 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, @@ -520,6 +529,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/interface/index_search_internal_ids.cc b/src/core/interface/index_search_internal_ids.cc index cfe83ffd36..17ad81af18 100644 --- a/src/core/interface/index_search_internal_ids.cc +++ b/src/core/interface/index_search_internal_ids.cc @@ -17,6 +17,7 @@ #include #include #include +#include "algorithm/flat/flat_streamer.h" namespace zvec::core_interface { @@ -59,6 +60,12 @@ int Index::search_internal_ids(const VectorData &vector_data, } const size_t topk = search_param->topk; + if (search_param->refiner_param && !search_param->bf_pks && + !search_param->is_linear) { + const int ret = _search_refine_internal_ids(vector_data, search_param, + output_ids, output_scores); + if (ret != core::IndexError_NotImplemented) return ret; + } if (RequiresSearchFallback(search_param)) { thread_local SearchResult result; const int ret = search(vector_data, search_param, &result); @@ -119,6 +126,106 @@ int Index::search_internal_ids(const VectorData &vector_data, return ret; } +int Index::_search_refine_internal_ids( + const VectorData &vector_data, + const BaseIndexQueryParam::Pointer &search_param, int64_t *output_ids, + float *output_scores) { + const auto &reference = search_param->refiner_param->reference_index; + // Reject unsupported shapes before preparing a context or consuming a + // filter: the ordinary search must still receive the original request. + 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 auto &quantizer = reference->param_.quantizer_param; + if ((reference->param_.metric_type != MetricType::kL2sq && + reference->param_.metric_type != MetricType::kInnerProduct) || + (quantizer && quantizer->type != QuantizerType::kNone && + 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); + if (ret != 0) { + context->reset(); + return ret; + } + const int coarse_topk = _get_coarse_search_topk(search_param); + if (coarse_topk < 0) { + context->reset(); + return coarse_topk; + } + context->set_topk(coarse_topk); + context->set_fetch_vector(false); + + // Only capacity is retained across calls, never pointers into an index. + // This path invokes no user filter/group callbacks and is non-reentrant + // within a thread, as is the existing candidate/context scratch. + thread_local std::string coarse_query_storage; + thread_local std::string reference_query_storage; + thread_local std::vector keys; + const void *query = nullptr; + core::IndexQueryMeta query_meta; + ret = _prepare_dense_query(vector_data, &coarse_query_storage, &query, + &query_meta); + 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 { + // Pooled Flat contexts may have served a different query/index. The + // radius is applied to coarse candidates, not the reference metric. + reference_context->set_fetch_vector(false); + reference_context->reset_filter(); + reference_context->reset_threshold(); + _set_group_by_on_context(search_param, reference_context); + 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, search_param->topk, + query_meta, reference_context); + if (ret == core::IndexError_NotImplemented) { + // The coarse search has already completed. Reuse its candidates + // with the ordinary refiner instead of repeating the entire search. + reference_context->reset(); + std::vector> candidates(1); + candidates[0].swap(keys); + thread_local SearchResult result; + ret = reference->_refine_dense_candidates(vector_data, search_param, + candidates, &result); + candidates[0].swap(keys); + if (ret == 0) { + CopySearchResult(result, search_param->topk, output_ids, + output_scores); + } + } else if (ret == 0 && output_scores) { + ret = reference->_normalize_buffer_scores( + vector_data, output_ids, output_scores, search_param->topk); + } + } + reference_context->reset(); + } + } + context->reset(); + return ret; +} + int Index::_normalize_buffer_scores(const VectorData &vector_data, const int64_t *output_ids, float *output_scores, size_t count) { diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index aa14c3b465..eb2093c22d 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -206,6 +206,9 @@ class ZVEC_CORE_API Index { int _normalize_buffer_scores(const VectorData &query, const int64_t *output_ids, float *output_scores, size_t count); + int _search_refine_internal_ids( + const VectorData &query, const BaseIndexQueryParam::Pointer &search_param, + int64_t *output_ids, float *output_scores); int _collect_dense_result(const VectorData &query, const core::IndexQueryMeta &query_meta, const BaseIndexQueryParam::Pointer &search_param, diff --git a/tests/core/algorithm/flat/flat_streamer_test.cc b/tests/core/algorithm/flat/flat_streamer_test.cc index 3044bbf4b6..3f4e70bee0 100644 --- a/tests/core/algorithm/flat/flat_streamer_test.cc +++ b/tests/core/algorithm/flat/flat_streamer_test.cc @@ -2007,6 +2007,110 @@ TEST_F(FlatStreamerTest, TestInt8WithRotate) { EXPECT_EQ(kTopk, linear_ctx->result().size()); } +TEST_F(FlatStreamerTest, TurboCandidateDirectOutputMatchesOrdinarySearch) { + constexpr size_t kCount = 32; + constexpr size_t kTopk = 10; + for (const char *name : {"Fp16Quantizer", "Fp32Quantizer"}) { + for (const char *metric : {"SquaredEuclidean", "InnerProduct"}) { + for (size_t dimension : {128U, 960U}) { + SCOPED_TRACE(name); + SCOPED_TRACE(metric); + SCOPED_TRACE(dimension); + IndexMeta raw_meta(IndexMeta::DT_FP32, dimension); + raw_meta.set_metric(metric, 0, Params()); + auto quantizer = IndexFactory::CreateQuantizer(name); + ASSERT_NE(nullptr, quantizer); + ASSERT_EQ(0, quantizer->init(raw_meta, Params())); + IndexMeta meta = quantizer->meta(); + meta.set_quantizer(name, 0, Params()); + Params params; + params.set(PARAM_FLAT_USE_CONTIGUOUS_MEMORY, true); + auto streamer = IndexFactory::CreateStreamer("FlatStreamer"); + ASSERT_NE(nullptr, streamer); + ASSERT_EQ(0, streamer->init(meta, params, quantizer)); + auto storage = IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_NE(nullptr, storage); + ASSERT_EQ(0, storage->init(Params())); + const std::string path = + dir_ + "direct_" + name + metric + std::to_string(dimension); + ASSERT_EQ(0, storage->open(path, true)); + ASSERT_EQ(0, streamer->open(storage)); + auto context = streamer->create_context(); + ASSERT_NE(nullptr, context); + std::vector query(dimension, 0.371f); + std::string query_code; + IndexQueryMeta qmeta; + ASSERT_EQ( + 0, quantizer->quantize( + query.data(), IndexQueryMeta(IndexMeta::DT_FP32, dimension), + &query_code, &qmeta)); + std::vector row(dimension); + std::string code(quantizer->quantized_datapoint_vector_length(), '\0'); + for (size_t id = 0; id < kCount; ++id) { + // Duplicate pairs also check the ordinary path's tie ordering. + for (size_t d = 0; d < dimension; ++d) { + row[d] = float(((id / 2) * 37 + d * 13) % 251) / 129.3f; + } + quantizer->quantize_data(row.data(), code.data()); + ASSERT_EQ( + 0, streamer->add_with_id_impl(id, code.data(), qmeta, context)); + } + ASSERT_EQ(0, streamer->flush(0)); + ASSERT_EQ(0, streamer->close()); + ASSERT_EQ(0, streamer->open(storage)); + auto *flat = dynamic_cast *>(streamer.get()); + ASSERT_NE(nullptr, flat); + ASSERT_NE(nullptr, flat->quantizer()); + const auto *entity = + dynamic_cast(&flat->entity()); + ASSERT_NE(nullptr, entity); + ASSERT_TRUE(entity->is_contiguous()); + context = streamer->create_context(); + ASSERT_NE(nullptr, context); + context->set_topk(kTopk); + for (size_t count : {11U, 14U, 26U, 32U}) { + SCOPED_TRACE(count); + std::vector keys(count); + for (size_t i = 0; i < count; ++i) keys[i] = (i * 7) % kCount; + ASSERT_EQ(0, streamer->search_bf_by_p_keys_impl( + query_code.data(), {keys}, qmeta, 1, context)); + const auto expected = context->result(); + ASSERT_EQ(kTopk, expected.size()); + for (bool return_scores : {false, true, false}) { + std::vector ids(kTopk + 1, -2); + std::vector scores(kTopk + 1, -123.0f); + // Call the direct operation itself: an Index-level fallback + // must not make this regression test pass accidentally. + ASSERT_EQ(0, flat->search_by_p_keys_fast( + query_code.data(), keys, ids.data(), + return_scores ? scores.data() : nullptr, kTopk, + qmeta, context)); + for (size_t rank = 0; rank < kTopk; ++rank) { + EXPECT_EQ(static_cast(expected[rank].key()), ids[rank]); + if (return_scores) { + EXPECT_FLOAT_EQ(expected[rank].score(), scores[rank]); + } + } + EXPECT_EQ(-2, ids[kTopk]); + EXPECT_FLOAT_EQ(-123.0f, scores[kTopk]); + } + keys[0] = kCount; + std::vector ids(kTopk, -2); + std::vector scores(kTopk, -123.0f); + EXPECT_EQ(IndexError_NotImplemented, + flat->search_by_p_keys_fast(query_code.data(), keys, + ids.data(), scores.data(), + kTopk, qmeta, context)); + EXPECT_EQ(std::vector(kTopk, -2), ids); + EXPECT_EQ(std::vector(kTopk, -123.0f), scores); + } + ASSERT_EQ(0, streamer->close()); + ASSERT_EQ(0, storage->close()); + } + } + } +} + // The streamer computes distances with the turbo quantizer. TEST_F(FlatStreamerTest, TestTurboQuantizerAddAndSearch) { const size_t raw_dim = 24; diff --git a/tests/core/interface/index_interface_test.cc b/tests/core/interface/index_interface_test.cc index bc779a4393..31287e54a5 100644 --- a/tests/core/interface/index_interface_test.cc +++ b/tests/core/interface/index_interface_test.cc @@ -42,6 +42,7 @@ #include "algorithm/cluster/cluster_params.h" #include "algorithm/hnsw/hnsw_entity.h" #include "algorithm/hnsw/hnsw_params.h" +#include "algorithm/hnsw/hnsw_streamer.h" #include "algorithm/ivf/ivf_params.h" #include "algorithm/vamana/vamana_entity.h" #include "algorithm/vamana/vamana_streamer.h" @@ -91,6 +92,47 @@ class TestableVamanaIndex : public VamanaIndex { } }; +class CountingHnswStreamer : public zvec::core::HnswStreamer { + public: + using HnswStreamer::init; + + int search_candidates_impl( + const void *query, const zvec::core::IndexQueryMeta &meta, + std::vector &keys, Context::Pointer &context, + std::vector *scores = nullptr) const override { + ++candidate_search_count; + return HnswStreamer::search_candidates_impl(query, meta, keys, context, + scores); + } + + mutable size_t candidate_search_count{0}; +}; + +class CountingHNSWIndex : public HNSWIndex { + public: + int init_for_test(const BaseIndexParam ¶m) { + return init(param); + } + + size_t candidate_search_count() const { + return std::static_pointer_cast(streamer_) + ->candidate_search_count; + } + + protected: + int create_and_init_streamer(const BaseIndexParam ¶m) override { + int ret = HNSWIndex::create_and_init_streamer(param); + if (ret != 0) return ret; + auto streamer = std::make_shared(); + ret = turbo_quantizer_ + ? streamer->init(proxima_index_meta_, proxima_index_params_, + turbo_quantizer_) + : streamer->init(proxima_index_meta_, proxima_index_params_); + if (ret == 0) streamer_ = std::move(streamer); + return ret; + } +}; + } // namespace TEST(IndexInterface, GraphRefineKeepsLegacyDefaultCandidateCount) { @@ -3027,6 +3069,24 @@ TEST(IndexInterface, VamanaNativeRefineMatchesExplicitCandidates) { EXPECT_EQ(expected.reverted_vector_list_, actual.reverted_vector_list_); } + // Default refine scale uses graph-specific query parameters to determine + // the candidate count. Validate their type before accessing those fields. + auto wrong_refiner = std::make_shared(); + wrong_refiner->reference_index = fine; + auto wrong_param = HNSWQueryParamBuilder() + .with_topk(kTopk) + .with_refiner_param(wrong_refiner) + .build(); + std::array ids{}; + const VectorData invalid_query{DenseVector{vectors[7].data()}}; + SearchResult invalid_result; + const int expected_error = + coarse->search(invalid_query, wrong_param, &invalid_result); + ASSERT_NE(0, expected_error); + EXPECT_EQ(expected_error, + coarse->search_internal_ids(invalid_query, wrong_param, ids.data(), + nullptr)); + ASSERT_EQ(0, coarse->close()); ASSERT_EQ(0, fine->close()); zvec::test_util::RemoveTestFiles(coarse_path); @@ -5059,12 +5119,46 @@ TEST(IndexInterface, RefineInternalIdsMatchesSearchAcrossStorageModes) { for (size_t rank = 0; rank < scores.size(); ++rank) { EXPECT_FLOAT_EQ(expected.doc_list_[rank].score(), scores[rank]); } - // Fewer coarse candidates than k takes the unchanged public fallback - // and must fill the remainder instead of exposing stale output. + auto expect_current_param_matches = [&]() { + SearchResult current_expected; + ASSERT_EQ(0, coarse->search(query, param, ¤t_expected)); + const size_t topk = param->topk; + const size_t count = current_expected.doc_list_.size(); + ASSERT_LE(count, topk); + // Alternate score requests and keep a guard beyond the caller's + // output range, including when a query returns fewer than topk. + for (bool return_scores : {true, false, true}) { + std::vector current_ids(topk + 1, -2); + std::vector current_scores(topk + 1, -123.0f); + ASSERT_EQ(0, coarse->search_internal_ids( + query, param, current_ids.data(), + return_scores ? current_scores.data() : nullptr)); + for (size_t rank = 0; rank < count; ++rank) { + EXPECT_EQ( + static_cast(current_expected.doc_list_[rank].key()), + current_ids[rank]); + if (return_scores) { + EXPECT_FLOAT_EQ(current_expected.doc_list_[rank].score(), + current_scores[rank]); + } + } + for (size_t rank = count; rank < topk; ++rank) { + EXPECT_EQ(-1, current_ids[rank]); + if (return_scores) { + EXPECT_TRUE(std::isnan(current_scores[rank])); + } + } + EXPECT_EQ(-2, current_ids[topk]); + EXPECT_FLOAT_EQ(-123.0f, current_scores[topk]); + } + }; + // Fewer coarse candidates than k use the ordinary refiner, which must + // fill the remainder instead of exposing stale output. param->radius = 0.5f; ASSERT_EQ( 0, coarse->search_internal_ids(query, param, ids.data(), nullptr)); EXPECT_EQ((std::array{{0, -1, -1}}), ids); + expect_current_param_matches(); param->radius = 0.0f; param->filter = std::make_shared(); param->filter->set([](uint64_t key) { return key >= 2; }); @@ -5074,12 +5168,25 @@ TEST(IndexInterface, RefineInternalIdsMatchesSearchAcrossStorageModes) { std::array filtered{{ids[0], ids[1]}}; std::sort(filtered.begin(), filtered.end()); EXPECT_EQ((std::array{{0, 1}}), filtered); + expect_current_param_matches(); param->filter.reset(); ASSERT_EQ( 0, coarse->search_internal_ids(query, param, ids.data(), nullptr)); if (!tied) { EXPECT_EQ((std::array{{5, 4, 3}}), ids); } + // Grow beyond the document count, then shrink on the same thread. + // Neither candidate scratch nor a preceding fallback may leak into + // the next refined query or overwrite an earlier caller's result. + for (uint32_t topk : {11U, 1U, 3U}) { + param->topk = topk; + expect_current_param_matches(); + } + for (size_t rank = 0; rank < scores.size(); ++rank) { + EXPECT_EQ(static_cast(expected.doc_list_[rank].key()), + scored_ids[rank]); + EXPECT_FLOAT_EQ(expected.doc_list_[rank].score(), scores[rank]); + } ASSERT_EQ(0, fine->close()); zvec::test_util::RemoveTestFiles(fine_path); } @@ -5089,6 +5196,108 @@ TEST(IndexInterface, RefineInternalIdsMatchesSearchAcrossStorageModes) { zvec::test_util::RemoveTestFiles(coarse_path); } +TEST(IndexInterface, RefineInternalIdsFallbackReusesCoarseCandidates) { + constexpr uint32_t kDimension = 128; + constexpr uint32_t kCoarseCount = 8; + const std::string coarse_path = "refine_once_coarse.index"; + const std::string fine_path = "refine_once_fine.index"; + zvec::test_util::RemoveTestFiles(coarse_path); + auto coarse_param = HNSWIndexParamBuilder() + .with_metric_type(MetricType::kL2sq) + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .with_m(8) + .with_ef_construction(32) + .build(); + auto coarse = std::make_shared(); + ASSERT_EQ(0, coarse->init_for_test(*coarse_param)); + ASSERT_EQ( + 0, coarse->open(coarse_path, {StorageOptions::StorageType::kMMAP, true})); + std::vector values(kDimension, 0.0f); + for (uint32_t key = 0; key < kCoarseCount; ++key) { + values[0] = static_cast(key); + ASSERT_EQ(0, coarse->add(VectorData{DenseVector{values.data()}}, key)); + } + const std::vector query_values(kDimension, 0.0f); + const VectorData query{DenseVector{query_values.data()}}; + struct Case { + const char *name; + bool contiguous; + uint32_t reference_count; + uint32_t topk; + }; + const std::array cases{{ + {"non_contiguous", false, kCoarseCount, 3}, + {"too_few_candidates", true, kCoarseCount, kCoarseCount + 1}, + {"missing_reference_keys", true, 2, 3}, + {"direct_output", true, kCoarseCount, 3}, + }}; + for (const auto &test_case : cases) { + SCOPED_TRACE(test_case.name); + zvec::test_util::RemoveTestFiles(fine_path); + auto fine_param = FlatIndexParamBuilder() + .with_metric_type(MetricType::kL2sq) + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .with_use_contiguous_memory(test_case.contiguous) + .build(); + auto fine = IndexFactory::CreateAndInitIndex(*fine_param); + ASSERT_TRUE(fine); + ASSERT_EQ( + 0, fine->open(fine_path, {StorageOptions::StorageType::kMMAP, true})); + for (uint32_t key = 0; key < test_case.reference_count; ++key) { + values[0] = static_cast(kCoarseCount - key); + ASSERT_EQ(0, fine->add(VectorData{DenseVector{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 = HNSWQueryParamBuilder() + .with_topk(test_case.topk) + .with_ef_search(32) + .with_refiner_param(refiner) + .build(); + SearchResult expected; + ASSERT_EQ(0, coarse->search(query, param, &expected)); + ASSERT_EQ(std::min(test_case.reference_count, test_case.topk), + expected.doc_list_.size()); + for (bool return_scores : {true, false, true}) { + const size_t count_before = coarse->candidate_search_count(); + std::vector ids(test_case.topk + 1, -2); + std::vector scores(test_case.topk + 1, -123.0f); + ASSERT_EQ(0, coarse->search_internal_ids( + query, param, ids.data(), + return_scores ? scores.data() : nullptr)); + // Falling back after candidate generation must only rerun refinement, + // never the coarse search that produced those candidates. + EXPECT_EQ(count_before + 1, coarse->candidate_search_count()); + for (size_t i = 0; i < expected.doc_list_.size(); ++i) { + EXPECT_EQ(static_cast(expected.doc_list_[i].key()), ids[i]); + if (return_scores) { + EXPECT_FLOAT_EQ(expected.doc_list_[i].score(), scores[i]); + } + } + for (size_t i = expected.doc_list_.size(); i < test_case.topk; ++i) { + EXPECT_EQ(-1, ids[i]); + if (return_scores) { + EXPECT_TRUE(std::isnan(scores[i])); + } + } + EXPECT_EQ(-2, ids.back()); + EXPECT_FLOAT_EQ(-123.0f, scores.back()); + } + ASSERT_EQ(0, fine->close()); + zvec::test_util::RemoveTestFiles(fine_path); + } + ASSERT_EQ(0, coarse->close()); + zvec::test_util::RemoveTestFiles(coarse_path); +} + TEST(IndexInterface, RefineInternalIdsInnerProductMatchesSearch) { constexpr uint32_t kDimension = 128; const std::string coarse_path = "docids_ip_coarse.index";