This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new fb86c9a3 [core] Search primary-key vector buckets and files
concurrently (#556)
fb86c9a3 is described below
commit fb86c9a3fcfa0d6d18bf18b29a8fcad0df551e97
Author: Junrui Lee <[email protected]>
AuthorDate: Wed Jul 22 18:20:05 2026 +0800
[core] Search primary-key vector buckets and files concurrently (#556)
---
crates/paimon/src/spec/core_options.rs | 58 +++
crates/paimon/src/table/pk_vector_orchestrator.rs | 500 +++++++++++++++++++---
crates/paimon/src/table/vector_search_builder.rs | 4 +
crates/paimon/src/vindex/pkvector/bucket.rs | 463 +++++++++++++++++---
4 files changed, 895 insertions(+), 130 deletions(-)
diff --git a/crates/paimon/src/spec/core_options.rs
b/crates/paimon/src/spec/core_options.rs
index fd4809bc..dd545f19 100644
--- a/crates/paimon/src/spec/core_options.rs
+++ b/crates/paimon/src/spec/core_options.rs
@@ -24,6 +24,7 @@ const DATA_EVOLUTION_ENABLED_OPTION: &str =
"data-evolution.enabled";
const GLOBAL_INDEX_ENABLED_OPTION: &str = "global-index.enabled";
const GLOBAL_INDEX_SEARCH_MODE_OPTION: &str = "global-index.search-mode";
const GLOBAL_INDEX_ROW_COUNT_PER_SHARD_OPTION: &str =
"global-index.row-count-per-shard";
+const GLOBAL_INDEX_THREAD_NUM_OPTION: &str = "global-index.thread-num";
const GLOBAL_INDEX_COLUMN_UPDATE_ACTION_OPTION: &str =
"global-index.column-update-action";
const SORTED_INDEX_RECORDS_PER_RANGE_OPTION: &str =
"sorted-index.records-per-range";
const BTREE_INDEX_FALLBACK_SCAN_MAX_SIZE_OPTION: &str =
"btree-index.fallback-scan-max-size";
@@ -114,6 +115,7 @@ const DEFAULT_READ_BATCH_SIZE: usize = 1024;
const DYNAMIC_BUCKET_TARGET_ROW_NUM_OPTION: &str =
"dynamic-bucket.target-row-num";
const DEFAULT_DYNAMIC_BUCKET_TARGET_ROW_NUM: i64 = 200_000;
const DEFAULT_GLOBAL_INDEX_ROW_COUNT_PER_SHARD: i64 = 100_000;
+const DEFAULT_GLOBAL_INDEX_THREAD_NUM: i64 = 32;
const DEFAULT_GLOBAL_INDEX_FALLBACK_SCAN_MAX_SIZE: i64 = 256 * 1024 * 1024;
const BLOB_AS_DESCRIPTOR_OPTION: &str = "blob-as-descriptor";
pub(crate) const BLOB_FIELD_OPTION: &str = "blob-field";
@@ -574,6 +576,28 @@ impl<'a> CoreOptions<'a> {
Ok(value)
}
+ /// Maximum number of concurrent tasks for global-index I/O, mirroring Java
+ /// `CoreOptions.GLOBAL_INDEX_THREAD_NUM` (key `global-index.thread-num`,
+ /// default 32). Used as the fan-out limit for the primary-key vector
search
+ /// (per-bucket and per-exact-file). A value of `1` reproduces strict
+ /// sequential execution. A non-positive value is a misconfiguration and
fails
+ /// loud rather than being silently clamped.
+ pub fn global_index_thread_num(&self) -> crate::Result<usize> {
+ let value = self
+ .parse_i64_option(GLOBAL_INDEX_THREAD_NUM_OPTION)?
+ .unwrap_or(DEFAULT_GLOBAL_INDEX_THREAD_NUM);
+ if value <= 0 {
+ return Err(crate::Error::DataInvalid {
+ message: format!(
+ "Option '{}' must be greater than 0, got: {}",
+ GLOBAL_INDEX_THREAD_NUM_OPTION, value
+ ),
+ source: None,
+ });
+ }
+ Ok(value as usize)
+ }
+
pub fn sorted_index_records_per_range(&self) -> crate::Result<i64> {
let value = self
.parse_i64_option(SORTED_INDEX_RECORDS_PER_RANGE_OPTION)?
@@ -1235,6 +1259,7 @@ mod tests {
core_options.global_index_row_count_per_shard().unwrap(),
100_000
);
+ assert_eq!(core_options.global_index_thread_num().unwrap(), 32);
assert_eq!(
core_options.sorted_index_records_per_range().unwrap(),
100_000
@@ -1365,6 +1390,39 @@ mod tests {
}
}
+ #[test]
+ fn test_global_index_thread_num_default_and_custom() {
+ // Default mirrors Java (32) when unset.
+ let empty = HashMap::new();
+ let core = CoreOptions::new(&empty);
+ assert_eq!(core.global_index_thread_num().unwrap(), 32);
+
+ // Explicit value is read back verbatim.
+ let options =
+ HashMap::from([(GLOBAL_INDEX_THREAD_NUM_OPTION.to_string(),
"8".to_string())]);
+ let core = CoreOptions::new(&options);
+ assert_eq!(core.global_index_thread_num().unwrap(), 8);
+ }
+
+ #[test]
+ fn test_global_index_thread_num_rejects_invalid_values() {
+ // Non-positive values are a misconfiguration and must fail loud (never
+ // clamped to 1); an unparsable value is likewise rejected.
+ for value in ["0", "-1", "abc"] {
+ let options = HashMap::from([(
+ GLOBAL_INDEX_THREAD_NUM_OPTION.to_string(),
+ value.to_string(),
+ )]);
+ let core = CoreOptions::new(&options);
+
+ let err = core
+ .global_index_thread_num()
+ .expect_err("invalid thread-num should fail");
+ assert!(matches!(err, crate::Error::DataInvalid { message, .. }
+ if message.contains(GLOBAL_INDEX_THREAD_NUM_OPTION)));
+ }
+ }
+
#[test]
fn test_sorted_index_records_per_range_rejects_invalid_values() {
for value in ["0", "-1", "abc"] {
diff --git a/crates/paimon/src/table/pk_vector_orchestrator.rs
b/crates/paimon/src/table/pk_vector_orchestrator.rs
index 572fa641..99ab5108 100644
--- a/crates/paimon/src/table/pk_vector_orchestrator.rs
+++ b/crates/paimon/src/table/pk_vector_orchestrator.rs
@@ -28,6 +28,9 @@ use std::sync::Arc;
use roaring::RoaringTreemap;
+use futures::stream::{self, StreamExt, TryStreamExt};
+use tokio::sync::Semaphore;
+
use crate::deletion_vector::DeletionVector;
use crate::spec::BinaryRow;
use crate::table::data_file_reader::DataFileReader;
@@ -391,6 +394,7 @@ impl PkVectorOrchestrator {
search_options: &HashMap<String, String>,
skip_exact_fallback: bool,
residual_by_split: Option<&[HashMap<String, RoaringTreemap>]>,
+ concurrency: usize,
) -> crate::Result<OrchestratorSearchResult> {
let mut results = self
.search_candidates_batch(
@@ -404,6 +408,7 @@ impl PkVectorOrchestrator {
search_options,
skip_exact_fallback,
residual_by_split,
+ concurrency,
)
.await?;
debug_assert_eq!(results.len(), 1);
@@ -422,8 +427,20 @@ impl PkVectorOrchestrator {
/// The residual allow-list depends only on the filter and the plan, not
the
/// query vector, so the SAME `residual_by_split` slice is shared across
every
/// query. Input-shape validation (positive limits, non-empty query,
residual
- /// count) is applied per query / once as appropriate. Kept a sequential
- /// per-split loop.
+ /// count) is applied per query / once as appropriate.
+ ///
+ /// `concurrency` is the global fan-out limit (Java
`GLOBAL_INDEX_THREAD_NUM`):
+ /// `1` runs the buckets and their files strictly sequentially, larger
values fan
+ /// them out with `buffer_unordered`. To match Java's single shared
+ /// `GlobalIndexReadThreadPool`, a single [`Semaphore`] budget of
`concurrency`
+ /// permits is shared across BOTH the per-bucket and per-exact-file
fan-outs and
+ /// acquired only around leaf exact-file I/O, so total in-flight exact-file
+ /// searches are capped at `concurrency` overall — not `concurrency` per
bucket,
+ /// which would allow up to `concurrency * concurrency`. Each bucket's
per-query
+ /// results feed per-query cross-bucket global Top-K heaps, which are
+ /// order-independent, so the output does not depend on which bucket
finished
+ /// first; results are collected and merged into the correct per-query
slot after
+ /// the parallel stage.
#[allow(clippy::too_many_arguments)]
#[allow(clippy::type_complexity)]
pub(crate) async fn search_candidates_batch(
@@ -448,6 +465,7 @@ impl PkVectorOrchestrator {
search_options: &HashMap<String, String>,
skip_exact_fallback: bool,
residual_by_split: Option<&[HashMap<String, RoaringTreemap>]>,
+ concurrency: usize,
) -> crate::Result<Vec<OrchestratorSearchResult>> {
// Eager input-shape validation (Java checkArgument parity).
if queries.is_empty() {
@@ -478,68 +496,121 @@ impl PkVectorOrchestrator {
(0..queries.len()).map(|_| Vec::new()).collect();
let mut exact_candidates: Vec<Vec<PkVectorCandidate>> =
(0..queries.len()).map(|_| Vec::new()).collect();
- for (split_index, split) in splits.iter().enumerate() {
- let dvs = build_bucket_dv_map(&self.reader, split).await?;
- // Adapt the split-scoped search closure to bucket_search's
per-file
- // closure by binding the current split index/split. The coercion
helper
- // ties the produced future's borrow to the arguments, which
closure
- // inference cannot express on its own.
- let bucket_search_closure = as_bucket_exact_file_search(
- |file: &BucketActiveFile,
- queries: &[&[f32]],
- metric: VectorSearchMetric,
- exact_limit: usize,
- is_excluded: &(dyn Fn(i64) -> bool + Sync)|
- -> ExactFileSearchFuture<'_> {
- exact_file_search(
- split_index,
- split,
- file,
- queries,
- metric,
- exact_limit,
- is_excluded,
- )
- },
- );
- let residual_ranges = residual_by_split.map(|per_split|
&per_split[split_index]);
- let per_query = bucket_search_batch(
- ann_searcher,
- &split.ann_segments,
- &split.active_files,
- &dvs,
- &bucket_search_closure,
- queries,
- metric,
- indexed_limit,
- limit,
- search_options,
- skip_exact_fallback,
- residual_ranges,
- )
- .await?;
- if per_query.len() != queries.len() {
- return Err(data_invalid(format!(
- "bucket search returned {} result lists for {} queries",
- per_query.len(),
- queries.len()
- )));
+
+ // One shared concurrency budget for the WHOLE search, mirroring
Java's single
+ // `GlobalIndexReadThreadPool`: the per-bucket and per-exact-file
fan-outs draw
+ // slots from the SAME N permits, so total in-flight exact-file I/O is
capped at
+ // N across all buckets and files (not N per bucket, which would allow
N*N).
+ // Only leaf exact-file work acquires a permit; bucket orchestration
never holds
+ // one, so it cannot starve leaf work. `concurrency <= 1` takes the
strictly
+ // sequential path at both levels and needs no budget.
+ let search_budget = (concurrency > 1).then(||
Arc::new(Semaphore::new(concurrency)));
+
+ // One lazy future per bucket. Each builds its own DV map + per-file
search
+ // closure, searches all queries against the bucket, and returns
per-query
+ // (indexed, exact) candidate lists already tagged with the bucket's
+ // partition/bucket/split_index. The futures are not polled until
driven
+ // below, so the sequential branch observes buckets in strict split
order.
+ let per_bucket = splits.iter().enumerate().map(|(split_index, split)| {
+ let search_budget = search_budget.clone();
+ async move {
+ let dvs = build_bucket_dv_map(&self.reader, split).await?;
+ // Adapt the split-scoped search closure to bucket_search's
per-file
+ // closure by binding the current split index/split. The
coercion helper
+ // ties the produced future's borrow to the arguments, which
closure
+ // inference cannot express on its own.
+ let bucket_search_closure = as_bucket_exact_file_search(
+ |file: &BucketActiveFile,
+ queries: &[&[f32]],
+ metric: VectorSearchMetric,
+ exact_limit: usize,
+ is_excluded: &(dyn Fn(i64) -> bool + Sync)|
+ -> ExactFileSearchFuture<'_> {
+ exact_file_search(
+ split_index,
+ split,
+ file,
+ queries,
+ metric,
+ exact_limit,
+ is_excluded,
+ )
+ },
+ );
+ let residual_ranges = residual_by_split.map(|per_split|
&per_split[split_index]);
+ let per_query = bucket_search_batch(
+ ann_searcher,
+ &split.ann_segments,
+ &split.active_files,
+ &dvs,
+ &bucket_search_closure,
+ queries,
+ metric,
+ indexed_limit,
+ limit,
+ search_options,
+ skip_exact_fallback,
+ residual_ranges,
+ concurrency,
+ search_budget,
+ )
+ .await?;
+ if per_query.len() != queries.len() {
+ return Err(data_invalid(format!(
+ "bucket search returned {} result lists for {}
queries",
+ per_query.len(),
+ queries.len()
+ )));
+ }
+ let tag = |PkVectorSearchResult {
+ data_file_name,
+ row_position,
+ distance,
+ }: PkVectorSearchResult| PkVectorCandidate {
+ split_index,
+ partition: split.data_split.partition().clone(),
+ bucket: split.data_split.bucket(),
+ data_file_name,
+ row_position,
+ distance,
+ };
+ let tagged: Vec<(Vec<PkVectorCandidate>,
Vec<PkVectorCandidate>)> = per_query
+ .into_iter()
+ .map(|result| {
+ (
+ result.indexed.into_iter().map(&tag).collect(),
+ result.exact.into_iter().map(&tag).collect(),
+ )
+ })
+ .collect();
+ Ok::<_, crate::Error>(tagged)
}
- let tag = |PkVectorSearchResult {
- data_file_name,
- row_position,
- distance,
- }: PkVectorSearchResult| PkVectorCandidate {
- split_index,
- partition: split.data_split.partition().clone(),
- bucket: split.data_split.bucket(),
- data_file_name,
- row_position,
- distance,
+ });
+
+ // Drive the per-bucket futures. `concurrency == 1` uses a strictly
+ // sequential loop so buckets are searched in split order; larger
values fan
+ // them out with `buffer_unordered`. Either way each bucket's
per-query lists
+ // are collected and only then folded into the per-query candidate
+ // accumulators, and the final per-query `global_top_k` is
order-independent
+ // (deterministic `candidate_cmp`), so the result does not depend on
bucket
+ // completion order.
+ let collected: Vec<Vec<(Vec<PkVectorCandidate>,
Vec<PkVectorCandidate>)>> =
+ if concurrency <= 1 {
+ let mut out = Vec::with_capacity(splits.len());
+ for fut in per_bucket {
+ out.push(fut.await?);
+ }
+ out
+ } else {
+ stream::iter(per_bucket)
+ .buffer_unordered(concurrency)
+ .try_collect::<Vec<_>>()
+ .await?
};
- for (query_index, result) in per_query.into_iter().enumerate() {
-
indexed_candidates[query_index].extend(result.indexed.into_iter().map(&tag));
-
exact_candidates[query_index].extend(result.exact.into_iter().map(&tag));
+ for tagged in collected {
+ for (query_index, (indexed, exact)) in
tagged.into_iter().enumerate() {
+ indexed_candidates[query_index].extend(indexed);
+ exact_candidates[query_index].extend(exact);
}
}
@@ -1240,7 +1311,7 @@ mod e2e_tests {
);
let result = orch
.search_candidates(
- splits, query, metric, limit, limit, ann, &wrapped, opts,
false, None,
+ splits, query, metric, limit, limit, ann, &wrapped, opts,
false, None, 1,
)
.await?;
// Merge the two bounded lists into the best-first survivors the
@@ -1277,6 +1348,7 @@ mod e2e_tests {
&opts,
false,
None,
+ 1,
)
.await
.map(|_| ())
@@ -1303,6 +1375,7 @@ mod e2e_tests {
&opts,
false,
None,
+ 1,
)
.await
.map(|_| ())
@@ -1647,6 +1720,7 @@ mod e2e_tests {
&opts,
false,
None,
+ 1,
)
.await
.unwrap();
@@ -1739,6 +1813,7 @@ mod e2e_tests {
&opts,
false,
Some(&residual_by_split),
+ 1,
)
.await
.unwrap();
@@ -1793,6 +1868,7 @@ mod e2e_tests {
&opts,
false,
Some(&residual_by_split),
+ 1,
)
.await
.map(|_| ())
@@ -1836,6 +1912,7 @@ mod e2e_tests {
&opts,
true,
None,
+ 1,
)
.await
.unwrap();
@@ -1950,6 +2027,7 @@ mod e2e_tests {
&opts,
false,
None,
+ 1,
)
.await
.unwrap();
@@ -1966,6 +2044,7 @@ mod e2e_tests {
&opts,
false,
None,
+ 1,
)
.await
.unwrap();
@@ -2027,6 +2106,7 @@ mod e2e_tests {
&opts,
false,
None,
+ 1,
)
.await
.unwrap();
@@ -2038,4 +2118,298 @@ mod e2e_tests {
assert_eq!(m1.len(), 1);
assert_eq!(m1[0].row_position, 0);
}
+
+ #[tokio::test]
+ async fn search_candidates_bucket_order_is_sequential_at_concurrency_one()
{
+ // At concurrency == 1 the per-bucket loop must search buckets in
strict
+ // split order. A split-scoped recording closure captures the visited
+ // split_index sequence; it must match the split order
deterministically.
+ let table_path = "memory:/pkvo_bucket_order";
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let mut splits = Vec::new();
+ for bucket in 0..4i32 {
+ let bucket_path = format!("{table_path}/bucket-{bucket}");
+ let meta = write_file(&file_io, &bucket_path, "d.mosaic",
vec![bucket]).await;
+ splits.push(PkVectorSearchSplit {
+ data_split: DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(BinaryRow::new(0))
+ .with_bucket(bucket)
+ .with_bucket_path(bucket_path)
+ .with_total_buckets(4)
+ .with_data_files(vec![meta])
+ .build()
+ .unwrap(),
+ ann_segments: Vec::new(),
+ active_files: vec![active("d.mosaic", 1)],
+ });
+ }
+ let visited = std::sync::Mutex::new(Vec::<usize>::new());
+ let factory = as_split_search(
+ |split_index: usize,
+ _: &PkVectorSearchSplit,
+ file: &BucketActiveFile,
+ queries: &[&[f32]],
+ metric: VectorSearchMetric,
+ exact_limit: usize,
+ is_excluded: &(dyn Fn(i64) -> bool + Sync)|
+ -> ExactFileSearchFuture<'_> {
+ visited.lock().unwrap().push(split_index);
+ let file_name = file.file_name.clone();
+ let query = queries[0].to_vec();
+ Box::pin(async move {
+ let mut reader = ArrayReader::new(2, vec![Some(vec![1.0,
0.0])]);
+ Ok(vec![exact_search(
+ &file_name,
+ &mut reader,
+ &query,
+ metric,
+ exact_limit,
+ is_excluded,
+ )?])
+ })
+ },
+ );
+ let opts = HashMap::new();
+ PkVectorOrchestrator::new(make_reader(file_io, table_path))
+ .search_candidates(
+ &splits,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 8,
+ 8,
+ None,
+ &factory,
+ &opts,
+ false,
+ None,
+ 1,
+ )
+ .await
+ .unwrap();
+ assert_eq!(
+ visited.lock().unwrap().as_slice(),
+ &[0usize, 1, 2, 3],
+ "concurrency == 1 must search buckets in strict split order"
+ );
+ }
+
+ #[tokio::test]
+ async fn search_candidates_serial_equals_parallel_with_ties() {
+ // Multiple buckets whose exact candidates share a distance (a tie
decided by
+ // the cross-bucket candidate_cmp) plus NaN-distance candidates that
must sort
+ // last. Running at concurrency == 1 and concurrency > 1 must yield
the same
+ // deterministic best-first survivors regardless of bucket completion
order.
+ let table_path = "memory:/pkvo_serial_parallel";
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let negative_nan = f32::from_bits(0xffc00000);
+ let mut splits = Vec::new();
+ for bucket in 0..4i32 {
+ let bucket_path = format!("{table_path}/bucket-{bucket}");
+ let meta = write_file(&file_io, &bucket_path, "d.mosaic",
vec![bucket, bucket]).await;
+ splits.push(PkVectorSearchSplit {
+ data_split: DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(BinaryRow::new(0))
+ .with_bucket(bucket)
+ .with_bucket_path(bucket_path)
+ .with_total_buckets(4)
+ .with_data_files(vec![meta])
+ .build()
+ .unwrap(),
+ ann_segments: Vec::new(),
+ active_files: vec![active("d.mosaic", 2)],
+ });
+ }
+ // Each bucket's exact file yields the same two candidates: pos0
distance 1.0
+ // (tied across buckets) and pos1 a NaN distance (must sort last).
Lower
+ // split_index buckets yield more, so they complete last under
+ // buffer_unordered.
+ let factory = as_split_search(
+ move |split_index: usize,
+ _: &PkVectorSearchSplit,
+ file: &BucketActiveFile,
+ _: &[&[f32]],
+ _: VectorSearchMetric,
+ _: usize,
+ _: &(dyn Fn(i64) -> bool + Sync)|
+ -> ExactFileSearchFuture<'_> {
+ let file_name = file.file_name.clone();
+ let yields = (4 - split_index) * 2;
+ Box::pin(async move {
+ for _ in 0..yields {
+ tokio::task::yield_now().await;
+ }
+ Ok(vec![vec![
+ PkVectorSearchResult {
+ data_file_name: file_name.clone(),
+ row_position: 0,
+ distance: 1.0,
+ },
+ PkVectorSearchResult {
+ data_file_name: file_name,
+ row_position: 1,
+ distance: negative_nan,
+ },
+ ]])
+ })
+ },
+ );
+ let opts = HashMap::new();
+ let serial_result =
PkVectorOrchestrator::new(make_reader(file_io.clone(), table_path))
+ .search_candidates(
+ &splits,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 8,
+ 8,
+ None,
+ &factory,
+ &opts,
+ false,
+ None,
+ 1,
+ )
+ .await
+ .unwrap();
+ let serial: Vec<(i32, i64, bool)> =
+ merge_candidates(serial_result.indexed, serial_result.exact, 8)
+ .iter()
+ .map(|c| (c.bucket, c.row_position, c.distance.is_nan()))
+ .collect();
+ let parallel_result = PkVectorOrchestrator::new(make_reader(file_io,
table_path))
+ .search_candidates(
+ &splits,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 8,
+ 8,
+ None,
+ &factory,
+ &opts,
+ false,
+ None,
+ 4,
+ )
+ .await
+ .unwrap();
+ let parallel: Vec<(i32, i64, bool)> =
+ merge_candidates(parallel_result.indexed, parallel_result.exact, 8)
+ .iter()
+ .map(|c| (c.bucket, c.row_position, c.distance.is_nan()))
+ .collect();
+ // Deterministic order: four tied distance-1.0 candidates by bucket
asc, then
+ // four NaN candidates by bucket asc.
+ assert_eq!(
+ serial,
+ vec![
+ (0, 0, false),
+ (1, 0, false),
+ (2, 0, false),
+ (3, 0, false),
+ (0, 1, true),
+ (1, 1, true),
+ (2, 1, true),
+ (3, 1, true),
+ ]
+ );
+ assert_eq!(
+ parallel, serial,
+ "parallel survivors must equal serial survivors"
+ );
+ }
+
+ #[tokio::test]
+ async fn search_candidates_peak_concurrency_capped_across_buckets() {
+ // Reproduces the reviewer's scenario: 2 buckets x 2 exact files with
+ // concurrency = 2. The per-bucket and per-exact-file fan-outs draw
from ONE
+ // shared budget, so at most `concurrency` (2) exact-file searches run
at once
+ // across ALL buckets. Before the shared budget each level capped at N
+ // independently, so 2 buckets x 2 files = 4 file searches ran
simultaneously.
+ use std::sync::atomic::{AtomicUsize, Ordering};
+ use std::sync::Arc;
+
+ let table_path = "memory:/pkvo_peak_concurrency";
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let mut splits = Vec::new();
+ for bucket in 0..2i32 {
+ let bucket_path = format!("{table_path}/bucket-{bucket}");
+ let meta_a = write_file(&file_io, &bucket_path, "a.mosaic",
vec![bucket]).await;
+ let meta_b = write_file(&file_io, &bucket_path, "b.mosaic",
vec![bucket]).await;
+ splits.push(PkVectorSearchSplit {
+ data_split: DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(BinaryRow::new(0))
+ .with_bucket(bucket)
+ .with_bucket_path(bucket_path)
+ .with_total_buckets(2)
+ .with_data_files(vec![meta_a, meta_b])
+ .build()
+ .unwrap(),
+ ann_segments: Vec::new(),
+ active_files: vec![active("a.mosaic", 1), active("b.mosaic",
1)],
+ });
+ }
+
+ // Shared (current, peak) in-flight counters. Each exact-file search
increments
+ // on entry (after acquiring its budget slot), yields so overlapping
searches
+ // are observable, then decrements; `peak` is the max simultaneous
count.
+ let counters = Arc::new((AtomicUsize::new(0), AtomicUsize::new(0)));
+ let counters_in_closure = counters.clone();
+ let factory = as_split_search(
+ move |_: usize,
+ _: &PkVectorSearchSplit,
+ file: &BucketActiveFile,
+ _: &[&[f32]],
+ _: VectorSearchMetric,
+ _: usize,
+ _: &(dyn Fn(i64) -> bool + Sync)|
+ -> ExactFileSearchFuture<'_> {
+ let counters = counters_in_closure.clone();
+ let file_name = file.file_name.clone();
+ Box::pin(async move {
+ let current = counters.0.fetch_add(1, Ordering::SeqCst) +
1;
+ counters.1.fetch_max(current, Ordering::SeqCst);
+ for _ in 0..8 {
+ tokio::task::yield_now().await;
+ }
+ counters.0.fetch_sub(1, Ordering::SeqCst);
+ Ok(vec![vec![PkVectorSearchResult {
+ data_file_name: file_name,
+ row_position: 0,
+ distance: 1.0,
+ }]])
+ })
+ },
+ );
+
+ let opts = HashMap::new();
+ PkVectorOrchestrator::new(make_reader(file_io, table_path))
+ .search_candidates(
+ &splits,
+ &[0.0],
+ VectorSearchMetric::L2,
+ 8,
+ 8,
+ None,
+ &factory,
+ &opts,
+ false,
+ None,
+ 2,
+ )
+ .await
+ .unwrap();
+
+ let peak = counters.1.load(Ordering::SeqCst);
+ assert!(
+ peak <= 2,
+ "shared budget must cap concurrent exact-file searches at
concurrency (2); \
+ observed peak {peak} (independent per-level fan-out would reach
4)"
+ );
+ assert!(
+ peak >= 2,
+ "test must actually exercise cross-bucket overlap; observed peak
{peak}"
+ );
+ }
}
diff --git a/crates/paimon/src/table/vector_search_builder.rs
b/crates/paimon/src/table/vector_search_builder.rs
index f14c08e9..9f88558e 100644
--- a/crates/paimon/src/table/vector_search_builder.rs
+++ b/crates/paimon/src/table/vector_search_builder.rs
@@ -629,6 +629,9 @@ async fn plan_and_search_pk_candidates_batch(
// `primary_key_vector_distance_metric` returns a validated name; re-parse
into
// the enum for the numeric semantics.
let metric =
VectorSearchMetric::parse(&core.primary_key_vector_distance_metric(pk_col)?)?;
+ // Fan-out limit for the per-bucket and per-exact-file search (Java
+ // `GLOBAL_INDEX_THREAD_NUM`); `1` reproduces strictly sequential
execution.
+ let concurrency = core.global_index_thread_num()?;
let index_type = core.primary_key_vector_index_type(pk_col)?;
let field_id = find_field_id_by_name(table.schema().fields(),
pk_col).ok_or_else(|| {
crate::Error::DataInvalid {
@@ -862,6 +865,7 @@ async fn plan_and_search_pk_candidates_batch(
&search_options,
skip_exact_fallback,
residual_by_split.as_deref(),
+ concurrency,
)
.await?;
diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs
b/crates/paimon/src/vindex/pkvector/bucket.rs
index 56900e07..b946918f 100644
--- a/crates/paimon/src/vindex/pkvector/bucket.rs
+++ b/crates/paimon/src/vindex/pkvector/bucket.rs
@@ -20,6 +20,8 @@ use std::collections::{BinaryHeap, HashMap, HashSet};
use std::sync::Arc;
use futures::future::BoxFuture;
+use futures::stream::{self, StreamExt, TryStreamExt};
+use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use super::ann::PkVectorAnnSearcher;
use super::data_invalid;
@@ -117,6 +119,64 @@ fn add_candidate(heap: &mut BinaryHeap<WorstFirst>,
candidate: PkVectorSearchRes
}
}
+/// Extract the sole per-query list from a single-query exact-file search
result,
+/// failing loud if the closure returned no lists.
+fn single_query_result(
+ per_query: Vec<Vec<PkVectorSearchResult>>,
+) -> crate::Result<Vec<PkVectorSearchResult>> {
+ per_query
+ .into_iter()
+ .next()
+ .ok_or_else(|| data_invalid("exact file search returned no per-query
results"))
+}
+
+/// Verify a multi-query exact-file search returned exactly one list per query,
+/// failing loud on an arity mismatch.
+fn validate_per_query_len(
+ per_query: Vec<Vec<PkVectorSearchResult>>,
+ expected: usize,
+) -> crate::Result<Vec<Vec<PkVectorSearchResult>>> {
+ if per_query.len() != expected {
+ return Err(data_invalid(format!(
+ "exact file search returned {} result lists for {} queries",
+ per_query.len(),
+ expected
+ )));
+ }
+ Ok(per_query)
+}
+
+/// Build the per-position exclusion predicate for one uncovered exact file: a
+/// physical position is excluded if the deletion vector marks it deleted, or
+/// (when a residual allow-list is present) if it is not in the allow-list.
Folds
+/// the residual ∩ DV exclusion the exact search applies per row. The returned
+/// closure borrows the allow-list for its lifetime.
+fn position_excluder(
+ dv: Option<Arc<DeletionVector>>,
+ residual_allowed: Option<&roaring::RoaringTreemap>,
+) -> impl Fn(i64) -> bool + Sync + '_ {
+ move |position: i64| -> bool {
+ let dv_deleted = match &dv {
+ Some(dv) => u64::try_from(position)
+ .map(|p| dv.is_deleted(p))
+ .unwrap_or(false),
+ None => false,
+ };
+ if dv_deleted {
+ return true;
+ }
+ match residual_allowed {
+ // No residual restriction: the row is allowed.
+ None => false,
+ // Residual present: exclude positions outside the allow-list.
+ Some(allowed) => match u64::try_from(position) {
+ Ok(p) => !allowed.contains(p),
+ Err(_) => true,
+ },
+ }
+ }
+}
+
/// Active data files whose rows are already covered by an ANN segment's source
/// metadata, matched by both file name AND row count. The bucket exact
fallback
/// skips these files (an ANN segment already covers their rows), so they never
@@ -145,6 +205,37 @@ pub(crate) fn covered_source_files(
covered
}
+/// Acquire one slot from the shared global-index search concurrency budget,
if a
+/// budget is set. The returned guard holds the slot until it is dropped, so
the
+/// caller must keep it alive for the duration of the leaf exact-file I/O it
gates.
+///
+/// A `None` budget means the leaf runs ungated (no cap) — the function does
not
+/// require any particular `concurrency` value; the orchestrator simply passes
+/// `None` on the strictly sequential `concurrency <= 1` path (which needs no
+/// gating). A `Some` budget is a single [`Semaphore`] shared across every
bucket
+/// and every exact file of one search, so total in-flight exact-file I/O is
capped
+/// at N regardless of how many buckets and files fan out — mirroring Java's
single
+/// shared `GlobalIndexReadThreadPool`. Only leaf exact-file work acquires a
permit;
+/// bucket orchestration never holds one, so it cannot starve leaf work (the
async
+/// analogue of Java's "start from the caller" note in
+/// `PrimaryKeyVectorRead.searchBuckets`).
+async fn acquire_search_permit(
+ budget: &Option<Arc<Semaphore>>,
+) -> crate::Result<Option<OwnedSemaphorePermit>> {
+ match budget {
+ Some(semaphore) => {
+ let permit = semaphore.clone().acquire_owned().await.map_err(|e| {
+ crate::Error::UnexpectedError {
+ message: "global-index search concurrency budget was
closed".to_string(),
+ source: Some(Box::new(e)),
+ }
+ })?;
+ Ok(Some(permit))
+ }
+ None => Ok(None),
+ }
+}
+
/// Separately bounded approximate-index and exact-fallback candidates for one
/// bucket. The approximate list may be over-fetched (for later exact
reranking)
/// while the exact-fallback list stays bounded to the caller's final limit.
@@ -189,6 +280,8 @@ pub(crate) async fn bucket_search(
search_options: &HashMap<String, String>,
skip_exact_fallback: bool,
residual_ranges: Option<&HashMap<String, roaring::RoaringTreemap>>,
+ concurrency: usize,
+ search_budget: Option<Arc<Semaphore>>,
) -> crate::Result<BucketSearchResult> {
if indexed_limit == 0 {
return Err(data_invalid("vector search limit must be positive"));
@@ -260,6 +353,9 @@ pub(crate) async fn bucket_search(
// skips them, so the lazy exact-reader factory is never invoked for those
files.
let covered = covered_source_files(ann_segments, active_files);
+ // The ANN searcher is synchronous CPU work (no `.await`), so iterating
segments
+ // sequentially is intentional: fanning it out would need
`spawn_blocking`. Only
+ // the exact-fallback file reads below (which are async I/O) are
parallelized.
for segment in ann_segments {
// An active ANN source with a mismatched row count is corruption (the
// ordinal-to-position mapping would be wrong). An inactive source (no
@@ -292,6 +388,11 @@ pub(crate) async fn bucket_search(
}
if !skip_exact_fallback {
+ // Collect the eligible uncovered files (active-file order preserved)
with
+ // their per-position exclusion predicate. A file with no
residual-allowed
+ // rows is skipped without reading.
+ #[allow(clippy::type_complexity)]
+ let mut tasks: Vec<(&BucketActiveFile, Box<dyn Fn(i64) -> bool +
Sync>)> = Vec::new();
for file in active_files {
if covered.contains(&file.file_name) {
continue;
@@ -307,41 +408,44 @@ pub(crate) async fn bucket_search(
None => None,
};
let dv = deletion_vectors.get(&file.file_name).cloned();
- let is_excluded = move |position: i64| -> bool {
- let dv_deleted = match &dv {
- Some(dv) => u64::try_from(position)
- .map(|p| dv.is_deleted(p))
- .unwrap_or(false),
- None => false,
- };
- if dv_deleted {
- return true;
- }
- match residual_allowed {
- // No residual restriction: the row is allowed.
- None => false,
- // Residual present: exclude positions outside the
allow-list.
- Some(allowed) => match u64::try_from(position) {
- Ok(p) => !allowed.contains(p),
- Err(_) => true,
- },
+ tasks.push((file, Box::new(position_excluder(dv,
residual_allowed))));
+ }
+
+ // Search each uncovered file for its exact Top-K. The caller passes a
+ // single-query slice and each search returns one per-query list. The
+ // per-file results feed the bounded heap, which is order-independent,
so
+ // the merge does not depend on which file finished first.
`concurrency == 1`
+ // takes a plain sequential loop so the file visit order is strictly
+ // deterministic; larger values fan the file searches out with
+ // `buffer_unordered`, each acquiring one slot of the shared
`search_budget`
+ // so total in-flight exact-file I/O across all buckets is capped at N.
+ let queries: [&[f32]; 1] = [query];
+ let per_file: Vec<Vec<PkVectorSearchResult>> = if concurrency <= 1 {
+ let mut out = Vec::with_capacity(tasks.len());
+ for (file, is_excluded) in &tasks {
+ let per_query =
+ exact_file_search(file, &queries, metric, exact_limit,
is_excluded.as_ref())
+ .await?;
+ out.push(single_query_result(per_query)?);
+ }
+ out
+ } else {
+ stream::iter(tasks.iter().map(|(file, is_excluded)| {
+ let queries = &queries;
+ let budget = search_budget.clone();
+ async move {
+ let _permit = acquire_search_permit(&budget).await?;
+ let per_query =
+ exact_file_search(file, queries, metric, exact_limit,
is_excluded.as_ref())
+ .await?;
+ single_query_result(per_query)
}
- };
- // Search this one file for its per-query exact Top-K. The caller
passes
- // a single-query slice and reads the sole returned list.
- let queries: [&[f32]; 1] = [query];
- let per_query = exact_file_search(
- file,
- &queries,
- metric,
- exact_limit,
- &is_excluded as &(dyn Fn(i64) -> bool + Sync),
- )
- .await?;
- let results = per_query
- .into_iter()
- .next()
- .ok_or_else(|| data_invalid("exact file search returned no
per-query results"))?;
+ }))
+ .buffer_unordered(concurrency)
+ .try_collect::<Vec<_>>()
+ .await?
+ };
+ for results in per_file {
for result in results {
add_candidate(&mut exact_heap, result, exact_limit);
}
@@ -389,6 +493,8 @@ pub(crate) async fn bucket_search_batch(
search_options: &HashMap<String, String>,
skip_exact_fallback: bool,
residual_ranges: Option<&HashMap<String, roaring::RoaringTreemap>>,
+ concurrency: usize,
+ search_budget: Option<Arc<Semaphore>>,
) -> crate::Result<Vec<BucketSearchResult>> {
if queries.is_empty() {
return Err(data_invalid("vector search requires at least one query"));
@@ -410,6 +516,8 @@ pub(crate) async fn bucket_search_batch(
search_options,
skip_exact_fallback,
residual_ranges,
+ concurrency,
+ search_budget,
)
.await?;
return Ok(vec![single]);
@@ -489,6 +597,9 @@ pub(crate) async fn bucket_search_batch(
files_by_name.keys().map(|name| name.to_string()).collect();
let covered = covered_source_files(ann_segments, active_files);
+ // The ANN searcher is synchronous CPU work (no `.await`), so iterating
segments
+ // sequentially is intentional: fanning it out would need
`spawn_blocking`. Only
+ // the exact-fallback file reads below (which are async I/O) are
parallelized.
for segment in ann_segments {
for source in segment.source_meta.source_files() {
if let Some(active) = files_by_name.get(source.file_name()) {
@@ -528,6 +639,11 @@ pub(crate) async fn bucket_search_batch(
}
if !skip_exact_fallback {
+ // Eligible uncovered files (active-file order preserved) with their
+ // per-position exclusion predicate; a file with no residual-allowed
rows is
+ // skipped without reading.
+ #[allow(clippy::type_complexity)]
+ let mut tasks: Vec<(&BucketActiveFile, Box<dyn Fn(i64) -> bool +
Sync>)> = Vec::new();
for file in active_files {
if covered.contains(&file.file_name) {
continue;
@@ -540,40 +656,41 @@ pub(crate) async fn bucket_search_batch(
None => None,
};
let dv = deletion_vectors.get(&file.file_name).cloned();
- let is_excluded = move |position: i64| -> bool {
- let dv_deleted = match &dv {
- Some(dv) => u64::try_from(position)
- .map(|p| dv.is_deleted(p))
- .unwrap_or(false),
- None => false,
- };
- if dv_deleted {
- return true;
- }
- match residual_allowed {
- None => false,
- Some(allowed) => match u64::try_from(position) {
- Ok(p) => !allowed.contains(p),
- Err(_) => true,
- },
- }
- };
- // One shared stream per file scores every query into its own heap.
- let per_query = exact_file_search(
- file,
- queries,
- metric,
- exact_limit,
- &is_excluded as &(dyn Fn(i64) -> bool + Sync),
- )
- .await?;
- if per_query.len() != queries.len() {
- return Err(data_invalid(format!(
- "exact file search returned {} result lists for {}
queries",
- per_query.len(),
- queries.len()
- )));
+ tasks.push((file, Box::new(position_excluder(dv,
residual_allowed))));
+ }
+
+ // One shared stream per file scores every query into its own per-query
+ // list. The per-file lists feed per-query bounded heaps
(order-independent),
+ // so the fan-in does not depend on which file finished first: collect
every
+ // file's per-query result, then merge. `concurrency == 1` uses a
strictly
+ // sequential loop; larger values fan the file searches out with
+ // `buffer_unordered`, each acquiring one slot of the shared
`search_budget`
+ // so total in-flight exact-file I/O across all buckets is capped at N.
+ let per_file: Vec<Vec<Vec<PkVectorSearchResult>>> = if concurrency <=
1 {
+ let mut out = Vec::with_capacity(tasks.len());
+ for (file, is_excluded) in &tasks {
+ let per_query =
+ exact_file_search(file, queries, metric, exact_limit,
is_excluded.as_ref())
+ .await?;
+ out.push(validate_per_query_len(per_query, queries.len())?);
}
+ out
+ } else {
+ stream::iter(tasks.iter().map(|(file, is_excluded)| {
+ let budget = search_budget.clone();
+ async move {
+ let _permit = acquire_search_permit(&budget).await?;
+ let per_query =
+ exact_file_search(file, queries, metric, exact_limit,
is_excluded.as_ref())
+ .await?;
+ validate_per_query_len(per_query, queries.len())
+ }
+ }))
+ .buffer_unordered(concurrency)
+ .try_collect::<Vec<_>>()
+ .await?
+ };
+ for per_query in per_file {
for (results, heap) in
per_query.into_iter().zip(exact_heaps.iter_mut()) {
for result in results {
add_candidate(heap, result, exact_limit);
@@ -777,6 +894,8 @@ mod tests {
&opts,
false,
None,
+ 1,
+ None,
)
.await
.unwrap();
@@ -807,6 +926,8 @@ mod tests {
&HashMap::new(),
false,
None,
+ 1,
+ None,
)
.await
.unwrap_err();
@@ -849,6 +970,8 @@ mod tests {
&HashMap::new(),
false,
None,
+ 1,
+ None,
)
.await
.unwrap();
@@ -904,6 +1027,8 @@ mod tests {
&HashMap::new(),
false,
None,
+ 1,
+ None,
)
.await
.unwrap();
@@ -967,6 +1092,8 @@ mod tests {
&HashMap::new(),
false,
None,
+ 1,
+ None,
)
.await
.unwrap();
@@ -1042,6 +1169,8 @@ mod tests {
&HashMap::new(),
false,
None,
+ 1,
+ None,
)
.await
.unwrap();
@@ -1084,6 +1213,8 @@ mod tests {
&HashMap::new(),
false,
None,
+ 1,
+ None,
)
.await
.unwrap_err();
@@ -1110,6 +1241,8 @@ mod tests {
&HashMap::new(),
false,
None,
+ 1,
+ None,
)
.await
.unwrap_err();
@@ -1158,6 +1291,8 @@ mod tests {
&HashMap::new(),
false,
None,
+ 1,
+ None,
)
.await
.unwrap();
@@ -1194,6 +1329,8 @@ mod tests {
&HashMap::new(),
false,
None,
+ 1,
+ None,
)
.await
.unwrap_err();
@@ -1221,6 +1358,8 @@ mod tests {
&HashMap::new(),
true, // skip_exact_fallback
None,
+ 1,
+ None,
)
.await
.unwrap();
@@ -1257,6 +1396,8 @@ mod tests {
&HashMap::new(),
false,
None,
+ 1,
+ None,
)
.await
.unwrap_err();
@@ -1295,6 +1436,8 @@ mod tests {
&HashMap::new(),
false,
None,
+ 1,
+ None,
)
.await
.unwrap_err();
@@ -1322,6 +1465,8 @@ mod tests {
&HashMap::new(),
false,
None,
+ 1,
+ None,
)
.await
.unwrap_err();
@@ -1390,6 +1535,8 @@ mod tests {
&HashMap::new(),
false,
Some(&residual),
+ 1,
+ None,
)
.await
.unwrap();
@@ -1458,6 +1605,8 @@ mod tests {
&HashMap::new(),
false,
Some(&residual),
+ 1,
+ None,
)
.await
.unwrap();
@@ -1503,6 +1652,8 @@ mod tests {
&HashMap::new(),
false,
Some(&residual),
+ 1,
+ None,
)
.await
.unwrap();
@@ -1539,6 +1690,8 @@ mod tests {
&HashMap::new(),
false,
Some(&residual),
+ 1,
+ None,
)
.await
.unwrap();
@@ -1582,6 +1735,8 @@ mod tests {
&HashMap::new(),
false,
None,
+ 1,
+ None,
)
.await
.unwrap_err();
@@ -1684,6 +1839,8 @@ mod tests {
&opts,
false,
None,
+ 1,
+ None,
)
.await
.unwrap()
@@ -1708,6 +1865,8 @@ mod tests {
&opts,
false,
None,
+ 1,
+ None,
)
.await
.unwrap();
@@ -1742,6 +1901,8 @@ mod tests {
&HashMap::new(),
false,
None,
+ 1,
+ None,
)
.await
.unwrap();
@@ -1800,6 +1961,8 @@ mod tests {
&HashMap::new(),
false,
None,
+ 1,
+ None,
)
.await
.unwrap();
@@ -1842,6 +2005,8 @@ mod tests {
&HashMap::new(),
false,
None,
+ 1,
+ None,
)
.await
.unwrap_err();
@@ -1851,4 +2016,168 @@ mod tests {
"the exact-file search closure must not be invoked for a malformed
batch"
);
}
+
+ #[tokio::test]
+ async fn test_exact_file_search_order_is_sequential_at_concurrency_one() {
+ // At concurrency == 1 the per-exact-file loop must visit uncovered
files in
+ // strict active-file (submission) order. A recording closure captures
the
+ // visit order; it must match the active-file order deterministically.
+ let calls = std::sync::Mutex::new(Vec::<String>::new());
+ let factory = as_search(
+ |file: &BucketActiveFile,
+ queries: &[&[f32]],
+ metric: VectorSearchMetric,
+ exact_limit: usize,
+ is_excluded: &(dyn Fn(i64) -> bool + Sync)|
+ -> ExactFileSearchFuture<'_> {
+ calls.lock().unwrap().push(file.file_name.clone());
+ let file_name = file.file_name.clone();
+ let query = queries[0].to_vec();
+ Box::pin(async move {
+ let mut reader = ArrayReader::new(2, vec![Some(vec![1.0,
0.0])]);
+ Ok(vec![exact_search(
+ &file_name,
+ &mut reader,
+ &query,
+ metric,
+ exact_limit,
+ is_excluded,
+ )?])
+ })
+ },
+ );
+ let out = bucket_search(
+ None,
+ &[],
+ &[
+ active("data-a", 1),
+ active("data-b", 1),
+ active("data-c", 1),
+ active("data-d", 1),
+ ],
+ &HashMap::new(),
+ &factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 8,
+ 8,
+ &HashMap::new(),
+ false,
+ None,
+ 1,
+ None,
+ )
+ .await
+ .unwrap();
+ assert_eq!(out.exact.len(), 4);
+ assert_eq!(
+ calls.lock().unwrap().as_slice(),
+ &[
+ "data-a".to_string(),
+ "data-b".to_string(),
+ "data-c".to_string(),
+ "data-d".to_string(),
+ ],
+ "concurrency == 1 must visit files in strict active-file order"
+ );
+ }
+
+ #[tokio::test]
+ async fn
test_parallel_exact_files_tie_and_nan_rank_deterministically_out_of_order() {
+ // Two uncovered files whose candidates share a distance (a tie
decided by
+ // the BEST_FIRST file/position tie-break) plus a NaN-distance
candidate that
+ // must sort LAST. The futures complete OUT OF ORDER: the
alphabetically
+ // first file (data-a) yields more times before returning, so under
+ // buffer_unordered(concurrency > 1) data-b completes first. The
bounded-heap
+ // merge is order-independent, so the final best-first ranking must
still be
+ // the deterministic BEST_FIRST order (identical to a serial run).
+ let negative_nan = f32::from_bits(0xffc00000);
+ assert!(negative_nan.is_nan());
+ let factory = as_search(
+ move |file: &BucketActiveFile,
+ queries: &[&[f32]],
+ _metric: VectorSearchMetric,
+ exact_limit: usize,
+ _is_excluded: &(dyn Fn(i64) -> bool + Sync)|
+ -> ExactFileSearchFuture<'_> {
+ let name = file.file_name.clone();
+ // data-a is submitted first but yields more, so it finishes
last.
+ let yields = if name == "data-a" { 8 } else { 0 };
+ let _ = (queries, exact_limit);
+ Box::pin(async move {
+ for _ in 0..yields {
+ tokio::task::yield_now().await;
+ }
+ // Each file returns two candidates: one tied distance 1.0
and one
+ // NaN distance that must never outrank a finite candidate.
+ let results = vec![
+ PkVectorSearchResult {
+ data_file_name: name.clone(),
+ row_position: 0,
+ distance: 1.0,
+ },
+ PkVectorSearchResult {
+ data_file_name: name.clone(),
+ row_position: 1,
+ distance: negative_nan,
+ },
+ ];
+ Ok(vec![results])
+ })
+ },
+ );
+
+ let run = |concurrency: usize| {
+ let factory = &factory;
+ async move {
+ let out = bucket_search(
+ None,
+ &[],
+ &[active("data-a", 2), active("data-b", 2)],
+ &HashMap::new(),
+ factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 8,
+ 8,
+ &HashMap::new(),
+ false,
+ None,
+ concurrency,
+ None,
+ )
+ .await
+ .unwrap();
+ out.exact
+ .iter()
+ .map(|r| {
+ (
+ r.data_file_name.clone(),
+ r.row_position,
+ r.distance.is_nan(),
+ )
+ })
+ .collect::<Vec<_>>()
+ }
+ };
+
+ let serial = run(1).await;
+ let parallel = run(4).await;
+ // Serial == parallel, and the deterministic BEST_FIRST order is: tied
+ // distance-1.0 candidates first (data-a before data-b by file name),
then the
+ // NaN-distance candidates last (again data-a before data-b).
+ assert_eq!(
+ serial,
+ vec![
+ ("data-a".to_string(), 0, false),
+ ("data-b".to_string(), 0, false),
+ ("data-a".to_string(), 1, true),
+ ("data-b".to_string(), 1, true),
+ ]
+ );
+ assert_eq!(
+ parallel, serial,
+ "parallel ranking must equal serial ranking"
+ );
+ }
}
