diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index 281f77bb..1b216598 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -65,6 +65,7 @@ use arrow_select::interleave::interleave_record_batch; use futures::{stream, TryStreamExt}; use paimon_vindex_core::distance::MetricType; use paimon_vindex_core::index::VectorIndexReader as VIndexReader; +use paimon_vindex_core::io::SeekRead; use roaring::RoaringTreemap; use std::cmp::Ordering; use std::collections::{BinaryHeap, HashMap, HashSet}; @@ -98,11 +99,30 @@ impl VectorIndexBackend { } } -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)), +async fn execute_vindex_searches( + io_meta: GlobalIndexIOMeta, + options: HashMap, + vector_searches: Vec, + source: S, + file_name: String, + shard_concurrency: usize, +) -> 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 || { + let mut reader = VindexVectorGlobalIndexReader::new(io_meta, options) + .with_batch_shard_concurrency(shard_concurrency); + reader + .visit_batch_vector_search(&vector_searches, |_| Ok(source)) + .map_err(|e| crate::Error::DataInvalid { + message: format!("Failed to read vindex index file '{}': {}", file_name, e), + source: Some(Box::new(e)), + }) }) + .await } pub struct VectorSearchBuilder<'a> { @@ -839,7 +859,8 @@ async fn plan_and_search_pk_candidates_batch( reader.visit_batch_vector_search(searches, |_| Ok(Cursor::new(data))) } VectorIndexBackend::Vindex => { - let mut reader = VindexVectorGlobalIndexReader::new(io_meta, options.clone()); + let mut reader = VindexVectorGlobalIndexReader::new(io_meta, options.clone()) + .with_batch_shard_concurrency(concurrency); reader.visit_batch_vector_search(searches, |_| Ok(Cursor::new(data))) } } @@ -1424,10 +1445,21 @@ async fn evaluate_batch_vector_search( let mut merged = vec![SearchResult::empty(); vector_searches.len()]; if !vector_entries.is_empty() { let concurrency = core_options.global_index_thread_num()?; + if concurrency > tokio::sync::Semaphore::MAX_PERMITS { + return Err(crate::Error::DataInvalid { + message: format!( + "Global index thread count must not exceed {}", + tokio::sync::Semaphore::MAX_PERMITS + ), + source: None, + }); + } ensure_global_index_executor_capacity(concurrency); + let range_read_permits = Arc::new(tokio::sync::Semaphore::new(concurrency)); let futures: Vec<_> = vector_entries .into_iter() .map(|entry| { + let range_read_permits = Arc::clone(&range_read_permits); let global_meta = entry.index_file.global_index_meta.as_ref().unwrap(); let backend = VectorIndexBackend::from_index_type(&entry.index_file.index_type) .expect("filtered vector index type"); @@ -1482,64 +1514,62 @@ async fn evaluate_batch_vector_search( .await? } VectorIndexBackend::Vindex => { - if vector_searches.len() > 1 { - let data = input.read().await.map_err(|e| { - crate::Error::DataInvalid { - message: format!( - "Failed to read vindex index file '{}': {}", - file_name, e - ), - source: None, - } - })?; - execute_global_index( - "vindex global-index batch search task failed", - move || { - let mut reader = VindexVectorGlobalIndexReader::new( - io_meta, options, - ); - reader.visit_batch_vector_search(&vector_searches, |_| { - Ok(Cursor::new(data)) - }) - }, - ) - .await? - } else { - let file_reader = input.reader().await.map_err(|e| { - crate::Error::DataInvalid { - message: format!( - "Failed to open vindex file '{}' for range reads: {}", - file_name, e - ), - source: None, - } - })?; - let source = VindexFileReader::new( - Arc::new(file_reader), - current_tokio_runtime_handle()?, - file_size, - file_name.clone(), - ); - execute_global_index( - "vindex global-index search task failed", - move || { - let mut reader = VindexVectorGlobalIndexReader::new( - io_meta, options, - ); - reader - .visit_batch_vector_search(&vector_searches, |_| { - Ok(source) - }) - .map_err(|e| crate::Error::DataInvalid { - message: format!( - "Failed to read vindex index file '{}': {}", - file_name, e - ), - source: Some(Box::new(e)), - }) - }, - ) - .await? + match tokio::runtime::Handle::try_current() { + Ok(runtime) => { + let file_reader = input.reader().await.map_err(|e| { + crate::Error::DataInvalid { + message: format!( + "Failed to open vindex file '{}' for range reads: {}", + file_name, e + ), + source: None, + } + })?; + let source = VindexFileReader::new_with_permits( + Arc::new(file_reader), + runtime, + range_read_permits, + file_size, + file_name.clone(), + ); + execute_vindex_searches( + io_meta, + options, + vector_searches, + source, + file_name, + concurrency, + ) + .await? + } + Err(_) if query_count > 1 => { + let data = input.read().await.map_err(|e| { + crate::Error::DataInvalid { + message: format!( + "Failed to read vindex index file '{}': {}", + file_name, e + ), + source: None, + } + })?; + execute_vindex_searches( + io_meta, + options, + vector_searches, + Cursor::new(data), + file_name, + concurrency, + ) + .await? + } + Err(error) => { + return Err(crate::Error::UnexpectedError { + message: + "Vector index range reader requires a Tokio runtime" + .to_string(), + source: Some(Box::new(error)), + }); + } } } }; @@ -3627,6 +3657,46 @@ mod tests { }); } + #[test] + fn test_batch_vindex_outside_tokio_uses_buffered_fallback() { + futures::executor::block_on(async { + let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); + let index = build_vindex_segment_bytes("l2"); + file_io + .new_output("memory:///test_table/index/test.idx") + .unwrap() + .write(bytes::Bytes::from(index.clone())) + .await + .unwrap(); + let fields = vec![make_field(2, "embedding")]; + let searches = vec![ + VectorSearch::new(vec![1.0, 0.0], 2, "embedding".to_string()).unwrap(), + VectorSearch::new(vec![0.0, 1.0], 2, "embedding".to_string()).unwrap(), + ]; + let options = HashMap::new(); + let mut entry = make_lumina_entry("test.idx", IVF_FLAT_IDENTIFIER, FileKind::Add, 2); + entry.index_file.file_size = index.len() as i64; + entry.index_file.row_count = 3; + entry + .index_file + .global_index_meta + .as_mut() + .unwrap() + .row_range_end = 2; + + let results = evaluate_batch_vector_search( + eval_context(&file_io, &options, &fields, None), + &[entry], + &searches, + ) + .await + .expect("batch vindex search should fall back to buffered I/O outside Tokio"); + + assert_eq!(results.len(), searches.len()); + assert!(results.iter().all(|result| !result.is_empty())); + }); + } + #[tokio::test] async fn test_execute_fails_closed_when_query_auth_enabled() { let table = crate::table::query_auth_table(); diff --git a/crates/paimon/src/vindex/range_reader.rs b/crates/paimon/src/vindex/range_reader.rs index 1e66472c..07431749 100644 --- a/crates/paimon/src/vindex/range_reader.rs +++ b/crates/paimon/src/vindex/range_reader.rs @@ -67,16 +67,33 @@ pub(crate) struct VindexFileReader { } impl VindexFileReader { + #[cfg(test)] pub(crate) fn new( reader: Arc, runtime: tokio::runtime::Handle, file_size: u64, path: String, + ) -> Self { + Self::new_with_permits( + reader, + runtime, + Arc::new(tokio::sync::Semaphore::new(RANGE_READ_CONCURRENCY)), + file_size, + path, + ) + } + + pub(crate) fn new_with_permits( + reader: Arc, + runtime: tokio::runtime::Handle, + permits: Arc, + file_size: u64, + path: String, ) -> Self { Self { reader, runtime, - permits: Arc::new(tokio::sync::Semaphore::new(RANGE_READ_CONCURRENCY)), + permits, file_size, path, scalar_cache: None, @@ -288,6 +305,7 @@ mod tests { use super::*; use crate::io::FileIO; use async_trait::async_trait; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Mutex; use std::time::Duration; @@ -324,6 +342,12 @@ mod tests { runtime_id: Mutex>, } + struct ConcurrencyTrackingRead { + data: Bytes, + active: AtomicUsize, + max_active: AtomicUsize, + } + #[async_trait] impl FileRead for RuntimeTrackingRead { async fn read(&self, range: Range) -> crate::Result { @@ -332,6 +356,17 @@ mod tests { } } + #[async_trait] + impl FileRead for ConcurrencyTrackingRead { + async fn read(&self, range: Range) -> crate::Result { + let active = self.active.fetch_add(1, Ordering::SeqCst) + 1; + self.max_active.fetch_max(active, Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(25)).await; + self.active.fetch_sub(1, Ordering::SeqCst); + Ok(self.data.slice(range.start as usize..range.end as usize)) + } + } + #[async_trait] impl FileRead for TrackingRead { async fn read(&self, range: Range) -> crate::Result { @@ -557,6 +592,47 @@ mod tests { assert_eq!(cloned.read_capabilities(), SeekReadCapabilities::default()); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn shared_permits_bound_reads_across_independent_readers() { + let data = Bytes::from(vec![8u8; 1024]); + let tracking = Arc::new(ConcurrencyTrackingRead { + data: data.clone(), + active: AtomicUsize::new(0), + max_active: AtomicUsize::new(0), + }); + let permits = Arc::new(tokio::sync::Semaphore::new(1)); + let make_reader = |path: &str| { + let source: Arc = tracking.clone(); + VindexFileReader::new_with_permits( + source, + tokio::runtime::Handle::current(), + Arc::clone(&permits), + data.len() as u64, + path.to_string(), + ) + }; + let mut first_reader = make_reader("first.index"); + let mut second_reader = make_reader("second.index"); + assert!(Arc::ptr_eq(&first_reader.permits, &second_reader.permits)); + + let first = tokio::task::spawn_blocking(move || { + let mut output = [0u8; 128]; + first_reader + .pread(&mut [ReadRequest::new(0, &mut output)]) + .unwrap(); + }); + let second = tokio::task::spawn_blocking(move || { + let mut output = [0u8; 128]; + second_reader + .pread(&mut [ReadRequest::new(128, &mut output)]) + .unwrap(); + }); + first.await.unwrap(); + second.await.unwrap(); + + assert_eq!(tracking.max_active.load(Ordering::SeqCst), 1); + } + #[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 ecc39a57..3dc9dda1 100644 --- a/crates/paimon/src/vindex/reader.rs +++ b/crates/paimon/src/vindex/reader.rs @@ -15,6 +15,7 @@ // 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::{ @@ -27,6 +28,7 @@ use std::io; const DEFAULT_NPROBE: usize = 16; const NPROBE_PARAMETER: &str = "ivf.nprobe"; +const NATIVE_BATCH_OPERATION_WORKING_SET_BYTES: usize = 64 * 1024 * 1024; trait ErasedSeekRead: Send { fn pread_erased(&mut self, ranges: &mut [ReadRequest<'_>]) -> io::Result<()>; @@ -76,6 +78,7 @@ impl SeekRead for VindexInput { pub struct VindexVectorGlobalIndexReader { io_meta: GlobalIndexIOMeta, options: HashMap, + batch_shard_concurrency: Option, reader: Option>, metadata: Option, } @@ -85,17 +88,23 @@ impl VindexVectorGlobalIndexReader { Self { io_meta, options, + batch_shard_concurrency: None, reader: None, metadata: None, } } + pub(crate) fn with_batch_shard_concurrency(mut self, concurrency: usize) -> Self { + self.batch_shard_concurrency = Some(concurrency.max(1)); + self + } + pub fn visit_vector_search( &mut self, vector_search: &VectorSearch, stream_fn: impl FnOnce(&str) -> crate::Result, ) -> crate::Result>> { - self.ensure_loaded(stream_fn)?; + self.ensure_loaded(stream_fn, |_| Ok(()))?; self.search(vector_search) } @@ -104,11 +113,56 @@ impl VindexVectorGlobalIndexReader { vector_searches: &[VectorSearch], stream_fn: impl FnOnce(&str) -> crate::Result, ) -> crate::Result>>> { - self.ensure_loaded(stream_fn)?; - vector_searches - .iter() - .map(|vector_search| self.search(vector_search)) - .collect() + self.ensure_loaded(stream_fn, |_| Ok(()))?; + self.search_batch(vector_searches) + } + + #[cfg(test)] + pub(crate) fn load( + &mut self, + stream_fn: impl FnOnce(&str) -> crate::Result, + ) -> crate::Result<()> { + self.ensure_loaded(stream_fn, |_| Ok(())) + } + + pub(crate) fn metadata(&self) -> crate::Result<&VectorIndexMetadata> { + self.metadata + .as_ref() + .ok_or_else(|| crate::Error::DataInvalid { + message: "vindex metadata not initialized".to_string(), + source: None, + }) + } + + pub(crate) fn search_batch( + &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() + .ok_or_else(|| crate::Error::DataInvalid { + message: "vindex reader not initialized".to_string(), + source: None, + })?; + let metadata = self + .metadata + .as_ref() + .ok_or_else(|| crate::Error::DataInvalid { + message: "vindex metadata not initialized".to_string(), + source: None, + })?; + search_batch_vindex( + reader, + metadata, + &self.options, + vector_searches, + shard_concurrency, + ) } fn search(&mut self, vector_search: &VectorSearch) -> crate::Result>> { @@ -130,12 +184,38 @@ impl VindexVectorGlobalIndexReader { search_vindex(reader, metadata, &self.options, vector_search) } - fn ensure_loaded( + fn ensure_loaded( &mut self, stream_fn: impl FnOnce(&str) -> crate::Result, - ) -> crate::Result<()> { + validate: F, + ) -> crate::Result<()> + where + S: SeekRead + 'static, + F: FnOnce(&VectorIndexMetadata) -> crate::Result<()>, + { + self.ensure_loaded_with_optimizer(stream_fn, validate, |reader| { + reader + .optimize_for_search() + .map_err(|e| crate::Error::DataInvalid { + message: format!("Failed to optimize paimon-vindex-core reader: {}", e), + source: Some(Box::new(e)), + }) + }) + } + + fn ensure_loaded_with_optimizer( + &mut self, + stream_fn: impl FnOnce(&str) -> crate::Result, + validate: F, + optimize: O, + ) -> crate::Result<()> + where + S: SeekRead + 'static, + F: FnOnce(&VectorIndexMetadata) -> crate::Result<()>, + O: FnOnce(&mut VIndexReader) -> crate::Result<()>, + { if self.reader.is_some() { - return Ok(()); + return validate(self.metadata()?); } let source = stream_fn(&self.io_meta.file_path)?; @@ -146,12 +226,8 @@ impl VindexVectorGlobalIndexReader { } })?; let metadata = reader.metadata(); - reader - .optimize_for_search() - .map_err(|e| crate::Error::DataInvalid { - message: format!("Failed to optimize paimon-vindex-core reader: {}", e), - source: Some(Box::new(e)), - })?; + validate(&metadata)?; + optimize(&mut reader)?; self.reader = Some(reader); self.metadata = Some(metadata); @@ -165,12 +241,35 @@ fn search_vindex( options: &HashMap, vector_search: &VectorSearch, ) -> crate::Result>> { - let expected_dim = metadata.dimension; - if vector_search.vector.len() != expected_dim { + let Some(prepared) = prepare_search(metadata, options, vector_search)? else { + return Ok(None); + }; + let (labels, distances) = execute_scalar_search(reader, vector_search, &prepared)?; + let id_to_scores = collect_results(&labels, &distances, prepared.top_k, metadata.metric); + if id_to_scores.is_empty() { + return Ok(None); + } + + Ok(Some(id_to_scores)) +} + +#[derive(Clone, PartialEq, Eq)] +struct PreparedSearch { + top_k: usize, + nprobe: usize, + filter_bytes: Option>, +} + +fn prepare_search( + metadata: &VectorIndexMetadata, + options: &HashMap, + vector_search: &VectorSearch, +) -> crate::Result> { + if vector_search.vector.len() != metadata.dimension { return Err(crate::Error::DataInvalid { message: format!( "Query vector dimension mismatch: index expects {}, but got {}", - expected_dim, + metadata.dimension, vector_search.vector.len() ), source: None, @@ -178,48 +277,195 @@ fn search_vindex( } let count = usize::try_from(metadata.total_vectors).unwrap_or(0); - let effective_k = std::cmp::min(vector_search.limit, count); - if effective_k == 0 { + let mut top_k = vector_search.limit.min(count); + if top_k == 0 { return Ok(None); } - let nprobe = int_parameter(options, NPROBE_PARAMETER, DEFAULT_NPROBE)?; - let params = VectorSearchParams::new(effective_k, nprobe); - let (labels, distances) = if let Some(include_ids) = &vector_search.include_row_ids { + let filter_bytes = if let Some(include_ids) = &vector_search.include_row_ids { if include_ids.is_empty() { return Ok(None); } - let ek = std::cmp::min(effective_k, include_ids.len() as usize); - let params = VectorSearchParams::new(params.top_k.min(ek), nprobe); - let mut filter_bytes = Vec::new(); + top_k = top_k.min(include_ids.len() as usize); + let mut bytes = Vec::new(); include_ids - .serialize_into(&mut filter_bytes) + .serialize_into(&mut bytes) .map_err(|e| crate::Error::DataInvalid { message: format!("Failed to serialize vector search row-id filter: {}", e), source: Some(Box::new(e)), })?; - reader - .search_with_roaring_filter(&vector_search.vector, params, &filter_bytes) - .map_err(|e| crate::Error::DataInvalid { - message: format!("paimon-vindex-core filtered search failed: {}", e), - source: Some(Box::new(e)), - })? + Some(bytes) } else { - reader - .search(&vector_search.vector, params) + None + }; + + Ok(Some(PreparedSearch { + top_k, + nprobe, + filter_bytes, + })) +} + +fn execute_scalar_search( + reader: &mut VIndexReader, + vector_search: &VectorSearch, + prepared: &PreparedSearch, +) -> crate::Result<(Vec, Vec)> { + let params = VectorSearchParams::new(prepared.top_k, prepared.nprobe); + match &prepared.filter_bytes { + Some(filter) => reader + .search_with_roaring_filter(&vector_search.vector, params, filter) .map_err(|e| crate::Error::DataInvalid { - message: format!("paimon-vindex-core search failed: {}", e), + message: format!("paimon-vindex-core filtered search failed: {}", e), source: Some(Box::new(e)), - })? - }; + }), + None => { + reader + .search(&vector_search.vector, params) + .map_err(|e| crate::Error::DataInvalid { + message: format!("paimon-vindex-core search failed: {}", e), + source: Some(Box::new(e)), + }) + } + } +} - let id_to_scores = collect_results(&labels, &distances, effective_k, metadata.metric); - if id_to_scores.is_empty() { - return Ok(None); +fn search_batch_vindex( + reader: &mut VIndexReader, + metadata: &VectorIndexMetadata, + options: &HashMap, + vector_searches: &[VectorSearch], + shard_concurrency: usize, +) -> crate::Result>>> { + let mut results: Vec>> = + (0..vector_searches.len()).map(|_| None).collect(); + let mut groups: Vec<(PreparedSearch, Vec)> = Vec::new(); + + for (index, search) in vector_searches.iter().enumerate() { + let Some(prepared) = prepare_search(metadata, options, search)? else { + continue; + }; + if let Some((_, indices)) = groups.iter_mut().find(|(key, _)| key == &prepared) { + indices.push(index); + } else { + groups.push((prepared, vec![index])); + } } - Ok(Some(id_to_scores)) + for (prepared, indices) in groups { + let chunk_size = native_batch_chunk_size(metadata, &prepared, shard_concurrency); + for indices in indices.chunks(chunk_size) { + if indices.len() == 1 { + let index = indices[0]; + let (labels, distances) = + execute_scalar_search(reader, &vector_searches[index], &prepared)?; + let map = collect_results(&labels, &distances, prepared.top_k, metadata.metric); + if !map.is_empty() { + results[index] = Some(map); + } + continue; + } + + let mut queries = Vec::with_capacity(indices.len() * metadata.dimension); + for &index in indices { + queries.extend_from_slice(&vector_searches[index].vector); + } + let params = VectorSearchParams::new(prepared.top_k, prepared.nprobe); + let (labels, distances) = match &prepared.filter_bytes { + Some(filter) => reader + .search_batch_with_roaring_filter(&queries, indices.len(), params, filter) + .map_err(|e| crate::Error::DataInvalid { + message: format!("paimon-vindex-core filtered batch search failed: {}", e), + source: Some(Box::new(e)), + })?, + None => reader + .search_batch(&queries, indices.len(), params) + .map_err(|e| crate::Error::DataInvalid { + message: format!("paimon-vindex-core batch search failed: {}", e), + source: Some(Box::new(e)), + })?, + }; + let expected = indices.len() * prepared.top_k; + if labels.len() != expected || distances.len() != expected { + return Err(crate::Error::DataInvalid { + message: format!( + "paimon-vindex-core batch search returned labels/distances of length {}/{}, expected {expected}", + labels.len(), + distances.len() + ), + source: None, + }); + } + for (query_index, &result_index) in indices.iter().enumerate() { + let start = query_index * prepared.top_k; + let end = start + prepared.top_k; + let map = collect_results( + &labels[start..end], + &distances[start..end], + prepared.top_k, + metadata.metric, + ); + if !map.is_empty() { + results[result_index] = Some(map); + } + } + } + } + + Ok(results) +} + +fn native_batch_chunk_size( + metadata: &VectorIndexMetadata, + prepared: &PreparedSearch, + shard_concurrency: usize, +) -> usize { + let per_shard_budget = NATIVE_BATCH_OPERATION_WORKING_SET_BYTES + .checked_div(shard_concurrency.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); + query_budget + .checked_div(native_batch_query_working_set_bytes(metadata, prepared)) + .unwrap_or(0) + .max(1) +} + +fn native_batch_query_working_set_bytes( + metadata: &VectorIndexMetadata, + prepared: &PreparedSearch, +) -> usize { + let query_vectors = metadata + .dimension + .saturating_mul(std::mem::size_of::() * 2); + let centroid_products = metadata.nlist.saturating_mul(std::mem::size_of::()); + let probe_results = prepared + .nprobe + .min(metadata.nlist) + .saturating_mul(std::mem::size_of::() + std::mem::size_of::()); + let top_k_results = prepared.top_k.saturating_mul( + std::mem::size_of::() + std::mem::size_of::() + std::mem::size_of::<(f32, i64)>(), + ); + let pq_tables = match (metadata.pq_m, metadata.pq_bits) { + (Some(m), Some(bits)) => 1usize + .checked_shl(bits as u32) + .unwrap_or(usize::MAX) + .saturating_mul(m) + .saturating_mul(std::mem::size_of::()), + _ => 0, + }; + + query_vectors + .saturating_add(centroid_products) + .saturating_add(probe_results) + .saturating_add(top_k_results) + .saturating_add(pq_tables) + .saturating_add(256) + .max(1) } fn collect_results( @@ -305,6 +551,8 @@ mod tests { use bytes::Bytes; use paimon_vindex_core::index::{VectorIndexConfig, VectorIndexTrainer, VectorIndexWriter}; use paimon_vindex_core::io::{PosWriter, SeekReadCapabilities}; + use std::cell::Cell; + use std::io::Cursor; use std::ops::Range; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; @@ -385,17 +633,65 @@ mod tests { Bytes::from(output) } - fn query() -> VectorSearch { + fn query_for_cluster(cluster: usize, limit: usize) -> VectorSearch { + let cluster = cluster as f32 * 100.0; VectorSearch::new( (0..TEST_DIMENSION) - .map(|dimension| dimension as f32 * 0.01) + .map(|dimension| cluster + dimension as f32 * 0.01) .collect(), - 10, + limit, "embedding".to_string(), ) .unwrap() } + fn query() -> VectorSearch { + query_for_cluster(0, 10) + } + + async fn tracked_batch_search( + index: Bytes, + query_count: usize, + ) -> (Vec>>, usize) { + tracked_batch_search_with_options( + index, + query_count, + HashMap::from([(NPROBE_PARAMETER.to_string(), "1".to_string())]), + 32, + ) + .await + } + + async fn tracked_batch_search_with_options( + index: Bytes, + query_count: usize, + options: HashMap, + shard_concurrency: usize, + ) -> (Vec>>, usize) { + let tracking = TrackingIndexRead::new(index.clone()); + let source: Arc = tracking.clone(); + let runtime = tokio::runtime::Handle::current(); + let results = tokio::task::spawn_blocking(move || { + let source = VindexFileReader::new( + source, + runtime, + index.len() as u64, + "batch.index".to_string(), + ); + let io_meta = + 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); + reader + .visit_batch_vector_search(&searches, |_| Ok(source)) + .unwrap() + }) + .await + .unwrap(); + (results, tracking.bytes_read.load(Ordering::SeqCst)) + } + #[test] fn test_convert_distance_to_score() { assert_eq!(convert_distance_to_score(0.0, MetricType::L2), 1.0); @@ -434,6 +730,43 @@ mod tests { assert!(!result.contains_key(&3)); } + #[test] + fn native_batch_chunk_size_tracks_working_set_inputs() { + let base_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 base_prepared = PreparedSearch { + top_k: 10, + nprobe: 16, + filter_bytes: None, + }; + let base = native_batch_chunk_size(&base_metadata, &base_prepared, 32); + + 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); + + 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); + + 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(&base_metadata, &base_prepared, 64) < base); + } + #[test] fn test_int_parameter() { let mut options = HashMap::new(); @@ -500,4 +833,296 @@ mod tests { "range search unexpectedly read the entire index" ); } + + #[test] + fn mixed_batch_matches_scalar_searches_and_preserves_order() { + let index = build_ivf_flat_index(); + let options = HashMap::from([(NPROBE_PARAMETER.to_string(), "1".to_string())]); + let mut include_row_ids = roaring::RoaringTreemap::new(); + include_row_ids.insert(0); + include_row_ids.insert(16); + include_row_ids.insert(32); + let searches = vec![ + query_for_cluster(0, 4), + query_for_cluster(5, 1), + query_for_cluster(15, 4), + query_for_cluster(0, 10).with_include_row_ids(include_row_ids), + ]; + + let scalar_meta = + GlobalIndexIOMeta::new("scalar.index".to_string(), index.len() as u64, Vec::new()); + let mut scalar_reader = VindexVectorGlobalIndexReader::new(scalar_meta, options.clone()); + scalar_reader + .load(|_| Ok(Cursor::new(index.clone()))) + .unwrap(); + let expected: Vec<_> = searches + .iter() + .map(|search| scalar_reader.search(search).unwrap()) + .collect(); + + let batch_meta = + GlobalIndexIOMeta::new("batch.index".to_string(), index.len() as u64, Vec::new()); + let mut batch_reader = VindexVectorGlobalIndexReader::new(batch_meta, options); + let actual = batch_reader + .visit_batch_vector_search(&searches, |_| Ok(Cursor::new(index))) + .unwrap(); + + assert_eq!(actual, expected); + assert_ne!( + actual[0], actual[2], + "interleaved batch groups lost query order" + ); + assert_eq!(actual[1].as_ref().map(HashMap::len), Some(1)); + assert_eq!(actual[3].as_ref().map(HashMap::len), Some(3)); + } + + #[test] + fn metadata_validation_runs_before_optimization() { + let index = build_ivf_flat_index(); + let io_meta = GlobalIndexIOMeta::new( + "validated.index".to_string(), + index.len() as u64, + Vec::new(), + ); + let mut reader = VindexVectorGlobalIndexReader::new(io_meta, HashMap::new()); + let optimized = Cell::new(false); + + let error = reader + .ensure_loaded_with_optimizer( + |_| Ok(Cursor::new(index)), + |metadata| { + assert_eq!(metadata.dimension, TEST_DIMENSION); + Err(crate::Error::DataInvalid { + message: "rejected test metric".to_string(), + source: None, + }) + }, + |_| { + optimized.set(true); + Ok(()) + }, + ) + .unwrap_err(); + + assert!(error.to_string().contains("rejected test metric")); + assert!(!optimized.get()); + assert!(reader.metadata().is_err()); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn batch_search_reuses_probed_lists_and_avoids_full_file_read() { + let index = build_ivf_flat_index(); + let batch_index = index.clone(); + let options = HashMap::from([(NPROBE_PARAMETER.to_string(), "1".to_string())]); + let search = query(); + + let scalar_tracking = TrackingIndexRead::new(index.clone()); + let scalar_source: Arc = scalar_tracking.clone(); + let scalar_options = options.clone(); + let scalar_search = search.clone(); + let scalar_runtime = tokio::runtime::Handle::current(); + let scalar_results = tokio::task::spawn_blocking(move || { + let source = VindexFileReader::new( + scalar_source, + scalar_runtime, + index.len() as u64, + "scalar.index".to_string(), + ); + let io_meta = + GlobalIndexIOMeta::new("scalar.index".to_string(), index.len() as u64, Vec::new()); + let mut reader = VindexVectorGlobalIndexReader::new(io_meta, scalar_options); + reader.load(|_| Ok(source)).unwrap(); + vec![ + reader.search(&scalar_search).unwrap(), + reader.search(&scalar_search).unwrap(), + ] + }) + .await + .unwrap(); + + let batch_tracking = TrackingIndexRead::new(batch_index.clone()); + let batch_source: Arc = batch_tracking.clone(); + let batch_options = options; + let batch_search = search.clone(); + let batch_runtime = tokio::runtime::Handle::current(); + let batch_results = tokio::task::spawn_blocking(move || { + let source = VindexFileReader::new( + batch_source, + batch_runtime, + batch_index.len() as u64, + "batch.index".to_string(), + ); + let io_meta = GlobalIndexIOMeta::new( + "batch.index".to_string(), + batch_index.len() as u64, + Vec::new(), + ); + let mut reader = VindexVectorGlobalIndexReader::new(io_meta, batch_options); + reader + .visit_batch_vector_search(&[batch_search.clone(), batch_search], |_| Ok(source)) + .unwrap() + }) + .await + .unwrap(); + + assert_eq!(batch_results, scalar_results); + let scalar_bytes = scalar_tracking.bytes_read.load(Ordering::SeqCst); + let batch_bytes = batch_tracking.bytes_read.load(Ordering::SeqCst); + assert!( + batch_bytes < scalar_bytes, + "batch should read a shared probed list once: batch={batch_bytes}, scalar={scalar_bytes}" + ); + assert!( + batch_bytes < batch_tracking.data.len() / 2, + "nprobe=1 should read substantially less than the full index: read={batch_bytes}, file={} ", + batch_tracking.data.len() + ); + assert!(batch_tracking + .ranges() + .iter() + .all(|range| { range.start != 0 || range.end != batch_tracking.data.len() as u64 })); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn filtered_batch_matches_scalar_searches_and_reuses_probed_lists() { + let index = build_ivf_flat_index(); + let batch_index = index.clone(); + let options = HashMap::from([(NPROBE_PARAMETER.to_string(), "1".to_string())]); + let mut include_row_ids = roaring::RoaringTreemap::new(); + for row_id in (0..256).step_by(16) { + include_row_ids.insert(row_id); + } + let searches = vec![ + query().with_include_row_ids(include_row_ids.clone()), + query().with_include_row_ids(include_row_ids), + ]; + + let scalar_tracking = TrackingIndexRead::new(index.clone()); + let scalar_source: Arc = scalar_tracking.clone(); + let scalar_options = options.clone(); + let scalar_searches = searches.clone(); + let scalar_runtime = tokio::runtime::Handle::current(); + let scalar_results = tokio::task::spawn_blocking(move || { + let source = VindexFileReader::new( + scalar_source, + scalar_runtime, + index.len() as u64, + "scalar-filtered.index".to_string(), + ); + let io_meta = GlobalIndexIOMeta::new( + "scalar-filtered.index".to_string(), + index.len() as u64, + Vec::new(), + ); + let mut reader = VindexVectorGlobalIndexReader::new(io_meta, scalar_options); + reader.load(|_| Ok(source)).unwrap(); + scalar_searches + .iter() + .map(|search| reader.search(search).unwrap()) + .collect::>() + }) + .await + .unwrap(); + + let batch_tracking = TrackingIndexRead::new(batch_index.clone()); + let batch_source: Arc = batch_tracking.clone(); + let batch_runtime = tokio::runtime::Handle::current(); + let batch_results = tokio::task::spawn_blocking(move || { + let source = VindexFileReader::new( + batch_source, + batch_runtime, + batch_index.len() as u64, + "batch-filtered.index".to_string(), + ); + let io_meta = GlobalIndexIOMeta::new( + "batch-filtered.index".to_string(), + batch_index.len() as u64, + Vec::new(), + ); + let mut reader = VindexVectorGlobalIndexReader::new(io_meta, options); + reader + .visit_batch_vector_search(&searches, |_| Ok(source)) + .unwrap() + }) + .await + .unwrap(); + + assert_eq!(batch_results, scalar_results); + let scalar_bytes = scalar_tracking.bytes_read.load(Ordering::SeqCst); + let batch_bytes = batch_tracking.bytes_read.load(Ordering::SeqCst); + assert!( + batch_bytes < scalar_bytes, + "filtered batch should read a shared probed list once: batch={batch_bytes}, scalar={scalar_bytes}" + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn large_homogeneous_batch_reuses_lists_across_previous_boundary() { + let previous_batch_boundary = 16; + let index = build_ivf_flat_index(); + let (within_limit_results, within_limit_bytes) = + tracked_batch_search(index.clone(), previous_batch_boundary).await; + let (over_limit_results, over_limit_bytes) = + tracked_batch_search(index, previous_batch_boundary + 1).await; + + assert_eq!(within_limit_results.len(), previous_batch_boundary); + assert_eq!(over_limit_results.len(), previous_batch_boundary + 1); + assert!(within_limit_results + .iter() + .all(|result| result == &within_limit_results[0])); + assert!(over_limit_results + .iter() + .all(|result| result == &over_limit_results[0])); + assert_eq!(over_limit_results[0], within_limit_results[0]); + assert_eq!( + over_limit_bytes, within_limit_bytes, + "compatible queries should reuse the same probed list across the former 16-query boundary" + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn homogeneous_batch_chunks_at_working_set_boundary() { + let shard_concurrency = 4096; + let metadata = VectorIndexMetadata { + index_type: paimon_vindex_core::index::IndexType::IvfFlat, + dimension: TEST_DIMENSION, + nlist: 16, + metric: MetricType::L2, + total_vectors: 8192, + pq_m: None, + pq_bits: None, + rq_bits: None, + diskann: None, + }; + let options = HashMap::from([ + (NPROBE_PARAMETER.to_string(), "1".to_string()), + ("global-index.thread-num".to_string(), "1".to_string()), + ]); + let prepared = prepare_search(&metadata, &options, &query()) + .unwrap() + .unwrap(); + let chunk_size = native_batch_chunk_size(&metadata, &prepared, shard_concurrency); + assert!(chunk_size > 16); + + let index = build_ivf_flat_index(); + let (within_results, within_bytes) = tracked_batch_search_with_options( + index.clone(), + chunk_size, + options.clone(), + shard_concurrency, + ) + .await; + let (over_results, over_bytes) = + tracked_batch_search_with_options(index, chunk_size + 1, options, shard_concurrency) + .await; + + assert_eq!(within_results.len(), chunk_size); + assert_eq!(over_results.len(), chunk_size + 1); + assert!(within_results + .iter() + .all(|result| result == &within_results[0])); + assert!(over_results.iter().all(|result| result == &over_results[0])); + assert_eq!(over_results[0], within_results[0]); + assert!(over_bytes > within_bytes); + } }