From b2dd1fd57ea170c46abe2e66659f2ac3b86a0e41 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Sat, 1 Aug 2026 12:49:13 +0800 Subject: [PATCH] perf(vector-search): execute PK-vector leaves concurrently --- crates/paimon/src/spec/core_options.rs | 40 + .../src/table/pk_vector_orchestrator.rs | 60 +- .../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 e33cd09b6..1039800d1 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 2278b0d25..1a816bd9a 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>, 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>, exact_file_search: &(dyn for<'s, 'a> Fn( usize, &'s PkVectorSearchSplit, @@ -497,14 +490,11 @@ impl PkVectorOrchestrator { let mut exact_candidates: Vec> = (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::>() - .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, } impl PkVectorAnnSearcher for FakeAnn { + fn load_segment( + &self, + _segment: &BucketAnnSegment, + ) -> futures::future::BoxFuture<'static, crate::Result> { + 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>, 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 b2d675973..f1e22b20e 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( +async fn execute_vindex_searches( io_meta: GlobalIndexIOMeta, options: HashMap, vector_searches: Vec, source: S, file_name: String, shard_concurrency: usize, + guard: G, ) -> crate::Result>>> { 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( .await } +fn current_tokio_runtime_handle() -> crate::Result { + 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, @@ -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 = 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>> { - 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>, +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 76b8a185d..173b5d7d8 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; @@ -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 = OnceLock::new(); +static PROCESS_GLOBAL_CAPACITY: AtomicUsize = AtomicUsize::new(0); + struct ExecutorState { receiver: Receiver, 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> { + 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( + panic_context: &'static str, + guard: G, + task: F, +) -> crate::Result +where + T: Send + 'static, + F: FnOnce() -> crate::Result + Send + 'static, + G: Send + 'static, +{ + execute_global_index(panic_context, move || { + let _guard = guard; + task() + }) + .await +} + async fn execute_on( 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 ab50196a9..76cf6aa49 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, deletion_vectors: &HashMap>, residual_ranges: Option<&HashMap>, @@ -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, deletion_vectors: &HashMap>, residual_ranges: Option<&HashMap>, @@ -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` 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>; + + /// 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> { + 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>, ) -> crate::Result>>; + #[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, + deletion_vectors: &HashMap>, + search_options: &HashMap, + residual_ranges: Option<&HashMap>, + ) -> crate::Result>> { + 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> { 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, + deletion_vectors: &HashMap>, + search_options: &HashMap, + residual_ranges: Option<&HashMap>, + ) -> crate::Result> { + 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>>> + dyn Fn( + &BucketAnnSegment, + Bytes, + &[VectorSearch], + ) -> crate::Result>>> + + Send + + Sync, +>; + +pub(crate) type SourceBatchScorer = Box< + dyn Fn( + &BucketAnnSegment, + AnnSegmentSource, + &[VectorSearch], + ) -> crate::Result>>> + 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 BoxFuture<'static, crate::Result> + 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> + 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> { + 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> { + (self.loader)(segment) + } + fn search_batch( &self, segment: &BucketAnnSegment, + segment_bytes: Bytes, + queries: &[&[f32]], + metric: VectorSearchMetric, + limit: usize, + active_source_files: &HashSet, + deletion_vectors: &HashMap>, + search_options: &HashMap, + residual_ranges: Option<&HashMap>, + ) -> crate::Result>> { + 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 { @@ -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>>> = 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 4666d4037..c52a72e59 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, @@ -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>, + /// 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, +} + +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` + /// 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) -> 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 { + 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, + _process_global: Option>, +} + +async fn acquire_owned(sem: Arc) -> crate::Result { + 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>, -) -> crate::Result> { + budget: &Option, +) -> crate::Result> { 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, } +/// 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>), + Exact(Vec>), +} + /// 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>, ann_segments: &[BucketAnnSegment], active_files: &[BucketActiveFile], deletion_vectors: &HashMap>, @@ -281,7 +383,7 @@ pub(crate) async fn bucket_search( skip_exact_fallback: bool, residual_ranges: Option<&HashMap>, concurrency: usize, - search_budget: Option>, + search_budget: Option, ) -> crate::Result { 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 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 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> = 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>> = 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::>() - .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>, ann_segments: &[BucketAnnSegment], active_files: &[BucketActiveFile], deletion_vectors: &HashMap>, @@ -494,7 +660,7 @@ pub(crate) async fn bucket_search_batch( skip_exact_fallback: bool, residual_ranges: Option<&HashMap>, concurrency: usize, - search_budget: Option>, + search_budget: Option, ) -> crate::Result> { 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>> = + 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 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 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>> = 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>> = 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::>() - .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> { + 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, } impl PkVectorAnnSearcher for FakeAnnSearcher { + fn load_segment( + &self, + segment: &BucketAnnSegment, + ) -> futures::future::BoxFuture<'static, crate::Result> { + 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, + peak: std::sync::Arc, + } + impl PkVectorAnnSearcher for PeakProbeAnn { + fn load_segment( + &self, + segment: &BucketAnnSegment, + ) -> futures::future::BoxFuture<'static, crate::Result> { + empty_ann_loader(segment) + } + fn search_batch( + &self, + _segment: &BucketAnnSegment, + _segment_bytes: Bytes, + queries: &[&[f32]], + _metric: VectorSearchMetric, + _limit: usize, + _active_source_files: &HashSet, + _dvs: &HashMap>, + _opts: &HashMap, + _residual_ranges: Option<&HashMap>, + ) -> crate::Result>> { + 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, + Vec, + std::sync::Arc, + std::sync::Arc, + ) { + 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 = 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 = 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 = 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 = 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 = 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> { + empty_ann_loader(segment) + } + fn search_batch( + &self, + _segment: &BucketAnnSegment, + _segment_bytes: Bytes, + _queries: &[&[f32]], + _metric: VectorSearchMetric, + _limit: usize, + _active_source_files: &HashSet, + _dvs: &HashMap>, + _opts: &HashMap, + _residual_ranges: Option<&HashMap>, + ) -> crate::Result>> { + 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 = 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, + fail_files: HashSet, + slow_files: HashSet, + } + impl PkVectorAnnSearcher for DrainProbeAnn { + fn load_segment( + &self, + segment: &BucketAnnSegment, + ) -> futures::future::BoxFuture<'static, crate::Result> { + empty_ann_loader(segment) + } + fn search_batch( + &self, + segment: &BucketAnnSegment, + _segment_bytes: Bytes, + queries: &[&[f32]], + _metric: VectorSearchMetric, + _limit: usize, + _active_source_files: &HashSet, + _dvs: &HashMap>, + _opts: &HashMap, + _residual_ranges: Option<&HashMap>, + ) -> crate::Result>> { + 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 = 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, + peak: std::sync::Arc, + } + impl PkVectorAnnSearcher for SharedPeakAnn { + fn load_segment( + &self, + segment: &BucketAnnSegment, + ) -> futures::future::BoxFuture<'static, crate::Result> { + empty_ann_loader(segment) + } + fn search_batch( + &self, + _segment: &BucketAnnSegment, + _segment_bytes: Bytes, + queries: &[&[f32]], + _metric: VectorSearchMetric, + _limit: usize, + _active_source_files: &HashSet, + _dvs: &HashMap>, + _opts: &HashMap, + _residual_ranges: Option<&HashMap>, + ) -> crate::Result>> { + 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 = + 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>>, + overlapped: std::sync::Arc, + } + impl PkVectorAnnSearcher for HandshakeAnn { + fn load_segment( + &self, + segment: &BucketAnnSegment, + ) -> futures::future::BoxFuture<'static, crate::Result> { + empty_ann_loader(segment) + } + fn search_batch( + &self, + _segment: &BucketAnnSegment, + _segment_bytes: Bytes, + queries: &[&[f32]], + _metric: VectorSearchMetric, + _limit: usize, + _active_source_files: &HashSet, + _dvs: &HashMap>, + _opts: &HashMap, + _residual_ranges: Option<&HashMap>, + ) -> crate::Result>> { + // 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 = 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>>, + } + impl PkVectorAnnSearcher for RecordingLoaderAnn { + fn load_segment( + &self, + segment: &BucketAnnSegment, + ) -> futures::future::BoxFuture<'static, crate::Result> { + 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, + _dvs: &HashMap>, + _opts: &HashMap, + _residual_ranges: Option<&HashMap>, + ) -> crate::Result>> { + // 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 = 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 3dc9dda1b..21603f17f 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( + &mut self, + stream_fn: impl FnOnce(&str) -> crate::Result, + 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()