diff --git a/crates/paimon/src/table/pk_vector_data_file_reader.rs b/crates/paimon/src/table/pk_vector_data_file_reader.rs index 33067b6a..d445f153 100644 --- a/crates/paimon/src/table/pk_vector_data_file_reader.rs +++ b/crates/paimon/src/table/pk_vector_data_file_reader.rs @@ -15,15 +15,20 @@ // specific language governing permissions and limitations // under the License. -//! Exact sequential vector reader over one data file's vector column. Mirrors -//! Java `org.apache.paimon.index.pkvector.PkVectorDataFileReader`. +//! Streaming exact vector search over one data file's vector column. Mirrors +//! Java `org.apache.paimon.index.pkvector.PkVectorDataFileReader` + +//! `PkVectorExactSearcher`. //! -//! The factory projects the single vector column, reads the whole file in -//! physical order, and preloads it into memory as `Vec>>` -//! (a NULL row is `None`). Deletion vectors are deliberately NOT applied here: -//! physical position must stay in lockstep with the segment ordinal so the -//! bucket search can address rows by position. The returned reader then serves -//! vectors from memory one physical row at a time. +//! The factory projects the single vector column and, per uncovered file, +//! streams the column one Arrow batch at a time — feeding each row into per-query +//! bounded Top-K heaps and dropping the batch — so peak memory is one batch plus +//! the heaps rather than the whole column. Deletion vectors and residual filters +//! are deliberately NOT applied by the stream itself: physical position must stay +//! in lockstep with the segment ordinal, so exclusion is folded in via the +//! caller-supplied `is_excluded(position)` predicate. A NULL row is not scored but +//! still advances the physical position. + +use std::collections::BinaryHeap; use arrow_array::{Array, FixedSizeListArray, Float32Array, ListArray}; use futures::TryStreamExt; @@ -32,7 +37,9 @@ use crate::spec::{DataField, DataType}; use crate::table::data_file_reader::DataFileReader; use crate::table::source::DataSplit; use crate::vindex::pkvector::bucket::BucketActiveFile; -use crate::vindex::pkvector::reader::PkVectorReader; +use crate::vindex::pkvector::exact::{drain_best_first, push_bounded, validate_query, WorstFirst}; +use crate::vindex::pkvector::metric::VectorSearchMetric; +use crate::vindex::pkvector::result::PkVectorSearchResult; fn data_invalid(message: impl Into) -> crate::Error { crate::Error::DataInvalid { @@ -41,7 +48,8 @@ fn data_invalid(message: impl Into) -> crate::Error { } } -/// Builds an exact [`PkVectorReader`] over one data file's vector column. +/// Runs a streaming exact [`PkVectorSearchResult`] search over one data file's +/// vector column. /// /// `reader` is configured (via [`DataFileReader::with_read_type`]) to project /// only the vector column, so each read returns a single-column batch. Mirrors @@ -79,14 +87,33 @@ impl DataFilePkVectorReaderFactory { }) } - /// Preload the whole vector column of `file` into memory and return a - /// sequential reader over it. `file` must name a data file present in this - /// factory's split. The drained row count is checked against the file's - /// `DataFileMeta.row_count`. - pub(crate) async fn create( + /// Stream the vector column of `file` one Arrow batch at a time and return one + /// bounded, BEST_FIRST Top-K list per query (outer index aligned to `queries`). + /// `file` must name a data file present in this factory's split. + /// + /// All queries are validated (dimension + finite) BEFORE the file stream is + /// opened. Each surviving physical position (not NULL, not `is_excluded`) is + /// scored against every query into that query's bounded heap; a NULL row is + /// skipped but still advances the position so the position stays in lockstep + /// with `is_excluded`. The drained row count is checked against the file's + /// `DataFileMeta.row_count` (both truncation and overrun fail loud). + pub(crate) async fn search_file( &self, file: &BucketActiveFile, - ) -> crate::Result> { + queries: &[&[f32]], + metric: VectorSearchMetric, + exact_limit: usize, + is_excluded: &(dyn Fn(i64) -> bool + Sync), + ) -> crate::Result>> { + if exact_limit == 0 { + return Err(data_invalid("vector search limit must be positive")); + } + // Validate every query before opening the stream (validate-before-POLL); + // a malformed query fails loud before any file I/O. + for query in queries { + validate_query(query, self.dimension)?; + } + let file_meta = self .data_split .data_files() @@ -110,34 +137,58 @@ impl DataFilePkVectorReaderFactory { None, )?; - let mut vectors: Vec>> = Vec::new(); + let mut heaps: Vec> = (0..queries.len()) + .map(|_| BinaryHeap::with_capacity(exact_limit + 1)) + .collect(); + // One reused buffer per batch; a NULL row leaves it untouched (and is not + // scored). `position` is the monotonic physical row counter across batches. + let mut batch_vectors: Vec>> = Vec::new(); + let mut position: i64 = 0; while let Some(batch) = stream.try_next().await? { + batch_vectors.clear(); append_batch_vectors( &batch, self.vector_field.name(), self.dimension, - &mut vectors, + &mut batch_vectors, )?; + for entry in &batch_vectors { + let pos = position; + position += 1; + if pos >= row_count { + return Err(data_invalid( + "data file produced more rows than DataFileMeta.row_count", + )); + } + let Some(vector) = entry else { + continue; // NULL row: not scored, position already advanced. + }; + if is_excluded(pos) { + continue; + } + for (query, heap) in queries.iter().zip(heaps.iter_mut()) { + let candidate = PkVectorSearchResult { + data_file_name: file.file_name.clone(), + row_position: pos, + distance: metric.compute_distance(query, vector), + }; + push_bounded(heap, candidate, exact_limit); + } + } } - let drained = vectors.len() as i64; - if drained > row_count { + if position > row_count { return Err(data_invalid( "data file produced more rows than DataFileMeta.row_count", )); } - if drained < row_count { + if position < row_count { return Err(data_invalid( "data file ended before DataFileMeta.row_count", )); } - Ok(Box::new(DataFilePkVectorReader { - dimension: self.dimension, - row_count, - vectors, - position: 0, - })) + Ok(heaps.into_iter().map(drain_best_first).collect()) } } @@ -219,97 +270,6 @@ pub(crate) fn append_batch_vectors( Ok(()) } -/// In-memory sequential reader over one file's preloaded vector column. Each -/// [`read_next_vector`](PkVectorReader::read_next_vector) advances exactly one -/// physical row; a NULL row returns `false` but still advances the position. -struct DataFilePkVectorReader { - dimension: usize, - row_count: i64, - /// Preloaded whole-file column in physical order; `None` = NULL row. - vectors: Vec>>, - position: usize, -} - -impl PkVectorReader for DataFilePkVectorReader { - fn dimension(&self) -> usize { - self.dimension - } - - fn row_count(&self) -> i64 { - self.row_count - } - - fn read_next_vector(&mut self, reuse: &mut [f32]) -> crate::Result { - if reuse.len() != self.dimension { - return Err(data_invalid(format!( - "reuse buffer length {} does not match vector dimension {}", - reuse.len(), - self.dimension - ))); - } - if self.position as i64 >= self.row_count { - return Err(data_invalid("read past row count")); - } - let entry = &self.vectors[self.position]; - self.position += 1; - match entry { - Some(vector) => { - reuse.copy_from_slice(vector); - Ok(true) - } - None => Ok(false), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn null_row_consumes_ordinal() { - let mut r = DataFilePkVectorReader { - dimension: 2, - row_count: 3, - vectors: vec![Some(vec![1.0, 2.0]), None, Some(vec![3.0, 4.0])], - position: 0, - }; - let mut buf = [0.0f32; 2]; - assert!(r.read_next_vector(&mut buf).unwrap()); - assert_eq!(buf, [1.0, 2.0]); - assert!(!r.read_next_vector(&mut buf).unwrap()); // null: false, ordinal advanced - assert!(r.read_next_vector(&mut buf).unwrap()); - assert_eq!(buf, [3.0, 4.0]); - assert_eq!(r.row_count(), 3); - assert_eq!(r.dimension(), 2); - } - - #[test] - fn read_past_row_count_errors() { - let mut r = DataFilePkVectorReader { - dimension: 1, - row_count: 1, - vectors: vec![Some(vec![1.0])], - position: 0, - }; - let mut buf = [0.0f32; 1]; - assert!(r.read_next_vector(&mut buf).unwrap()); - assert!(r.read_next_vector(&mut buf).is_err()); // past row_count - } - - #[test] - fn reuse_len_mismatch_errors() { - let mut r = DataFilePkVectorReader { - dimension: 2, - row_count: 1, - vectors: vec![Some(vec![1.0, 2.0])], - position: 0, - }; - let mut buf = [0.0f32; 1]; - assert!(r.read_next_vector(&mut buf).is_err()); - } -} - #[cfg(test)] mod integration_tests { use super::*; @@ -384,12 +344,16 @@ mod integration_tests { ); } - /// Write a FixedSizeList vector column - /// (`[1,2]`, NULL, `[3,4]`) as a parquet data file, build the factory over - /// its split, preload via `create`, and assert the whole-file sequential - /// read plus the "file not in split" error. - #[tokio::test] - async fn create_preloads_and_reads_whole_file() { + /// Build a FixedSizeList vector column from `rows` (`None` = NULL + /// row), write it as one parquet data file across `batches` write calls, and + /// return a factory over its split plus the file name. `stated_row_count` is + /// what the `DataFileMeta` claims (usually the true row count, but a test can + /// pass a wrong value to exercise the row-count guard). + async fn build_factory( + rows: &[Option>], + stated_row_count: i64, + table_path: &str, + ) -> (DataFilePkVectorReaderFactory, String) { let field = vector_field(); let read_fields = vec![field.clone()]; let arrow_schema = build_target_arrow_schema(&read_fields).unwrap(); @@ -397,20 +361,24 @@ mod integration_tests { let mut builder = FixedSizeListBuilder::new(Float32Builder::new(), 2).with_field(Arc::new( ArrowField::new("element", ArrowDataType::Float32, true), )); - builder.values().append_value(1.0); - builder.values().append_value(2.0); - builder.append(true); - builder.values().append_value(0.0); - builder.values().append_value(0.0); - builder.append(false); // NULL vector row - builder.values().append_value(3.0); - builder.values().append_value(4.0); - builder.append(true); + for row in rows { + match row { + Some(v) => { + builder.values().append_value(v[0]); + builder.values().append_value(v[1]); + builder.append(true); + } + None => { + builder.values().append_value(0.0); + builder.values().append_value(0.0); + builder.append(false); + } + } + } let vec_array = builder.finish(); let batch = RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(vec_array)]).unwrap(); let file_io = FileIOBuilder::new("memory").build().unwrap(); - let table_path = "memory:/pk_vector_data_file_reader"; let bucket_path = format!("{table_path}/bucket-0"); let file_name = "part-0.parquet"; let file_path = format!("{bucket_path}/{file_name}"); @@ -440,7 +408,7 @@ mod integration_tests { .with_data_files(vec![data_file( file_name, file_size as i64, - 3, + stated_row_count, table_schema_id, )]) .build() @@ -455,36 +423,162 @@ mod integration_tests { read_fields.clone(), Vec::new(), ); + let factory = DataFilePkVectorReaderFactory::new(reader, data_split, field).unwrap(); + (factory, file_name.to_string()) + } - let factory = - DataFilePkVectorReaderFactory::new(reader, data_split, field.clone()).unwrap(); + /// The streaming per-file search must produce candidates byte-identical to + /// the reference `exact_search` over an in-memory `ArrayReader` of the same + /// data, including a NULL row and a residual/DV exclusion. + #[tokio::test] + async fn search_file_matches_exact_search_reference() { + use crate::vindex::pkvector::exact::exact_search; + use crate::vindex::pkvector::reader::test_support::ArrayReader; - let present = BucketActiveFile { - file_name: file_name.to_string(), + let rows = vec![ + Some(vec![3.0, 0.0]), + None, + Some(vec![1.0, 0.0]), + Some(vec![2.0, 0.0]), + ]; + let (factory, file_name) = + build_factory(&rows, rows.len() as i64, "memory:/pkvdfr_equiv").await; + let active = BucketActiveFile { + file_name: file_name.clone(), + row_count: rows.len() as i64, + }; + // Exclude physical position 2 (residual/DV fold): mirrors the closure the + // bucket search passes in. + let is_excluded = |pos: i64| pos == 2; + let query = [0.0f32, 0.0]; + + let streamed = factory + .search_file(&active, &[&query], VectorSearchMetric::L2, 2, &is_excluded) + .await + .unwrap(); + + let mut ref_reader = ArrayReader::new(2, rows.clone()); + let reference = exact_search( + &file_name, + &mut ref_reader, + &query, + VectorSearchMetric::L2, + 2, + &is_excluded, + ) + .unwrap(); + + assert_eq!(streamed.len(), 1, "one query in, one result list out"); + assert_eq!(streamed[0], reference); + } + + /// Same streaming-vs-reference equivalence, but with more scorable rows than + /// `exact_limit` so the bounded heap's eviction branch is exercised on both + /// paths (the shared `push_bounded` must evict identically). + #[tokio::test] + async fn search_file_matches_exact_search_reference_with_eviction() { + use crate::vindex::pkvector::exact::exact_search; + use crate::vindex::pkvector::reader::test_support::ArrayReader; + + // Five scorable rows, no NULL/exclusion; keep only the 2 closest to [0,0]. + let rows = vec![ + Some(vec![4.0, 0.0]), + Some(vec![1.0, 0.0]), + Some(vec![3.0, 0.0]), + Some(vec![2.0, 0.0]), + Some(vec![5.0, 0.0]), + ]; + let (factory, file_name) = + build_factory(&rows, rows.len() as i64, "memory:/pkvdfr_evict").await; + let active = BucketActiveFile { + file_name: file_name.clone(), + row_count: rows.len() as i64, + }; + let query = [0.0f32, 0.0]; + + let streamed = factory + .search_file(&active, &[&query], VectorSearchMetric::L2, 2, &|_| false) + .await + .unwrap(); + + let mut ref_reader = ArrayReader::new(2, rows.clone()); + let reference = exact_search( + &file_name, + &mut ref_reader, + &query, + VectorSearchMetric::L2, + 2, + &|_| false, + ) + .unwrap(); + + assert_eq!(streamed[0], reference); + // The two closest are positions 1 ([1,0]) then 3 ([2,0]), best-first. + assert_eq!(streamed[0].len(), 2, "bounded to exact_limit"); + assert_eq!(streamed[0][0].row_position, 1); + assert_eq!(streamed[0][1].row_position, 3); + } + + /// A `DataFileMeta.row_count` larger than the file's real row count means the + /// stream ends early; the search must fail loud rather than return a short + /// result. + #[tokio::test] + async fn search_file_fails_loud_on_row_count_truncation() { + let rows = vec![Some(vec![1.0, 0.0]), Some(vec![2.0, 0.0])]; + // Claim 3 rows but only write 2. + let (factory, file_name) = build_factory(&rows, 3, "memory:/pkvdfr_trunc").await; + let active = BucketActiveFile { + file_name, row_count: 3, }; - let mut pk_reader = factory.create(&present).await.unwrap(); - assert_eq!(pk_reader.dimension(), 2); - assert_eq!(pk_reader.row_count(), 3); - - let mut buf = [0.0f32; 2]; - assert!(pk_reader.read_next_vector(&mut buf).unwrap()); - assert_eq!(buf, [1.0, 2.0]); - assert!(!pk_reader.read_next_vector(&mut buf).unwrap()); // NULL row - assert!(pk_reader.read_next_vector(&mut buf).unwrap()); - assert_eq!(buf, [3.0, 4.0]); - assert!(pk_reader.read_next_vector(&mut buf).is_err()); // past row count - - // A file name absent from the split is rejected as invalid. + let query = [0.0f32, 0.0]; + let err = factory + .search_file(&active, &[&query], VectorSearchMetric::L2, 2, &|_| false) + .await + .expect_err("row-count truncation must fail loud"); + assert!(err.to_string().contains("ended before"), "got: {err}"); + } + + /// A malformed query (wrong dimension / non-finite element) fails loud, and a + /// file name absent from the split is rejected as invalid. + #[tokio::test] + async fn search_file_validates_query_and_rejects_absent_file() { + let rows = vec![Some(vec![1.0, 2.0]), None, Some(vec![3.0, 4.0])]; + let (factory, file_name) = + build_factory(&rows, rows.len() as i64, "memory:/pkvdfr_validate").await; + let present = BucketActiveFile { + file_name: file_name.clone(), + row_count: rows.len() as i64, + }; + + // Wrong dimension. + let bad_dim = [1.0f32]; + let err = factory + .search_file(&present, &[&bad_dim], VectorSearchMetric::L2, 2, &|_| false) + .await + .expect_err("dimension mismatch must fail loud"); + assert!(err.to_string().contains("dimension"), "got: {err}"); + + // Non-finite element. + let bad_finite = [f32::NAN, 0.0]; + let err = factory + .search_file(&present, &[&bad_finite], VectorSearchMetric::L2, 2, &|_| { + false + }) + .await + .expect_err("non-finite query must fail loud"); + assert!(err.to_string().contains("finite"), "got: {err}"); + + // Absent file. let missing = BucketActiveFile { file_name: "absent.parquet".to_string(), row_count: 3, }; + let query = [0.0f32, 0.0]; let err = factory - .create(&missing) + .search_file(&missing, &[&query], VectorSearchMetric::L2, 2, &|_| false) .await - .err() - .expect("absent file must be rejected"); + .expect_err("absent file must be rejected"); assert!(matches!(err, crate::Error::DataInvalid { .. })); } } diff --git a/crates/paimon/src/table/pk_vector_orchestrator.rs b/crates/paimon/src/table/pk_vector_orchestrator.rs index 358211a5..3e6a09fa 100644 --- a/crates/paimon/src/table/pk_vector_orchestrator.rs +++ b/crates/paimon/src/table/pk_vector_orchestrator.rs @@ -35,7 +35,7 @@ use crate::table::pk_vector_indexed_split_read::PkVectorIndexedSplit; use crate::table::source::{DataSplit, DataSplitBuilder, RowRange}; use crate::vindex::pkvector::ann::PkVectorAnnSearcher; use crate::vindex::pkvector::bucket::{ - bucket_search, BucketActiveFile, BucketAnnSegment, ExactReaderFuture, + bucket_search, BucketActiveFile, BucketAnnSegment, ExactFileSearchFuture, }; use crate::vindex::pkvector::metric::{java_float_compare, VectorSearchMetric}; use crate::vindex::pkvector::result::PkVectorSearchResult; @@ -47,29 +47,44 @@ fn data_invalid(message: impl Into) -> crate::Error { } } -/// Coerce a closure into the higher-ranked per-file exact-reader factory shape so -/// its returned future borrows for exactly the file argument's lifetime. Closure -/// return types cannot express this borrow through inference alone, so the bound -/// is supplied here. -fn as_bucket_factory(f: F) -> F +/// Coerce a closure into the higher-ranked split-scoped exact-file search shape +/// expected by [`PkVectorOrchestrator::search_candidates`], binding the returned +/// future's borrow to the arguments' lifetime. Callers building a search closure +/// use this so the higher-ranked bound is supplied where inference cannot. The +/// closure is `Fn + Send + Sync` (it is called concurrently by the parallel +/// search), not `FnMut`. +#[allow(clippy::type_complexity)] +pub(crate) fn as_split_exact_file_search(f: F) -> F where - F: for<'f> FnMut(&'f BucketActiveFile) -> ExactReaderFuture<'f> + Send, + F: for<'s, 'a> Fn( + usize, + &'s PkVectorSearchSplit, + &'a BucketActiveFile, + &'a [&'a [f32]], + VectorSearchMetric, + usize, + &'a (dyn Fn(i64) -> bool + Sync), + ) -> ExactFileSearchFuture<'a> + + Send + + Sync, { f } -/// Coerce a closure into the split-scoped exact-reader factory shape expected by -/// [`PkVectorOrchestrator::search_candidates`], binding the returned future's -/// borrow to the file argument's lifetime. Callers building a factory closure use -/// this so the higher-ranked bound is supplied where inference cannot. -pub(crate) fn as_split_exact_reader_factory(f: F) -> F +/// Coerce a closure into the per-file exact-file search shape `bucket_search` +/// expects, supplying the higher-ranked bound closure inference cannot express. +#[allow(clippy::type_complexity)] +fn as_bucket_exact_file_search(f: F) -> F where - F: for<'s, 'f> FnMut( + F: for<'a> Fn( + &'a BucketActiveFile, + &'a [&'a [f32]], + VectorSearchMetric, usize, - &'s PkVectorSearchSplit, - &'f BucketActiveFile, - ) -> ExactReaderFuture<'f> - + Send, + &'a (dyn Fn(i64) -> bool + Sync), + ) -> ExactFileSearchFuture<'a> + + Send + + Sync, { f } @@ -356,12 +371,17 @@ impl PkVectorOrchestrator { limit: usize, indexed_limit: usize, ann_searcher: Option<&dyn PkVectorAnnSearcher>, - exact_reader_factory: &mut (dyn for<'s, 'f> FnMut( + exact_file_search: &(dyn for<'s, 'a> Fn( usize, &'s PkVectorSearchSplit, - &'f BucketActiveFile, - ) -> ExactReaderFuture<'f> - + Send), + &'a BucketActiveFile, + &'a [&'a [f32]], + VectorSearchMetric, + usize, + &'a (dyn Fn(i64) -> bool + Sync), + ) -> ExactFileSearchFuture<'a> + + Send + + Sync), search_options: &HashMap, skip_exact_fallback: bool, residual_by_split: Option<&[HashMap]>, @@ -389,19 +409,35 @@ impl PkVectorOrchestrator { let mut exact_candidates: Vec = Vec::new(); for (split_index, split) in splits.iter().enumerate() { let dvs = build_bucket_dv_map(&self.reader, split).await?; - // Wrap the split-scoped factory into bucket_search's per-file signature. - // The coercion helper ties the produced future's borrow to the file - // argument, which closure inference cannot express on its own. - let mut bucket_factory = as_bucket_factory(|file: &BucketActiveFile| { - exact_reader_factory(split_index, split, file) - }); + // Adapt the split-scoped search closure to bucket_search's per-file + // closure by binding the current split index/split. The coercion helper + // ties the produced future's borrow to the arguments, which closure + // inference cannot express on its own. + let bucket_search_closure = as_bucket_exact_file_search( + |file: &BucketActiveFile, + queries: &[&[f32]], + metric: VectorSearchMetric, + exact_limit: usize, + is_excluded: &(dyn Fn(i64) -> bool + Sync)| + -> ExactFileSearchFuture<'_> { + exact_file_search( + split_index, + split, + file, + queries, + metric, + exact_limit, + is_excluded, + ) + }, + ); let residual_ranges = residual_by_split.map(|per_split| &per_split[split_index]); let result = bucket_search( ann_searcher, &split.ann_segments, &split.active_files, &dvs, - &mut bucket_factory, + &bucket_search_closure, query, metric, indexed_limit, @@ -729,8 +765,8 @@ mod e2e_tests { use crate::table::pk_vector_position_read::{PKEY_VECTOR_POSITION_COLUMN, SEARCH_SCORE_COLUMN}; use crate::table::schema_manager::SchemaManager; use crate::table::source::DeletionFile; + use crate::vindex::pkvector::exact::exact_search; use crate::vindex::pkvector::reader::test_support::ArrayReader; - use crate::vindex::pkvector::reader::PkVectorReader; use arrow_array::{Array, Float32Array, Int32Array, Int64Array, RecordBatch}; use bytes::Bytes; use futures::TryStreamExt; @@ -897,30 +933,114 @@ mod e2e_tests { } } - /// Coerce a closure into the higher-ranked split-scoped exact-reader factory - /// shape so its returned future borrows for exactly the file argument's - /// lifetime. Closure return types cannot express this borrow through inference - /// alone, so the bound is supplied here. - fn as_split_factory(f: F) -> F + /// Coerce a closure into the split-scoped exact-file search shape, supplying + /// the higher-ranked bound closure inference cannot express. The closure is + /// `Fn + Send + Sync`. + #[allow(clippy::type_complexity)] + fn as_split_search(f: F) -> F where - F: for<'s, 'f> FnMut( + F: for<'s, 'a> Fn( usize, &'s PkVectorSearchSplit, - &'f BucketActiveFile, - ) -> ExactReaderFuture<'f> - + Send, + &'a BucketActiveFile, + &'a [&'a [f32]], + VectorSearchMetric, + usize, + &'a (dyn Fn(i64) -> bool + Sync), + ) -> ExactFileSearchFuture<'a> + + Send + + Sync, { - as_split_exact_reader_factory(f) + as_split_exact_file_search(f) + } + + /// A split-scoped search closure that must never be invoked. + #[allow(clippy::type_complexity)] + fn unreachable_split_search() -> impl for<'s, 'a> Fn( + usize, + &'s PkVectorSearchSplit, + &'a BucketActiveFile, + &'a [&'a [f32]], + VectorSearchMetric, + usize, + &'a (dyn Fn(i64) -> bool + Sync), + ) -> ExactFileSearchFuture<'a> + + Send + + Sync { + as_split_search( + |_: usize, + _: &PkVectorSearchSplit, + _: &BucketActiveFile, + _: &[&[f32]], + _: VectorSearchMetric, + _: usize, + _: &(dyn Fn(i64) -> bool + Sync)| + -> ExactFileSearchFuture<'_> { + Box::pin(async { unreachable!("closure must not be invoked in this test") }) + }, + ) } - /// Same coercion for the per-file factory used by `materialize_via_splits`. - fn as_file_factory(f: F) -> F + #[allow(clippy::type_complexity)] + fn as_file_search(f: F) -> F where - F: for<'f> FnMut(&'f BucketActiveFile) -> ExactReaderFuture<'f> + Send, + F: for<'a> Fn( + &'a BucketActiveFile, + &'a [&'a [f32]], + VectorSearchMetric, + usize, + &'a (dyn Fn(i64) -> bool + Sync), + ) -> ExactFileSearchFuture<'a> + + Send + + Sync, { f } + /// Build a per-file search closure that runs the reference `exact_search` over + /// an in-memory `ArrayReader` whose vectors come from `vectors_for(file_name)`. + /// The caller wires a single query, so the returned outer `Vec` has one element. + #[allow(clippy::type_complexity)] + fn file_array_search( + vectors_for: fn(&str) -> Vec>>, + ) -> impl for<'a> Fn( + &'a BucketActiveFile, + &'a [&'a [f32]], + VectorSearchMetric, + usize, + &'a (dyn Fn(i64) -> bool + Sync), + ) -> ExactFileSearchFuture<'a> + + Send + + Sync { + as_file_search( + move |file: &BucketActiveFile, + queries: &[&[f32]], + metric: VectorSearchMetric, + exact_limit: usize, + is_excluded: &(dyn Fn(i64) -> bool + Sync)| + -> ExactFileSearchFuture<'_> { + let vectors = vectors_for(&file.file_name); + let dimension = vectors + .first() + .and_then(|v| v.as_ref()) + .map_or(2, |v| v.len()); + let file_name = file.file_name.clone(); + let query = queries[0].to_vec(); + Box::pin(async move { + let mut reader = ArrayReader::new(dimension, vectors); + Ok(vec![exact_search( + &file_name, + &mut reader, + &query, + metric, + exact_limit, + is_excluded, + )?]) + }) + }, + ) + } + fn column_by_name<'a>(batch: &'a RecordBatch, name: &str) -> Option<&'a Arc> { batch .schema() @@ -1000,6 +1120,7 @@ mod e2e_tests { /// materialization path the production best-first read reorders on top of; the /// tests below drive it directly through its `pub(crate)` components. #[allow(clippy::too_many_arguments)] + #[allow(clippy::type_complexity)] async fn materialize_via_splits( reader: DataFileReader, splits: &[PkVectorSearchSplit], @@ -1007,26 +1128,35 @@ mod e2e_tests { metric: VectorSearchMetric, limit: usize, ann: Option<&dyn PkVectorAnnSearcher>, - factory: &mut (dyn for<'f> FnMut(&'f BucketActiveFile) -> ExactReaderFuture<'f> + Send), + search: &(dyn for<'a> Fn( + &'a BucketActiveFile, + &'a [&'a [f32]], + VectorSearchMetric, + usize, + &'a (dyn Fn(i64) -> bool + Sync), + ) -> ExactFileSearchFuture<'a> + + Send + + Sync), opts: &HashMap, ) -> crate::Result> { let orch = PkVectorOrchestrator::new(reader.clone()); - // Wrap the per-file factory into the split-scoped shape search_candidates + // Wrap the per-file search into the split-scoped shape search_candidates // expects; the split index/split are unused here. - let mut wrapped = - as_split_factory(|_: usize, _: &PkVectorSearchSplit, f: &BucketActiveFile| factory(f)); + let wrapped = as_split_search( + |_: usize, + _: &PkVectorSearchSplit, + file: &BucketActiveFile, + queries: &[&[f32]], + metric: VectorSearchMetric, + exact_limit: usize, + is_excluded: &(dyn Fn(i64) -> bool + Sync)| + -> ExactFileSearchFuture<'_> { + search(file, queries, metric, exact_limit, is_excluded) + }, + ); let result = orch .search_candidates( - splits, - query, - metric, - limit, - limit, - ann, - &mut wrapped, - opts, - false, - None, + splits, query, metric, limit, limit, ann, &wrapped, opts, false, None, ) .await?; // Merge the two bounded lists into the best-first survivors the @@ -1049,11 +1179,7 @@ mod e2e_tests { let file_io = FileIOBuilder::new("memory").build().unwrap(); let reader = make_reader(file_io, "memory:/pkvo_zero"); let splits: Vec = Vec::new(); - let mut factory = as_split_factory( - |_: usize, _: &PkVectorSearchSplit, _: &BucketActiveFile| -> ExactReaderFuture<'_> { - Box::pin(async { unreachable!("no bucket search on eager-rejected input") }) - }, - ); + let factory = unreachable_split_search(); let opts = HashMap::new(); let err = PkVectorOrchestrator::new(reader) .search_candidates( @@ -1063,7 +1189,7 @@ mod e2e_tests { 0, 0, None, - &mut factory, + &factory, &opts, false, None, @@ -1079,11 +1205,7 @@ mod e2e_tests { let file_io = FileIOBuilder::new("memory").build().unwrap(); let reader = make_reader(file_io, "memory:/pkvo_empty_query"); let splits: Vec = Vec::new(); - let mut factory = as_split_factory( - |_: usize, _: &PkVectorSearchSplit, _: &BucketActiveFile| -> ExactReaderFuture<'_> { - Box::pin(async { unreachable!("no bucket search on eager-rejected input") }) - }, - ); + let factory = unreachable_split_search(); let opts = HashMap::new(); let err = PkVectorOrchestrator::new(reader) .search_candidates( @@ -1093,7 +1215,7 @@ mod e2e_tests { 5, 5, None, - &mut factory, + &factory, &opts, false, None, @@ -1138,14 +1260,9 @@ mod e2e_tests { }], }; // Exact fallback scans only "exact.mosaic": pos0 {1,0} d=1.0, pos1 {2,0} d=4.0. - let mut factory = as_file_factory(|f: &BucketActiveFile| -> ExactReaderFuture<'_> { - let vectors = match f.file_name.as_str() { - "exact.mosaic" => vec![Some(vec![1.0, 0.0]), Some(vec![2.0, 0.0])], - other => panic!("unexpected exact scan of covered file {other}"), - }; - Box::pin(async move { - Ok(Box::new(ArrayReader::new(2, vectors)) as Box) - }) + let factory = file_array_search(|file_name| match file_name { + "exact.mosaic" => vec![Some(vec![1.0, 0.0]), Some(vec![2.0, 0.0])], + other => panic!("unexpected exact scan of covered file {other}"), }); let opts = HashMap::new(); let batches = materialize_via_splits( @@ -1155,7 +1272,7 @@ mod e2e_tests { VectorSearchMetric::L2, 3, Some(&ann), - &mut factory, + &factory, &opts, ) .await @@ -1221,23 +1338,18 @@ mod e2e_tests { // b0: x = 1,4,6 -> d = 1,16,36. b1: x = 2,3,5 -> d = 4,9,25. // Global best 3: d1 (b0 pos0 id10), d4 (b1 pos0 id20), d9 (b1 pos1 id21). - let mut factory = as_file_factory(|f: &BucketActiveFile| -> ExactReaderFuture<'_> { - let vectors = match f.file_name.as_str() { - "b0.mosaic" => vec![ - Some(vec![1.0, 0.0]), - Some(vec![4.0, 0.0]), - Some(vec![6.0, 0.0]), - ], - "b1.mosaic" => vec![ - Some(vec![2.0, 0.0]), - Some(vec![3.0, 0.0]), - Some(vec![5.0, 0.0]), - ], - other => panic!("unexpected file {other}"), - }; - Box::pin(async move { - Ok(Box::new(ArrayReader::new(2, vectors)) as Box) - }) + let factory = file_array_search(|file_name| match file_name { + "b0.mosaic" => vec![ + Some(vec![1.0, 0.0]), + Some(vec![4.0, 0.0]), + Some(vec![6.0, 0.0]), + ], + "b1.mosaic" => vec![ + Some(vec![2.0, 0.0]), + Some(vec![3.0, 0.0]), + Some(vec![5.0, 0.0]), + ], + other => panic!("unexpected file {other}"), }); let opts = HashMap::new(); let batches = materialize_via_splits( @@ -1247,7 +1359,7 @@ mod e2e_tests { VectorSearchMetric::L2, 3, None, - &mut factory, + &factory, &opts, ) .await @@ -1292,19 +1404,14 @@ mod e2e_tests { }; // pos0 {1,0} d=1, pos1 {2,0} d=4 (DELETED), pos2 {3,0} d=9, pos3 {0,0} d=0. - let mut factory = as_file_factory(|f: &BucketActiveFile| -> ExactReaderFuture<'_> { - let vectors = match f.file_name.as_str() { - "d.mosaic" => vec![ - Some(vec![1.0, 0.0]), - Some(vec![2.0, 0.0]), - Some(vec![3.0, 0.0]), - Some(vec![0.0, 0.0]), - ], - other => panic!("unexpected file {other}"), - }; - Box::pin(async move { - Ok(Box::new(ArrayReader::new(2, vectors)) as Box) - }) + let factory = file_array_search(|file_name| match file_name { + "d.mosaic" => vec![ + Some(vec![1.0, 0.0]), + Some(vec![2.0, 0.0]), + Some(vec![3.0, 0.0]), + Some(vec![0.0, 0.0]), + ], + other => panic!("unexpected file {other}"), }); let opts = HashMap::new(); let batches = materialize_via_splits( @@ -1314,7 +1421,7 @@ mod e2e_tests { VectorSearchMetric::L2, 4, None, - &mut factory, + &factory, &opts, ) .await @@ -1355,18 +1462,13 @@ mod e2e_tests { }; // pos0 {3,0} d=9, pos1 {1,0} d=1, pos2 {2,0} d=4. Best-first = [1,2,0]. - let mut factory = as_file_factory(|f: &BucketActiveFile| -> ExactReaderFuture<'_> { - let vectors = match f.file_name.as_str() { - "o.mosaic" => vec![ - Some(vec![3.0, 0.0]), - Some(vec![1.0, 0.0]), - Some(vec![2.0, 0.0]), - ], - other => panic!("unexpected file {other}"), - }; - Box::pin(async move { - Ok(Box::new(ArrayReader::new(2, vectors)) as Box) - }) + let factory = file_array_search(|file_name| match file_name { + "o.mosaic" => vec![ + Some(vec![3.0, 0.0]), + Some(vec![1.0, 0.0]), + Some(vec![2.0, 0.0]), + ], + other => panic!("unexpected file {other}"), }); let opts = HashMap::new(); let batches = materialize_via_splits( @@ -1376,7 +1478,7 @@ mod e2e_tests { VectorSearchMetric::L2, 3, None, - &mut factory, + &factory, &opts, ) .await @@ -1416,18 +1518,35 @@ mod e2e_tests { active_files: vec![active("c.mosaic", 3)], }; // pos0 {3,0} d=9, pos1 {1,0} d=1, pos2 {2,0} d=4. - let mut factory = as_split_factory( - |_: usize, _: &PkVectorSearchSplit, f: &BucketActiveFile| -> ExactReaderFuture<'_> { - assert_eq!(f.file_name, "c.mosaic"); - Box::pin(async { - Ok(Box::new(ArrayReader::new( + let factory = as_split_search( + |_: usize, + _: &PkVectorSearchSplit, + file: &BucketActiveFile, + queries: &[&[f32]], + metric: VectorSearchMetric, + exact_limit: usize, + is_excluded: &(dyn Fn(i64) -> bool + Sync)| + -> ExactFileSearchFuture<'_> { + assert_eq!(file.file_name, "c.mosaic"); + let file_name = file.file_name.clone(); + let query = queries[0].to_vec(); + Box::pin(async move { + let mut reader = ArrayReader::new( 2, vec![ Some(vec![3.0, 0.0]), Some(vec![1.0, 0.0]), Some(vec![2.0, 0.0]), ], - )) as Box) + ); + Ok(vec![exact_search( + &file_name, + &mut reader, + &query, + metric, + exact_limit, + is_excluded, + )?]) }) }, ); @@ -1440,7 +1559,7 @@ mod e2e_tests { 2, 2, None, - &mut factory, + &factory, &opts, false, None, @@ -1485,18 +1604,35 @@ mod e2e_tests { active_files: vec![active("r.mosaic", 3)], }; // pos0 {3,0} d=9, pos1 {1,0} d=1, pos2 {2,0} d=4. - let mut factory = as_split_factory( - |_: usize, _: &PkVectorSearchSplit, f: &BucketActiveFile| -> ExactReaderFuture<'_> { - assert_eq!(f.file_name, "r.mosaic"); - Box::pin(async { - Ok(Box::new(ArrayReader::new( + let factory = as_split_search( + |_: usize, + _: &PkVectorSearchSplit, + file: &BucketActiveFile, + queries: &[&[f32]], + metric: VectorSearchMetric, + exact_limit: usize, + is_excluded: &(dyn Fn(i64) -> bool + Sync)| + -> ExactFileSearchFuture<'_> { + assert_eq!(file.file_name, "r.mosaic"); + let file_name = file.file_name.clone(); + let query = queries[0].to_vec(); + Box::pin(async move { + let mut reader = ArrayReader::new( 2, vec![ Some(vec![3.0, 0.0]), Some(vec![1.0, 0.0]), Some(vec![2.0, 0.0]), ], - )) as Box) + ); + Ok(vec![exact_search( + &file_name, + &mut reader, + &query, + metric, + exact_limit, + is_excluded, + )?]) }) }, ); @@ -1515,7 +1651,7 @@ mod e2e_tests { 3, 3, None, - &mut factory, + &factory, &opts, false, Some(&residual_by_split), @@ -1556,11 +1692,7 @@ mod e2e_tests { ann_segments: Vec::new(), active_files: vec![active("m.mosaic", 2)], }; - let mut factory = as_split_factory( - |_: usize, _: &PkVectorSearchSplit, _: &BucketActiveFile| -> ExactReaderFuture<'_> { - Box::pin(async { unreachable!("length guard must fire before any bucket search") }) - }, - ); + let factory = unreachable_split_search(); // Two residual maps for a single split. let residual_by_split: Vec> = vec![HashMap::new(), HashMap::new()]; @@ -1573,7 +1705,7 @@ mod e2e_tests { 3, 3, None, - &mut factory, + &factory, &opts, false, Some(&residual_by_split), @@ -1606,11 +1738,7 @@ mod e2e_tests { ann_segments: Vec::new(), active_files: vec![active("f.mosaic", 2)], }; - let mut factory = as_split_factory( - |_: usize, _: &PkVectorSearchSplit, _: &BucketActiveFile| -> ExactReaderFuture<'_> { - Box::pin(async { unreachable!("fast mode must not read exact") }) - }, - ); + let factory = unreachable_split_search(); let opts = HashMap::new(); let result = PkVectorOrchestrator::new(make_reader(file_io, table_path)) .search_candidates( @@ -1620,7 +1748,7 @@ mod e2e_tests { 2, 2, None, - &mut factory, + &factory, &opts, true, None, diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index 9f50ee46..6f4cb27a 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -34,8 +34,8 @@ use crate::table::pk_vector_data_file_reader::{ }; use crate::table::pk_vector_indexed_split_read::{expand_ranges, PkVectorIndexedSplitRead}; use crate::table::pk_vector_orchestrator::{ - as_split_exact_reader_factory, build_indexed_splits, merge_candidates, - OrchestratorSearchResult, PkVectorCandidate, PkVectorOrchestrator, PkVectorSearchSplit, + as_split_exact_file_search, build_indexed_splits, merge_candidates, OrchestratorSearchResult, + PkVectorCandidate, PkVectorOrchestrator, PkVectorSearchSplit, }; use crate::table::pk_vector_position_read::{ PkVectorPositionRead, PKEY_VECTOR_POSITION_COLUMN, SEARCH_SCORE_COLUMN, @@ -49,7 +49,7 @@ use crate::table::{ use crate::vector_search::{GlobalIndexIOMeta, SearchResult, VectorSearch}; use crate::vindex::is_vindex_index_type; use crate::vindex::pkvector::ann::VindexAnnSearcher; -use crate::vindex::pkvector::bucket::{BucketActiveFile, BucketAnnSegment, ExactReaderFuture}; +use crate::vindex::pkvector::bucket::{BucketActiveFile, BucketAnnSegment, ExactFileSearchFuture}; use crate::vindex::pkvector::metric::VectorSearchMetric; use crate::vindex::reader::VindexVectorGlobalIndexReader; use arrow_array::{Array, FixedSizeListArray, Float32Array, Int64Array, ListArray, RecordBatch}; @@ -577,18 +577,23 @@ impl<'a> VectorSearchBuilder<'a> { None => None, }; - // Build the exact-fallback vector reader on demand: the kernel calls this - // only for a file it actually searches (uncovered by ANN, residual-allowed, - // and only when the search mode is not FAST). Everything the future needs is + // Build the exact-fallback search on demand: the kernel calls this only + // for a file it actually searches (uncovered by ANN, residual-allowed, and + // only when the search mode is not FAST). Everything the future needs is // cloned/owned up front so it borrows neither the split nor the file across - // the await. + // the await. The search streams the file's vector column one Arrow batch at + // a time into per-query bounded heaps (the caller passes a single query). let reader_for_factory = reader.clone(); let vector_field_for_factory = vector_field.clone(); - let mut factory = as_split_exact_reader_factory( + let factory = as_split_exact_file_search( move |_split_index: usize, split: &PkVectorSearchSplit, - file: &BucketActiveFile| - -> ExactReaderFuture<'_> { + file: &BucketActiveFile, + queries: &[&[f32]], + metric: VectorSearchMetric, + exact_limit: usize, + is_excluded: &(dyn Fn(i64) -> bool + Sync)| + -> ExactFileSearchFuture<'_> { let reader = reader_for_factory.clone(); let vector_field = vector_field_for_factory.clone(); let data_split = split.data_split.clone(); @@ -596,10 +601,15 @@ impl<'a> VectorSearchBuilder<'a> { file_name: file.file_name.clone(), row_count: file.row_count, }; + let owned_queries: Vec> = queries.iter().map(|q| q.to_vec()).collect(); Box::pin(async move { let factory = DataFilePkVectorReaderFactory::new(reader, data_split, vector_field)?; - factory.create(&active).await + let query_refs: Vec<&[f32]> = + owned_queries.iter().map(|q| q.as_slice()).collect(); + factory + .search_file(&active, &query_refs, metric, exact_limit, is_excluded) + .await }) }, ); @@ -612,7 +622,7 @@ impl<'a> VectorSearchBuilder<'a> { limit, indexed_limit, Some(&ann_searcher), - &mut factory, + &factory, &search_options, skip_exact_fallback, residual_by_split.as_deref(), diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs b/crates/paimon/src/vindex/pkvector/bucket.rs index f3cfc35f..3b2f7fb1 100644 --- a/crates/paimon/src/vindex/pkvector/bucket.rs +++ b/crates/paimon/src/vindex/pkvector/bucket.rs @@ -23,16 +23,19 @@ use futures::future::BoxFuture; use super::ann::PkVectorAnnSearcher; use super::data_invalid; -use super::exact::exact_search; use super::metric::{java_float_compare, VectorSearchMetric}; -use super::reader::PkVectorReader; use super::result::PkVectorSearchResult; use crate::deletion_vector::DeletionVector; use crate::spec::PkVectorSourceMeta; -/// Future returned by the exact-fallback reader factory: builds the sequential -/// vector reader for one data file on demand. -pub(crate) type ExactReaderFuture<'a> = BoxFuture<'a, crate::Result>>; +/// 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 +/// passed to the closure). The table-layer implementation streams the file's +/// vector column one Arrow batch at a time, so peak memory is one batch plus the +/// bounded heaps, not the whole column. `is_excluded(position)` folds residual + +/// deletion-vector exclusion (built by `bucket_search`, which owns those inputs). +pub(crate) type ExactFileSearchFuture<'a> = + BoxFuture<'a, crate::Result>>>; /// One ANN segment to be searched by the bucket kernel. `source_meta` resolves /// segment ordinals back to physical `(data file, position)` and drives live-row @@ -164,13 +167,21 @@ pub(crate) struct BucketSearchResult { /// with an empty set) has no allowed rows and produces no candidates. Mirrors Java /// `rowRangesByFile`. #[allow(clippy::too_many_arguments)] +#[allow(clippy::type_complexity)] pub(crate) async fn bucket_search( ann_searcher: Option<&dyn PkVectorAnnSearcher>, ann_segments: &[BucketAnnSegment], active_files: &[BucketActiveFile], deletion_vectors: &HashMap>, - exact_reader_factory: &mut (dyn for<'a> FnMut(&'a BucketActiveFile) -> ExactReaderFuture<'a> - + Send), + exact_file_search: &(dyn for<'a> Fn( + &'a BucketActiveFile, + &'a [&'a [f32]], + VectorSearchMetric, + usize, + &'a (dyn Fn(i64) -> bool + Sync), + ) -> ExactFileSearchFuture<'a> + + Send + + Sync), query: &[f32], metric: VectorSearchMetric, indexed_limit: usize, @@ -185,6 +196,15 @@ pub(crate) async fn bucket_search( if exact_limit == 0 { return Err(data_invalid("vector search limit must be positive")); } + // Validate-before-OPEN: a non-finite query element fails loud before any + // exact-file search closure is invoked (so before any file stream is opened). + // The per-file closure additionally validates the query dimension before it + // polls its stream (validate-before-POLL). + if let Some(i) = query.iter().position(|v| !v.is_finite()) { + return Err(data_invalid(format!( + "query vector element at position {i} must be finite" + ))); + } let mut files_by_name: HashMap<&str, &BucketActiveFile> = HashMap::new(); for file in active_files { @@ -307,15 +327,22 @@ pub(crate) async fn bucket_search( }, } }; - let mut reader = exact_reader_factory(file).await?; - for result in exact_search( - &file.file_name, - reader.as_mut(), - query, + // Search this one file for its per-query exact Top-K. The caller passes + // a single-query slice and reads the sole returned list. + let queries: [&[f32]; 1] = [query]; + let per_query = exact_file_search( + file, + &queries, metric, exact_limit, - &is_excluded, - )? { + &is_excluded as &(dyn Fn(i64) -> bool + Sync), + ) + .await?; + let results = per_query + .into_iter() + .next() + .ok_or_else(|| data_invalid("exact file search returned no per-query results"))?; + for result in results { add_candidate(&mut exact_heap, result, exact_limit); } } @@ -333,6 +360,7 @@ mod tests { use super::*; use crate::spec::PkVectorSourceFile; use crate::vindex::pkvector::ann::PkVectorAnnSearcher; + use crate::vindex::pkvector::exact::exact_search; use crate::vindex::pkvector::reader::test_support::ArrayReader; use roaring::RoaringBitmap; @@ -354,17 +382,93 @@ mod tests { } } - /// Coerce a closure into the higher-ranked exact-reader factory shape so its - /// returned future borrows for exactly the argument's lifetime. Closure return + /// Coerce a closure into the higher-ranked per-file exact-search shape so its + /// returned future borrows for exactly the argument lifetime. Closure return /// types cannot express this borrow through inference alone, so the bound is - /// supplied here. - fn as_factory(f: F) -> F + /// supplied here. The closure is `Fn + Send + Sync` (called concurrently by + /// the parallel search), not `FnMut`. + fn as_search(f: F) -> F where - F: for<'a> FnMut(&'a BucketActiveFile) -> ExactReaderFuture<'a> + Send, + F: for<'a> Fn( + &'a BucketActiveFile, + &'a [&'a [f32]], + VectorSearchMetric, + usize, + &'a (dyn Fn(i64) -> bool + Sync), + ) -> ExactFileSearchFuture<'a> + + Send + + Sync, { f } + /// Build a per-file exact-search closure that runs the reference `exact_search` + /// over an in-memory `ArrayReader` of `vectors`, returning one per-query list. + /// The caller wires a single query, so the returned outer `Vec` has one element. + #[allow(clippy::type_complexity)] + fn array_search( + vectors: Vec>>, + ) -> impl for<'a> Fn( + &'a BucketActiveFile, + &'a [&'a [f32]], + VectorSearchMetric, + usize, + &'a (dyn Fn(i64) -> bool + Sync), + ) -> ExactFileSearchFuture<'a> + + Send + + Sync { + let vectors = Arc::new(vectors); + as_search( + move |file: &BucketActiveFile, + queries: &[&[f32]], + metric: VectorSearchMetric, + exact_limit: usize, + is_excluded: &(dyn Fn(i64) -> bool + Sync)| + -> ExactFileSearchFuture<'_> { + let dimension = vectors + .first() + .and_then(|v| v.as_ref()) + .map_or(0, |v| v.len()); + let owned: Vec>> = (*vectors).clone(); + let file_name = file.file_name.clone(); + let query: Vec = queries[0].to_vec(); + Box::pin(async move { + let mut reader = ArrayReader::new(dimension, owned); + let results = exact_search( + &file_name, + &mut reader, + &query, + metric, + exact_limit, + is_excluded, + )?; + Ok(vec![results]) + }) + }, + ) + } + + /// A per-file exact-search closure that must never be invoked. + #[allow(clippy::type_complexity)] + fn unreachable_search() -> impl for<'a> Fn( + &'a BucketActiveFile, + &'a [&'a [f32]], + VectorSearchMetric, + usize, + &'a (dyn Fn(i64) -> bool + Sync), + ) -> ExactFileSearchFuture<'a> + + Send + + Sync { + as_search( + |_: &BucketActiveFile, + _: &[&[f32]], + _: VectorSearchMetric, + _: usize, + _: &(dyn Fn(i64) -> bool + Sync)| + -> ExactFileSearchFuture<'_> { Box::pin(async { unreachable!() }) }, + ) + } + /// Fake ANN searcher returning preset results and recording calls. struct FakeAnnSearcher { result: Vec, @@ -412,17 +516,11 @@ mod tests { ], }; // Exact file has three rows; query nearest is position 0. - let mut factory = as_factory(|_f: &BucketActiveFile| -> ExactReaderFuture<'_> { - let reader = ArrayReader::new( - 2, - vec![ - Some(vec![1.0, 0.0]), - Some(vec![9.0, 0.0]), - Some(vec![8.0, 0.0]), - ], - ); - Box::pin(async move { Ok(Box::new(reader) as Box) }) - }); + let factory = array_search(vec![ + Some(vec![1.0, 0.0]), + Some(vec![9.0, 0.0]), + Some(vec![8.0, 0.0]), + ]); let active_files = vec![active("ann.mosaic", 3), active("exact.mosaic", 3)]; let dvs: HashMap> = HashMap::new(); let opts = HashMap::new(); @@ -432,7 +530,7 @@ mod tests { &[segment], &active_files, &dvs, - &mut factory, + &factory, &[1.0, 0.0], VectorSearchMetric::L2, 3, // indexed_limit @@ -456,15 +554,13 @@ mod tests { #[tokio::test] async fn test_rejects_non_positive_limit() { - let mut factory = as_factory(|_: &BucketActiveFile| -> ExactReaderFuture<'_> { - Box::pin(async { unreachable!() }) - }); + let factory = unreachable_search(); let err = bucket_search( None, &[], &[], &HashMap::new(), - &mut factory, + &factory, &[0.0, 0.0], VectorSearchMetric::L2, 0, @@ -500,15 +596,13 @@ mod tests { hit("data-1", 1), ], }; - let mut factory = as_factory(|_: &BucketActiveFile| -> ExactReaderFuture<'_> { - Box::pin(async { unreachable!() }) - }); + let factory = unreachable_search(); let out = bucket_search( Some(&ann), &[segment], &[active("data-1", 3)], &HashMap::new(), - &mut factory, + &factory, &[0.0, 0.0], VectorSearchMetric::L2, 3, @@ -557,15 +651,13 @@ mod tests { }, ], }; - let mut factory = as_factory(|_: &BucketActiveFile| -> ExactReaderFuture<'_> { - Box::pin(async { unreachable!() }) - }); + let factory = unreachable_search(); let out = bucket_search( Some(&ann), &[segment], &[active("data-1", 2)], &HashMap::new(), - &mut factory, + &factory, &[0.0, 0.0], VectorSearchMetric::L2, 1, @@ -598,22 +690,37 @@ mod tests { }], }; let calls = std::sync::Mutex::new(Vec::::new()); - let mut factory = as_factory(|f: &BucketActiveFile| -> ExactReaderFuture<'_> { - calls.lock().unwrap().push(f.file_name.clone()); - // data-2 vectors: pos0 {1,0} dist 1.0, pos1 {3,0} dist 9.0 - Box::pin(async { - Ok(Box::new(ArrayReader::new( - 2, - vec![Some(vec![1.0, 0.0]), Some(vec![3.0, 0.0])], - )) as Box) - }) - }); + let factory = as_search( + |file: &BucketActiveFile, + queries: &[&[f32]], + metric: VectorSearchMetric, + exact_limit: usize, + is_excluded: &(dyn Fn(i64) -> bool + Sync)| + -> ExactFileSearchFuture<'_> { + calls.lock().unwrap().push(file.file_name.clone()); + // data-2 vectors: pos0 {1,0} dist 1.0, pos1 {3,0} dist 9.0 + let file_name = file.file_name.clone(); + let query = queries[0].to_vec(); + Box::pin(async move { + let mut reader = + ArrayReader::new(2, vec![Some(vec![1.0, 0.0]), Some(vec![3.0, 0.0])]); + Ok(vec![exact_search( + &file_name, + &mut reader, + &query, + metric, + exact_limit, + is_excluded, + )?]) + }) + }, + ); let out = bucket_search( Some(&ann), &[segment], &[active("data-1", 2), active("data-2", 2)], &HashMap::new(), - &mut factory, + &factory, &[0.0, 0.0], VectorSearchMetric::L2, 2, @@ -650,17 +757,34 @@ mod tests { async fn test_exact_fallback_merges_files_and_applies_deletion_vectors() { // No ANN. data-1 pos0 {0,0} deleted; remaining candidates merge across files. let calls = std::sync::Mutex::new(0usize); - let mut factory = as_factory(|f: &BucketActiveFile| -> ExactReaderFuture<'_> { - *calls.lock().unwrap() += 1; - let vectors = match f.file_name.as_str() { - "data-1" => vec![Some(vec![0.0, 0.0]), Some(vec![2.0, 0.0])], - "data-2" => vec![Some(vec![1.0, 0.0]), None], - _ => unreachable!(), - }; - Box::pin(async move { - Ok(Box::new(ArrayReader::new(2, vectors)) as Box) - }) - }); + let factory = as_search( + |file: &BucketActiveFile, + queries: &[&[f32]], + metric: VectorSearchMetric, + exact_limit: usize, + is_excluded: &(dyn Fn(i64) -> bool + Sync)| + -> ExactFileSearchFuture<'_> { + *calls.lock().unwrap() += 1; + let vectors = match file.file_name.as_str() { + "data-1" => vec![Some(vec![0.0, 0.0]), Some(vec![2.0, 0.0])], + "data-2" => vec![Some(vec![1.0, 0.0]), None], + _ => unreachable!(), + }; + let file_name = file.file_name.clone(); + let query = queries[0].to_vec(); + Box::pin(async move { + let mut reader = ArrayReader::new(2, vectors); + Ok(vec![exact_search( + &file_name, + &mut reader, + &query, + metric, + exact_limit, + is_excluded, + )?]) + }) + }, + ); let mut dvs: HashMap> = HashMap::new(); let mut bm = RoaringBitmap::new(); bm.insert(0); // data-1 position 0 deleted @@ -671,7 +795,7 @@ mod tests { &[], &[active("data-1", 2), active("data-2", 2)], &dvs, - &mut factory, + &factory, &[0.0, 0.0], VectorSearchMetric::L2, 2, @@ -707,15 +831,13 @@ mod tests { #[tokio::test] async fn test_rejects_duplicate_active_file_name() { - let mut factory = as_factory(|_: &BucketActiveFile| -> ExactReaderFuture<'_> { - Box::pin(async { unreachable!() }) - }); + let factory = unreachable_search(); let err = bucket_search( None, &[], &[active("dup", 1), active("dup", 1)], &HashMap::new(), - &mut factory, + &factory, &[0.0, 0.0], VectorSearchMetric::L2, 1, @@ -735,15 +857,13 @@ mod tests { // Segment references data-1 with 2 rows, but the active file has 3 rows. // An active source with a mismatched row count is still a hard error. let segment = BucketAnnSegment::for_test(meta(&[("data-1", 2)])); - let mut factory = as_factory(|_: &BucketActiveFile| -> ExactReaderFuture<'_> { - Box::pin(async { unreachable!() }) - }); + let factory = unreachable_search(); let err = bucket_search( Some(&ann), &[segment], &[active("data-1", 3)], &HashMap::new(), - &mut factory, + &factory, &[0.0, 0.0], VectorSearchMetric::L2, 1, @@ -775,16 +895,23 @@ mod tests { }], }; let calls = std::sync::Mutex::new(Vec::::new()); - let mut factory = as_factory(|f: &BucketActiveFile| -> ExactReaderFuture<'_> { - calls.lock().unwrap().push(f.file_name.clone()); - Box::pin(async { unreachable!("only data-1 is active and it is ANN-covered") }) - }); + let factory = as_search( + |file: &BucketActiveFile, + _: &[&[f32]], + _: VectorSearchMetric, + _: usize, + _: &(dyn Fn(i64) -> bool + Sync)| + -> ExactFileSearchFuture<'_> { + calls.lock().unwrap().push(file.file_name.clone()); + Box::pin(async { unreachable!("only data-1 is active and it is ANN-covered") }) + }, + ); let out = bucket_search( Some(&ann), &[segment], &[active("data-1", 2)], &HashMap::new(), - &mut factory, + &factory, &[0.0, 0.0], VectorSearchMetric::L2, 2, @@ -814,15 +941,13 @@ mod tests { #[tokio::test] async fn test_rejects_segments_without_ann_searcher() { let segment = BucketAnnSegment::for_test(meta(&[("data-1", 2)])); - let mut factory = as_factory(|_: &BucketActiveFile| -> ExactReaderFuture<'_> { - Box::pin(async { unreachable!() }) - }); + let factory = unreachable_search(); let err = bucket_search( None, &[segment], &[active("data-1", 2)], &HashMap::new(), - &mut factory, + &factory, &[0.0, 0.0], VectorSearchMetric::L2, 1, @@ -843,15 +968,13 @@ mod tests { async fn test_skip_exact_fallback_does_not_call_factory() { // No ANN segments, two active files. With skip_exact_fallback = true the // factory must never be called and the result is empty. - let mut factory = as_factory(|_: &BucketActiveFile| -> ExactReaderFuture<'_> { - Box::pin(async { unreachable!() }) - }); + let factory = unreachable_search(); let out = bucket_search( None, &[], &[active("data-1", 2), active("data-2", 2)], &HashMap::new(), - &mut factory, + &factory, &[0.0, 0.0], VectorSearchMetric::L2, 2, @@ -881,15 +1004,13 @@ mod tests { index_meta: vec![4, 5, 6], }; let ann = FakeAnnSearcher { result: vec![] }; - let mut factory = as_factory(|_: &BucketActiveFile| -> ExactReaderFuture<'_> { - Box::pin(async { unreachable!() }) - }); + let factory = unreachable_search(); let err = bucket_search( Some(&ann), &[seg1, seg2], &[active("data-1", 2), active("data-2", 2)], &HashMap::new(), - &mut factory, + &factory, &[0.0, 0.0], VectorSearchMetric::L2, 1, @@ -921,15 +1042,13 @@ mod tests { index_meta: vec![4, 5, 6], }; let ann = FakeAnnSearcher { result: vec![] }; - let mut factory = as_factory(|_: &BucketActiveFile| -> ExactReaderFuture<'_> { - Box::pin(async { unreachable!() }) - }); + let factory = unreachable_search(); let err = bucket_search( Some(&ann), &[seg1, seg2], &[active("data-1", 2), active("data-2", 2)], &HashMap::new(), - &mut factory, + &factory, &[0.0, 0.0], VectorSearchMetric::L2, 1, @@ -950,15 +1069,13 @@ mod tests { #[tokio::test] async fn test_negative_active_row_count_rejected() { - let mut factory = as_factory(|_: &BucketActiveFile| -> ExactReaderFuture<'_> { - Box::pin(async { unreachable!() }) - }); + let factory = unreachable_search(); let err = bucket_search( None, &[], &[active("data-1", -1)], &HashMap::new(), - &mut factory, + &factory, &[0.0, 0.0], VectorSearchMetric::L2, 1, @@ -1014,18 +1131,11 @@ mod tests { // No ANN. data-1 has 3 rows: pos0 {1,0} dist 1.0, pos1 {2,0} dist 4.0, // pos2 {3,0} dist 9.0. residual allows only {0, 2} -> pos1 excluded even // though it is not deletion-vector deleted. - let mut factory = as_factory(|_: &BucketActiveFile| -> ExactReaderFuture<'_> { - Box::pin(async { - Ok(Box::new(ArrayReader::new( - 2, - vec![ - Some(vec![1.0, 0.0]), - Some(vec![2.0, 0.0]), - Some(vec![3.0, 0.0]), - ], - )) as Box) - }) - }); + let factory = array_search(vec![ + Some(vec![1.0, 0.0]), + Some(vec![2.0, 0.0]), + Some(vec![3.0, 0.0]), + ]); let mut residual: HashMap = HashMap::new(); residual.insert("data-1".into(), treemap(&[0, 2])); let out = bucket_search( @@ -1033,7 +1143,7 @@ mod tests { &[], &[active("data-1", 3)], &HashMap::new(), - &mut factory, + &factory, &[0.0, 0.0], VectorSearchMetric::L2, 5, @@ -1070,15 +1180,30 @@ mod tests { // residual covers only data-1; data-2 has no entry -> no allowed rows, so // data-2 is skipped entirely (its factory reader is never built). let calls = std::sync::Mutex::new(Vec::::new()); - let mut factory = as_factory(|f: &BucketActiveFile| -> ExactReaderFuture<'_> { - calls.lock().unwrap().push(f.file_name.clone()); - Box::pin(async { - Ok(Box::new(ArrayReader::new( - 2, - vec![Some(vec![1.0, 0.0]), Some(vec![2.0, 0.0])], - )) as Box) - }) - }); + let factory = as_search( + |file: &BucketActiveFile, + queries: &[&[f32]], + metric: VectorSearchMetric, + exact_limit: usize, + is_excluded: &(dyn Fn(i64) -> bool + Sync)| + -> ExactFileSearchFuture<'_> { + calls.lock().unwrap().push(file.file_name.clone()); + let file_name = file.file_name.clone(); + let query = queries[0].to_vec(); + Box::pin(async move { + let mut reader = + ArrayReader::new(2, vec![Some(vec![1.0, 0.0]), Some(vec![2.0, 0.0])]); + Ok(vec![exact_search( + &file_name, + &mut reader, + &query, + metric, + exact_limit, + is_excluded, + )?]) + }) + }, + ); let mut residual: HashMap = HashMap::new(); residual.insert("data-1".into(), treemap(&[0, 1])); let out = bucket_search( @@ -1086,7 +1211,7 @@ mod tests { &[], &[active("data-1", 2), active("data-2", 2)], &HashMap::new(), - &mut factory, + &factory, &[0.0, 0.0], VectorSearchMetric::L2, 5, @@ -1111,10 +1236,19 @@ mod tests { // data-1 has an entry but it is empty -> no allowed rows, skipped without // reading. Mirrors a file with no residual matches. let calls = std::sync::Mutex::new(0usize); - let mut factory = as_factory(|_: &BucketActiveFile| -> ExactReaderFuture<'_> { - *calls.lock().unwrap() += 1; - Box::pin(async { unreachable!("data-1 has an empty allow set and must not be read") }) - }); + let factory = as_search( + |_: &BucketActiveFile, + _: &[&[f32]], + _: VectorSearchMetric, + _: usize, + _: &(dyn Fn(i64) -> bool + Sync)| + -> ExactFileSearchFuture<'_> { + *calls.lock().unwrap() += 1; + Box::pin(async { + unreachable!("data-1 has an empty allow set and must not be read") + }) + }, + ); let mut residual: HashMap = HashMap::new(); residual.insert("data-1".into(), treemap(&[])); let out = bucket_search( @@ -1122,7 +1256,7 @@ mod tests { &[], &[active("data-1", 3)], &HashMap::new(), - &mut factory, + &factory, &[0.0, 0.0], VectorSearchMetric::L2, 5, @@ -1142,18 +1276,11 @@ mod tests { async fn test_exact_residual_intersects_with_deletion_vector() { // residual allows {0, 1, 2} but the deletion vector deletes pos0; the // surviving candidates are the residual-allowed AND not-deleted rows. - let mut factory = as_factory(|_: &BucketActiveFile| -> ExactReaderFuture<'_> { - Box::pin(async { - Ok(Box::new(ArrayReader::new( - 2, - vec![ - Some(vec![1.0, 0.0]), - Some(vec![2.0, 0.0]), - Some(vec![3.0, 0.0]), - ], - )) as Box) - }) - }); + let factory = array_search(vec![ + Some(vec![1.0, 0.0]), + Some(vec![2.0, 0.0]), + Some(vec![3.0, 0.0]), + ]); let mut dvs: HashMap> = HashMap::new(); let mut bm = RoaringBitmap::new(); bm.insert(0); // pos0 deleted @@ -1165,7 +1292,7 @@ mod tests { &[], &[active("data-1", 3)], &dvs, - &mut factory, + &factory, &[0.0, 0.0], VectorSearchMetric::L2, 5, @@ -1185,4 +1312,44 @@ mod tests { vec![1, 2] ); } + + #[tokio::test] + async fn test_non_finite_query_fails_before_opening_any_file() { + // A non-finite query element must fail loud before the exact-file search + // closure is ever invoked (validate-before-OPEN). The recording closure + // flips a flag if called; the flag must stay false. + let called = std::sync::atomic::AtomicBool::new(false); + let factory = as_search( + |_: &BucketActiveFile, + _: &[&[f32]], + _: VectorSearchMetric, + _: usize, + _: &(dyn Fn(i64) -> bool + Sync)| + -> ExactFileSearchFuture<'_> { + called.store(true, std::sync::atomic::Ordering::SeqCst); + Box::pin(async { unreachable!("closure must not run for a malformed query") }) + }, + ); + let err = bucket_search( + None, + &[], + &[active("data-1", 2)], + &HashMap::new(), + &factory, + &[f32::NAN, 0.0], + VectorSearchMetric::L2, + 2, + 2, + &HashMap::new(), + false, + None, + ) + .await + .unwrap_err(); + assert!(err.to_string().contains("finite"), "got: {err}"); + assert!( + !called.load(std::sync::atomic::Ordering::SeqCst), + "the exact-file search closure must not be invoked for a malformed query" + ); + } } diff --git a/crates/paimon/src/vindex/pkvector/exact.rs b/crates/paimon/src/vindex/pkvector/exact.rs index e5d043c6..80333b58 100644 --- a/crates/paimon/src/vindex/pkvector/exact.rs +++ b/crates/paimon/src/vindex/pkvector/exact.rs @@ -19,7 +19,10 @@ use std::cmp::Ordering; use std::collections::BinaryHeap; use super::data_invalid; -use super::metric::{java_float_compare, VectorSearchMetric}; +use super::metric::java_float_compare; +#[cfg(test)] +use super::metric::VectorSearchMetric; +#[cfg(test)] use super::reader::PkVectorReader; use super::result::PkVectorSearchResult; @@ -28,7 +31,10 @@ use super::result::PkVectorSearchResult; /// top therefore evicts the least-wanted candidate. Uses `java_float_compare` /// for a deterministic total order over f32 that ranks NaN distances as worst /// (largest), so a NaN is evicted before any finite candidate (no panic). -struct WorstFirst(PkVectorSearchResult); +/// +/// Ordered by distance then row position only (a single data file, so the file +/// name is constant across all its candidates). +pub(crate) struct WorstFirst(PkVectorSearchResult); impl PartialEq for WorstFirst { fn eq(&self, other: &Self) -> bool { @@ -50,15 +56,74 @@ impl Ord for WorstFirst { /// True if `candidate` ranks strictly better (BEST_FIRST) than the current /// worst-on-heap `weakest`: smaller distance, ties broken by smaller position. -fn is_better_than(candidate: &PkVectorSearchResult, weakest: &PkVectorSearchResult) -> bool { +pub(crate) fn is_better_than( + candidate: &PkVectorSearchResult, + weakest: &PkVectorSearchResult, +) -> bool { java_float_compare(candidate.distance, weakest.distance) .then_with(|| candidate.row_position.cmp(&weakest.row_position)) == Ordering::Less } +/// Add `candidate` to a bounded (size `limit`) BEST_FIRST Top-K max-heap over one +/// file's candidates: push if under capacity, else replace the current worst iff +/// the candidate beats it. `O(log limit)` per call. The single shared push step +/// so the whole-file `exact_search` and the streaming per-file search cannot +/// drift in how they bound their heaps. +pub(crate) fn push_bounded( + heap: &mut BinaryHeap, + candidate: PkVectorSearchResult, + limit: usize, +) { + if heap.len() < limit { + heap.push(WorstFirst(candidate)); + } else if heap + .peek() + .is_some_and(|worst| is_better_than(&candidate, &worst.0)) + { + heap.pop(); + heap.push(WorstFirst(candidate)); + } +} + +/// Drain a bounded Top-K heap into a BEST_FIRST-sorted result list: distance +/// ASC, then row_position ASC (single file, so data_file_name is constant). +pub(crate) fn drain_best_first(heap: BinaryHeap) -> Vec { + let mut results: Vec = heap.into_iter().map(|w| w.0).collect(); + results.sort_by(|a, b| { + java_float_compare(a.distance, b.distance).then_with(|| a.row_position.cmp(&b.row_position)) + }); + results +} + +/// Validate one query vector against the index dimension: length must match and +/// every element must be finite. Shared by the whole-file `exact_search` and the +/// table-layer streaming per-file search so both reject the same malformed +/// queries with the same messages (before any read is performed). +pub(crate) fn validate_query(query: &[f32], dimension: usize) -> crate::Result<()> { + if query.len() != dimension { + return Err(data_invalid(format!( + "query vector dimension does not match: index expects {}, got {}", + dimension, + query.len() + ))); + } + if let Some(i) = query.iter().position(|v| !v.is_finite()) { + return Err(data_invalid(format!( + "query vector element at position {i} must be finite" + ))); + } + Ok(()) +} + /// Exact Top-K over one sequential physical-row vector source. Mirrors Java /// `PkVectorExactSearcher.search`. Results are sorted BEST_FIRST: distance ASC, /// then row_position ASC (single file, so data_file_name is constant). +/// +/// No longer the production read path (the table layer streams the vector column +/// one Arrow batch at a time via the per-file search closure); kept as a tested +/// reference the streaming search must stay byte-identical to. +#[cfg(test)] pub(crate) fn exact_search( data_file_name: &str, reader: &mut dyn PkVectorReader, @@ -67,21 +132,10 @@ pub(crate) fn exact_search( limit: usize, is_excluded: &dyn Fn(i64) -> bool, ) -> crate::Result> { - if query.len() != reader.dimension() { - return Err(data_invalid(format!( - "query vector dimension does not match: index expects {}, got {}", - reader.dimension(), - query.len() - ))); - } + validate_query(query, reader.dimension())?; if limit == 0 { return Err(data_invalid("vector search limit must be positive")); } - if let Some(i) = query.iter().position(|v| !v.is_finite()) { - return Err(data_invalid(format!( - "query vector element at position {i} must be finite" - ))); - } let row_count = reader.row_count(); if row_count < 0 { return Err(data_invalid(format!( @@ -101,22 +155,10 @@ pub(crate) fn exact_search( row_position: position, distance: metric.compute_distance(query, &reuse), }; - if heap.len() < limit { - heap.push(WorstFirst(candidate)); - } else if heap - .peek() - .is_some_and(|worst| is_better_than(&candidate, &worst.0)) - { - heap.pop(); - heap.push(WorstFirst(candidate)); - } + push_bounded(&mut heap, candidate, limit); } - let mut results: Vec = heap.into_iter().map(|w| w.0).collect(); - results.sort_by(|a, b| { - java_float_compare(a.distance, b.distance).then_with(|| a.row_position.cmp(&b.row_position)) - }); - Ok(results) + Ok(drain_best_first(heap)) } #[cfg(test)] diff --git a/crates/paimon/src/vindex/pkvector/reader.rs b/crates/paimon/src/vindex/pkvector/reader.rs index f0f7d0aa..1e0fe8b3 100644 --- a/crates/paimon/src/vindex/pkvector/reader.rs +++ b/crates/paimon/src/vindex/pkvector/reader.rs @@ -21,6 +21,10 @@ /// `Send` so a boxed reader 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). +/// +/// Retained as a test-only reference source for the exact search: production now +/// streams the vector column one Arrow batch at a time in the table layer. +#[cfg(test)] pub(crate) trait PkVectorReader: Send { fn dimension(&self) -> usize;