perf(search): output refined internal IDs directly

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)):

src/core/algorithm/flat/flat_streamer.cc

@@ -449,6 +449,37 @@ int FlatStreamer<BATCH_SIZE>::search_bf_by_p_keys_impl(
   return 0;
 }
 
+template <size_t BATCH_SIZE>
+int FlatStreamer<BATCH_SIZE>::search_by_p_keys_fast(
+    const void *query, const std::vector<uint64_t> &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<FlatStreamerContext<BATCH_SIZE> *>(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<uint32_t>(quantizer_->dim()) ||
+        qmeta.quantize_type() != static_cast<uint32_t>(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 <size_t BATCH_SIZE>
 int FlatStreamer<BATCH_SIZE>::group_by_search_impl(
     const void *query, const IndexQueryMeta &qmeta, uint32_t count,
@@ -545,6 +576,14 @@ int FlatStreamer<BATCH_SIZE>::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<uint64_t> &, int64_t *, float *, size_t,
+    const IndexQueryMeta &, Context::UPointer &) const;
+template int FlatStreamer<32>::search_by_p_keys_fast(
+    const void *, const std::vector<uint64_t> &, 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>);

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<uint64_t> &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;
 

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<uint64_t> & /*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<uint64_t> &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<size_t>(std::numeric_limits<int>::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<const char *>(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<int>(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<uint64_t, float> 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<float>::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<int64_t>(documents[i].first);
+    if (output_scores) output_scores[i] = documents[i].second;
+  }
+  return 0;
+}
+
 int FlatContiguousStreamerEntity::build_contiguous_memory() {
   degrade_to_mmap();
 

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<uint64_t> &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<uint64_t> &keys,
+                            int64_t *output_ids, float *output_scores,
+                            size_t topk,
+                            FlatSearchScratch *scratch) const override;
+
   bool is_contiguous() const {
     return !!load_contiguous_storage();
   }

src/core/interface/index_search_internal_ids.cc

@@ -17,6 +17,7 @@
 #include <turbo/quantizer/quantizer.h>
 #include <zvec/core/framework/index_error.h>
 #include <zvec/core/interface/index.h>
+#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<const core::FlatStreamer<32> *>(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<uint64_t> 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<std::vector<uint64_t>> 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) {

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,