From cb84b15863bb061073cefe19f7187d2c9a55fc3c Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 14 Aug 2026 10:10:44 +0800 Subject: [PATCH 1/4] perf(vindex): decouple native batch chunk concurrency --- .../paimon/src/table/vector_search_builder.rs | 39 ++++++++++-- crates/paimon/src/vindex/range_reader.rs | 36 +++++++++++ crates/paimon/src/vindex/reader.rs | 59 ++++++++++--------- 3 files changed, 101 insertions(+), 33 deletions(-) diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index 1295db93..f5825309 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -106,7 +106,7 @@ async fn execute_vindex_searches( vector_searches: Vec, source: S, file_name: String, - shard_concurrency: usize, + index_parallelism: usize, guard: G, ) -> crate::Result>>> { let panic_context = if vector_searches.len() > 1 { @@ -116,7 +116,7 @@ async fn execute_vindex_searches( }; execute_global_index_with_guard(panic_context, guard, move || { let mut reader = VindexVectorGlobalIndexReader::new(io_meta, options) - .with_batch_shard_concurrency(shard_concurrency); + .with_batch_index_parallelism(index_parallelism); reader .visit_batch_vector_search(&vector_searches, |_| Ok(source)) .map_err(|e| crate::Error::DataInvalid { @@ -134,6 +134,10 @@ fn current_tokio_runtime_handle() -> crate::Result { }) } +fn vindex_index_parallelism(entry_count: usize, max_concurrency: usize) -> usize { + entry_count.min(max_concurrency).max(1) +} + pub struct VectorSearchBuilder<'a> { table: &'a Table, vector_column: Option, @@ -821,6 +825,16 @@ async fn plan_and_search_pk_candidates_batch( source: None, } })?; + let batch_index_parallelism = match backend { + VectorIndexBackend::Vindex => vindex_index_parallelism( + plan.splits + .iter() + .map(|split| split.ann_segments.len()) + .sum(), + concurrency, + ), + VectorIndexBackend::Lumina => 1, + }; // Production data-file reader, mirroring `table_read.rs::new_data_file_reader` // but projecting only the vector column with no predicates. @@ -906,7 +920,7 @@ async fn plan_and_search_pk_candidates_batch( } (VectorIndexBackend::Vindex, AnnSegmentSource::Vindex(source)) => { let mut reader = VindexVectorGlobalIndexReader::new(io_meta, options.clone()) - .with_batch_shard_concurrency(concurrency); + .with_batch_index_parallelism(batch_index_parallelism); reader.load_validated( |_| Ok(source), |metadata| { @@ -1523,6 +1537,13 @@ async fn evaluate_batch_vector_search( } ensure_global_index_executor_capacity(concurrency); let range_read_permits = Arc::new(tokio::sync::Semaphore::new(concurrency)); + let batch_index_parallelism = vindex_index_parallelism( + vector_entries + .iter() + .filter(|entry| is_vindex_index_type(&entry.index_file.index_type)) + .count(), + concurrency, + ); let futures: Vec<_> = vector_entries .into_iter() .map(|entry| { @@ -1607,7 +1628,7 @@ async fn evaluate_batch_vector_search( vector_searches, source, file_name, - concurrency, + batch_index_parallelism, permit, ) .await? @@ -1628,7 +1649,7 @@ async fn evaluate_batch_vector_search( vector_searches, Cursor::new(data), file_name, - concurrency, + batch_index_parallelism, permit, ) .await? @@ -3084,6 +3105,14 @@ mod tests { VectorSearchMetric::L2.distance_to_score(distance) } + #[test] + fn vindex_batch_parallelism_tracks_active_entries() { + assert_eq!(vindex_index_parallelism(1, 1), 1); + assert_eq!(vindex_index_parallelism(1, 64), 1); + assert_eq!(vindex_index_parallelism(8, 4), 4); + assert_eq!(vindex_index_parallelism(4, 8), 4); + } + fn make_field(id: i32, name: &str) -> DataField { DataField::new(id, name.to_string(), DataType::Int(IntType::default())) } diff --git a/crates/paimon/src/vindex/range_reader.rs b/crates/paimon/src/vindex/range_reader.rs index 07431749..f9654aa7 100644 --- a/crates/paimon/src/vindex/range_reader.rs +++ b/crates/paimon/src/vindex/range_reader.rs @@ -633,6 +633,42 @@ mod tests { assert_eq!(tracking.max_active.load(Ordering::SeqCst), 1); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn read_many_caps_each_batch_at_32_ranges() { + let range_count = RANGE_READ_CONCURRENCY + 1; + let stride = RANGE_COALESCE_GAP + 1; + let data = Bytes::from(vec![8u8; range_count * stride as usize]); + let tracking = Arc::new(ConcurrencyTrackingRead { + data: data.clone(), + active: AtomicUsize::new(0), + max_active: AtomicUsize::new(0), + }); + let source: Arc = tracking.clone(); + let mut reader = VindexFileReader::new( + source, + tokio::runtime::Handle::current(), + data.len() as u64, + "index".to_string(), + ); + + tokio::task::spawn_blocking(move || { + let mut buffers = vec![[0u8; 1]; range_count]; + let mut requests = buffers + .iter_mut() + .enumerate() + .map(|(index, buffer)| ReadRequest::new(index as u64 * stride, buffer)) + .collect::>(); + reader.pread(&mut requests).unwrap(); + }) + .await + .unwrap(); + + assert_eq!( + tracking.max_active.load(Ordering::SeqCst), + RANGE_READ_CONCURRENCY + ); + } + #[test] fn local_fs_read_completes_with_one_host_blocking_thread() { let temp_dir = tempfile::tempdir().unwrap(); diff --git a/crates/paimon/src/vindex/reader.rs b/crates/paimon/src/vindex/reader.rs index 21603f17..74166174 100644 --- a/crates/paimon/src/vindex/reader.rs +++ b/crates/paimon/src/vindex/reader.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use crate::spec::CoreOptions; use crate::vector_search::{GlobalIndexIOMeta, VectorSearch}; use paimon_vindex_core::distance::MetricType; use paimon_vindex_core::index::{ @@ -78,7 +77,7 @@ impl SeekRead for VindexInput { pub struct VindexVectorGlobalIndexReader { io_meta: GlobalIndexIOMeta, options: HashMap, - batch_shard_concurrency: Option, + batch_index_parallelism: usize, reader: Option>, metadata: Option, } @@ -88,14 +87,14 @@ impl VindexVectorGlobalIndexReader { Self { io_meta, options, - batch_shard_concurrency: None, + batch_index_parallelism: 1, reader: None, metadata: None, } } - pub(crate) fn with_batch_shard_concurrency(mut self, concurrency: usize) -> Self { - self.batch_shard_concurrency = Some(concurrency.max(1)); + pub(crate) fn with_batch_index_parallelism(mut self, parallelism: usize) -> Self { + self.batch_index_parallelism = parallelism.max(1); self } @@ -150,10 +149,6 @@ impl VindexVectorGlobalIndexReader { &mut self, vector_searches: &[VectorSearch], ) -> crate::Result>>> { - let shard_concurrency = match self.batch_shard_concurrency { - Some(concurrency) => concurrency, - None => CoreOptions::new(&self.options).global_index_thread_num()?, - }; let reader = self .reader .as_mut() @@ -173,7 +168,7 @@ impl VindexVectorGlobalIndexReader { metadata, &self.options, vector_searches, - shard_concurrency, + self.batch_index_parallelism, ) } @@ -348,7 +343,7 @@ fn search_batch_vindex( metadata: &VectorIndexMetadata, options: &HashMap, vector_searches: &[VectorSearch], - shard_concurrency: usize, + index_parallelism: usize, ) -> crate::Result>>> { let mut results: Vec>> = (0..vector_searches.len()).map(|_| None).collect(); @@ -366,7 +361,7 @@ fn search_batch_vindex( } for (prepared, indices) in groups { - let chunk_size = native_batch_chunk_size(metadata, &prepared, shard_concurrency); + let chunk_size = native_batch_chunk_size(metadata, &prepared, index_parallelism); for indices in indices.chunks(chunk_size) { if indices.len() == 1 { let index = indices[0]; @@ -431,16 +426,16 @@ fn search_batch_vindex( fn native_batch_chunk_size( metadata: &VectorIndexMetadata, prepared: &PreparedSearch, - shard_concurrency: usize, + index_parallelism: usize, ) -> usize { - let per_shard_budget = NATIVE_BATCH_OPERATION_WORKING_SET_BYTES - .checked_div(shard_concurrency.max(1)) + let per_index_budget = NATIVE_BATCH_OPERATION_WORKING_SET_BYTES + .checked_div(index_parallelism.max(1)) .unwrap_or(0); let filter_bytes = prepared .filter_bytes .as_ref() .map_or(0, |filter| filter.len().saturating_mul(2)); - let query_budget = per_shard_budget.saturating_sub(filter_bytes); + let query_budget = per_index_budget.saturating_sub(filter_bytes); query_budget .checked_div(native_batch_query_working_set_bytes(metadata, prepared)) .unwrap_or(0) @@ -669,7 +664,7 @@ mod tests { index, query_count, HashMap::from([(NPROBE_PARAMETER.to_string(), "1".to_string())]), - 32, + 1, ) .await } @@ -678,7 +673,7 @@ mod tests { index: Bytes, query_count: usize, options: HashMap, - shard_concurrency: usize, + index_parallelism: usize, ) -> (Vec>>, usize) { let tracking = TrackingIndexRead::new(index.clone()); let source: Arc = tracking.clone(); @@ -694,7 +689,7 @@ mod tests { GlobalIndexIOMeta::new("batch.index".to_string(), index.len() as u64, Vec::new()); let searches = vec![query(); query_count]; let mut reader = VindexVectorGlobalIndexReader::new(io_meta, options) - .with_batch_shard_concurrency(shard_concurrency); + .with_batch_index_parallelism(index_parallelism); reader .visit_batch_vector_search(&searches, |_| Ok(source)) .unwrap() @@ -760,23 +755,31 @@ mod tests { nprobe: 16, filter_bytes: None, }; - let base = native_batch_chunk_size(&base_metadata, &base_prepared, 32); + let index_parallelism = 32; + let base = native_batch_chunk_size(&base_metadata, &base_prepared, index_parallelism); let mut larger_index = base_metadata.clone(); larger_index.dimension *= 2; larger_index.nlist *= 2; - assert!(native_batch_chunk_size(&larger_index, &base_prepared, 32) < base); + assert!(native_batch_chunk_size(&larger_index, &base_prepared, index_parallelism) < base); let mut larger_top_k = base_prepared.clone(); larger_top_k.top_k *= 4; - assert!(native_batch_chunk_size(&base_metadata, &larger_top_k, 32) < base); + assert!(native_batch_chunk_size(&base_metadata, &larger_top_k, index_parallelism) < base); let mut pq_metadata = base_metadata.clone(); pq_metadata.pq_m = Some(64); pq_metadata.pq_bits = Some(8); - assert!(native_batch_chunk_size(&pq_metadata, &base_prepared, 32) < base); + assert!(native_batch_chunk_size(&pq_metadata, &base_prepared, index_parallelism) < base); - assert!(native_batch_chunk_size(&base_metadata, &base_prepared, 64) < base); + let per_index_working_set = base.saturating_mul(native_batch_query_working_set_bytes( + &base_metadata, + &base_prepared, + )); + assert!( + per_index_working_set.saturating_mul(index_parallelism) + <= NATIVE_BATCH_OPERATION_WORKING_SET_BYTES + ); } #[test] @@ -1094,7 +1097,7 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn homogeneous_batch_chunks_at_working_set_boundary() { - let shard_concurrency = 4096; + let index_parallelism = 4096; let metadata = VectorIndexMetadata { index_type: paimon_vindex_core::index::IndexType::IvfFlat, dimension: TEST_DIMENSION, @@ -1113,7 +1116,7 @@ mod tests { let prepared = prepare_search(&metadata, &options, &query()) .unwrap() .unwrap(); - let chunk_size = native_batch_chunk_size(&metadata, &prepared, shard_concurrency); + let chunk_size = native_batch_chunk_size(&metadata, &prepared, index_parallelism); assert!(chunk_size > 16); let index = build_ivf_flat_index(); @@ -1121,11 +1124,11 @@ mod tests { index.clone(), chunk_size, options.clone(), - shard_concurrency, + index_parallelism, ) .await; let (over_results, over_bytes) = - tracked_batch_search_with_options(index, chunk_size + 1, options, shard_concurrency) + tracked_batch_search_with_options(index, chunk_size + 1, options, index_parallelism) .await; assert_eq!(within_results.len(), chunk_size); From 202cb019df5194420d038d035ed9d1cab851e00b Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 14 Aug 2026 11:42:02 +0800 Subject: [PATCH 2/4] fix(vindex): bound native batch memory process-wide --- crates/paimon/src/vindex/range_reader.rs | 2 +- crates/paimon/src/vindex/reader.rs | 54 +++++++++++++++++++++--- 2 files changed, 50 insertions(+), 6 deletions(-) diff --git a/crates/paimon/src/vindex/range_reader.rs b/crates/paimon/src/vindex/range_reader.rs index f9654aa7..d9c5e158 100644 --- a/crates/paimon/src/vindex/range_reader.rs +++ b/crates/paimon/src/vindex/range_reader.rs @@ -636,7 +636,7 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn read_many_caps_each_batch_at_32_ranges() { let range_count = RANGE_READ_CONCURRENCY + 1; - let stride = RANGE_COALESCE_GAP + 1; + let stride = RANGE_COALESCE_GAP + 2; let data = Bytes::from(vec![8u8; range_count * stride as usize]); let tracking = Arc::new(ConcurrencyTrackingRead { data: data.clone(), diff --git a/crates/paimon/src/vindex/reader.rs b/crates/paimon/src/vindex/reader.rs index 74166174..5ac6fb44 100644 --- a/crates/paimon/src/vindex/reader.rs +++ b/crates/paimon/src/vindex/reader.rs @@ -24,10 +24,48 @@ use paimon_vindex_core::io::{ReadRequest, SeekRead, SeekReadCapabilities}; use std::collections::BinaryHeap; use std::collections::HashMap; use std::io; +use std::sync::{Condvar, Mutex}; const DEFAULT_NPROBE: usize = 16; const NPROBE_PARAMETER: &str = "ivf.nprobe"; -const NATIVE_BATCH_OPERATION_WORKING_SET_BYTES: usize = 64 * 1024 * 1024; +const NATIVE_BATCH_PROCESS_WORKING_SET_BYTES: usize = 64 * 1024 * 1024; +// Native searches run on dedicated executor threads, so blocking here does not block async I/O. +static NATIVE_BATCH_AVAILABLE_BYTES: Mutex = + Mutex::new(NATIVE_BATCH_PROCESS_WORKING_SET_BYTES); +static NATIVE_BATCH_MEMORY_AVAILABLE: Condvar = Condvar::new(); + +struct NativeBatchMemoryPermit { + bytes: usize, +} + +impl Drop for NativeBatchMemoryPermit { + fn drop(&mut self) { + let mut available = NATIVE_BATCH_AVAILABLE_BYTES + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + *available += self.bytes; + drop(available); + NATIVE_BATCH_MEMORY_AVAILABLE.notify_all(); + } +} + +fn acquire_native_batch_memory(index_parallelism: usize) -> NativeBatchMemoryPermit { + let bytes = native_batch_memory_reservation(index_parallelism); + let mut available = NATIVE_BATCH_AVAILABLE_BYTES + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + while *available < bytes { + available = NATIVE_BATCH_MEMORY_AVAILABLE + .wait(available) + .unwrap_or_else(|poisoned| poisoned.into_inner()); + } + *available -= bytes; + NativeBatchMemoryPermit { bytes } +} + +fn native_batch_memory_reservation(index_parallelism: usize) -> usize { + NATIVE_BATCH_PROCESS_WORKING_SET_BYTES / index_parallelism.max(1) +} trait ErasedSeekRead: Send { fn pread_erased(&mut self, ranges: &mut [ReadRequest<'_>]) -> io::Result<()>; @@ -362,6 +400,8 @@ fn search_batch_vindex( for (prepared, indices) in groups { let chunk_size = native_batch_chunk_size(metadata, &prepared, index_parallelism); + let _memory_permit = (chunk_size > 1 && indices.len() > 1) + .then(|| acquire_native_batch_memory(index_parallelism)); for indices in indices.chunks(chunk_size) { if indices.len() == 1 { let index = indices[0]; @@ -428,9 +468,7 @@ fn native_batch_chunk_size( prepared: &PreparedSearch, index_parallelism: usize, ) -> usize { - let per_index_budget = NATIVE_BATCH_OPERATION_WORKING_SET_BYTES - .checked_div(index_parallelism.max(1)) - .unwrap_or(0); + let per_index_budget = native_batch_memory_reservation(index_parallelism); let filter_bytes = prepared .filter_bytes .as_ref() @@ -778,8 +816,14 @@ mod tests { )); assert!( per_index_working_set.saturating_mul(index_parallelism) - <= NATIVE_BATCH_OPERATION_WORKING_SET_BYTES + <= NATIVE_BATCH_PROCESS_WORKING_SET_BYTES ); + for parallelism in [1, 2, 3, 32, 64] { + assert!( + native_batch_memory_reservation(parallelism).saturating_mul(parallelism) + <= NATIVE_BATCH_PROCESS_WORKING_SET_BYTES + ); + } } #[test] From 145331545ae8dbe0988587404ef3cc4abe442083 Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 14 Aug 2026 13:22:30 +0800 Subject: [PATCH 3/4] perf(vindex): reserve native batch memory per chunk --- crates/paimon/src/vindex/reader.rs | 148 ++++++++++++++++++++++++----- 1 file changed, 123 insertions(+), 25 deletions(-) diff --git a/crates/paimon/src/vindex/reader.rs b/crates/paimon/src/vindex/reader.rs index 5ac6fb44..41e163ce 100644 --- a/crates/paimon/src/vindex/reader.rs +++ b/crates/paimon/src/vindex/reader.rs @@ -30,37 +30,58 @@ const DEFAULT_NPROBE: usize = 16; const NPROBE_PARAMETER: &str = "ivf.nprobe"; const NATIVE_BATCH_PROCESS_WORKING_SET_BYTES: usize = 64 * 1024 * 1024; // Native searches run on dedicated executor threads, so blocking here does not block async I/O. -static NATIVE_BATCH_AVAILABLE_BYTES: Mutex = - Mutex::new(NATIVE_BATCH_PROCESS_WORKING_SET_BYTES); -static NATIVE_BATCH_MEMORY_AVAILABLE: Condvar = Condvar::new(); +static NATIVE_BATCH_MEMORY_POOL: NativeBatchMemoryPool = + NativeBatchMemoryPool::new(NATIVE_BATCH_PROCESS_WORKING_SET_BYTES); -struct NativeBatchMemoryPermit { +struct NativeBatchMemoryPool { + available_bytes: Mutex, + memory_available: Condvar, +} + +impl NativeBatchMemoryPool { + const fn new(bytes: usize) -> Self { + Self { + available_bytes: Mutex::new(bytes), + memory_available: Condvar::new(), + } + } + + fn acquire(&self, bytes: usize) -> NativeBatchMemoryPermit<'_> { + let mut available = self + .available_bytes + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + while *available < bytes { + available = self + .memory_available + .wait(available) + .unwrap_or_else(|poisoned| poisoned.into_inner()); + } + *available -= bytes; + NativeBatchMemoryPermit { pool: self, bytes } + } +} + +struct NativeBatchMemoryPermit<'a> { + pool: &'a NativeBatchMemoryPool, bytes: usize, } -impl Drop for NativeBatchMemoryPermit { +impl Drop for NativeBatchMemoryPermit<'_> { fn drop(&mut self) { - let mut available = NATIVE_BATCH_AVAILABLE_BYTES + let mut available = self + .pool + .available_bytes .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); *available += self.bytes; drop(available); - NATIVE_BATCH_MEMORY_AVAILABLE.notify_all(); + self.pool.memory_available.notify_all(); } } -fn acquire_native_batch_memory(index_parallelism: usize) -> NativeBatchMemoryPermit { - let bytes = native_batch_memory_reservation(index_parallelism); - let mut available = NATIVE_BATCH_AVAILABLE_BYTES - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - while *available < bytes { - available = NATIVE_BATCH_MEMORY_AVAILABLE - .wait(available) - .unwrap_or_else(|poisoned| poisoned.into_inner()); - } - *available -= bytes; - NativeBatchMemoryPermit { bytes } +fn acquire_native_batch_memory(bytes: usize) -> NativeBatchMemoryPermit<'static> { + NATIVE_BATCH_MEMORY_POOL.acquire(bytes) } fn native_batch_memory_reservation(index_parallelism: usize) -> usize { @@ -400,8 +421,6 @@ fn search_batch_vindex( for (prepared, indices) in groups { let chunk_size = native_batch_chunk_size(metadata, &prepared, index_parallelism); - let _memory_permit = (chunk_size > 1 && indices.len() > 1) - .then(|| acquire_native_batch_memory(index_parallelism)); for indices in indices.chunks(chunk_size) { if indices.len() == 1 { let index = indices[0]; @@ -414,6 +433,10 @@ fn search_batch_vindex( continue; } + let reservation = + native_batch_chunk_working_set_bytes(metadata, &prepared, indices.len()); + debug_assert!(reservation <= native_batch_memory_reservation(index_parallelism)); + let _memory_permit = acquire_native_batch_memory(reservation); let mut queries = Vec::with_capacity(indices.len() * metadata.dimension); for &index in indices { queries.extend_from_slice(&vector_searches[index].vector); @@ -469,10 +492,7 @@ fn native_batch_chunk_size( index_parallelism: usize, ) -> usize { let per_index_budget = native_batch_memory_reservation(index_parallelism); - let filter_bytes = prepared - .filter_bytes - .as_ref() - .map_or(0, |filter| filter.len().saturating_mul(2)); + let filter_bytes = native_batch_filter_working_set_bytes(prepared); let query_budget = per_index_budget.saturating_sub(filter_bytes); query_budget .checked_div(native_batch_query_working_set_bytes(metadata, prepared)) @@ -480,6 +500,23 @@ fn native_batch_chunk_size( .max(1) } +fn native_batch_chunk_working_set_bytes( + metadata: &VectorIndexMetadata, + prepared: &PreparedSearch, + query_count: usize, +) -> usize { + native_batch_filter_working_set_bytes(prepared).saturating_add( + query_count.saturating_mul(native_batch_query_working_set_bytes(metadata, prepared)), + ) +} + +fn native_batch_filter_working_set_bytes(prepared: &PreparedSearch) -> usize { + prepared + .filter_bytes + .as_ref() + .map_or(0, |filter| filter.len().saturating_mul(2)) +} + fn native_batch_query_working_set_bytes( metadata: &VectorIndexMetadata, prepared: &PreparedSearch, @@ -826,6 +863,67 @@ mod tests { } } + #[test] + fn native_batch_chunk_reservation_tracks_actual_chunk() { + let metadata = VectorIndexMetadata { + index_type: paimon_vindex_core::index::IndexType::IvfFlat, + dimension: 128, + nlist: 256, + metric: MetricType::L2, + total_vectors: 8192, + pq_m: None, + pq_bits: None, + rq_bits: None, + diskann: None, + }; + let prepared = PreparedSearch { + top_k: 10, + nprobe: 16, + filter_bytes: Some(vec![0; 128]), + }; + let chunk_size = native_batch_chunk_size(&metadata, &prepared, 1); + let full_chunk = native_batch_chunk_working_set_bytes(&metadata, &prepared, chunk_size); + let final_chunk = native_batch_chunk_working_set_bytes(&metadata, &prepared, 2); + + assert!(chunk_size > 2); + assert!(full_chunk <= native_batch_memory_reservation(1)); + assert!(final_chunk < full_chunk); + } + + #[test] + fn native_batch_memory_pool_admits_only_available_bytes() { + let pool = NativeBatchMemoryPool::new(64); + let large = pool.acquire(48); + + std::thread::scope(|scope| { + let (fits_tx, fits_rx) = std::sync::mpsc::channel(); + scope.spawn(|| { + let _permit = pool.acquire(16); + fits_tx.send(()).unwrap(); + }); + fits_rx + .recv_timeout(std::time::Duration::from_secs(1)) + .expect("reservation fitting the available bytes should not wait"); + + let (blocked_tx, blocked_rx) = std::sync::mpsc::channel(); + scope.spawn(|| { + let _permit = pool.acquire(17); + blocked_tx.send(()).unwrap(); + }); + assert!( + blocked_rx + .recv_timeout(std::time::Duration::from_millis(50)) + .is_err(), + "reservation exceeding the available bytes should wait" + ); + + drop(large); + blocked_rx + .recv_timeout(std::time::Duration::from_secs(1)) + .expect("waiting reservation should proceed after bytes are released"); + }); + } + #[test] fn test_int_parameter() { let mut options = HashMap::new(); From 6d26ae36eb7f42e5c22321a12bba5a1a90289d40 Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 14 Aug 2026 13:36:22 +0800 Subject: [PATCH 4/4] test(vindex): move senders into scoped threads --- crates/paimon/src/vindex/reader.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/crates/paimon/src/vindex/reader.rs b/crates/paimon/src/vindex/reader.rs index 41e163ce..2a45902c 100644 --- a/crates/paimon/src/vindex/reader.rs +++ b/crates/paimon/src/vindex/reader.rs @@ -896,8 +896,9 @@ mod tests { let large = pool.acquire(48); std::thread::scope(|scope| { + let pool = &pool; let (fits_tx, fits_rx) = std::sync::mpsc::channel(); - scope.spawn(|| { + scope.spawn(move || { let _permit = pool.acquire(16); fits_tx.send(()).unwrap(); }); @@ -906,7 +907,7 @@ mod tests { .expect("reservation fitting the available bytes should not wait"); let (blocked_tx, blocked_rx) = std::sync::mpsc::channel(); - scope.spawn(|| { + scope.spawn(move || { let _permit = pool.acquire(17); blocked_tx.send(()).unwrap(); });