feat: support custom batch query distances

docs/docs/en/src/api/search.md

@@ -41,6 +41,44 @@ enum class SearchMode {
 | `limited_size_` | `int64_t` | `-1` | Cap on range results; `-1` means no limit. |
 | `params_str_` | `std::string` | `""` | Algorithm-specific search params as JSON (e.g. `ef_search`). |
 
+### Custom query distance callback
+
+`distance_batch_func_` optionally supplies query-to-vector scores from application code. It receives
+stable external vector IDs, so a closure can keep the query and query-specific state outside VSAG:
+
+```cpp
+request.distance_batch_size_ = 32;
+request.distance_batch_func_ = [query](const int64_t* ids, uint64_t count, float* scores) {
+    for (uint64_t i = 0; i < count; ++i) {
+        scores[i] = Score(query, ids[i]);
+    }
+};
+```
+
+| Field | Type | Default | Meaning |
+|-------|------|---------|---------|
+| `distance_batch_func_` | `SearchDistanceBatchFunc` | `nullptr` | Fills one score for each input external ID, in input order. Smaller finite scores are better. |
+| `distance_batch_size_` | `uint64_t` | `1` | Maximum IDs passed to one callback invocation. Set this to the scorer's efficient batch width for batched inference or external data access. Must be positive when a callback is set. |
+
+The callback is request-scoped and is not serialized. It may capture the query. `hgraph` permits a
+null `query_` in callback mode, while `ivf` still requires one query vector for bucket routing. The
+callback must be stable for one request, thread-safe if the caller enables parallel execution,
+non-throwing, and must write only finite scores.
+
+`brute_force` uses the callback for exact KNN and range search. `hgraph` supports KNN only; the
+callback drives graph traversal and final ordering, while the graph remains built with its configured
+built-in metric. Therefore recall under the callback score is not guaranteed. HGraph may score
+filtered traversal nodes to preserve graph connectivity, but filtered nodes are never returned.
+Callback HGraph does not support parallel search or `brute_force_threshold`; those configurations
+are rejected rather than silently changing search semantics.
+
+`ivf` supports callback KNN with a non-null single query vector. IVF uses its configured built-in
+metric to select `scan_buckets_count` buckets, then applies the callback to candidates in those
+buckets. The callback controls candidate ranking but cannot recover vectors outside the selected
+buckets; increase `scan_buckets_count` to improve recall. Callback IVF does not support range search,
+`disable_bucket_scan`, bucket-graph search, or parallel search. Those configurations are rejected.
+Reordering is automatically disabled in callback mode. Other indexes do not support the callback.
+
 ### IVF bucket routing
 
 IVF accepts `{"ivf":{"scan_buckets_count":N,"disable_bucket_scan":true}}` through

docs/docs/zh/src/api/search.md

@@ -41,6 +41,39 @@ enum class SearchMode {
 | `limited_size_` | `int64_t` | `-1` | 范围结果的上限；`-1` 表示不限。 |
 | `params_str_` | `std::string` | `""` | 算法特有的搜索参数 JSON（如 `ef_search`）。 |
 
+### 自定义查询距离回调
+
+`distance_batch_func_` 可选地让应用提供 query 到向量的分数。它接收稳定的外部向量 ID，
+因此闭包可以在 VSAG 之外持有 query 和该 query 的上下文：
+
+```cpp
+request.distance_batch_size_ = 32;
+request.distance_batch_func_ = [query](const int64_t* ids, uint64_t count, float* scores) {
+    for (uint64_t i = 0; i < count; ++i) {
+        scores[i] = Score(query, ids[i]);
+    }
+};
+```
+
+| 字段 | 类型 | 默认值 | 含义 |
+|------|------|--------|------|
+| `distance_batch_func_` | `SearchDistanceBatchFunc` | `nullptr` | 按输入外部 ID 的顺序填写对应分数。分数越小越好，且必须是有限值。 |
+| `distance_batch_size_` | `uint64_t` | `1` | 单次回调的最大 ID 数。批量推理或外部数据访问应设置为 scorer 的高效批大小；设置回调时必须为正。 |
+
+回调仅属于当前请求，不参与序列化。它可以捕获 query。`hgraph` 的回调模式允许 `query_` 为 null，
+但 `ivf` 仍需要一个查询向量用于桶路由。回调在一次请求内必须稳定；若调用方启用并行执行则必须
+线程安全；不得抛异常，且只能写入有限分数。
+
+`brute_force` 用该回调执行精确 KNN 和范围搜索。`hgraph` 仅支持 KNN：回调驱动图遍历和最终
+排序，但图仍由配置的内置 metric 构建，因此无法保证该回调分数下的 recall。HGraph 可能为保持
+图连通性而计算被过滤的遍历节点，但不会返回这些节点。
+回调 HGraph 不支持并行检索或 `brute_force_threshold`；这些配置会被拒绝，而不会静默改变检索语义。
+
+`ivf` 支持带回调的 KNN，但必须提供一个非空的单查询向量。IVF 先以配置的内置 metric 选出
+`scan_buckets_count` 个桶，再对这些桶中的候选调用回调。回调决定候选排序，但无法召回未被选中的
+桶内向量；增大 `scan_buckets_count` 可提高 recall。回调 IVF 不支持范围搜索、`disable_bucket_scan`、
+bucket graph 或并行搜索，这些配置会被拒绝。回调模式会自动关闭精排。其他索引不支持该回调。
+
 ### IVF 桶路由
 
 IVF 可通过 `params_str_` 接收

include/vsag/search_request.h

@@ -15,6 +15,7 @@
 #pragma once
 
 #include <cstdint>
+#include <functional>
 #include <string>
 #include <vector>
 
@@ -26,6 +27,9 @@
 
 namespace vsag {
 
+using SearchDistanceBatchFunc =
+    std::function<void(const int64_t* ids, uint64_t count, float* distances)>;
+
 enum class SearchMode {
     KNN_SEARCH = 1,
     RANGE_SEARCH = 2,
@@ -83,6 +87,25 @@ class SearchRequest {
      */
     std::string params_str_{};
 
+    /**
+     * @brief Optional request-scoped callback for custom query scoring.
+     *
+     * Receives stable external IDs and writes one lower-is-better score per ID.
+     * When non-null, supported indexes use this callback for search traversal and
+     * result ordering instead of their built-in vector metric. Graph traversal can
+     * score filtered IDs to keep the graph connected, but filtered IDs are never
+     * returned as results.
+     */
+    SearchDistanceBatchFunc distance_batch_func_{nullptr};
+
+    /**
+     * @brief Maximum number of IDs submitted to distance_batch_func_ per invocation.
+     *
+     * Defaults to 1 for scalar scorers. Set this to the scorer's efficient batch
+     * width when it benefits from batched inference or data access.
+     */
+    uint64_t distance_batch_size_{1};
+
     // for attribute filter
     /**
      * @brief Flag to enable attribute-based filtering during search

src/algorithm/bruteforce/bruteforce.cpp

@@ -15,7 +15,9 @@
 #include "bruteforce.h"
 
 #include <atomic>
+#include <cmath>
 #include <cstring>
+#include <exception>
 #include <mutex>
 #include <new>
 #include <optional>
@@ -367,15 +369,30 @@ DatasetPtr
 BruteForce::SearchWithRequest(const SearchRequest& request) const {
     std::shared_lock read_lock(this->global_mutex_);
 
-    auto computer = this->make_search_computer(request.query_);
+    const bool use_custom_distance = request.distance_batch_func_ != nullptr;
+    if (use_custom_distance) {
+        CHECK_ARGUMENT(request.distance_batch_size_ > 0,
+                       "distance_batch_size must be greater than 0");
+    }
+
+    ComputerInterfacePtr computer = nullptr;
+    if (not use_custom_distance) {
+        computer = this->make_search_computer(request.query_);
+    }
 
     bool is_range = (request.mode_ == SearchMode::RANGE_SEARCH);
     if (is_range) {
-        if (not is_multi_vector_) {
+        if (use_custom_distance) {
+            CHECK_ARGUMENT(std::isfinite(request.radius_), "radius must be finite");
+            CHECK_ARGUMENT(request.radius_ >= 0.0F, "radius must be non-negative");
+            CHECK_ARGUMENT(request.limited_size_ != 0, "limited_size must not be 0");
+        } else if (not is_multi_vector_) {
             this->validate_range_args(request.query_, request.radius_, request.limited_size_);
         }
     } else {
-        if (not is_multi_vector_) {
+        if (use_custom_distance) {
+            CHECK_ARGUMENT(request.topk_ > 0, "topk must be greater than 0");
+        } else if (not is_multi_vector_) {
             this->validate_knn_args(request.query_, request.topk_);
         }
     }
@@ -434,7 +451,13 @@ BruteForce::SearchWithRequest(const SearchRequest& request) const {
         for (const auto& pair : label_to_inner_id) {
             float dist = 0.0F;
             const auto inner_id = pair.second;
-            this->inner_codes_->Query(&dist, computer, &inner_id, 1);
+            if (use_custom_distance) {
+                const auto label = this->label_table_->GetLabelById(inner_id);
+                request.distance_batch_func_(&label, 1, &dist);
+                CHECK_ARGUMENT(std::isfinite(dist), "distance callback must return finite scores");
+            } else {
+                this->inner_codes_->Query(&dist, computer, &inner_id, 1);
+            }
             reasoning_ctx->SetTrueDistance(inner_id, dist);
         }
     }
@@ -450,30 +473,71 @@ BruteForce::SearchWithRequest(const SearchRequest& request) const {
 
     auto search_func = [&](InnerIdType start, InnerIdType end, const DistHeapPtr& cur_heap) {
         uint32_t dist_cmp_local = 0;
+        std::vector<InnerIdType> custom_inner_ids;
+        std::vector<int64_t> custom_labels;
+        std::vector<float> custom_dists;
+        if (use_custom_distance) {
+            const uint64_t batch_capacity =
+                std::min<uint64_t>(request.distance_batch_size_, end - start);
+            custom_inner_ids.reserve(batch_capacity);
+            custom_labels.reserve(batch_capacity);
+            custom_dists.resize(batch_capacity);
+        }
+
+        auto flush_custom_batch = [&]() {
+            if (custom_inner_ids.empty()) {
+                return;
+            }
+            request.distance_batch_func_(
+                custom_labels.data(), custom_labels.size(), custom_dists.data());
+            for (uint64_t j = 0; j < custom_inner_ids.size(); ++j) {
+                const float dist = custom_dists[j];
+                CHECK_ARGUMENT(std::isfinite(dist), "distance callback must return finite scores");
+                if (reasoning != nullptr) {
+                    reasoning->RecordVisit(custom_inner_ids[j], dist, 0);
+                }
+                if (not is_range || dist <= radius) {
+                    cur_heap->Push(dist, custom_inner_ids[j]);
+                }
+            }
+            dist_cmp_local += static_cast<uint32_t>(custom_inner_ids.size());
+            custom_inner_ids.clear();
+            custom_labels.clear();
+        };
+
         for (InnerIdType i = start; i < end; ++i) {
-            float dist = 0.0F;
             if (attr_filter != nullptr and not attr_filter->CheckValid(i)) {
                 if (reasoning != nullptr) {
                     reasoning->RecordFilterReject(i);
                 }
                 continue;
             }
             if (ft == nullptr or ft->CheckValid(i)) {
-                inner_codes_->Query(&dist, computer, &i, 1);
-                ++dist_cmp_local;
-                if (reasoning != nullptr) {
-                    reasoning->RecordVisit(i, dist, 0);
-                }
-                if (is_range and dist > radius) {
-                    continue;
+                if (use_custom_distance) {
+                    custom_inner_ids.push_back(i);
+                    custom_labels.push_back(this->label_table_->GetLabelById(i));
+                    if (custom_inner_ids.size() == request.distance_batch_size_) {
+                        flush_custom_batch();
+                    }
+                } else {
+                    float dist = 0.0F;
+                    inner_codes_->Query(&dist, computer, &i, 1);
+                    ++dist_cmp_local;
+                    if (reasoning != nullptr) {
+                        reasoning->RecordVisit(i, dist, 0);
+                    }
+                    if (is_range and dist > radius) {
+                        continue;
+                    }
+                    cur_heap->Push(dist, i);
                 }
-                cur_heap->Push(dist, i);
             } else {
                 if (reasoning != nullptr) {
                     reasoning->RecordFilterReject(i);
                 }
             }
         }
+        flush_custom_batch();
         dist_cmp.fetch_add(dist_cmp_local, std::memory_order_relaxed);
     };
 
@@ -491,11 +555,24 @@ BruteForce::SearchWithRequest(const SearchRequest& request) const {
         for (auto i = 0; i < parallel_count; ++i) {
             auto start = i * chunk_size;
             auto end = std::min(start + chunk_size, count);
+            if (start >= end) {
+                continue;
+            }
             auto future = this->thread_pool_->GeneralEnqueue(search_func, start, end, heaps[i]);
             futures.emplace_back(std::move(future));
         }
+        std::exception_ptr first_error = nullptr;
         for (auto& future : futures) {
-            future.get();
+            try {
+                future.get();
+            } catch (...) {
+                if (first_error == nullptr) {
+                    first_error = std::current_exception();
+                }
+            }
+        }
+        if (first_error != nullptr) {
+            std::rethrow_exception(first_error);
         }
         heap = heaps[0];
         for (auto i = 1; i < parallel_count; ++i) {