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 7d12424d perf(vector-search): execute PK-vector leaves concurrently
(#647)
7d12424d is described below
commit 7d12424da183e6c80c898cf6c4a140ffe4e0562c
Author: Junrui Lee <[email protected]>
AuthorDate: Tue Aug 4 23:45:00 2026 +0800
perf(vector-search): execute PK-vector leaves concurrently (#647)
---
crates/paimon/src/spec/core_options.rs | 40 +
crates/paimon/src/table/pk_vector_orchestrator.rs | 60 +-
crates/paimon/src/table/vector_search_builder.rs | 295 +++--
crates/paimon/src/vindex/executor.rs | 136 ++-
crates/paimon/src/vindex/pkvector/ann.rs | 414 +++++--
crates/paimon/src/vindex/pkvector/bucket.rs | 1298 ++++++++++++++++++---
crates/paimon/src/vindex/reader.rs | 12 +
7 files changed, 1799 insertions(+), 456 deletions(-)
diff --git a/crates/paimon/src/spec/core_options.rs
b/crates/paimon/src/spec/core_options.rs
index e33cd09b..1039800d 100644
--- a/crates/paimon/src/spec/core_options.rs
+++ b/crates/paimon/src/spec/core_options.rs
@@ -123,6 +123,15 @@ const DYNAMIC_BUCKET_TARGET_ROW_NUM_OPTION: &str =
"dynamic-bucket.target-row-nu
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 MAX_GLOBAL_INDEX_THREAD_NUM: i64 = {
+ let tokio_max = (usize::MAX >> 3) as u64;
+ let i32_max = i32::MAX as u64;
+ if tokio_max < i32_max {
+ tokio_max as i64
+ } else {
+ i32_max as i64
+ }
+};
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";
@@ -664,6 +673,15 @@ impl<'a> CoreOptions<'a> {
source: None,
});
}
+ if value > MAX_GLOBAL_INDEX_THREAD_NUM {
+ return Err(crate::Error::DataInvalid {
+ message: format!(
+ "Option '{}' must not exceed {}, got: {}",
+ GLOBAL_INDEX_THREAD_NUM_OPTION,
MAX_GLOBAL_INDEX_THREAD_NUM, value
+ ),
+ source: None,
+ });
+ }
Ok(value as usize)
}
@@ -1578,6 +1596,28 @@ mod tests {
}
}
+ #[test]
+ fn test_global_index_thread_num_rejects_values_above_max() {
+ assert!(MAX_GLOBAL_INDEX_THREAD_NUM as usize <=
tokio::sync::Semaphore::MAX_PERMITS);
+
+ let too_big = (MAX_GLOBAL_INDEX_THREAD_NUM + 1).to_string();
+ let options =
HashMap::from([(GLOBAL_INDEX_THREAD_NUM_OPTION.to_string(), too_big)]);
+ let err = CoreOptions::new(&options)
+ .global_index_thread_num()
+ .expect_err("thread-num above maximum should fail");
+ assert!(matches!(err, crate::Error::DataInvalid { message, .. }
+ if message.contains("must not exceed")));
+
+ let at_max = HashMap::from([(
+ GLOBAL_INDEX_THREAD_NUM_OPTION.to_string(),
+ MAX_GLOBAL_INDEX_THREAD_NUM.to_string(),
+ )]);
+ assert_eq!(
+ CoreOptions::new(&at_max).global_index_thread_num().unwrap(),
+ MAX_GLOBAL_INDEX_THREAD_NUM as usize
+ );
+ }
+
#[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 2278b0d2..1a816bd9 100644
--- a/crates/paimon/src/table/pk_vector_orchestrator.rs
+++ b/crates/paimon/src/table/pk_vector_orchestrator.rs
@@ -28,17 +28,15 @@ 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;
use crate::table::pk_vector_indexed_split_read::PkVectorIndexedSplit;
use crate::table::source::{DataSplit, DataSplitBuilder, RowRange};
+use crate::vindex::executor::drain_indexed_jobs;
use crate::vindex::pkvector::ann::PkVectorAnnSearcher;
use crate::vindex::pkvector::bucket::{
- bucket_search_batch, BucketActiveFile, BucketAnnSegment,
ExactFileSearchFuture,
+ bucket_search_batch, BucketActiveFile, BucketAnnSegment,
ExactFileSearchFuture, SearchBudget,
};
use crate::vindex::pkvector::metric::{java_float_compare, VectorSearchMetric};
use crate::vindex::pkvector::result::PkVectorSearchResult;
@@ -379,7 +377,7 @@ impl PkVectorOrchestrator {
metric: VectorSearchMetric,
limit: usize,
indexed_limit: usize,
- ann_searcher: Option<&dyn PkVectorAnnSearcher>,
+ ann_searcher: Option<Arc<dyn PkVectorAnnSearcher>>,
exact_file_search: &(dyn for<'s, 'a> Fn(
usize,
&'s PkVectorSearchSplit,
@@ -430,17 +428,12 @@ impl PkVectorOrchestrator {
/// 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.
+ /// `1` runs buckets and their leaves strictly sequentially, while larger
values
+ /// fan them out concurrently. A shared [`SearchBudget`] gates every ANN
segment
+ /// and exact-file leaf across all buckets, so a query cannot multiply the
limit
+ /// at each nesting level. The process-global side of the budget also
bounds
+ /// concurrent work across separate queries. Each query's final
cross-bucket
+ /// Top-K merge is order-independent, so completion order does not affect
output.
#[allow(clippy::too_many_arguments)]
#[allow(clippy::type_complexity)]
pub(crate) async fn search_candidates_batch(
@@ -450,7 +443,7 @@ impl PkVectorOrchestrator {
metric: VectorSearchMetric,
limit: usize,
indexed_limit: usize,
- ann_searcher: Option<&dyn PkVectorAnnSearcher>,
+ ann_searcher: Option<Arc<dyn PkVectorAnnSearcher>>,
exact_file_search: &(dyn for<'s, 'a> Fn(
usize,
&'s PkVectorSearchSplit,
@@ -497,14 +490,11 @@ impl PkVectorOrchestrator {
let mut exact_candidates: Vec<Vec<PkVectorCandidate>> =
(0..queries.len()).map(|_| Vec::new()).collect();
- // 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 shared budget for the whole search. Every ANN and exact-file
leaf
+ // draws from the same query-local permits and process-global pool;
bucket
+ // orchestration itself never holds a permit, so it cannot starve leaf
work.
+ // The budget also applies at concurrency 1 to preserve process-wide
bounds.
+ let search_budget = Some(SearchBudget::production(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
@@ -513,6 +503,7 @@ impl PkVectorOrchestrator {
// 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();
+ let ann_searcher = ann_searcher.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
@@ -589,7 +580,7 @@ impl PkVectorOrchestrator {
// 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
+ // them out through `drain_indexed_jobs`. 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
@@ -602,10 +593,7 @@ impl PkVectorOrchestrator {
}
out
} else {
- stream::iter(per_bucket)
- .buffer_unordered(concurrency)
- .try_collect::<Vec<_>>()
- .await?
+ drain_indexed_jobs(per_bucket, concurrency).await?
};
for tagged in collected {
for (query_index, (indexed, exact)) in
tagged.into_iter().enumerate() {
@@ -1256,9 +1244,17 @@ mod e2e_tests {
hits: Vec<PkVectorSearchResult>,
}
impl PkVectorAnnSearcher for FakeAnn {
+ fn load_segment(
+ &self,
+ _segment: &BucketAnnSegment,
+ ) -> futures::future::BoxFuture<'static, crate::Result<Bytes>> {
+ Box::pin(async { Ok(Bytes::new()) })
+ }
+
fn search_batch(
&self,
_segment: &BucketAnnSegment,
+ _segment_bytes: Bytes,
queries: &[&[f32]],
_metric: VectorSearchMetric,
_limit: usize,
@@ -1283,7 +1279,7 @@ mod e2e_tests {
query: &[f32],
metric: VectorSearchMetric,
limit: usize,
- ann: Option<&dyn PkVectorAnnSearcher>,
+ ann: Option<Arc<dyn PkVectorAnnSearcher>>,
search: &(dyn for<'a> Fn(
&'a BucketActiveFile,
&'a [&'a [f32]],
@@ -1429,7 +1425,7 @@ mod e2e_tests {
&[0.0, 0.0],
VectorSearchMetric::L2,
3,
- Some(&ann),
+ Some(Arc::new(ann)),
&factory,
&opts,
)
diff --git a/crates/paimon/src/table/vector_search_builder.rs
b/crates/paimon/src/table/vector_search_builder.rs
index b2d67597..f1e22b20 100644
--- a/crates/paimon/src/table/vector_search_builder.rs
+++ b/crates/paimon/src/table/vector_search_builder.rs
@@ -51,9 +51,10 @@ use crate::table::{
};
use crate::vector_search::{GlobalIndexIOMeta, SearchResult, VectorSearch};
use crate::vindex::executor::{
- drain_indexed_jobs, ensure_global_index_executor_capacity,
execute_global_index,
+ acquire_process_global_search_permit, drain_indexed_jobs,
+ ensure_global_index_executor_capacity, execute_global_index_with_guard,
};
-use crate::vindex::pkvector::ann::VindexAnnSearcher;
+use crate::vindex::pkvector::ann::{AnnSegmentSource, PkVectorAnnSearcher,
VindexAnnSearcher};
use crate::vindex::pkvector::bucket::{BucketActiveFile, BucketAnnSegment,
ExactFileSearchFuture};
use crate::vindex::pkvector::exact::validate_query;
use crate::vindex::pkvector::metric::VectorSearchMetric;
@@ -99,20 +100,21 @@ impl VectorIndexBackend {
}
}
-async fn execute_vindex_searches<S: SeekRead + 'static>(
+async fn execute_vindex_searches<S: SeekRead + 'static, G: Send + 'static>(
io_meta: GlobalIndexIOMeta,
options: HashMap<String, String>,
vector_searches: Vec<VectorSearch>,
source: S,
file_name: String,
shard_concurrency: usize,
+ guard: G,
) -> crate::Result<Vec<Option<HashMap<u64, f32>>>> {
let panic_context = if vector_searches.len() > 1 {
"vindex global-index batch search task failed"
} else {
"vindex global-index search task failed"
};
- execute_global_index(panic_context, move || {
+ execute_global_index_with_guard(panic_context, guard, move || {
let mut reader = VindexVectorGlobalIndexReader::new(io_meta, options)
.with_batch_shard_concurrency(shard_concurrency);
reader
@@ -125,6 +127,13 @@ async fn execute_vindex_searches<S: SeekRead + 'static>(
.await
}
+fn current_tokio_runtime_handle() -> crate::Result<tokio::runtime::Handle> {
+ tokio::runtime::Handle::try_current().map_err(|error|
crate::Error::UnexpectedError {
+ message: "Vector index range reader requires a Tokio
runtime".to_string(),
+ source: Some(Box::new(error)),
+ })
+}
+
pub struct VectorSearchBuilder<'a> {
table: &'a Table,
vector_column: Option<String>,
@@ -667,7 +676,7 @@ pub(crate) fn ensure_no_reserved_read_columns(fields:
&[DataField]) -> crate::Re
}
/// Batch PK-vector search core shared by the single and batch builders: plan
ONE
-/// per-bucket split set, segment preload, ANN scorer, exact-fallback search
+/// per-bucket split set, lazy segment loader, ANN scorer, exact-fallback
search
/// closure, and residual allow-list (all query-independent), then run
/// `search_candidates_batch` ONCE so N queries share the opened readers. Per
/// query, the approximate candidates are exact-reranked (when a refine factor
is
@@ -732,7 +741,7 @@ 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
+ // Fan-out limit for bucket orchestration plus ANN and exact-file leaves
(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)?;
@@ -824,14 +833,9 @@ async fn plan_and_search_pk_candidates_batch(
Vec::new(),
);
- // Real ANN scorer: preload each segment's bytes (keyed by resolved,
globally
- // unique path) and drive the vindex reader from memory. The reader is
opened
- // once per segment and every query in the batch is searched against it,
- // mirroring the shared-reader batch search.
- let segment_bytes = preload_segment_bytes(table.file_io(),
&plan.splits).await?;
- // Fail loud on a config/segment metric mismatch before scoring, mirroring
Java
- // `PkVectorAnnSegmentSearcher.search`.
- verify_pk_vector_segment_metrics(&plan.splits, &segment_bytes, metric,
backend)?;
+ // Real ANN scorer + loader. Each segment source is opened lazily inside
its
+ // bucket leaf and dropped after scoring. Lumina keeps its buffered-byte
path;
+ // vindex remains range-backed and reads only metadata and probed lists.
let options = {
let mut o = table.schema().options().clone();
o.extend(query_options.clone());
@@ -839,34 +843,97 @@ async fn plan_and_search_pk_candidates_batch(
};
let search_options = options.clone();
let field_name = pk_col.to_string();
- let scorer: crate::vindex::pkvector::ann::BatchScorer = Box::new(
- move |segment: &BucketAnnSegment, searches: &[VectorSearch]| {
- let data = segment_bytes
- .get(&segment.path)
- .ok_or_else(|| crate::Error::DataInvalid {
- message: "missing preloaded ANN bytes for
segment".to_string(),
- source: None,
- })?
- .clone();
+
+ let loader_io = table.file_io().clone();
+ let loader_range_read_permits =
Arc::new(tokio::sync::Semaphore::new(concurrency));
+ let loader: crate::vindex::pkvector::ann::SourceSegmentLoader = Box::new(
+ move |segment: &BucketAnnSegment| {
+ let io = loader_io.clone();
+ let range_read_permits = Arc::clone(&loader_range_read_permits);
+ let path = segment.path.clone();
+ let file_size = segment.file_size;
+ Box::pin(async move {
+ let input = io.new_input(&path)?;
+ match backend {
+ VectorIndexBackend::Lumina => input
+ .read()
+ .await
+ .map(AnnSegmentSource::Buffered)
+ .map_err(|error| crate::Error::DataInvalid {
+ message: format!("failed to read ANN index file
'{path}': {error}"),
+ source: None,
+ }),
+ VectorIndexBackend::Vindex => {
+ let file_reader =
+ input
+ .reader()
+ .await
+ .map_err(|error| crate::Error::DataInvalid {
+ message: format!(
+ "failed to open ANN index file
'{path}' for range reads: {error}"
+ ),
+ source: None,
+ })?;
+ Ok(AnnSegmentSource::Vindex(
+ VindexFileReader::new_with_permits(
+ Arc::new(file_reader),
+ current_tokio_runtime_handle()?,
+ range_read_permits,
+ file_size,
+ path,
+ ),
+ ))
+ }
+ }
+ })
+ },
+ );
+
+ let scorer: crate::vindex::pkvector::ann::SourceBatchScorer = Box::new(
+ move |segment: &BucketAnnSegment, source: AnnSegmentSource, searches:
&[VectorSearch]| {
let io_meta = GlobalIndexIOMeta::new(
segment.path.clone(),
segment.file_size,
segment.index_meta.clone(),
);
- match backend {
- VectorIndexBackend::Lumina => {
+ match (backend, source) {
+ (VectorIndexBackend::Lumina, AnnSegmentSource::Buffered(data))
=> {
+ let lumina_metric =
+
LuminaIndexMeta::deserialize(&segment.index_meta)?.metric()?;
+ verify_segment_metric(metric,
VectorSearchMetric::from_lumina(lumina_metric))?;
let mut reader =
LuminaVectorGlobalIndexReader::new(io_meta, options.clone());
reader.visit_batch_vector_search(searches, |_|
Ok(Cursor::new(data)))
}
- VectorIndexBackend::Vindex => {
+ (VectorIndexBackend::Vindex, AnnSegmentSource::Vindex(source))
=> {
let mut reader =
VindexVectorGlobalIndexReader::new(io_meta, options.clone())
.with_batch_shard_concurrency(concurrency);
- reader.visit_batch_vector_search(searches, |_|
Ok(Cursor::new(data)))
+ reader.load_validated(
+ |_| Ok(source),
+ |metadata| {
+ verify_segment_metric(
+ metric,
+
VectorSearchMetric::from_vindex(metadata.metric),
+ )
+ },
+ )?;
+ reader.search_batch(searches)
+ }
+ (VectorIndexBackend::Lumina, AnnSegmentSource::Vindex(_))
+ | (VectorIndexBackend::Vindex, AnnSegmentSource::Buffered(_))
=> {
+ Err(crate::Error::DataInvalid {
+ message: format!(
+ "ANN segment '{}' was loaded with the wrong
backend source",
+ segment.path
+ ),
+ source: None,
+ })
}
}
},
);
- let ann_searcher = VindexAnnSearcher::new(field_name, scorer);
+ let ann_searcher: Arc<dyn PkVectorAnnSearcher> =
Arc::new(VindexAnnSearcher::new_with_source(
+ field_name, scorer, loader,
+ ));
// Residual (post-recall) filtering: for each candidate file, re-read its
// physical rows and keep the positions whose rows satisfy the filter. The
@@ -983,7 +1050,7 @@ async fn plan_and_search_pk_candidates_batch(
metric,
limit,
indexed_limit,
- Some(&ann_searcher),
+ Some(ann_searcher),
&factory,
&search_options,
skip_exact_fallback,
@@ -1253,7 +1320,7 @@ impl<'a> BatchVectorSearchBuilder<'a> {
// Vec.
let read_type = self.resolve_materialize_read_type()?;
- // One shared plan / segment preload / residual across all N queries;
the
+ // One shared plan / lazy segment loader / residual across all N
queries; the
// per-query candidate lists come back in strict input order. Any query
// error (or a shared-plan error) propagates here, so no partial Vec is
// returned.
@@ -1484,6 +1551,7 @@ async fn evaluate_batch_vector_search(
options.extend(search_options.clone());
let input = evaluation.file_io.new_input(&path);
async move {
+ let permit =
acquire_process_global_search_permit(concurrency).await?;
let input = input?;
let query_count = vector_searches.len();
let io_meta =
@@ -1501,8 +1569,9 @@ async fn evaluate_batch_vector_search(
source: None,
}
})?;
- execute_global_index(
+ execute_global_index_with_guard(
"Lumina global-index batch search task failed",
+ permit,
move || {
let mut reader =
LuminaVectorGlobalIndexReader::new(io_meta, options);
@@ -1539,6 +1608,7 @@ async fn evaluate_batch_vector_search(
source,
file_name,
concurrency,
+ permit,
)
.await?
}
@@ -1559,6 +1629,7 @@ async fn evaluate_batch_vector_search(
Cursor::new(data),
file_name,
concurrency,
+ permit,
)
.await?
}
@@ -1752,87 +1823,19 @@ async fn residual_positions_by_file(
Ok(out)
}
-/// Preload every ANN segment's bytes into a map keyed by the resolved
(globally
-/// unique) segment path. The scorer closure reads from this map so the vindex
-/// reader is driven from memory without per-search IO.
-async fn preload_segment_bytes(
- file_io: &FileIO,
- splits: &[PkVectorSearchSplit],
-) -> crate::Result<HashMap<String, Vec<u8>>> {
- let mut out = HashMap::new();
- for split in splits {
- for segment in &split.ann_segments {
- if out.contains_key(&segment.path) {
- continue;
- }
- let input = file_io.new_input(&segment.path)?;
- let bytes = input.read().await.map_err(|e|
crate::Error::DataInvalid {
- message: format!("failed to read ANN index file '{}': {e}",
segment.path),
- source: None,
- })?;
- out.insert(segment.path.clone(), bytes.to_vec());
- }
- }
- Ok(out)
-}
-
-/// Fail loud when an ANN segment was trained with a metric other than the
-/// configured one, mirroring the search-time `checkArgument` in Java
-/// `PkVectorAnnSegmentSearcher.search`. Opens each distinct segment's
preloaded
-/// bytes once and compares its trained metric against `configured`.
-fn verify_pk_vector_segment_metrics(
- splits: &[PkVectorSearchSplit],
- segment_bytes: &HashMap<String, Vec<u8>>,
+fn verify_segment_metric(
configured: VectorSearchMetric,
- backend: VectorIndexBackend,
+ segment_metric: VectorSearchMetric,
) -> crate::Result<()> {
- let mut checked: HashSet<&str> = HashSet::new();
- for split in splits {
- for segment in &split.ann_segments {
- if !checked.insert(segment.path.as_str()) {
- continue;
- }
- let segment_metric = match backend {
- VectorIndexBackend::Lumina => {
- // Lumina records its metric in the serialized index
metadata
- // (`index_meta`), not in the segment file bytes.
- let lumina_metric =
-
LuminaIndexMeta::deserialize(&segment.index_meta)?.metric()?;
- VectorSearchMetric::from_lumina(lumina_metric)
- }
- VectorIndexBackend::Vindex => {
- let bytes = segment_bytes.get(&segment.path).ok_or_else(||
{
- crate::Error::DataInvalid {
- message: format!(
- "missing preloaded ANN bytes for segment '{}'",
- segment.path
- ),
- source: None,
- }
- })?;
- let reader =
VIndexReader::open(Cursor::new(bytes.clone())).map_err(|e| {
- crate::Error::DataInvalid {
- message: format!(
- "failed to open ANN index file '{}' for metric
check: {e}",
- segment.path
- ),
- source: Some(Box::new(e)),
- }
- })?;
- VectorSearchMetric::from_vindex(reader.metadata().metric)
- }
- };
- if segment_metric != configured {
- return Err(crate::Error::DataInvalid {
- message: format!(
- "ANN segment metric {} does not match configured
metric {}",
- segment_metric.as_str(),
- configured.as_str()
- ),
- source: None,
- });
- }
- }
+ if segment_metric != configured {
+ return Err(crate::Error::DataInvalid {
+ message: format!(
+ "ANN segment metric {} does not match configured metric {}",
+ segment_metric.as_str(),
+ configured.as_str()
+ ),
+ source: None,
+ });
}
Ok(())
}
@@ -4430,20 +4433,6 @@ mod tests {
bytes
}
- /// A `PkVectorSearchSplit` carrying a single ANN segment addressed by
`path`.
- fn pk_split_with_segment(path: &str) -> PkVectorSearchSplit {
- let mut split = pk_search_split(0, vec![pk_data_file("file-a", 3,
Some(0))]);
- let source_meta = crate::spec::PrimaryKeyIndexSourceMeta::new(
- 1,
-
vec![crate::spec::PrimaryKeyIndexSourceFile::new("file-a".to_string(),
3).unwrap()],
- )
- .unwrap();
- let mut segment = BucketAnnSegment::for_test(source_meta);
- segment.path = path.to_string();
- split.ann_segments = vec![segment];
- split
- }
-
fn pk_split_with_lumina_segment(path: &str, metric: &str) ->
PkVectorSearchSplit {
let mut split = pk_search_split(0, vec![pk_data_file("file-a", 3,
Some(0))]);
let source_meta = crate::spec::PrimaryKeyIndexSourceMeta::new(
@@ -4465,31 +4454,35 @@ mod tests {
}
#[test]
- fn verify_pk_vector_segment_metrics_accepts_matching_lumina_metric() {
+ fn verify_segment_metric_accepts_matching_lumina_metric() {
// Lumina segment metadata says cosine; configured cosine => Ok. No
segment
// file bytes are needed on the Lumina path.
- let splits = vec![pk_split_with_lumina_segment("seg-lumina",
"cosine")];
- let segment_bytes = HashMap::new();
- verify_pk_vector_segment_metrics(
- &splits,
- &segment_bytes,
+ let split = pk_split_with_lumina_segment("seg-lumina", "cosine");
+ let segment = &split.ann_segments[0];
+ let lumina_metric = LuminaIndexMeta::deserialize(&segment.index_meta)
+ .unwrap()
+ .metric()
+ .unwrap();
+ verify_segment_metric(
VectorSearchMetric::Cosine,
- VectorIndexBackend::Lumina,
+ VectorSearchMetric::from_lumina(lumina_metric),
)
.expect("matching lumina metric must pass");
}
#[test]
- fn verify_pk_vector_segment_metrics_rejects_mismatched_lumina_metric() {
+ fn verify_segment_metric_rejects_mismatched_lumina_metric() {
// Lumina segment metadata says l2; configured inner_product => fail
loud,
// naming both metrics.
- let splits = vec![pk_split_with_lumina_segment("seg-lumina", "l2")];
- let segment_bytes = HashMap::new();
- let err = verify_pk_vector_segment_metrics(
- &splits,
- &segment_bytes,
+ let split = pk_split_with_lumina_segment("seg-lumina", "l2");
+ let segment = &split.ann_segments[0];
+ let lumina_metric = LuminaIndexMeta::deserialize(&segment.index_meta)
+ .unwrap()
+ .metric()
+ .unwrap();
+ let err = verify_segment_metric(
VectorSearchMetric::InnerProduct,
- VectorIndexBackend::Lumina,
+ VectorSearchMetric::from_lumina(lumina_metric),
)
.expect_err("mismatched lumina metric must fail loud");
assert!(
@@ -4520,31 +4513,25 @@ mod tests {
}
#[test]
- fn verify_pk_vector_segment_metrics_accepts_matching_metric() {
+ fn verify_segment_metric_accepts_matching_vindex_metric() {
// Real IVF segment trained with L2; configured metric L2 => Ok.
- let bytes = build_vindex_segment_bytes("l2");
- let splits = vec![pk_split_with_segment("seg-l2")];
- let segment_bytes = HashMap::from([("seg-l2".to_string(), bytes)]);
- verify_pk_vector_segment_metrics(
- &splits,
- &segment_bytes,
+ let bytes = bytes::Bytes::from(build_vindex_segment_bytes("l2"));
+ let reader = VIndexReader::open(Cursor::new(bytes)).unwrap();
+ verify_segment_metric(
VectorSearchMetric::L2,
- VectorIndexBackend::Vindex,
+ VectorSearchMetric::from_vindex(reader.metadata().metric),
)
.expect("matching metric must pass");
}
#[test]
- fn verify_pk_vector_segment_metrics_rejects_mismatched_metric() {
+ fn verify_segment_metric_rejects_mismatched_vindex_metric() {
// Real IVF segment trained with L2; configured metric Cosine => fail
loud.
- let bytes = build_vindex_segment_bytes("l2");
- let splits = vec![pk_split_with_segment("seg-l2")];
- let segment_bytes = HashMap::from([("seg-l2".to_string(), bytes)]);
- let err = verify_pk_vector_segment_metrics(
- &splits,
- &segment_bytes,
+ let bytes = bytes::Bytes::from(build_vindex_segment_bytes("l2"));
+ let reader = VIndexReader::open(Cursor::new(bytes)).unwrap();
+ let err = verify_segment_metric(
VectorSearchMetric::Cosine,
- VectorIndexBackend::Vindex,
+ VectorSearchMetric::from_vindex(reader.metadata().metric),
)
.expect_err("mismatched metric must fail loud");
assert!(
diff --git a/crates/paimon/src/vindex/executor.rs
b/crates/paimon/src/vindex/executor.rs
index 76b8a185..173b5d7d 100644
--- a/crates/paimon/src/vindex/executor.rs
+++ b/crates/paimon/src/vindex/executor.rs
@@ -22,7 +22,7 @@ use std::panic::{catch_unwind, AssertUnwindSafe};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::Duration;
-use tokio::sync::oneshot;
+use tokio::sync::{oneshot, Semaphore};
type Job = Box<dyn FnOnce() + Send + 'static>;
@@ -30,6 +30,9 @@ const DEFAULT_IO_BOUND_WORKERS: usize = 32;
const IO_BOUND_WORKERS_PER_CPU: usize = 4;
const WORKER_IDLE_TIMEOUT: Duration = Duration::from_secs(60);
+static PROCESS_GLOBAL_SEMAPHORE: OnceLock<Semaphore> = OnceLock::new();
+static PROCESS_GLOBAL_CAPACITY: AtomicUsize = AtomicUsize::new(0);
+
struct ExecutorState {
receiver: Receiver<Job>,
max_workers: AtomicUsize,
@@ -249,6 +252,64 @@ pub(crate) fn
ensure_global_index_executor_capacity(max_workers: usize) {
global_executor().ensure_capacity(max_workers);
}
+/// Grow a shared semaphore monotonically. The successful CAS owns one disjoint
+/// capacity delta, so concurrent callers neither over-add permits nor shrink
the
+/// high-watermark established by an earlier, larger query configuration.
+fn grow_semaphore_capacity(
+ tracked_capacity: &AtomicUsize,
+ semaphore: &Semaphore,
+ capacity: usize,
+) -> usize {
+ loop {
+ let current = tracked_capacity.load(Ordering::SeqCst);
+ if capacity <= current {
+ return 0;
+ }
+ if tracked_capacity
+ .compare_exchange(current, capacity, Ordering::SeqCst,
Ordering::SeqCst)
+ .is_ok()
+ {
+ let added = capacity - current;
+ semaphore.add_permits(added);
+ return added;
+ }
+ }
+}
+
+/// Java's shared executor is initially sized to `availableProcessors()` and is
+/// only replaced by a larger pool. A smaller configured `thread-num` limits
one
+/// query through its per-query scheduler or semaphore, but does not shrink
+/// cross-query process capacity below the machine's available parallelism.
+fn effective_process_global_capacity(requested: usize) -> usize {
+ requested.max(default_worker_count()).max(1)
+}
+
+/// Return the process-global permit pool, growing both it and the dedicated
CPU
+/// executor to the same monotonic high-watermark. This caps aggregate work
across
+/// queries while preserving Java's available-CPU floor.
+fn process_global_semaphore(capacity: usize) -> &'static Semaphore {
+ let effective = effective_process_global_capacity(capacity);
+ ensure_global_index_executor_capacity(effective);
+ let semaphore = PROCESS_GLOBAL_SEMAPHORE.get_or_init(|| {
+ PROCESS_GLOBAL_CAPACITY.store(effective, Ordering::SeqCst);
+ Semaphore::new(effective)
+ });
+ grow_semaphore_capacity(&PROCESS_GLOBAL_CAPACITY, semaphore, effective);
+ semaphore
+}
+
+pub(crate) async fn acquire_process_global_search_permit(
+ capacity: usize,
+) -> crate::Result<tokio::sync::SemaphorePermit<'static>> {
+ process_global_semaphore(capacity)
+ .acquire()
+ .await
+ .map_err(|error| crate::Error::UnexpectedError {
+ message: "global-index process concurrency budget was
closed".to_string(),
+ source: Some(Box::new(error)),
+ })
+}
+
/// Runs bounded global-index jobs to completion, restores submission order,
and
/// returns the first error by submission index. Dedicated executor jobs
cannot be
/// interrupted after they start, so short-circuiting would only hide
background
@@ -282,6 +343,23 @@ where
execute_on(global_executor(), panic_context, task).await
}
+pub(crate) async fn execute_global_index_with_guard<T, F, G>(
+ panic_context: &'static str,
+ guard: G,
+ task: F,
+) -> crate::Result<T>
+where
+ T: Send + 'static,
+ F: FnOnce() -> crate::Result<T> + Send + 'static,
+ G: Send + 'static,
+{
+ execute_global_index(panic_context, move || {
+ let _guard = guard;
+ task()
+ })
+ .await
+}
+
async fn execute_on<T, F>(
executor: &GlobalIndexExecutor,
panic_context: &'static str,
@@ -499,6 +577,39 @@ mod tests {
.expect("submitted task stopped before reporting completion");
}
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn guarded_task_keeps_guard_after_waiter_cancellation() {
+ let semaphore = Arc::new(Semaphore::new(1));
+ let permit = semaphore.clone().acquire_owned().await.unwrap();
+ let (started_sender, started_receiver) = oneshot::channel();
+ let (release_sender, release_receiver) = std::sync::mpsc::channel();
+ let join = tokio::spawn(async move {
+ execute_global_index_with_guard("guarded task failed", permit,
move || {
+ let _ = started_sender.send(());
+ let _ = release_receiver.recv();
+ Ok(())
+ })
+ .await
+ });
+ started_receiver.await.unwrap();
+
+ join.abort();
+ assert!(join.await.unwrap_err().is_cancelled());
+ assert!(
+ tokio::time::timeout(Duration::from_millis(50),
semaphore.clone().acquire_owned())
+ .await
+ .is_err(),
+ "cancelling the waiter released the guard while the task was still
running"
+ );
+
+ release_sender.send(()).unwrap();
+ let _recovered_permit =
+ tokio::time::timeout(Duration::from_secs(1),
semaphore.acquire_owned())
+ .await
+ .expect("guard was not released after the task finished")
+ .unwrap();
+ }
+
#[test]
fn cancelling_started_search_keeps_current_thread_runtime_live() {
let runtime = tokio::runtime::Builder::new_current_thread()
@@ -600,4 +711,27 @@ mod tests {
);
assert!(!second_ran.load(Ordering::SeqCst));
}
+
+ #[test]
+ fn semaphore_capacity_is_monotonic_high_watermark() {
+ let tracked = AtomicUsize::new(0);
+ let semaphore = Semaphore::new(0);
+ assert_eq!(grow_semaphore_capacity(&tracked, &semaphore, 4), 4);
+ assert_eq!(semaphore.available_permits(), 4);
+ assert_eq!(tracked.load(Ordering::SeqCst), 4);
+ assert_eq!(grow_semaphore_capacity(&tracked, &semaphore, 7), 3);
+ assert_eq!(semaphore.available_permits(), 7);
+ assert_eq!(grow_semaphore_capacity(&tracked, &semaphore, 4), 0);
+ assert_eq!(grow_semaphore_capacity(&tracked, &semaphore, 7), 0);
+ assert_eq!(tracked.load(Ordering::SeqCst), 7);
+ assert_eq!(semaphore.available_permits(), 7);
+ }
+
+ #[test]
+ fn process_capacity_floors_at_available_parallelism() {
+ let cores = default_worker_count();
+ assert_eq!(effective_process_global_capacity(1), cores);
+ assert_eq!(effective_process_global_capacity(cores + 100), cores +
100);
+ assert!(effective_process_global_capacity(0) >= 1);
+ }
}
diff --git a/crates/paimon/src/vindex/pkvector/ann.rs
b/crates/paimon/src/vindex/pkvector/ann.rs
index ab50196a..76cf6aa4 100644
--- a/crates/paimon/src/vindex/pkvector/ann.rs
+++ b/crates/paimon/src/vindex/pkvector/ann.rs
@@ -18,13 +18,20 @@
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
+use bytes::Bytes;
+use futures::future::BoxFuture;
+
use super::bucket::BucketAnnSegment;
use super::data_invalid;
use super::metric::{java_float_compare, VectorSearchMetric};
use super::result::PkVectorSearchResult;
use crate::deletion_vector::DeletionVector;
-use crate::spec::{PrimaryKeyIndexSourceFile, PrimaryKeyIndexSourceMeta};
+use crate::spec::{
+ PrimaryKeyIndexSourceFile as PkVectorSourceFile,
+ PrimaryKeyIndexSourceMeta as PkVectorSourceMeta,
+};
use crate::vector_search::VectorSearch;
+use crate::vindex::range_reader::VindexFileReader;
/// Build the live-row-id mask for the ANN reader's `include_row_ids` filter,
in
/// segment-ordinal space (source files concatenated in order). Mirrors Java
@@ -46,7 +53,7 @@ use crate::vector_search::VectorSearch;
/// no deletion vector is relevant — nothing to mask. Otherwise returns the
masked
/// live ids.
pub(crate) fn build_live_row_ids(
- source_files: &[PrimaryKeyIndexSourceFile],
+ source_files: &[PkVectorSourceFile],
active_source_files: &HashSet<String>,
deletion_vectors: &HashMap<String, Arc<DeletionVector>>,
residual_ranges: Option<&HashMap<String, roaring::RoaringTreemap>>,
@@ -122,7 +129,7 @@ pub(crate) fn build_live_row_ids(
/// sorted BEST_FIRST.
pub(crate) fn map_ann_results(
scored: &[(u64, f32)],
- source_meta: &PrimaryKeyIndexSourceMeta,
+ source_meta: &PkVectorSourceMeta,
active_source_files: &HashSet<String>,
deletion_vectors: &HashMap<String, Arc<DeletionVector>>,
residual_ranges: Option<&HashMap<String, roaring::RoaringTreemap>>,
@@ -172,18 +179,37 @@ pub(crate) fn map_ann_results(
/// One ANN segment's search dependency for the bucket kernel. Bucket tests
fake
/// this (mirroring Java's mock of `PkVectorAnnSegmentSearcher`).
///
-/// `Send + Sync` so a `&dyn PkVectorAnnSearcher` can be held across the
`.await`
-/// points of the async search path (the returned future is spawned on a `Send`
-/// runtime by callers such as the DataFusion integration).
+/// `Send + Sync` so an `Arc<dyn PkVectorAnnSearcher>` can be cloned into
concurrent
+/// leaf futures and moved onto the dedicated global-index executor.
pub(crate) trait PkVectorAnnSearcher: Send + Sync {
+ /// Load one buffered ANN segment. Runs on the async side of the bucket
leaf,
+ /// BEFORE the blocking score, so the bytes are loaded lazily per segment
and
+ /// dropped after that leaf.
+ /// Returns a `'static` boxed future so it borrows nothing from
`self`/`segment`
+ /// past the await (production clones `FileIO` + the path into the future).
+ fn load_segment(&self, segment: &BucketAnnSegment) -> BoxFuture<'static,
crate::Result<Bytes>>;
+
+ /// Production source loader. Must NOT acquire a search-concurrency permit:
+ /// the bucket leaf already holds one across both this load and the
subsequent
+ /// score, so a second acquisition would deadlock at capacity 1.
+ fn load_segment_source(
+ &self,
+ segment: &BucketAnnSegment,
+ ) -> BoxFuture<'static, crate::Result<AnnSegmentSource>> {
+ let future = self.load_segment(segment);
+ Box::pin(async move { future.await.map(AnnSegmentSource::Buffered) })
+ }
+
/// Search one ANN segment for a batch of query vectors, returning one
/// BEST_FIRST result list per query (outer index aligned to `queries`).
The
/// live-row mask (residual ∩ DV) is query-independent, so it is built
once and
- /// shared across all queries; only the per-query scores differ.
+ /// shared across all queries; only the per-query scores differ. Buffered
callers
+ /// pass the bytes from `load_segment` by value so they cannot outlive the
leaf.
#[allow(clippy::too_many_arguments)]
fn search_batch(
&self,
segment: &BucketAnnSegment,
+ segment_bytes: Bytes,
queries: &[&[f32]],
metric: VectorSearchMetric,
limit: usize,
@@ -193,12 +219,45 @@ pub(crate) trait PkVectorAnnSearcher: Send + Sync {
residual_ranges: Option<&HashMap<String, roaring::RoaringTreemap>>,
) -> crate::Result<Vec<Vec<PkVectorSearchResult>>>;
+ #[allow(clippy::too_many_arguments)]
+ fn search_batch_source(
+ &self,
+ segment: &BucketAnnSegment,
+ segment_source: AnnSegmentSource,
+ queries: &[&[f32]],
+ metric: VectorSearchMetric,
+ limit: usize,
+ active_source_files: &HashSet<String>,
+ deletion_vectors: &HashMap<String, Arc<DeletionVector>>,
+ search_options: &HashMap<String, String>,
+ residual_ranges: Option<&HashMap<String, roaring::RoaringTreemap>>,
+ ) -> crate::Result<Vec<Vec<PkVectorSearchResult>>> {
+ match segment_source {
+ AnnSegmentSource::Buffered(bytes) => self.search_batch(
+ segment,
+ bytes,
+ queries,
+ metric,
+ limit,
+ active_source_files,
+ deletion_vectors,
+ search_options,
+ residual_ranges,
+ ),
+ AnnSegmentSource::Vindex(_) => Err(data_invalid(
+ "ANN searcher does not support a range-backed segment source",
+ )),
+ }
+ }
+
/// Single-query wrapper over `search_batch`: searches the one query and
/// returns its result list. Asserts the batch produced exactly one list.
+ #[allow(dead_code)]
#[allow(clippy::too_many_arguments)]
fn search(
&self,
segment: &BucketAnnSegment,
+ segment_bytes: Bytes,
query: &[f32],
metric: VectorSearchMetric,
limit: usize,
@@ -209,6 +268,40 @@ pub(crate) trait PkVectorAnnSearcher: Send + Sync {
) -> crate::Result<Vec<PkVectorSearchResult>> {
let mut results = self.search_batch(
segment,
+ segment_bytes,
+ &[query],
+ metric,
+ limit,
+ active_source_files,
+ deletion_vectors,
+ search_options,
+ residual_ranges,
+ )?;
+ if results.len() != 1 {
+ return Err(data_invalid(format!(
+ "ANN batch search returned {} result lists for a single query",
+ results.len()
+ )));
+ }
+ Ok(results.pop().expect("length checked to be 1"))
+ }
+
+ #[allow(clippy::too_many_arguments)]
+ fn search_source(
+ &self,
+ segment: &BucketAnnSegment,
+ segment_source: AnnSegmentSource,
+ query: &[f32],
+ metric: VectorSearchMetric,
+ limit: usize,
+ active_source_files: &HashSet<String>,
+ deletion_vectors: &HashMap<String, Arc<DeletionVector>>,
+ search_options: &HashMap<String, String>,
+ residual_ranges: Option<&HashMap<String, roaring::RoaringTreemap>>,
+ ) -> crate::Result<Vec<PkVectorSearchResult>> {
+ let mut results = self.search_batch_source(
+ segment,
+ segment_source,
&[query],
metric,
limit,
@@ -228,40 +321,144 @@ pub(crate) trait PkVectorAnnSearcher: Send + Sync {
}
/// Batch scorer seam: drives the underlying vindex ANN reader for a batch of
-/// searches over ONE segment, opening the reader once and searching each query
-/// against it (mirroring Java's shared-reader `visitBatchVectorSearch`).
Returns
+/// searches over ONE segment, opening the reader once and issuing one backend
batch
+/// search (mirroring Java's shared-reader `visitBatchVectorSearch`). Returns
/// one `ordinal -> score` map (higher-is-better) per input search, aligned to
the
/// `searches` slice. Any negative labels are skipped by the existing `vindex`
/// reader (`collect_results` drops `row_id < 0`), so this seam only ever
yields
/// non-negative `u64` ordinals — no signed-label handling is needed
downstream.
///
-/// The production scorer drives
`VindexVectorGlobalIndexReader::visit_batch_vector_search`
-/// with a segment's index bytes; tests inject a synthetic scorer. The
adapter's
+/// The production scorer drives one backend reader from a typed segment
source;
+/// tests inject a synthetic buffered scorer. The adapter's
/// own logic (live-row masking, ordinal mapping, deletion checks, ordering) is
/// exercised independently of the scorer.
+#[cfg(test)]
pub(crate) type BatchScorer = Box<
- dyn Fn(&BucketAnnSegment, &[VectorSearch]) ->
crate::Result<Vec<Option<HashMap<u64, f32>>>>
+ dyn Fn(
+ &BucketAnnSegment,
+ Bytes,
+ &[VectorSearch],
+ ) -> crate::Result<Vec<Option<HashMap<u64, f32>>>>
+ + Send
+ + Sync,
+>;
+
+pub(crate) type SourceBatchScorer = Box<
+ dyn Fn(
+ &BucketAnnSegment,
+ AnnSegmentSource,
+ &[VectorSearch],
+ ) -> crate::Result<Vec<Option<HashMap<u64, f32>>>>
+ Send
+ Sync,
>;
-/// Structural vindex-backed `PkVectorAnnSearcher`. Composes the pure helpers
-/// (`build_live_row_ids`, `map_ann_results`) around the batch scorer seam.
+pub(crate) enum AnnSegmentSource {
+ Buffered(Bytes),
+ Vindex(VindexFileReader),
+}
+
+/// Test-only buffered loader retained for simple fake searchers.
+#[cfg(test)]
+pub(crate) type SegmentLoader =
+ Box<dyn Fn(&BucketAnnSegment) -> BoxFuture<'static, crate::Result<Bytes>>
+ Send + Sync>;
+
+/// Production loader for either buffered Lumina data or a range-backed vindex
+/// source. It runs before the dedicated-executor score and must not acquire a
+/// second search permit; the bucket leaf already holds one across both phases.
+pub(crate) type SourceSegmentLoader = Box<
+ dyn Fn(&BucketAnnSegment) -> BoxFuture<'static,
crate::Result<AnnSegmentSource>> + Send + Sync,
+>;
+
+/// Structural ANN-backed `PkVectorAnnSearcher`. Composes the pure helpers
+/// (`build_live_row_ids`, `map_ann_results`) around the batch scorer seam, and
+/// carries the async segment loader so each segment source is opened lazily
in its
+/// own bucket leaf and dropped after scoring (no up-front all-segments map).
pub(crate) struct VindexAnnSearcher {
field_name: String,
- scorer: BatchScorer,
+ scorer: SourceBatchScorer,
+ loader: SourceSegmentLoader,
}
impl VindexAnnSearcher {
- pub(crate) fn new(field_name: String, scorer: BatchScorer) -> Self {
- Self { field_name, scorer }
+ #[cfg(test)]
+ pub(crate) fn new(field_name: String, scorer: BatchScorer, loader:
SegmentLoader) -> Self {
+ let source_scorer: SourceBatchScorer =
+ Box::new(move |segment, source, searches| match source {
+ AnnSegmentSource::Buffered(bytes) => scorer(segment, bytes,
searches),
+ AnnSegmentSource::Vindex(_) => Err(data_invalid(
+ "buffered ANN scorer received a range-backed segment
source",
+ )),
+ });
+ let source_loader: SourceSegmentLoader = Box::new(move |segment| {
+ let future = loader(segment);
+ Box::pin(async move { future.await.map(AnnSegmentSource::Buffered)
})
+ });
+ Self::new_with_source(field_name, source_scorer, source_loader)
+ }
+
+ pub(crate) fn new_with_source(
+ field_name: String,
+ scorer: SourceBatchScorer,
+ loader: SourceSegmentLoader,
+ ) -> Self {
+ Self {
+ field_name,
+ scorer,
+ loader,
+ }
}
}
impl PkVectorAnnSearcher for VindexAnnSearcher {
+ fn load_segment(&self, segment: &BucketAnnSegment) -> BoxFuture<'static,
crate::Result<Bytes>> {
+ let future = (self.loader)(segment);
+ Box::pin(async move {
+ match future.await? {
+ AnnSegmentSource::Buffered(bytes) => Ok(bytes),
+ AnnSegmentSource::Vindex(_) => Err(data_invalid(
+ "range-backed ANN segment cannot be converted to buffered
bytes",
+ )),
+ }
+ })
+ }
+
+ fn load_segment_source(
+ &self,
+ segment: &BucketAnnSegment,
+ ) -> BoxFuture<'static, crate::Result<AnnSegmentSource>> {
+ (self.loader)(segment)
+ }
+
fn search_batch(
&self,
segment: &BucketAnnSegment,
+ segment_bytes: Bytes,
+ queries: &[&[f32]],
+ metric: VectorSearchMetric,
+ limit: usize,
+ active_source_files: &HashSet<String>,
+ deletion_vectors: &HashMap<String, Arc<DeletionVector>>,
+ search_options: &HashMap<String, String>,
+ residual_ranges: Option<&HashMap<String, roaring::RoaringTreemap>>,
+ ) -> crate::Result<Vec<Vec<PkVectorSearchResult>>> {
+ self.search_batch_source(
+ segment,
+ AnnSegmentSource::Buffered(segment_bytes),
+ queries,
+ metric,
+ limit,
+ active_source_files,
+ deletion_vectors,
+ search_options,
+ residual_ranges,
+ )
+ }
+
+ fn search_batch_source(
+ &self,
+ segment: &BucketAnnSegment,
+ segment_source: AnnSegmentSource,
queries: &[&[f32]],
metric: VectorSearchMetric,
limit: usize,
@@ -292,7 +489,7 @@ impl PkVectorAnnSearcher for VindexAnnSearcher {
}
searches.push(search);
}
- let scored_batch = (self.scorer)(segment, &searches)?;
+ let scored_batch = (self.scorer)(segment, segment_source, &searches)?;
if scored_batch.len() != queries.len() {
return Err(data_invalid(format!(
"ANN batch scorer returned {} result maps for {} queries",
@@ -327,12 +524,24 @@ mod tests {
use super::*;
use roaring::RoaringBitmap;
- fn source_meta(files: &[(&str, i64)]) -> PrimaryKeyIndexSourceMeta {
+ /// A trivial loader returning empty bytes — the synthetic scorers below
ignore
+ /// their `segment_bytes` (they model behavior above physical index
decoding).
+ fn empty_loader() -> SegmentLoader {
+ Box::new(|_: &BucketAnnSegment| Box::pin(async { Ok(Bytes::new()) }))
+ }
+
+ /// Build a `VindexAnnSearcher` with a trivial (empty-bytes) loader, for
tests
+ /// that only exercise the scorer/adapter logic.
+ fn vindex_searcher(field: &str, scorer: BatchScorer) -> VindexAnnSearcher {
+ VindexAnnSearcher::new(field.to_string(), scorer, empty_loader())
+ }
+
+ fn source_meta(files: &[(&str, i64)]) -> PkVectorSourceMeta {
let files = files
.iter()
- .map(|(name, rows)|
PrimaryKeyIndexSourceFile::new((*name).to_string(), *rows).unwrap())
+ .map(|(name, rows)| PkVectorSourceFile::new((*name).to_string(),
*rows).unwrap())
.collect();
- PrimaryKeyIndexSourceMeta::new(1, files).unwrap()
+ PkVectorSourceMeta::new(1, files).unwrap()
}
fn dv(deleted: &[u32]) -> Arc<DeletionVector> {
@@ -349,7 +558,7 @@ mod tests {
#[test]
fn test_build_live_row_ids_none_when_all_active_and_no_relevant_dv() {
- let files = [PrimaryKeyIndexSourceFile::new("f0".into(), 3).unwrap()];
+ let files = [PkVectorSourceFile::new("f0".into(), 3).unwrap()];
let active = active_set(&["f0"]);
// All active + empty map -> None.
assert!(build_live_row_ids(&files, &active, &HashMap::new(), None)
@@ -368,8 +577,8 @@ mod tests {
// f0 rows 0..3 (global 0,1,2), f1 rows 0..2 (global 3,4). f1 is
inactive,
// so its whole ordinal range is masked out; f0 stays fully live. No
DV.
let files = vec![
- PrimaryKeyIndexSourceFile::new("f0".into(), 3).unwrap(),
- PrimaryKeyIndexSourceFile::new("f1".into(), 2).unwrap(),
+ PkVectorSourceFile::new("f0".into(), 3).unwrap(),
+ PkVectorSourceFile::new("f1".into(), 2).unwrap(),
];
let live = build_live_row_ids(&files, &active_set(&["f0"]),
&HashMap::new(), None)
.unwrap()
@@ -381,8 +590,8 @@ mod tests {
fn test_build_live_row_ids_masks_deleted_positions_with_file_offsets() {
// f0 rows 0..3 (global 0,1,2), f1 rows 0..2 (global 3,4).
let files = vec![
- PrimaryKeyIndexSourceFile::new("f0".into(), 3).unwrap(),
- PrimaryKeyIndexSourceFile::new("f1".into(), 2).unwrap(),
+ PkVectorSourceFile::new("f0".into(), 3).unwrap(),
+ PkVectorSourceFile::new("f1".into(), 2).unwrap(),
];
let mut dvs = HashMap::new();
dvs.insert("f0".to_string(), dv(&[1])); // deletes global 1
@@ -483,10 +692,10 @@ mod tests {
let seen_has_filter = Arc::new(Mutex::new(false));
let scorer_limit = Arc::clone(&seen_limit);
let scorer_has_filter = Arc::clone(&seen_has_filter);
- let searcher = VindexAnnSearcher::new(
- "embedding".to_string(),
+ let searcher = vindex_searcher(
+ "embedding",
Box::new(
- move |_segment: &BucketAnnSegment, searches: &[VectorSearch]| {
+ move |_segment: &BucketAnnSegment, _bytes: Bytes, searches:
&[VectorSearch]| {
let search = &searches[0];
*scorer_limit.lock().unwrap() = search.limit;
*scorer_has_filter.lock().unwrap() =
search.include_row_ids.is_some();
@@ -498,12 +707,15 @@ mod tests {
),
);
let segment = BucketAnnSegment::for_test({
- use crate::spec::{PrimaryKeyIndexSourceFile,
PrimaryKeyIndexSourceMeta};
- PrimaryKeyIndexSourceMeta::new(
+ use crate::spec::{
+ PrimaryKeyIndexSourceFile as PkVectorSourceFile,
+ PrimaryKeyIndexSourceMeta as PkVectorSourceMeta,
+ };
+ PkVectorSourceMeta::new(
1,
vec![
- PrimaryKeyIndexSourceFile::new("f0".into(), 3).unwrap(),
- PrimaryKeyIndexSourceFile::new("f1".into(), 5).unwrap(),
+ PkVectorSourceFile::new("f0".into(), 3).unwrap(),
+ PkVectorSourceFile::new("f1".into(), 5).unwrap(),
],
)
.unwrap()
@@ -513,6 +725,7 @@ mod tests {
let results = searcher
.search(
&segment,
+ Bytes::new(),
&[0.0, 0.0],
VectorSearchMetric::L2,
2,
@@ -534,23 +747,26 @@ mod tests {
#[test]
fn test_vindex_adapter_rejects_non_positive_limit() {
- let searcher = VindexAnnSearcher::new(
- "embedding".to_string(),
- Box::new(|_: &BucketAnnSegment, searches: &[VectorSearch]| {
- Ok(vec![None; searches.len()])
- }),
+ let searcher = vindex_searcher(
+ "embedding",
+ Box::new(
+ |_: &BucketAnnSegment, _bytes: Bytes, searches:
&[VectorSearch]| {
+ Ok(vec![None; searches.len()])
+ },
+ ),
);
let segment = BucketAnnSegment::for_test({
- use crate::spec::{PrimaryKeyIndexSourceFile,
PrimaryKeyIndexSourceMeta};
- PrimaryKeyIndexSourceMeta::new(
- 1,
- vec![PrimaryKeyIndexSourceFile::new("f0".into(), 1).unwrap()],
- )
- .unwrap()
+ use crate::spec::{
+ PrimaryKeyIndexSourceFile as PkVectorSourceFile,
+ PrimaryKeyIndexSourceMeta as PkVectorSourceMeta,
+ };
+ PkVectorSourceMeta::new(1,
vec![PkVectorSourceFile::new("f0".into(), 1).unwrap()])
+ .unwrap()
});
let err = searcher
.search(
&segment,
+ Bytes::new(),
&[0.0, 0.0],
VectorSearchMetric::L2,
0,
@@ -565,23 +781,26 @@ mod tests {
#[test]
fn test_vindex_adapter_empty_scorer_result_is_empty() {
- let searcher = VindexAnnSearcher::new(
- "embedding".to_string(),
- Box::new(|_: &BucketAnnSegment, searches: &[VectorSearch]| {
- Ok(vec![None; searches.len()])
- }),
+ let searcher = vindex_searcher(
+ "embedding",
+ Box::new(
+ |_: &BucketAnnSegment, _bytes: Bytes, searches:
&[VectorSearch]| {
+ Ok(vec![None; searches.len()])
+ },
+ ),
);
let segment = BucketAnnSegment::for_test({
- use crate::spec::{PrimaryKeyIndexSourceFile,
PrimaryKeyIndexSourceMeta};
- PrimaryKeyIndexSourceMeta::new(
- 1,
- vec![PrimaryKeyIndexSourceFile::new("f0".into(), 1).unwrap()],
- )
- .unwrap()
+ use crate::spec::{
+ PrimaryKeyIndexSourceFile as PkVectorSourceFile,
+ PrimaryKeyIndexSourceMeta as PkVectorSourceMeta,
+ };
+ PkVectorSourceMeta::new(1,
vec![PkVectorSourceFile::new("f0".into(), 1).unwrap()])
+ .unwrap()
});
let results = searcher
.search(
&segment,
+ Bytes::new(),
&[0.0, 0.0],
VectorSearchMetric::L2,
2,
@@ -609,8 +828,8 @@ mod tests {
// entry (empty allow). Result: f0 keeps {0} (1 is residual-allowed but
// deleted, 2 not residual-allowed); f1 contributes nothing.
let files = vec![
- PrimaryKeyIndexSourceFile::new("f0".into(), 3).unwrap(),
- PrimaryKeyIndexSourceFile::new("f1".into(), 2).unwrap(),
+ PkVectorSourceFile::new("f0".into(), 3).unwrap(),
+ PkVectorSourceFile::new("f1".into(), 2).unwrap(),
];
let mut dvs = HashMap::new();
dvs.insert("f0".to_string(), dv(&[1]));
@@ -627,8 +846,8 @@ mod tests {
// f0 rows global 0,1,2; f1 rows global 3,4. residual allows f0={2},
f1={1}.
// f1 physical pos 1 -> global 3 + 1 = 4. Result {2, 4}. No DV.
let files = vec![
- PrimaryKeyIndexSourceFile::new("f0".into(), 3).unwrap(),
- PrimaryKeyIndexSourceFile::new("f1".into(), 2).unwrap(),
+ PkVectorSourceFile::new("f0".into(), 3).unwrap(),
+ PkVectorSourceFile::new("f1".into(), 2).unwrap(),
];
let mut residual = HashMap::new();
residual.insert("f0".to_string(), treemap(&[2]));
@@ -648,7 +867,7 @@ mod tests {
fn
test_build_live_row_ids_residual_some_returns_mask_even_when_all_active_no_dv()
{
// All active, no DV: without residual this returns None. With a
residual
// present, a mask is always required.
- let files = [PrimaryKeyIndexSourceFile::new("f0".into(), 3).unwrap()];
+ let files = [PkVectorSourceFile::new("f0".into(), 3).unwrap()];
let mut residual = HashMap::new();
residual.insert("f0".to_string(), treemap(&[0, 2]));
let live = build_live_row_ids(
@@ -705,10 +924,10 @@ mod tests {
use std::sync::{Arc, Mutex};
let seen_rows: Arc<Mutex<Option<Vec<u64>>>> =
Arc::new(Mutex::new(None));
let scorer_rows = Arc::clone(&seen_rows);
- let searcher = VindexAnnSearcher::new(
- "embedding".to_string(),
+ let searcher = vindex_searcher(
+ "embedding",
Box::new(
- move |_segment: &BucketAnnSegment, searches: &[VectorSearch]| {
+ move |_segment: &BucketAnnSegment, _bytes: Bytes, searches:
&[VectorSearch]| {
*scorer_rows.lock().unwrap() = searches[0]
.include_row_ids
.as_ref()
@@ -723,6 +942,7 @@ mod tests {
searcher
.search(
&segment,
+ Bytes::new(),
&[0.0, 0.0],
VectorSearchMetric::L2,
2,
@@ -740,24 +960,27 @@ mod tests {
// The single-query `search` wrapper must return exactly what
// `search_batch(&[q])[0]` returns for the same inputs.
let make = || {
- VindexAnnSearcher::new(
- "embedding".to_string(),
- Box::new(|_: &BucketAnnSegment, searches: &[VectorSearch]| {
- let mut out = Vec::with_capacity(searches.len());
- for _ in searches {
- let mut scores = HashMap::new();
- scores.insert(3u64, 0.5f32); // -> (f1, 0)
- scores.insert(0u64, 0.25f32); // -> (f0, 0)
- out.push(Some(scores));
- }
- Ok(out)
- }),
+ vindex_searcher(
+ "embedding",
+ Box::new(
+ |_: &BucketAnnSegment, _bytes: Bytes, searches:
&[VectorSearch]| {
+ let mut out = Vec::with_capacity(searches.len());
+ for _ in searches {
+ let mut scores = HashMap::new();
+ scores.insert(3u64, 0.5f32); // -> (f1, 0)
+ scores.insert(0u64, 0.25f32); // -> (f0, 0)
+ out.push(Some(scores));
+ }
+ Ok(out)
+ },
+ ),
)
};
let meta = source_meta(&[("f0", 3), ("f1", 5)]);
let single = make()
.search(
&BucketAnnSegment::for_test(meta.clone()),
+ Bytes::new(),
&[0.0, 0.0],
VectorSearchMetric::L2,
2,
@@ -771,6 +994,7 @@ mod tests {
let batch = make()
.search_batch(
&BucketAnnSegment::for_test(meta),
+ Bytes::new(),
&[query],
VectorSearchMetric::L2,
2,
@@ -788,28 +1012,31 @@ mod tests {
fn test_search_batch_returns_independent_per_query_results() {
// Two queries route to different synthetic scores; each result list is
// mapped from that query's own scores, with a shared live-row mask.
- let searcher = VindexAnnSearcher::new(
- "embedding".to_string(),
- Box::new(|_: &BucketAnnSegment, searches: &[VectorSearch]| {
- let mut out = Vec::with_capacity(searches.len());
- for (i, _) in searches.iter().enumerate() {
- let mut scores = HashMap::new();
- // Query 0 -> ordinal 0 (f0,0); query 1 -> ordinal 3
(f1,0).
- if i == 0 {
- scores.insert(0u64, 0.5f32);
- } else {
- scores.insert(3u64, 0.5f32);
+ let searcher = vindex_searcher(
+ "embedding",
+ Box::new(
+ |_: &BucketAnnSegment, _bytes: Bytes, searches:
&[VectorSearch]| {
+ let mut out = Vec::with_capacity(searches.len());
+ for (i, _) in searches.iter().enumerate() {
+ let mut scores = HashMap::new();
+ // Query 0 -> ordinal 0 (f0,0); query 1 -> ordinal 3
(f1,0).
+ if i == 0 {
+ scores.insert(0u64, 0.5f32);
+ } else {
+ scores.insert(3u64, 0.5f32);
+ }
+ out.push(Some(scores));
}
- out.push(Some(scores));
- }
- Ok(out)
- }),
+ Ok(out)
+ },
+ ),
);
let q0: &[f32] = &[0.0, 0.0];
let q1: &[f32] = &[1.0, 1.0];
let results = searcher
.search_batch(
&BucketAnnSegment::for_test(source_meta(&[("f0", 3), ("f1",
5)])),
+ Bytes::new(),
&[q0, q1],
VectorSearchMetric::L2,
2,
@@ -828,9 +1055,9 @@ mod tests {
fn test_search_batch_fails_loud_on_result_count_mismatch() {
// A batch scorer that returns the wrong number of result maps is
corruption
// and must fail loud, not be silently padded/truncated.
- let searcher = VindexAnnSearcher::new(
- "embedding".to_string(),
- Box::new(|_: &BucketAnnSegment, _: &[VectorSearch]| {
+ let searcher = vindex_searcher(
+ "embedding",
+ Box::new(|_: &BucketAnnSegment, _: Bytes, _: &[VectorSearch]| {
// Only one map returned regardless of query count.
Ok(vec![None])
}),
@@ -840,6 +1067,7 @@ mod tests {
let err = searcher
.search_batch(
&BucketAnnSegment::for_test(source_meta(&[("f0", 3)])),
+ Bytes::new(),
&[q0, q1],
VectorSearchMetric::L2,
2,
diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs
b/crates/paimon/src/vindex/pkvector/bucket.rs
index 4666d403..c52a72e5 100644
--- a/crates/paimon/src/vindex/pkvector/bucket.rs
+++ b/crates/paimon/src/vindex/pkvector/bucket.rs
@@ -20,7 +20,6 @@ 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;
@@ -28,7 +27,10 @@ use super::data_invalid;
use super::metric::{java_float_compare, VectorSearchMetric};
use super::result::PkVectorSearchResult;
use crate::deletion_vector::DeletionVector;
-use crate::spec::PrimaryKeyIndexSourceMeta;
+use crate::spec::PrimaryKeyIndexSourceMeta as PkVectorSourceMeta;
+use crate::vindex::executor::{
+ acquire_process_global_search_permit, drain_indexed_jobs,
execute_global_index,
+};
/// Search one uncovered data file for its per-query exact Top-K. Returns one
/// bounded, BEST_FIRST list per query (outer index aligns to the `queries`
slice
@@ -43,9 +45,16 @@ pub(crate) type ExactFileSearchFuture<'a> =
/// segment ordinals back to physical `(data file, position)` and drives
live-row
/// masking; the remaining fields address the segment's index file for the ANN
/// scorer that reads it.
+///
+/// `Clone` so one segment's search inputs can be moved into a dedicated
global-index
+/// executor leaf (the ANN CPU search runs off the async worker). Cloning
copies only
+/// the segment's addressing metadata (path, size, `index_meta`, source meta).
Each
+/// leaf opens its own segment source lazily and drops it after scoring; vindex
+/// sources stay range-backed instead of holding the full index bytes.
+#[derive(Clone)]
pub(crate) struct BucketAnnSegment {
- pub source_meta: PrimaryKeyIndexSourceMeta,
- /// Resolved index-file path (globally unique; the scorer's preload key).
+ pub source_meta: PkVectorSourceMeta,
+ /// Resolved index-file path (globally unique; the key the loader reads
by).
pub path: String,
pub file_size: u64,
pub index_meta: Vec<u8>,
@@ -55,7 +64,7 @@ pub(crate) struct BucketAnnSegment {
impl BucketAnnSegment {
/// Build a segment with dummy index-file fields for tests that exercise
only
/// `source_meta`-driven logic.
- pub(crate) fn for_test(source_meta: PrimaryKeyIndexSourceMeta) -> Self {
+ pub(crate) fn for_test(source_meta: PkVectorSourceMeta) -> Self {
Self {
source_meta,
path: "seg".to_string(),
@@ -205,33 +214,114 @@ 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.
+/// Concurrency budget for one search's global-index leaf work. Mirrors Java's
+/// two-level shape: a shared process pool with an available-CPU floor, plus a
+/// per-query limit (`per_query`) that keeps a single query from occupying more
+/// than its configured `thread_num` slots.
+///
+/// - `production(n)` — the real read path: per-query semaphore of `n` permits
AND
+/// the shared process pool. Used even when `n <= 1`; concurrent queries may
+/// share the CPU-sized process pool, matching Java's cached executor plus
its
+/// per-caller `SemaphoredDelegatingExecutor`.
+/// - `per_query_only(sem)` / `shared_for_test(sem)` — tests only: an explicit
+/// per-query semaphore with NO process-global gating, so a shared static
cannot
+/// pollute a cap assertion.
///
-/// 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`).
+/// A `None` budget (passed only by low-level tests that opt out of gating)
runs the
+/// leaf ungated; production never passes `None`.
+#[derive(Clone)]
+pub(crate) struct SearchBudget {
+ /// Per-query permit source: caps ONE query's in-flight leaves at
`thread_num`.
+ per_query: Option<Arc<Semaphore>>,
+ /// When set, also draw from the process-global semaphore. Its effective
size
+ /// is at least the available CPU count and grows with larger requested
values,
+ /// matching Java's shared cached executor. `None` skips global gating
(tests).
+ process_global_capacity: Option<usize>,
+}
+
+impl SearchBudget {
+ /// The production budget: per-query cap of `concurrency` plus the shared
+ /// process pool. Applied even when `concurrency <= 1`.
+ pub(crate) fn production(concurrency: usize) -> Self {
+ Self {
+ per_query: Some(Arc::new(Semaphore::new(concurrency.max(1)))),
+ process_global_capacity: Some(concurrency.max(1)),
+ }
+ }
+
+ /// A per-query-only budget for tests that assert a single query's cap in
+ /// isolation. No process-global gating, so a static initialized by
another test
+ /// cannot change the observed peak.
+ #[cfg(test)]
+ pub(crate) fn per_query_only(permits: usize) -> Self {
+ Self {
+ per_query: Some(Arc::new(Semaphore::new(permits))),
+ process_global_capacity: None,
+ }
+ }
+
+ /// A budget backed by an explicit, caller-provided semaphore, with NO
+ /// process-global gating. Two concurrent searches given the SAME
`Arc<Semaphore>`
+ /// share one cap deterministically — the "deliberately shared" test seam
that
+ /// exercises cross-query capping (what the process-global static does in
+ /// production) without touching the shared static (so tests stay
isolated).
+ #[cfg(test)]
+ pub(crate) fn shared_for_test(semaphore: Arc<Semaphore>) -> Self {
+ Self {
+ per_query: Some(semaphore),
+ process_global_capacity: None,
+ }
+ }
+
+ /// Acquire one slot: per-query permit FIRST, then the process-global
permit.
+ /// This ordering prevents a job from holding a scarce global permit while
it
+ /// merely waits for its own query-local permit. Both guards are returned
and
+ /// must be kept alive for the duration of the leaf work (an ANN segment's
+ /// blocking CPU search or an exact file's async I/O).
+ async fn acquire(&self) -> crate::Result<SearchPermit> {
+ let per_query = match &self.per_query {
+ Some(sem) => Some(acquire_owned(sem.clone()).await?),
+ None => None,
+ };
+ let process_global = match self.process_global_capacity {
+ Some(cap) =>
Some(acquire_process_global_search_permit(cap).await?),
+ None => None,
+ };
+ Ok(SearchPermit {
+ _per_query: per_query,
+ _process_global: process_global,
+ })
+ }
+}
+
+/// Guard holding the acquired permits (per-query and/or process-global) until
it is
+/// dropped. Held for the whole leaf, including inside the dedicated executor
task
+/// (whose work continues if the awaiting future is dropped), so the budget
stays
+/// accurate.
+struct SearchPermit {
+ _per_query: Option<OwnedSemaphorePermit>,
+ _process_global: Option<tokio::sync::SemaphorePermit<'static>>,
+}
+
+async fn acquire_owned(sem: Arc<Semaphore>) ->
crate::Result<OwnedSemaphorePermit> {
+ sem.acquire_owned()
+ .await
+ .map_err(|e| crate::Error::UnexpectedError {
+ message: "global-index search concurrency budget was
closed".to_string(),
+ source: Some(Box::new(e)),
+ })
+}
+
+/// Acquire one slot from a search's concurrency budget, if one is set. The
returned
+/// guard holds the slot(s) until dropped, so the caller must keep it alive
for the
+/// duration of the leaf work it gates. A `None` budget (passed only by
low-level
+/// tests) runs the leaf ungated; production always supplies a `SearchBudget`,
+/// including on the `concurrency <= 1` sequential path.
async fn acquire_search_permit(
- budget: &Option<Arc<Semaphore>>,
-) -> crate::Result<Option<OwnedSemaphorePermit>> {
+ budget: &Option<SearchBudget>,
+) -> crate::Result<Option<SearchPermit>> {
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))
- }
+ Some(budget) => Ok(Some(budget.acquire().await?)),
None => Ok(None),
}
}
@@ -245,6 +335,18 @@ pub(crate) struct BucketSearchResult {
pub(crate) exact: Vec<PkVectorSearchResult>,
}
+/// One completed leaf of the merged per-bucket search: either an ANN segment's
+/// per-query hit lists (fold into the INDEXED heaps) or an uncovered exact
file's
+/// per-query hit lists (fold into the EXACT heaps). ANN and exact leaves run
in one
+/// combined concurrency window under the shared budget (mirroring Java's
single
+/// `allOf` over ANN + exact futures) rather than ANN-fully-then-exact; the
tag says
+/// which heap set each result feeds. Both carry one list per query (the
single-query
+/// path uses a 1-element outer vec).
+enum BucketLeaf {
+ Indexed(Vec<Vec<PkVectorSearchResult>>),
+ Exact(Vec<Vec<PkVectorSearchResult>>),
+}
+
/// ANN + exact data-file fallback search for one snapshot bucket. Mirrors Java
/// `org.apache.paimon.index.pkvector.PrimaryKeyVectorBucketSearch.search`.
///
@@ -260,7 +362,7 @@ pub(crate) struct BucketSearchResult {
#[allow(clippy::too_many_arguments)]
#[allow(clippy::type_complexity)]
pub(crate) async fn bucket_search(
- ann_searcher: Option<&dyn PkVectorAnnSearcher>,
+ ann_searcher: Option<Arc<dyn PkVectorAnnSearcher>>,
ann_segments: &[BucketAnnSegment],
active_files: &[BucketActiveFile],
deletion_vectors: &HashMap<String, Arc<DeletionVector>>,
@@ -281,7 +383,7 @@ pub(crate) async fn bucket_search(
skip_exact_fallback: bool,
residual_ranges: Option<&HashMap<String, roaring::RoaringTreemap>>,
concurrency: usize,
- search_budget: Option<Arc<Semaphore>>,
+ search_budget: Option<SearchBudget>,
) -> crate::Result<BucketSearchResult> {
if indexed_limit == 0 {
return Err(data_invalid("vector search limit must be positive"));
@@ -353,9 +455,20 @@ 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.
+ // ANN segment search. Each segment's scorer is synchronous CPU work (the
+ // faiss-like vindex search, which may itself use Rayon), so running it
inline
+ // would monopolize the async worker and serialize every segment across the
+ // whole read. Instead each segment is a dedicated-executor leaf gated by
the
+ // shared `search_budget`, mirroring Java's single
`GlobalIndexReadThreadPool`
+ // where every ANN segment is a pool task. `concurrency <= 1` still goes
through
+ // the dedicated executor (so a long CPU search never blocks the runtime)
but strictly
+ // one at a time.
+ //
+ // Structural validation of ALL segments happens before any leaf launches,
so a
+ // malformed later segment fails loud without half the segments having
already
+ // scored. Row-count validation runs BEFORE resolving the searcher so a
+ // corruption diagnostic is not masked by an "ANN search is not configured"
+ // error when both are wrong.
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
@@ -372,34 +485,54 @@ pub(crate) async fn bucket_search(
}
}
}
- let searcher = ann_searcher.ok_or_else(|| data_invalid("ANN search is
not configured"))?;
- for result in searcher.search(
- segment,
- query,
- metric,
- indexed_limit,
- &active_source_files,
- deletion_vectors,
- search_options,
- residual_ranges,
- )? {
- add_candidate(&mut indexed_heap, result, indexed_limit);
- }
}
+ let searcher = if ann_segments.is_empty() {
+ None
+ } else {
+ Some(
+ ann_searcher
+ .as_ref()
+ .ok_or_else(|| data_invalid("ANN search is not configured"))?
+ .clone(),
+ )
+ };
+
+ // Merged leaf stream: ANN segment searches and uncovered exact-file
searches run
+ // in ONE combined concurrency window under the shared budget — mirroring
Java's
+ // single `allOf` over the ANN and exact futures — instead of
ANN-fully-then-exact.
+ // Leaves are ordered ANN-first (segment order), then exact (active-file
order), so
+ // `drain_indexed_jobs` (which awaits ALL leaves and returns the
lowest-ordinal
+ // error) yields a deterministic error: first ANN error by segment order,
else
+ // first exact error by active-file order. At `concurrency <= 1` the
ANN-first
+ // order reproduces sequential ANN-then-exact execution. Each leaf is
tagged so its
+ // results fold into the correct heap; heaps are order-independent, so
completion
+ // order does not affect the final Top-K.
+ //
+ // Owned exact-query slice (one element for the single-query path) so
exact leaves
+ // borrow no locals across the await.
+ let queries: [&[f32]; 1] = [query];
+
+ // ANN leaves: shared owned inputs, one dedicated-executor scorer per
segment.
+ let ann_shared = searcher.as_ref().map(|searcher| {
+ (
+ searcher.clone(),
+ residual_ranges.map(|r| Arc::new(r.clone())),
+ Arc::new(active_source_files.clone()),
+ Arc::new(deletion_vectors.clone()),
+ Arc::new(search_options.clone()),
+ Arc::<[f32]>::from(query.to_vec()),
+ )
+ });
+ // Eligible uncovered exact files (active-file order) with their exclusion
+ // predicate; a file with no residual-allowed rows is skipped without
reading.
+ #[allow(clippy::type_complexity)]
+ let mut exact_tasks: Vec<(&BucketActiveFile, Box<dyn Fn(i64) -> bool +
Sync>)> = Vec::new();
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;
}
- // Residual allow-list: when present, only rows whose physical
position
- // passes the predicate may produce candidates. A file with no
entry (or
- // an empty entry) has no allowed rows, so it is skipped without
reading.
let residual_allowed: Option<&roaring::RoaringTreemap> = match
residual_ranges {
Some(ranges) => match ranges.get(&file.file_name) {
Some(allowed) if !allowed.is_empty() => Some(allowed),
@@ -408,46 +541,79 @@ pub(crate) async fn bucket_search(
None => None,
};
let dv = deletion_vectors.get(&file.file_name).cloned();
- tasks.push((file, Box::new(position_excluder(dv,
residual_allowed))));
+ exact_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)?);
+ // Build the ANN-first-then-exact leaf futures.
+ let mut leaves: Vec<BoxFuture<'_, crate::Result<BucketLeaf>>> = Vec::new();
+ if let Some((
+ searcher,
+ residual_arc,
+ active_source_files,
+ deletion_vectors,
+ search_options,
+ query_owned,
+ )) = ann_shared
+ {
+ for segment in ann_segments {
+ let searcher = searcher.clone();
+ let segment = segment.clone();
+ let active_source_files = active_source_files.clone();
+ let deletion_vectors = deletion_vectors.clone();
+ let search_options = search_options.clone();
+ let residual_arc = residual_arc.clone();
+ let query_owned = query_owned.clone();
+ let budget = search_budget.clone();
+ leaves.push(Box::pin(async move {
+ // One permit spans BOTH the async load and the dedicated
executor
+ // score. `load_segment_source` must not acquire a second
permit or
+ // it would deadlock at capacity 1.
+ let permit = acquire_search_permit(&budget).await?;
+ let segment_source =
searcher.load_segment_source(&segment).await?;
+ let hits = execute_global_index("ANN segment search task
failed", move || {
+ let _permit = permit;
+ searcher.search_source(
+ &segment,
+ segment_source,
+ &query_owned,
+ metric,
+ indexed_limit,
+ &active_source_files,
+ &deletion_vectors,
+ &search_options,
+ residual_arc.as_deref(),
+ )
+ })
+ .await?;
+ Ok(BucketLeaf::Indexed(vec![hits]))
+ }));
+ }
+ }
+ for (file, is_excluded) in &exact_tasks {
+ let queries = &queries;
+ let budget = search_budget.clone();
+ leaves.push(Box::pin(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?;
+ Ok(BucketLeaf::Exact(vec![single_query_result(per_query)?]))
+ }));
+ }
+
+ // Drive all leaves in one window; drain-all + lowest-ordinal error.
+ let outs = drain_indexed_jobs(leaves.into_iter(), concurrency).await?;
+ for leaf in outs {
+ match leaf {
+ BucketLeaf::Indexed(per_query) => {
+ for result in per_query.into_iter().next().unwrap_or_default()
{
+ add_candidate(&mut indexed_heap, result, indexed_limit);
+ }
}
- 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)
+ BucketLeaf::Exact(per_query) => {
+ for result in per_query.into_iter().next().unwrap_or_default()
{
+ add_candidate(&mut exact_heap, result, exact_limit);
}
- }))
- .buffer_unordered(concurrency)
- .try_collect::<Vec<_>>()
- .await?
- };
- for results in per_file {
- for result in results {
- add_candidate(&mut exact_heap, result, exact_limit);
}
}
}
@@ -473,7 +639,7 @@ pub(crate) async fn bucket_search(
#[allow(clippy::too_many_arguments)]
#[allow(clippy::type_complexity)]
pub(crate) async fn bucket_search_batch(
- ann_searcher: Option<&dyn PkVectorAnnSearcher>,
+ ann_searcher: Option<Arc<dyn PkVectorAnnSearcher>>,
ann_segments: &[BucketAnnSegment],
active_files: &[BucketActiveFile],
deletion_vectors: &HashMap<String, Arc<DeletionVector>>,
@@ -494,7 +660,7 @@ pub(crate) async fn bucket_search_batch(
skip_exact_fallback: bool,
residual_ranges: Option<&HashMap<String, roaring::RoaringTreemap>>,
concurrency: usize,
- search_budget: Option<Arc<Semaphore>>,
+ search_budget: Option<SearchBudget>,
) -> crate::Result<Vec<BucketSearchResult>> {
if queries.is_empty() {
return Err(data_invalid("vector search requires at least one query"));
@@ -504,7 +670,7 @@ pub(crate) async fn bucket_search_batch(
// principle differ).
if queries.len() == 1 {
let single = bucket_search(
- ann_searcher,
+ ann_searcher.clone(),
ann_segments,
active_files,
deletion_vectors,
@@ -597,9 +763,15 @@ 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.
+ // ANN segment search (multi-query). As in the single-query path, each
segment's
+ // synchronous scorer runs as a dedicated-executor leaf gated by the shared
+ // `search_budget` so segments across all buckets share one budget and
never
+ // monopolize an async worker. All structural validation runs before any
leaf
+ // launches. One shared reader per segment searches every query; the
per-query
+ // hits fan into their own bounded heaps after the parallel stage.
+ // Row-count validation runs BEFORE resolving the searcher so a corruption
+ // diagnostic is not masked by an "ANN search is not configured" error
when both
+ // are wrong (mirrors the single-query path and Java's up-front checks).
for segment in ann_segments {
for source in segment.source_meta.source_files() {
if let Some(active) = files_by_name.get(source.file_name()) {
@@ -611,39 +783,40 @@ pub(crate) async fn bucket_search_batch(
}
}
}
- let searcher = ann_searcher.ok_or_else(|| data_invalid("ANN search is
not configured"))?;
- // One shared reader per segment searches all queries; fan the
per-query
- // hits into their own bounded heaps.
- let per_query = searcher.search_batch(
- segment,
- queries,
- metric,
- indexed_limit,
- &active_source_files,
- deletion_vectors,
- search_options,
- residual_ranges,
- )?;
- if per_query.len() != queries.len() {
- return Err(data_invalid(format!(
- "ANN batch search returned {} result lists for {} queries",
- per_query.len(),
- queries.len()
- )));
- }
- for (results, heap) in
per_query.into_iter().zip(indexed_heaps.iter_mut()) {
- for result in results {
- add_candidate(heap, result, indexed_limit);
- }
- }
}
+ let searcher = if ann_segments.is_empty() {
+ None
+ } else {
+ Some(
+ ann_searcher
+ .as_ref()
+ .ok_or_else(|| data_invalid("ANN search is not configured"))?
+ .clone(),
+ )
+ };
+
+ // Merged leaf stream (see the single-query `bucket_search` for the
rationale):
+ // ANN segment searches and uncovered exact-file searches run in ONE
combined
+ // concurrency window under the shared budget, ANN-first then exact,
drained
+ // together with a deterministic lowest-ordinal error. Each leaf carries
one hit
+ // list per query; the fan-in zips per-query results into the per-query
heaps.
+ let query_count = queries.len();
+ let ann_shared = searcher.as_ref().map(|searcher| {
+ let queries_owned: Arc<Vec<Vec<f32>>> =
+ Arc::new(queries.iter().map(|q| q.to_vec()).collect());
+ (
+ searcher.clone(),
+ residual_ranges.map(|r| Arc::new(r.clone())),
+ Arc::new(active_source_files.clone()),
+ Arc::new(deletion_vectors.clone()),
+ Arc::new(search_options.clone()),
+ queries_owned,
+ )
+ });
+ #[allow(clippy::type_complexity)]
+ let mut exact_tasks: Vec<(&BucketActiveFile, Box<dyn Fn(i64) -> bool +
Sync>)> = Vec::new();
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;
@@ -656,44 +829,89 @@ pub(crate) async fn bucket_search_batch(
None => None,
};
let dv = deletion_vectors.get(&file.file_name).cloned();
- tasks.push((file, Box::new(position_excluder(dv,
residual_allowed))));
+ exact_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())
+ let mut leaves: Vec<BoxFuture<'_, crate::Result<BucketLeaf>>> = Vec::new();
+ if let Some((
+ searcher,
+ residual_arc,
+ active_source_files,
+ deletion_vectors,
+ search_options,
+ queries_owned,
+ )) = ann_shared
+ {
+ for segment in ann_segments {
+ let searcher = searcher.clone();
+ let segment = segment.clone();
+ let active_source_files = active_source_files.clone();
+ let deletion_vectors = deletion_vectors.clone();
+ let search_options = search_options.clone();
+ let residual_arc = residual_arc.clone();
+ let queries_owned = queries_owned.clone();
+ let budget = search_budget.clone();
+ leaves.push(Box::pin(async move {
+ let permit = acquire_search_permit(&budget).await?;
+ let segment_source =
searcher.load_segment_source(&segment).await?;
+ let per_query = execute_global_index("ANN segment search task
failed", move || {
+ let _permit = permit;
+ let query_refs: Vec<&[f32]> =
+ queries_owned.iter().map(|q| q.as_slice()).collect();
+ let per_query = searcher.search_batch_source(
+ &segment,
+ segment_source,
+ &query_refs,
+ metric,
+ indexed_limit,
+ &active_source_files,
+ &deletion_vectors,
+ &search_options,
+ residual_arc.as_deref(),
+ )?;
+ if per_query.len() != query_count {
+ return Err(data_invalid(format!(
+ "ANN batch search returned {} result lists for {}
queries",
+ per_query.len(),
+ query_count
+ )));
+ }
+ Ok(per_query)
+ })
+ .await?;
+ Ok(BucketLeaf::Indexed(per_query))
+ }));
+ }
+ }
+ for (file, is_excluded) in &exact_tasks {
+ let budget = search_budget.clone();
+ leaves.push(Box::pin(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?;
+ Ok(BucketLeaf::Exact(validate_per_query_len(
+ per_query,
+ query_count,
+ )?))
+ }));
+ }
+
+ let outs = drain_indexed_jobs(leaves.into_iter(), concurrency).await?;
+ for leaf in outs {
+ match leaf {
+ BucketLeaf::Indexed(per_query) => {
+ for (results, heap) in
per_query.into_iter().zip(indexed_heaps.iter_mut()) {
+ for result in results {
+ add_candidate(heap, result, indexed_limit);
+ }
}
- }))
- .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);
+ }
+ BucketLeaf::Exact(per_query) => {
+ for (results, heap) in
per_query.into_iter().zip(exact_heaps.iter_mut()) {
+ for result in results {
+ add_candidate(heap, result, exact_limit);
+ }
}
}
}
@@ -714,18 +932,27 @@ pub(crate) async fn bucket_search_batch(
#[cfg(test)]
mod tests {
use super::*;
- use crate::spec::PrimaryKeyIndexSourceFile;
+ use crate::spec::PrimaryKeyIndexSourceFile as PkVectorSourceFile;
use crate::vindex::pkvector::ann::PkVectorAnnSearcher;
use crate::vindex::pkvector::exact::exact_search;
use crate::vindex::pkvector::reader::test_support::ArrayReader;
+ use bytes::Bytes;
use roaring::RoaringBitmap;
- fn meta(files: &[(&str, i64)]) -> PrimaryKeyIndexSourceMeta {
- PrimaryKeyIndexSourceMeta::new(
+ /// Trivial ANN-segment loader for fakes: returns empty owned bytes (the
fakes
+ /// model behavior above physical index decoding and ignore
`segment_bytes`).
+ fn empty_ann_loader(
+ _segment: &BucketAnnSegment,
+ ) -> futures::future::BoxFuture<'static, crate::Result<Bytes>> {
+ Box::pin(async { Ok(Bytes::new()) })
+ }
+
+ fn meta(files: &[(&str, i64)]) -> PkVectorSourceMeta {
+ PkVectorSourceMeta::new(
1,
files
.iter()
- .map(|(n, r)| PrimaryKeyIndexSourceFile::new((*n).into(),
*r).unwrap())
+ .map(|(n, r)| PkVectorSourceFile::new((*n).into(),
*r).unwrap())
.collect(),
)
.unwrap()
@@ -826,13 +1053,21 @@ mod tests {
}
/// Fake ANN searcher returning preset results and recording calls.
+ #[derive(Clone)]
struct FakeAnnSearcher {
result: Vec<PkVectorSearchResult>,
}
impl PkVectorAnnSearcher for FakeAnnSearcher {
+ fn load_segment(
+ &self,
+ segment: &BucketAnnSegment,
+ ) -> futures::future::BoxFuture<'static, crate::Result<Bytes>> {
+ empty_ann_loader(segment)
+ }
fn search_batch(
&self,
_segment: &BucketAnnSegment,
+ _segment_bytes: Bytes,
queries: &[&[f32]],
_metric: VectorSearchMetric,
_limit: usize,
@@ -882,7 +1117,7 @@ mod tests {
let opts = HashMap::new();
let out = bucket_search(
- Some(&ann),
+ Some(Arc::new(ann.clone())),
&[segment],
&active_files,
&dvs,
@@ -958,7 +1193,7 @@ mod tests {
};
let factory = unreachable_search();
let out = bucket_search(
- Some(&ann),
+ Some(Arc::new(ann.clone())),
&[segment],
&[active("data-1", 3)],
&HashMap::new(),
@@ -1015,7 +1250,7 @@ mod tests {
};
let factory = unreachable_search();
let out = bucket_search(
- Some(&ann),
+ Some(Arc::new(ann.clone())),
&[segment],
&[active("data-1", 2)],
&HashMap::new(),
@@ -1080,7 +1315,7 @@ mod tests {
},
);
let out = bucket_search(
- Some(&ann),
+ Some(Arc::new(ann.clone())),
&[segment],
&[active("data-1", 2), active("data-2", 2)],
&HashMap::new(),
@@ -1229,7 +1464,7 @@ mod tests {
let segment = BucketAnnSegment::for_test(meta(&[("data-1", 2)]));
let factory = unreachable_search();
let err = bucket_search(
- Some(&ann),
+ Some(Arc::new(ann.clone())),
&[segment],
&[active("data-1", 3)],
&HashMap::new(),
@@ -1279,7 +1514,7 @@ mod tests {
},
);
let out = bucket_search(
- Some(&ann),
+ Some(Arc::new(ann.clone())),
&[segment],
&[active("data-1", 2)],
&HashMap::new(),
@@ -1384,7 +1619,7 @@ mod tests {
let ann = FakeAnnSearcher { result: vec![] };
let factory = unreachable_search();
let err = bucket_search(
- Some(&ann),
+ Some(Arc::new(ann.clone())),
&[seg1, seg2],
&[active("data-1", 2), active("data-2", 2)],
&HashMap::new(),
@@ -1424,7 +1659,7 @@ mod tests {
let ann = FakeAnnSearcher { result: vec![] };
let factory = unreachable_search();
let err = bucket_search(
- Some(&ann),
+ Some(Arc::new(ann.clone())),
&[seg1, seg2],
&[active("data-1", 2), active("data-2", 2)],
&HashMap::new(),
@@ -1827,7 +2062,7 @@ mod tests {
Some(vec![8.0, 0.0]),
]);
bucket_search(
- Some(&ann),
+ Some(Arc::new(ann.clone())),
&[BucketAnnSegment::for_test(meta(&[("ann.mosaic", 3)]))],
&active_files,
&dvs,
@@ -1853,7 +2088,7 @@ mod tests {
]);
let query_ref: &[f32] = &query;
let batch = bucket_search_batch(
- Some(&ann),
+ Some(Arc::new(ann.clone())),
&[segment],
&active_files,
&dvs,
@@ -2180,4 +2415,715 @@ mod tests {
"parallel ranking must equal serial ranking"
);
}
+
+ /// ANN searcher probe that records the peak number of `search_batch` calls
+ /// running simultaneously. Each call does a real blocking sleep so
overlapping
+ /// calls are observable; `peak` is the max concurrent count seen. Mirrors
the
+ /// blocking CPU nature of the real vindex scorer.
+ struct PeakProbeAnn {
+ inflight: std::sync::Arc<std::sync::atomic::AtomicUsize>,
+ peak: std::sync::Arc<std::sync::atomic::AtomicUsize>,
+ }
+ impl PkVectorAnnSearcher for PeakProbeAnn {
+ fn load_segment(
+ &self,
+ segment: &BucketAnnSegment,
+ ) -> futures::future::BoxFuture<'static, crate::Result<Bytes>> {
+ empty_ann_loader(segment)
+ }
+ fn search_batch(
+ &self,
+ _segment: &BucketAnnSegment,
+ _segment_bytes: Bytes,
+ queries: &[&[f32]],
+ _metric: VectorSearchMetric,
+ _limit: usize,
+ _active_source_files: &HashSet<String>,
+ _dvs: &HashMap<String, Arc<DeletionVector>>,
+ _opts: &HashMap<String, String>,
+ _residual_ranges: Option<&HashMap<String,
roaring::RoaringTreemap>>,
+ ) -> crate::Result<Vec<Vec<PkVectorSearchResult>>> {
+ use std::sync::atomic::Ordering::SeqCst;
+ let current = self.inflight.fetch_add(1, SeqCst) + 1;
+ self.peak.fetch_max(current, SeqCst);
+ std::thread::sleep(std::time::Duration::from_millis(50));
+ self.inflight.fetch_sub(1, SeqCst);
+ Ok(queries.iter().map(|_| Vec::new()).collect())
+ }
+ }
+
+ /// `n` ANN segments in one bucket, each with a distinct payload path and a
+ /// distinct (covered) source file. Returns the searcher + its peak
counter so a
+ /// test can assert how many segment searches ran at once.
+ fn n_segment_bucket(
+ n: usize,
+ ) -> (
+ Vec<BucketAnnSegment>,
+ Vec<BucketActiveFile>,
+ std::sync::Arc<PeakProbeAnn>,
+ std::sync::Arc<std::sync::atomic::AtomicUsize>,
+ ) {
+ let mut segments = Vec::with_capacity(n);
+ let mut actives = Vec::with_capacity(n);
+ for i in 0..n {
+ let src = format!("cov-{i}");
+ segments.push(BucketAnnSegment {
+ source_meta: meta(&[(src.as_str(), 2)]),
+ path: format!("seg-{i}"),
+ file_size: 0,
+ index_meta: Vec::new(),
+ });
+ actives.push(active(&src, 2));
+ }
+ let peak = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+ let probe = std::sync::Arc::new(PeakProbeAnn {
+ inflight:
std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)),
+ peak: peak.clone(),
+ });
+ (segments, actives, probe, peak)
+ }
+
+ #[tokio::test]
+ async fn ann_segments_search_in_parallel_at_concurrency_above_one() {
+ // Two ANN segments in one bucket. At concurrency 4 their (blocking)
searches
+ // must overlap: peak simultaneous ANN calls must exceed 1. A serial
ANN loop
+ // yields peak == 1 and fails this assertion.
+ let (segments, active_files, probe, peak) = n_segment_bucket(2);
+ let searcher: Arc<dyn PkVectorAnnSearcher> = probe;
+ let factory = unreachable_search();
+ let out = bucket_search(
+ Some(searcher),
+ &segments,
+ &active_files,
+ &HashMap::new(),
+ &factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 8,
+ 8,
+ &HashMap::new(),
+ false,
+ None,
+ 4,
+ Some(SearchBudget::per_query_only(4)),
+ )
+ .await
+ .unwrap();
+ assert!(out.indexed.is_empty());
+ assert!(
+ peak.load(std::sync::atomic::Ordering::SeqCst) >= 2,
+ "ANN segment searches must run in parallel at concurrency 4;
observed peak {}",
+ peak.load(std::sync::atomic::Ordering::SeqCst)
+ );
+ }
+
+ #[tokio::test]
+ async fn ann_segments_search_serially_at_concurrency_one() {
+ // At concurrency 1 the ANN segments must NOT overlap: peak
simultaneous
+ // calls is exactly 1 (still off the async worker on the dedicated
executor,
+ // but one at a time). Guards the one-worker budget without regressing
inline.
+ let (segments, active_files, probe, peak) = n_segment_bucket(3);
+ let searcher: Arc<dyn PkVectorAnnSearcher> = probe;
+ let factory = unreachable_search();
+ bucket_search(
+ Some(searcher),
+ &segments,
+ &active_files,
+ &HashMap::new(),
+ &factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 8,
+ 8,
+ &HashMap::new(),
+ false,
+ None,
+ 1,
+ None,
+ )
+ .await
+ .unwrap();
+ assert_eq!(
+ peak.load(std::sync::atomic::Ordering::SeqCst),
+ 1,
+ "concurrency == 1 must search ANN segments one at a time"
+ );
+ }
+
+ #[tokio::test]
+ async fn ann_segment_search_respects_shared_budget() {
+ // Four ANN segments, shared budget of 2: at most 2 segment searches
may run
+ // at once even though concurrency (fan-out width) is 4. Mirrors
Java's single
+ // shared pool sized to threadNum capping total in-flight leaf work.
+ let (segments, active_files, probe, peak) = n_segment_bucket(4);
+ let searcher: Arc<dyn PkVectorAnnSearcher> = probe;
+ let factory = unreachable_search();
+ bucket_search(
+ Some(searcher),
+ &segments,
+ &active_files,
+ &HashMap::new(),
+ &factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 8,
+ 8,
+ &HashMap::new(),
+ false,
+ None,
+ 4,
+ Some(SearchBudget::per_query_only(2)),
+ )
+ .await
+ .unwrap();
+ let observed = peak.load(std::sync::atomic::Ordering::SeqCst);
+ assert!(
+ observed <= 2,
+ "shared budget must cap concurrent ANN segment searches at 2;
observed {observed}"
+ );
+ assert!(
+ observed >= 2,
+ "test must actually exercise ANN overlap; observed {observed}"
+ );
+ }
+
+ #[tokio::test]
+ async fn malformed_ann_source_fails_before_any_segment_search_runs() {
+ // A segment whose ANN source row count disagrees with the active file
is
+ // corruption. The mismatch must be detected during up-front structural
+ // validation, BEFORE any segment search leaf is spawned — so the
probe's
+ // peak stays 0. Guards Codex's "validate all segments before
launching any
+ // job" requirement.
+ let (mut segments, _active_files, probe, peak) = n_segment_bucket(2);
+ // Break segment 1's source row count vs the active file (active says
2).
+ segments[1] = BucketAnnSegment {
+ source_meta: meta(&[("cov-1", 99)]),
+ path: "seg-1".to_string(),
+ file_size: 0,
+ index_meta: Vec::new(),
+ };
+ let searcher: Arc<dyn PkVectorAnnSearcher> = probe;
+ let factory = unreachable_search();
+ let err = bucket_search(
+ Some(searcher),
+ &segments,
+ &[active("cov-0", 2), active("cov-1", 2)],
+ &HashMap::new(),
+ &factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 8,
+ 8,
+ &HashMap::new(),
+ false,
+ None,
+ 4,
+ Some(SearchBudget::per_query_only(4)),
+ )
+ .await
+ .unwrap_err();
+ assert!(
+ err.to_string().contains("does not match"),
+ "expected a row-count mismatch error, got: {err}"
+ );
+ assert_eq!(
+ peak.load(std::sync::atomic::Ordering::SeqCst),
+ 0,
+ "no ANN segment search may run when structural validation fails"
+ );
+ }
+
+ #[tokio::test]
+ async fn batch_ann_segments_search_in_parallel_at_concurrency_above_one() {
+ // Multi-query (2 queries) routes through the `bucket_search_batch`
multi-query
+ // path (not the batch-of-one short-circuit), so this covers the batch
ANN
+ // parallel loop specifically. Two segments at concurrency 4 must
overlap.
+ let (segments, active_files, probe, peak) = n_segment_bucket(2);
+ let searcher: Arc<dyn PkVectorAnnSearcher> = probe;
+ let factory = unreachable_search();
+ let q0: &[f32] = &[0.0, 0.0];
+ let q1: &[f32] = &[1.0, 1.0];
+ let out = bucket_search_batch(
+ Some(searcher),
+ &segments,
+ &active_files,
+ &HashMap::new(),
+ &factory,
+ &[q0, q1],
+ VectorSearchMetric::L2,
+ 8,
+ 8,
+ &HashMap::new(),
+ false,
+ None,
+ 4,
+ Some(SearchBudget::per_query_only(4)),
+ )
+ .await
+ .unwrap();
+ assert_eq!(out.len(), 2, "one result list per query");
+ assert!(
+ peak.load(std::sync::atomic::Ordering::SeqCst) >= 2,
+ "batch ANN segment searches must run in parallel at concurrency 4;
observed peak {}",
+ peak.load(std::sync::atomic::Ordering::SeqCst)
+ );
+ }
+
+ /// ANN searcher whose `search_batch` panics, exercising the dedicated
+ /// executor's panic-to-error mapping.
+ struct PanicAnn;
+ impl PkVectorAnnSearcher for PanicAnn {
+ fn load_segment(
+ &self,
+ segment: &BucketAnnSegment,
+ ) -> futures::future::BoxFuture<'static, crate::Result<Bytes>> {
+ empty_ann_loader(segment)
+ }
+ fn search_batch(
+ &self,
+ _segment: &BucketAnnSegment,
+ _segment_bytes: Bytes,
+ _queries: &[&[f32]],
+ _metric: VectorSearchMetric,
+ _limit: usize,
+ _active_source_files: &HashSet<String>,
+ _dvs: &HashMap<String, Arc<DeletionVector>>,
+ _opts: &HashMap<String, String>,
+ _residual_ranges: Option<&HashMap<String,
roaring::RoaringTreemap>>,
+ ) -> crate::Result<Vec<Vec<PkVectorSearchResult>>> {
+ panic!("scorer panic to exercise JoinError mapping");
+ }
+ }
+
+ #[tokio::test]
+ async fn ann_segment_scorer_panic_maps_to_unexpected_error() {
+ // A panic inside the dedicated-executor ANN leaf must surface as a
mapped
+ // `UnexpectedError` ("ANN segment search task failed"), not abort the
runtime
+ // or hang. Uses the parallel path (concurrency 4).
+ let (segments, active_files, _probe, _peak) = n_segment_bucket(1);
+ let searcher: Arc<dyn PkVectorAnnSearcher> = Arc::new(PanicAnn);
+ let factory = unreachable_search();
+ let err = bucket_search(
+ Some(searcher),
+ &segments,
+ &active_files,
+ &HashMap::new(),
+ &factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 8,
+ 8,
+ &HashMap::new(),
+ false,
+ None,
+ 4,
+ Some(SearchBudget::per_query_only(4)),
+ )
+ .await
+ .unwrap_err();
+ assert!(
+ err.to_string().contains("ANN segment search task failed"),
+ "panic must map to UnexpectedError, got: {err}"
+ );
+ }
+
+ /// ANN searcher that counts every completed `search_batch` call and
returns an
+ /// `Err` (naming the segment's source file) for any segment whose source
file is
+ /// in `fail_files`. A configurable pre-return spin makes the erroring
segment
+ /// finish FIRST while the others are still running, so a short-circuiting
drain
+ /// would surface a non-deterministic / non-lowest-index error and skip
jobs.
+ struct DrainProbeAnn {
+ completed: std::sync::Arc<std::sync::atomic::AtomicUsize>,
+ fail_files: HashSet<String>,
+ slow_files: HashSet<String>,
+ }
+ impl PkVectorAnnSearcher for DrainProbeAnn {
+ fn load_segment(
+ &self,
+ segment: &BucketAnnSegment,
+ ) -> futures::future::BoxFuture<'static, crate::Result<Bytes>> {
+ empty_ann_loader(segment)
+ }
+ fn search_batch(
+ &self,
+ segment: &BucketAnnSegment,
+ _segment_bytes: Bytes,
+ queries: &[&[f32]],
+ _metric: VectorSearchMetric,
+ _limit: usize,
+ _active_source_files: &HashSet<String>,
+ _dvs: &HashMap<String, Arc<DeletionVector>>,
+ _opts: &HashMap<String, String>,
+ _residual_ranges: Option<&HashMap<String,
roaring::RoaringTreemap>>,
+ ) -> crate::Result<Vec<Vec<PkVectorSearchResult>>> {
+ let file = segment.source_meta.source_files()[0]
+ .file_name()
+ .to_string();
+ // Non-failing "slow" segments block so they are still in flight
when the
+ // fast-failing segment returns — proving the drain awaits them
anyway.
+ if self.slow_files.contains(&file) {
+ std::thread::sleep(std::time::Duration::from_millis(80));
+ }
+ self.completed
+ .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+ if self.fail_files.contains(&file) {
+ return Err(data_invalid(format!("boom in {file}")));
+ }
+ Ok(queries.iter().map(|_| Vec::new()).collect())
+ }
+ }
+
+ #[tokio::test]
+ async fn drain_awaits_all_jobs_and_returns_lowest_index_error() {
+ // 4 segments (seg-0..seg-3 over cov-0..cov-3). Segments 1 and 2
error; 0 and 3
+ // are slow. The drain must (a) run ALL 4 scorers to completion even
though
+ // seg-1 errors early, and (b) surface seg-1's error (lowest index),
not seg-2's
+ // and not whichever finished first.
+ let (segments, active_files, _probe, _peak) = n_segment_bucket(4);
+ let completed =
std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+ let searcher: Arc<dyn PkVectorAnnSearcher> = Arc::new(DrainProbeAnn {
+ completed: completed.clone(),
+ fail_files: HashSet::from(["cov-1".to_string(),
"cov-2".to_string()]),
+ slow_files: HashSet::from(["cov-0".to_string(),
"cov-3".to_string()]),
+ });
+ let factory = unreachable_search();
+ let err = bucket_search(
+ Some(searcher),
+ &segments,
+ &active_files,
+ &HashMap::new(),
+ &factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 8,
+ 8,
+ &HashMap::new(),
+ false,
+ None,
+ 4,
+ Some(SearchBudget::per_query_only(4)),
+ )
+ .await
+ .unwrap_err();
+ // (b) lowest-index erroring segment wins, deterministically.
+ assert!(
+ err.to_string().contains("boom in cov-1"),
+ "must surface the lowest-index error (cov-1), got: {err}"
+ );
+ // (a) every segment's scorer ran to completion — no job was skipped
by an
+ // early short-circuit.
+ assert_eq!(
+ completed.load(std::sync::atomic::Ordering::SeqCst),
+ 4,
+ "all 4 ANN segment jobs must run to completion before the error
returns"
+ );
+ }
+
+ /// ANN searcher that records peak concurrent calls into an EXTERNALLY
shared
+ /// counter, so two independent `bucket_search` invocations can observe
their
+ /// combined in-flight count. Each call blocks briefly so overlap is
observable.
+ struct SharedPeakAnn {
+ inflight: std::sync::Arc<std::sync::atomic::AtomicUsize>,
+ peak: std::sync::Arc<std::sync::atomic::AtomicUsize>,
+ }
+ impl PkVectorAnnSearcher for SharedPeakAnn {
+ fn load_segment(
+ &self,
+ segment: &BucketAnnSegment,
+ ) -> futures::future::BoxFuture<'static, crate::Result<Bytes>> {
+ empty_ann_loader(segment)
+ }
+ fn search_batch(
+ &self,
+ _segment: &BucketAnnSegment,
+ _segment_bytes: Bytes,
+ queries: &[&[f32]],
+ _metric: VectorSearchMetric,
+ _limit: usize,
+ _active_source_files: &HashSet<String>,
+ _dvs: &HashMap<String, Arc<DeletionVector>>,
+ _opts: &HashMap<String, String>,
+ _residual_ranges: Option<&HashMap<String,
roaring::RoaringTreemap>>,
+ ) -> crate::Result<Vec<Vec<PkVectorSearchResult>>> {
+ use std::sync::atomic::Ordering::SeqCst;
+ let current = self.inflight.fetch_add(1, SeqCst) + 1;
+ self.peak.fetch_max(current, SeqCst);
+ std::thread::sleep(std::time::Duration::from_millis(60));
+ self.inflight.fetch_sub(1, SeqCst);
+ Ok(queries.iter().map(|_| Vec::new()).collect())
+ }
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+ async fn shared_budget_caps_ann_across_concurrent_searches() {
+ // Two independent searches (simulating two concurrent PK-vector
queries), each
+ // with 2 ANN segments and concurrency 4, but sharing ONE budget
semaphore of
+ // 2 permits (the deliberately-shared test seam). Their COMBINED
in-flight ANN
+ // count must never exceed 2 — proving the shared budget caps ANN work
ACROSS
+ // queries, which is what the process-global static does in
production. Before
+ // the fix, each query built its own Semaphore, so the combined peak
could
+ // reach 2 (segments) x 2 (queries) = 4.
+ let inflight =
std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+ let peak = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+ let shared = Arc::new(Semaphore::new(2));
+
+ let run = |budget: SearchBudget| {
+ let inflight = inflight.clone();
+ let peak = peak.clone();
+ async move {
+ let (segments, active_files, _p, _pk) = n_segment_bucket(2);
+ let searcher: Arc<dyn PkVectorAnnSearcher> =
+ Arc::new(SharedPeakAnn { inflight, peak });
+ let factory = unreachable_search();
+ bucket_search(
+ Some(searcher),
+ &segments,
+ &active_files,
+ &HashMap::new(),
+ &factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 8,
+ 8,
+ &HashMap::new(),
+ false,
+ None,
+ 4,
+ Some(budget),
+ )
+ .await
+ .unwrap();
+ }
+ };
+
+ // Both searches share the SAME budget semaphore.
+ let a = run(SearchBudget::shared_for_test(shared.clone()));
+ let b = run(SearchBudget::shared_for_test(shared.clone()));
+ tokio::join!(a, b);
+
+ let observed = peak.load(std::sync::atomic::Ordering::SeqCst);
+ assert!(
+ observed <= 2,
+ "shared budget must cap ANN across concurrent searches at 2;
observed {observed}"
+ );
+ assert!(
+ observed >= 2,
+ "test must actually exercise cross-search overlap; observed
{observed}"
+ );
+ }
+
+ /// ANN searcher whose score performs a BOUNDED two-way handshake with a
matching
+ /// exact-file probe: it signals arrival, then waits (with a timeout) for
the
+ /// exact leaf to signal too. If both arrive within the timeout they were
in
+ /// flight simultaneously — deterministic overlap detection with no sleeps
and no
+ /// unbounded wait. If they cannot overlap (e.g. a two-phase scheduler),
the leaf
+ /// that runs first times out waiting for the peer and RETURNS (setting no
overlap
+ /// flag), so the test fails cleanly on the assertion rather than hanging
forever.
+ struct HandshakeAnn {
+ /// ANN sends here on arrival; the exact side receives it.
+ ann_here: std::sync::mpsc::SyncSender<()>,
+ /// ANN receives the exact side's arrival here.
+ exact_here:
std::sync::Arc<std::sync::Mutex<std::sync::mpsc::Receiver<()>>>,
+ overlapped: std::sync::Arc<std::sync::atomic::AtomicBool>,
+ }
+ impl PkVectorAnnSearcher for HandshakeAnn {
+ fn load_segment(
+ &self,
+ segment: &BucketAnnSegment,
+ ) -> futures::future::BoxFuture<'static, crate::Result<Bytes>> {
+ empty_ann_loader(segment)
+ }
+ fn search_batch(
+ &self,
+ _segment: &BucketAnnSegment,
+ _segment_bytes: Bytes,
+ queries: &[&[f32]],
+ _metric: VectorSearchMetric,
+ _limit: usize,
+ _active_source_files: &HashSet<String>,
+ _dvs: &HashMap<String, Arc<DeletionVector>>,
+ _opts: &HashMap<String, String>,
+ _residual_ranges: Option<&HashMap<String,
roaring::RoaringTreemap>>,
+ ) -> crate::Result<Vec<Vec<PkVectorSearchResult>>> {
+ // Runs on the blocking pool. Announce arrival, then wait
(bounded) for the
+ // exact leaf. Both arriving proves overlap; a timeout means no
overlap.
+ let _ = self.ann_here.try_send(());
+ let got = self
+ .exact_here
+ .lock()
+ .unwrap()
+ .recv_timeout(std::time::Duration::from_secs(5));
+ if got.is_ok() {
+ self.overlapped
+ .store(true, std::sync::atomic::Ordering::SeqCst);
+ }
+ Ok(queries.iter().map(|_| Vec::new()).collect())
+ }
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+ async fn ann_and_exact_leaves_overlap_within_one_bucket() {
+ // One bucket with one ANN segment (covers "cov") AND one uncovered
exact file
+ // ("ex"). The ANN leaf (blocking) and the exact leaf (via
spawn_blocking) do a
+ // bounded two-way handshake: each signals arrival and waits (5s cap)
for the
+ // other. Both arriving proves they overlap. A two-phase scheduler
(all ANN,
+ // then all exact) cannot get both in flight, so the leaf that runs
first
+ // times out waiting for the peer (which never co-runs) and
`overlapped` stays
+ // false — the assertion fails cleanly within ~5s, not by hanging.
+ let (ann_tx, ann_rx) = std::sync::mpsc::sync_channel::<()>(1);
+ let (exact_tx, exact_rx) = std::sync::mpsc::sync_channel::<()>(1);
+ let ann_rx = std::sync::Arc::new(std::sync::Mutex::new(ann_rx));
+ let exact_rx = std::sync::Arc::new(std::sync::Mutex::new(exact_rx));
+ let overlapped =
std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
+
+ let segment = BucketAnnSegment::for_test(meta(&[("cov", 2)]));
+ let searcher: Arc<dyn PkVectorAnnSearcher> = Arc::new(HandshakeAnn {
+ ann_here: ann_tx,
+ exact_here: exact_rx.clone(),
+ overlapped: overlapped.clone(),
+ });
+
+ // Exact-file probe: rendezvous on the same channels via
`spawn_blocking`
+ // (so it can block-wait alongside the ANN blocking leaf), also
bounded.
+ let ex_ann_rx = ann_rx.clone();
+ let factory = as_search(
+ move |file: &BucketActiveFile,
+ queries: &[&[f32]],
+ _: VectorSearchMetric,
+ _: usize,
+ _: &(dyn Fn(i64) -> bool + Sync)|
+ -> ExactFileSearchFuture<'_> {
+ let ex_ann_rx = ex_ann_rx.clone();
+ let exact_tx = exact_tx.clone();
+ let n = queries.len();
+ let _ = file;
+ Box::pin(async move {
+ tokio::task::spawn_blocking(move || {
+ let _ = exact_tx.try_send(());
+ let _ = ex_ann_rx
+ .lock()
+ .unwrap()
+ .recv_timeout(std::time::Duration::from_secs(5));
+ })
+ .await
+ .expect("exact handshake task");
+ Ok(vec![Vec::new(); n])
+ })
+ },
+ );
+
+ let out = bucket_search(
+ Some(searcher),
+ &[segment],
+ &[active("cov", 2), active("ex", 2)],
+ &HashMap::new(),
+ &factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 8,
+ 8,
+ &HashMap::new(),
+ false, // NOT fast: exact fallback runs
+ None,
+ 4,
+ Some(SearchBudget::per_query_only(4)),
+ )
+ .await
+ .unwrap();
+ assert!(out.indexed.is_empty() && out.exact.is_empty());
+ assert!(
+ overlapped.load(std::sync::atomic::Ordering::SeqCst),
+ "ANN and exact leaves of one bucket must overlap (both completed
the handshake)"
+ );
+ }
+
+ /// ANN searcher whose `load_segment` records which segment paths it
loaded and
+ /// hands the loaded bytes through to `search_batch`, which asserts it
received
+ /// exactly those bytes. Proves the leaf loads per-segment via the seam
(not from
+ /// an up-front map) and that the loaded bytes reach the scorer.
+ struct RecordingLoaderAnn {
+ loaded: std::sync::Arc<std::sync::Mutex<Vec<String>>>,
+ }
+ impl PkVectorAnnSearcher for RecordingLoaderAnn {
+ fn load_segment(
+ &self,
+ segment: &BucketAnnSegment,
+ ) -> futures::future::BoxFuture<'static, crate::Result<Bytes>> {
+ let loaded = self.loaded.clone();
+ let path = segment.path.clone();
+ // Bytes are the segment path itself, so `search_batch` can verify
the
+ // exact bytes from THIS segment's load arrived (not some other
segment's).
+ Box::pin(async move {
+ loaded.lock().unwrap().push(path.clone());
+ Ok(Bytes::from(path.into_bytes()))
+ })
+ }
+ fn search_batch(
+ &self,
+ segment: &BucketAnnSegment,
+ segment_bytes: Bytes,
+ queries: &[&[f32]],
+ _metric: VectorSearchMetric,
+ _limit: usize,
+ _active_source_files: &HashSet<String>,
+ _dvs: &HashMap<String, Arc<DeletionVector>>,
+ _opts: &HashMap<String, String>,
+ _residual_ranges: Option<&HashMap<String,
roaring::RoaringTreemap>>,
+ ) -> crate::Result<Vec<Vec<PkVectorSearchResult>>> {
+ // The bytes handed to the scorer must be exactly this segment's
loaded
+ // bytes (its path), proving load→score threads the right payload.
+ assert_eq!(segment_bytes.as_ref(), segment.path.as_bytes());
+ Ok(queries.iter().map(|_| Vec::new()).collect())
+ }
+ }
+
+ #[tokio::test]
+ async fn ann_segment_bytes_are_loaded_lazily_per_segment_via_the_seam() {
+ // Two ANN segments; the searcher's `load_segment` is invoked once per
segment
+ // during the bucket search (lazily, in each leaf), and the bytes it
returns
+ // are the ones passed to that segment's scorer. This is the fused
+ // load→score path (no up-front all-segments preload map).
+ let loaded = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
+ let searcher: Arc<dyn PkVectorAnnSearcher> =
Arc::new(RecordingLoaderAnn {
+ loaded: loaded.clone(),
+ });
+ let seg_a = BucketAnnSegment {
+ source_meta: meta(&[("cov-a", 2)]),
+ path: "seg-a".to_string(),
+ file_size: 0,
+ index_meta: Vec::new(),
+ };
+ let seg_b = BucketAnnSegment {
+ source_meta: meta(&[("cov-b", 2)]),
+ path: "seg-b".to_string(),
+ file_size: 0,
+ index_meta: Vec::new(),
+ };
+ let factory = unreachable_search();
+ bucket_search(
+ Some(searcher),
+ &[seg_a, seg_b],
+ &[active("cov-a", 2), active("cov-b", 2)],
+ &HashMap::new(),
+ &factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 8,
+ 8,
+ &HashMap::new(),
+ false,
+ None,
+ 4,
+ Some(SearchBudget::per_query_only(4)),
+ )
+ .await
+ .unwrap();
+ let mut got = loaded.lock().unwrap().clone();
+ got.sort();
+ assert_eq!(
+ got,
+ vec!["seg-a".to_string(), "seg-b".to_string()],
+ "each ANN segment must be loaded exactly once via load_segment
during the search"
+ );
+ }
}
diff --git a/crates/paimon/src/vindex/reader.rs
b/crates/paimon/src/vindex/reader.rs
index 3dc9dda1..21603f17 100644
--- a/crates/paimon/src/vindex/reader.rs
+++ b/crates/paimon/src/vindex/reader.rs
@@ -125,6 +125,18 @@ impl VindexVectorGlobalIndexReader {
self.ensure_loaded(stream_fn, |_| Ok(()))
}
+ pub(crate) fn load_validated<S, F>(
+ &mut self,
+ stream_fn: impl FnOnce(&str) -> crate::Result<S>,
+ validate: F,
+ ) -> crate::Result<()>
+ where
+ S: SeekRead + 'static,
+ F: FnOnce(&VectorIndexMetadata) -> crate::Result<()>,
+ {
+ self.ensure_loaded(stream_fn, validate)
+ }
+
pub(crate) fn metadata(&self) -> crate::Result<&VectorIndexMetadata> {
self.metadata
.as_ref()
