diff --git a/Cargo.lock b/Cargo.lock index 889ccd99..43181baf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4426,6 +4426,7 @@ dependencies = [ "hmac 0.12.1", "indexmap 2.14.0", "libloading 0.9.0", + "log", "lz4_flex 0.11.6", "md-5 0.10.6", "opendal-core", 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 d445f153..9fadaf2b 100644 --- a/crates/paimon/src/table/pk_vector_data_file_reader.rs +++ b/crates/paimon/src/table/pk_vector_data_file_reader.rs @@ -519,6 +519,60 @@ mod integration_tests { assert_eq!(streamed[0][1].row_position, 3); } + /// A multi-query `search_file` returns independent per-query Top-K lists: the + /// slot for a query in a batch is identical to that query searched alone (no + /// cross-query bleed), and a shared `is_excluded` applies to every query. + #[tokio::test] + async fn search_file_multi_query_returns_independent_per_query_top_k() { + let rows = vec![ + Some(vec![0.0, 0.0]), + Some(vec![1.0, 0.0]), + Some(vec![2.0, 0.0]), + Some(vec![3.0, 0.0]), + ]; + let (factory, file_name) = + build_factory(&rows, rows.len() as i64, "memory:/pkvdfr_multiquery").await; + let active = BucketActiveFile { + file_name, + row_count: rows.len() as i64, + }; + // Exclude physical position 1 for all queries (shared predicate). + let is_excluded = |pos: i64| pos == 1; + let q0 = [0.0f32, 0.0]; // nearest is pos 0 + let q1 = [3.0f32, 0.0]; // nearest is pos 3 + + let batch = factory + .search_file( + &active, + &[&q0, &q1], + VectorSearchMetric::L2, + 2, + &is_excluded, + ) + .await + .unwrap(); + assert_eq!(batch.len(), 2, "one result list per query"); + + // Each query searched alone must equal its slot in the batch. + let only_q0 = factory + .search_file(&active, &[&q0], VectorSearchMetric::L2, 2, &is_excluded) + .await + .unwrap(); + let only_q1 = factory + .search_file(&active, &[&q1], VectorSearchMetric::L2, 2, &is_excluded) + .await + .unwrap(); + assert_eq!(batch[0], only_q0[0]); + assert_eq!(batch[1], only_q1[0]); + + // Sanity: distinct nearest neighbours, and the excluded position is absent. + assert_eq!(batch[0][0].row_position, 0); + assert_eq!(batch[1][0].row_position, 3); + assert!(batch + .iter() + .all(|list| list.iter().all(|r| r.row_position != 1))); + } + /// 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. diff --git a/crates/paimon/src/table/pk_vector_indexed_split_read.rs b/crates/paimon/src/table/pk_vector_indexed_split_read.rs index b6383c99..c6158108 100644 --- a/crates/paimon/src/table/pk_vector_indexed_split_read.rs +++ b/crates/paimon/src/table/pk_vector_indexed_split_read.rs @@ -47,7 +47,7 @@ fn data_invalid(message: impl Into) -> crate::Error { /// on both ends, required strictly ascending and non-overlapping. Adjacent /// (touching) ranges are allowed and need not be coalesced by the producer. /// `scores`, when present, is aligned to the expanded-range order (ascending -/// position); `None` means no `_PKEY_VECTOR_SCORE` column in the output. +/// position); `None` means no `__paimon_search_score` column in the output. /// /// Deliberately NOT reusing `DataSplit.row_ranges`, whose ranges mean stable/global /// row ids on the append/data-evolution path. Not serialized. @@ -604,7 +604,7 @@ mod e2e_tests { #[tokio::test] async fn reads_ranges_with_position_column_no_scores() { // rows id=[10,11,12], ranges [0..=0, 2..=2] -> ids [10,12], positions [0,2]; - // no _ROW_ID leak, no _PKEY_VECTOR_SCORE. + // no _ROW_ID leak, no __paimon_search_score. let data = write_mosaic_single_group(&id_batch(vec![10, 11, 12])); let (reader, split) = build_reader_and_split("memory:/pkvisr_basic", &data, 3, &[]).await; let indexed = PkVectorIndexedSplit { diff --git a/crates/paimon/src/table/pk_vector_orchestrator.rs b/crates/paimon/src/table/pk_vector_orchestrator.rs index 3e6a09fa..572fa641 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, ExactFileSearchFuture, + bucket_search_batch, BucketActiveFile, BucketAnnSegment, ExactFileSearchFuture, }; use crate::vindex::pkvector::metric::{java_float_compare, VectorSearchMetric}; use crate::vindex::pkvector::result::PkVectorSearchResult; @@ -361,8 +361,14 @@ impl PkVectorOrchestrator { /// absent from its split's map (or mapped to an empty set) contributes no /// candidates. `None` applies no residual filtering. The slice must have the /// same length as `splits`. + /// + /// This is the single-query wrapper over + /// [`search_candidates_batch`](Self::search_candidates_batch): it searches the + /// one query and returns its sole result, so its output is byte-identical to + /// the batch-of-one path. #[allow(clippy::too_many_arguments)] #[allow(clippy::type_complexity)] + #[cfg_attr(not(test), allow(dead_code))] pub(crate) async fn search_candidates( &self, splits: &[PkVectorSearchSplit], @@ -386,15 +392,77 @@ impl PkVectorOrchestrator { skip_exact_fallback: bool, residual_by_split: Option<&[HashMap]>, ) -> crate::Result { + let mut results = self + .search_candidates_batch( + splits, + &[query], + metric, + limit, + indexed_limit, + ann_searcher, + exact_file_search, + search_options, + skip_exact_fallback, + residual_by_split, + ) + .await?; + debug_assert_eq!(results.len(), 1); + Ok(results.remove(0)) + } + + /// Multi-query variant of [`search_candidates`](Self::search_candidates): + /// share ONE per-bucket plan (splits, DV maps, opened readers) across all N + /// queries and return one [`OrchestratorSearchResult`] per query (outer index + /// aligned to `queries`). Per split, the bucket state is built once and + /// `bucket_search_batch` fans all queries over the shared readers into + /// per-query bounded heaps; after every split, each query's indexed/exact + /// lists get their own cross-bucket global Top-K. No query's candidates bleed + /// into another's (independent per-query heaps). + /// + /// The residual allow-list depends only on the filter and the plan, not the + /// query vector, so the SAME `residual_by_split` slice is shared across every + /// query. Input-shape validation (positive limits, non-empty query, residual + /// count) is applied per query / once as appropriate. Kept a sequential + /// per-split loop. + #[allow(clippy::too_many_arguments)] + #[allow(clippy::type_complexity)] + pub(crate) async fn search_candidates_batch( + &self, + splits: &[PkVectorSearchSplit], + queries: &[&[f32]], + metric: VectorSearchMetric, + limit: usize, + indexed_limit: usize, + ann_searcher: Option<&dyn PkVectorAnnSearcher>, + exact_file_search: &(dyn for<'s, 'a> Fn( + usize, + &'s PkVectorSearchSplit, + &'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]>, + ) -> crate::Result> { // Eager input-shape validation (Java checkArgument parity). + if queries.is_empty() { + return Err(data_invalid("vector search requires at least one query")); + } if limit == 0 { return Err(data_invalid("vector search limit must be positive")); } if indexed_limit == 0 { return Err(data_invalid("vector indexed search limit must be positive")); } - if query.is_empty() { - return Err(data_invalid("vector search query must not be empty")); + for query in queries { + if query.is_empty() { + return Err(data_invalid("vector search query must not be empty")); + } } if let Some(per_split) = residual_by_split { if per_split.len() != splits.len() { @@ -404,9 +472,12 @@ impl PkVectorOrchestrator { } } - // Eager per-bucket search -> tagged candidates, kept split by path. - let mut indexed_candidates: Vec = Vec::new(); - let mut exact_candidates: Vec = Vec::new(); + // Eager per-bucket search -> per-query tagged candidates, kept split by + // path. One inner Vec per query. + let mut indexed_candidates: Vec> = + (0..queries.len()).map(|_| Vec::new()).collect(); + let mut exact_candidates: Vec> = + (0..queries.len()).map(|_| Vec::new()).collect(); for (split_index, split) in splits.iter().enumerate() { let dvs = build_bucket_dv_map(&self.reader, split).await?; // Adapt the split-scoped search closure to bucket_search's per-file @@ -432,13 +503,13 @@ impl PkVectorOrchestrator { }, ); let residual_ranges = residual_by_split.map(|per_split| &per_split[split_index]); - let result = bucket_search( + let per_query = bucket_search_batch( ann_searcher, &split.ann_segments, &split.active_files, &dvs, &bucket_search_closure, - query, + queries, metric, indexed_limit, limit, @@ -447,6 +518,13 @@ impl PkVectorOrchestrator { residual_ranges, ) .await?; + if per_query.len() != queries.len() { + return Err(data_invalid(format!( + "bucket search returned {} result lists for {} queries", + per_query.len(), + queries.len() + ))); + } let tag = |PkVectorSearchResult { data_file_name, row_position, @@ -459,14 +537,20 @@ impl PkVectorOrchestrator { row_position, distance, }; - indexed_candidates.extend(result.indexed.into_iter().map(&tag)); - exact_candidates.extend(result.exact.into_iter().map(&tag)); + for (query_index, result) in per_query.into_iter().enumerate() { + indexed_candidates[query_index].extend(result.indexed.into_iter().map(&tag)); + exact_candidates[query_index].extend(result.exact.into_iter().map(&tag)); + } } - Ok(OrchestratorSearchResult { - indexed: global_top_k(indexed_candidates, indexed_limit), - exact: global_top_k(exact_candidates, limit), - }) + Ok(indexed_candidates + .into_iter() + .zip(exact_candidates) + .map(|(indexed, exact)| OrchestratorSearchResult { + indexed: global_top_k(indexed, indexed_limit), + exact: global_top_k(exact, limit), + }) + .collect()) } } @@ -1100,18 +1184,18 @@ mod e2e_tests { hits: Vec, } impl PkVectorAnnSearcher for FakeAnn { - fn search( + fn search_batch( &self, _segment: &BucketAnnSegment, - _query: &[f32], + queries: &[&[f32]], _metric: VectorSearchMetric, _limit: usize, _active_source_files: &HashSet, _dvs: &HashMap>, _opts: &HashMap, _residual_ranges: Option<&HashMap>, - ) -> crate::Result> { - Ok(self.hits.clone()) + ) -> crate::Result>> { + Ok(queries.iter().map(|_| self.hits.clone()).collect()) } } @@ -1771,4 +1855,187 @@ mod e2e_tests { let err = validate_row_position("data-1", 9, 3).unwrap_err(); assert!(err.to_string().contains("out of range") && err.to_string().contains("data-1")); } + + /// A split-scoped multi-query search closure running `exact_search` per query + /// over a fresh `ArrayReader` of `vectors`, returning one list per query. Used + /// by the batch orchestrator tests. + #[allow(clippy::type_complexity)] + fn split_array_search_batch( + vectors: Vec>>, + ) -> 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 { + let vectors = Arc::new(vectors); + as_split_search( + move |_: usize, + _: &PkVectorSearchSplit, + 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(2, |v| v.len()); + let file_name = file.file_name.clone(); + let owned_queries: Vec> = queries.iter().map(|q| q.to_vec()).collect(); + let vectors = Arc::clone(&vectors); + Box::pin(async move { + let mut out = Vec::with_capacity(owned_queries.len()); + for query in &owned_queries { + let mut reader = ArrayReader::new(dimension, (*vectors).clone()); + out.push(exact_search( + &file_name, + &mut reader, + query, + metric, + exact_limit, + is_excluded, + )?); + } + Ok(out) + }) + }, + ) + } + + #[tokio::test] + async fn search_candidates_batch_of_one_equals_single() { + // A batch-of-one must equal the single-query `search_candidates` exactly. + let table_path = "memory:/pkvo_batch_one"; + let bucket_path = format!("{table_path}/bucket-0"); + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let meta = write_file(&file_io, &bucket_path, "c.mosaic", vec![1, 2, 3]).await; + let make_split = || PkVectorSearchSplit { + data_split: DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(bucket_path.clone()) + .with_total_buckets(1) + .with_data_files(vec![meta.clone()]) + .build() + .unwrap(), + ann_segments: Vec::new(), + active_files: vec![active("c.mosaic", 3)], + }; + // pos0 {3,0} d=9, pos1 {1,0} d=1, pos2 {2,0} d=4. + let vectors = vec![ + Some(vec![3.0, 0.0]), + Some(vec![1.0, 0.0]), + Some(vec![2.0, 0.0]), + ]; + let query: &[f32] = &[0.0, 0.0]; + let opts = HashMap::new(); + + let single = PkVectorOrchestrator::new(make_reader(file_io.clone(), table_path)) + .search_candidates( + &[make_split()], + query, + VectorSearchMetric::L2, + 2, + 2, + None, + &split_array_search_batch(vectors.clone()), + &opts, + false, + None, + ) + .await + .unwrap(); + + let batch = PkVectorOrchestrator::new(make_reader(file_io, table_path)) + .search_candidates_batch( + &[make_split()], + &[query], + VectorSearchMetric::L2, + 2, + 2, + None, + &split_array_search_batch(vectors), + &opts, + false, + None, + ) + .await + .unwrap(); + assert_eq!(batch.len(), 1); + assert_eq!( + batch[0] + .exact + .iter() + .map(|c| (c.row_position, c.distance)) + .collect::>(), + single + .exact + .iter() + .map(|c| (c.row_position, c.distance)) + .collect::>() + ); + } + + #[tokio::test] + async fn search_candidates_batch_per_query_results_independent() { + // Two queries over one exact file; each query's Top-K comes from its own + // heap with no cross-query bleed. q0 nearest pos1 (x=1); q1 nearest pos0 + // (x=3). + let table_path = "memory:/pkvo_batch_indep"; + let bucket_path = format!("{table_path}/bucket-0"); + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let meta = write_file(&file_io, &bucket_path, "c.mosaic", vec![1, 2, 3]).await; + let split = PkVectorSearchSplit { + data_split: DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(bucket_path) + .with_total_buckets(1) + .with_data_files(vec![meta]) + .build() + .unwrap(), + ann_segments: Vec::new(), + active_files: vec![active("c.mosaic", 3)], + }; + // pos0 {3,0}, pos1 {1,0}, pos2 {2,0}. + let vectors = vec![ + Some(vec![3.0, 0.0]), + Some(vec![1.0, 0.0]), + Some(vec![2.0, 0.0]), + ]; + let q0: &[f32] = &[1.0, 0.0]; // nearest pos1 (dist 0) + let q1: &[f32] = &[3.0, 0.0]; // nearest pos0 (dist 0) + let opts = HashMap::new(); + let out = PkVectorOrchestrator::new(make_reader(file_io, table_path)) + .search_candidates_batch( + &[split], + &[q0, q1], + VectorSearchMetric::L2, + 1, + 1, + None, + &split_array_search_batch(vectors), + &opts, + false, + None, + ) + .await + .unwrap(); + assert_eq!(out.len(), 2); + let m0 = merge_candidates(out[0].indexed.clone(), out[0].exact.clone(), 1); + let m1 = merge_candidates(out[1].indexed.clone(), out[1].exact.clone(), 1); + assert_eq!(m0.len(), 1); + assert_eq!(m0[0].row_position, 1); + assert_eq!(m1.len(), 1); + assert_eq!(m1[0].row_position, 0); + } } diff --git a/crates/paimon/src/table/pk_vector_position_read.rs b/crates/paimon/src/table/pk_vector_position_read.rs index 5cb1f2b1..29446822 100644 --- a/crates/paimon/src/table/pk_vector_position_read.rs +++ b/crates/paimon/src/table/pk_vector_position_read.rs @@ -595,7 +595,7 @@ mod tests { ); assert!( column_by_name(batch, SEARCH_SCORE_COLUMN).is_none(), - "_PKEY_VECTOR_SCORE must be absent when no scores are supplied" + "__paimon_search_score must be absent when no scores are supplied" ); } } diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index 50f87722..f14c08e9 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -19,7 +19,9 @@ use crate::arrow::format::FilePredicates; use crate::arrow::residual::{evaluate_predicates_mask, widen_scan_fields}; use crate::io::FileIO; use crate::lumina::reader::LuminaVectorGlobalIndexReader; -use crate::lumina::{is_lumina_index_type, LuminaIndexMeta, LuminaVectorMetric}; +use crate::lumina::{ + is_lumina_index_type, LuminaIndexMeta, LuminaVectorIndexOptions, LuminaVectorMetric, +}; use crate::spec::{ BigIntType, CoreOptions, DataField, DataType, FileKind, GlobalIndexSearchMode, IndexFileMeta, IndexManifest, IndexManifestEntry, Predicate, ROW_ID_FIELD_ID, ROW_ID_FIELD_NAME, @@ -47,11 +49,12 @@ use crate::table::{ find_field_id_by_name, merge_row_ranges, ArrowRecordBatchStream, RowRange, 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, ExactFileSearchFuture}; +use crate::vindex::pkvector::exact::validate_query; use crate::vindex::pkvector::metric::VectorSearchMetric; use crate::vindex::reader::VindexVectorGlobalIndexReader; +use crate::vindex::{is_vindex_index_type, VindexVectorIndexOptions}; use arrow_array::{Array, FixedSizeListArray, Float32Array, Int64Array, ListArray, RecordBatch}; use arrow_select::interleave::interleave_record_batch; use futures::{stream, TryStreamExt}; @@ -106,6 +109,8 @@ pub struct BatchVectorSearchBuilder<'a> { query_vectors: Option>>, limit: Option, options: HashMap, + projection: Option>, + filter: Option, } impl<'a> VectorSearchBuilder<'a> { @@ -381,13 +386,10 @@ impl<'a> VectorSearchBuilder<'a> { Ok(Box::pin(stream::iter(output.into_iter().map(Ok)))) } - /// Shared PK-vector search core for both the search-only and search-and-read - /// paths: plan the per-bucket splits, verify the configured metric against each - /// ANN segment, build the real vindex ANN scorer and (outside FAST mode) the - /// exact-fallback readers, and run the orchestrator. Returns the best-first - /// candidates together with the plan and resolved metric so the caller can - /// either serialize them to a `SearchResult` or materialize their rows. An - /// empty plan yields empty candidates. + /// Single-query wrapper over + /// [`plan_and_search_pk_candidates_batch`]: plan once, search the one query, + /// and return its candidate list. Output is byte-identical to the batch-of-one + /// path. async fn plan_and_search_pk_candidates( &self, core: &CoreOptions<'_>, @@ -395,273 +397,18 @@ impl<'a> VectorSearchBuilder<'a> { query_vector: &[f32], limit: usize, ) -> crate::Result<(Vec, PkVectorScanPlan, VectorSearchMetric)> { - // Residual pre-filter guard, mirroring Java `PrimaryKeyVectorScan`. A data - // predicate set via `with_filter` is applied post-recall by re-reading - // each candidate file's physical rows (see below). That physical-position - // filtering only agrees with the bucket search when the table exposes - // physical rows directly: deletion vectors enabled and merge-on-read - // disabled. Under merge-on-read (or without deletion vectors) a read - // merges multiple key versions, so a scalar filter could retain a stale - // version whose live version does not match — a silent wrong-read. Reject - // such queries rather than answer them incorrectly. No filter → nothing to - // guard, so the search-only and read paths are unaffected. - let physical_row_read = - core.deletion_vectors_enabled() && !core.deletion_vectors_merge_on_read(); - if self.filter.is_some() && !physical_row_read { - return Err(crate::Error::DataInvalid { - message: - "primary-key vector pre-filter requires deletion vectors without merge-on-read" - .to_string(), - source: None, - }); - } - // `primary_key_vector_distance_metric` returns a validated name; re-parse - // into the enum for the numeric semantics. - let metric = VectorSearchMetric::parse(&core.primary_key_vector_distance_metric(pk_col)?)?; - let index_type = core.primary_key_vector_index_type(pk_col)?; - let field_id = - find_field_id_by_name(self.table.schema().fields(), pk_col).ok_or_else(|| { - crate::Error::DataInvalid { - message: format!("PK-vector column '{pk_col}' not found in schema"), - source: None, - } - })?; - let vector_field = self - .table - .schema() - .fields() - .iter() - .find(|f| f.name() == pk_col) - .cloned() - .ok_or_else(|| crate::Error::DataInvalid { - message: format!("PK-vector column '{pk_col}' not found in schema"), - source: None, - })?; - - let search_mode = core.global_index_search_mode()?; - let skip_exact_fallback = search_mode == GlobalIndexSearchMode::Fast; - - // Resolve the refine factor from the query options first, then fall back to - // the table options; a positive factor over-fetches indexed (approximate) - // candidates so the exact rerank below has a wider pool to reorder. Factor 0 - // (unset) leaves `indexed_limit == limit`, byte-identical to the no-rerank - // path. The two option maps are kept distinct (query options passed - // separately from table options) so a broad query key cannot be overridden - // by a more specific table key: query options take precedence as a whole. - // Resolved before planning so an invalid factor (e.g. a non-numeric value) - // fails loud regardless of whether the table currently has searchable data. - let refine_factor = configured_refine_factor( + let (mut candidates, plan, metric) = plan_and_search_pk_candidates_batch( + self.table, &self.options, - self.table.schema().options(), + self.filter.as_ref(), + core, pk_col, - &index_type, - )?; - let indexed_limit = indexed_search_limit(limit, refine_factor)?; - - let plan = PkVectorScan::new( - self.table, - field_id, - index_type.clone(), - self.filter.clone(), + &[query_vector], + limit, ) - .plan() .await?; - if plan.splits.is_empty() { - return Ok((Vec::new(), plan, metric)); - } - - // Resolve the vector index backend from the single configured index type. - // Java enforces one index type per PK table and Rust filters segments to - // it, so one backend serves every segment. Computed after the empty-plan - // return so an empty table never errors on an unrecognized type. - let backend = VectorIndexBackend::from_index_type(&index_type).ok_or_else(|| { - crate::Error::DataInvalid { - message: format!("unsupported PK vector index backend/type: '{index_type}'"), - source: None, - } - })?; - - // Production data-file reader, mirroring `table_read.rs::new_data_file_reader` - // but projecting only the vector column with no predicates. - let reader = DataFileReader::new( - self.table.file_io().clone(), - self.table.schema_manager().clone(), - self.table.schema().id(), - self.table.schema().fields().to_vec(), - vec![vector_field.clone()], - Vec::new(), - ); - - // Real ANN scorer: preload each segment's bytes (keyed by resolved, - // globally unique path) and drive the vindex reader from memory. - let segment_bytes = preload_segment_bytes(self.table.file_io(), &plan.splits).await?; - // Fail loud on a config/segment metric mismatch before scoring, mirroring - // Java `PkVectorAnnSegmentSearcher.search`. - verify_pk_vector_segment_metrics(&plan.splits, &segment_bytes, metric, backend)?; - let options = { - let mut o = self.table.schema().options().clone(); - o.extend(self.options.clone()); - o - }; - let search_options = options.clone(); - let field_name = pk_col.to_string(); - let scorer: crate::vindex::pkvector::ann::Scorer = - Box::new(move |segment: &BucketAnnSegment, search: &VectorSearch| { - let data = segment_bytes - .get(&segment.path) - .ok_or_else(|| crate::Error::DataInvalid { - message: "missing preloaded ANN bytes for segment".to_string(), - source: None, - })? - .clone(); - let io_meta = GlobalIndexIOMeta::new( - segment.path.clone(), - segment.file_size, - segment.index_meta.clone(), - ); - match backend { - VectorIndexBackend::Lumina => { - let mut reader = - LuminaVectorGlobalIndexReader::new(io_meta, options.clone()); - reader.visit_vector_search(search, |_| Ok(Cursor::new(data))) - } - VectorIndexBackend::Vindex => { - let mut reader = - VindexVectorGlobalIndexReader::new(io_meta, options.clone()); - reader.visit_vector_search(search, |_| Ok(Cursor::new(data))) - } - } - }); - let ann_searcher = VindexAnnSearcher::new(field_name, scorer); - - // Residual (post-recall) filtering: for each candidate file, re-read its - // physical rows and keep the positions whose rows satisfy the filter. The - // per-split allow-list is threaded into the bucket search so the residual - // folds into recall (best-first order and Top-K are preserved). Built only - // when a filter is set; otherwise `None` leaves the search unfiltered. The - // residual reader projects only the predicate columns and carries no - // pushdown; `residual_positions_by_file` recovers each surviving row's - // file-local physical position from its ordinal in the unfiltered scan (no - // `_ROW_ID`, no `first_row_id`). A file the allow-list leaves empty is - // skipped by the bucket search without opening an exact reader. - let residual_by_split: Option>> = match &self.filter { - Some(filter) => { - let file_predicates = FilePredicates { - predicates: vec![filter.clone()], - row_filter_factory: None, - file_fields: self.table.schema().fields().to_vec(), - }; - let residual_read_type = widen_scan_fields(&[], Some(&file_predicates)); - let residual_reader = DataFileReader::new( - self.table.file_io().clone(), - self.table.schema_manager().clone(), - self.table.schema().id(), - self.table.schema().fields().to_vec(), - residual_read_type, - Vec::new(), - ); - let mut per_split = Vec::with_capacity(plan.splits.len()); - for split in &plan.splits { - per_split.push( - residual_positions_by_file( - &residual_reader, - &split.data_split, - &split.active_files, - &file_predicates, - ) - .await?, - ); - } - Some(per_split) - } - None => None, - }; - - // 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 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 factory = as_split_exact_file_search( - move |_split_index: usize, - split: &PkVectorSearchSplit, - 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(); - let active = BucketActiveFile { - 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)?; - 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 - }) - }, - ); - - let search: OrchestratorSearchResult = PkVectorOrchestrator::new(reader) - .search_candidates( - &plan.splits, - query_vector, - metric, - limit, - indexed_limit, - Some(&ann_searcher), - &factory, - &search_options, - skip_exact_fallback, - residual_by_split.as_deref(), - ) - .await?; - - // Exact rerank of the approximate candidates when a refine factor is set; - // exact-fallback candidates are already exact and are not reranked. With no - // refine factor this is a plain merge, byte-identical to the no-rerank path. - let indexed = if refine_factor > 0 && !search.indexed.is_empty() { - // Vector-only reader (project just the vector field); the position read - // appends _PKEY_VECTOR_POSITION itself and injects _ROW_ID internally. - let rerank_reader = DataFileReader::new( - self.table.file_io().clone(), - self.table.schema_manager().clone(), - self.table.schema().id(), - self.table.schema().fields().to_vec(), - vec![vector_field.clone()], - Vec::new(), - ); - rerank_indexed_positional( - &rerank_reader, - search.indexed, - &plan.splits, - query_vector, - metric, - limit, - &vector_field, - ) - .await? - } else { - search.indexed - }; - // Merge the (possibly reranked) indexed list with the exact-fallback list - // back into one best-first list bounded to the caller's limit; the - // downstream materialization consumes a single ranked candidate list. - let candidates = merge_candidates(indexed, search.exact, limit); - - Ok((candidates, plan, metric)) + debug_assert_eq!(candidates.len(), 1); + Ok((candidates.remove(0), plan, metric)) } /// Materialize the best-first PK-vector search hits into Arrow rows. Mirrors @@ -691,10 +438,6 @@ impl<'a> VectorSearchBuilder<'a> { // Default (no `with_projection`) is every user table column. let read_type = self.resolve_materialize_read_type()?; - if candidates.is_empty() { - return Ok(Box::pin(stream::empty())); - } - // A separate, predicate-free materialization reader projecting the user // columns (the search reader projects only the vector column). Mirrors // `table_read.rs::new_data_file_reader` with an empty predicate list. @@ -707,6 +450,25 @@ impl<'a> VectorSearchBuilder<'a> { Vec::new(), ); + Self::materialize_candidates(candidates, &plan, metric, &materialize_reader).await + } + + /// Materialize one best-first candidate list into an Arrow stream, best-first, + /// with a `__paimon_search_score` column and `_PKEY_VECTOR_POSITION` stripped. + /// An empty candidate list yields an empty stream (never skipped) so a batch + /// caller preserves per-query arity. `materialize_reader` must project the + /// output columns (predicate-free). Both the single-query and batch read paths + /// use this so their materialization is identical. + async fn materialize_candidates( + candidates: Vec, + plan: &PkVectorScanPlan, + metric: VectorSearchMetric, + materialize_reader: &DataFileReader, + ) -> crate::Result { + if candidates.is_empty() { + return Ok(Box::pin(stream::empty())); + } + // Rank each candidate by its best-first position, then reduce the physical // materialization order back to best-first. The orchestrator emits rows in // ascending (partition, bucket, file, position); the rank map keyed by @@ -762,27 +524,17 @@ impl<'a> VectorSearchBuilder<'a> { /// names resolved via `resolve_projected_fields`. Rejects reserved metadata /// names and `_ROW_ID` so a user cannot request a hidden column. fn resolve_materialize_read_type(&self) -> crate::Result> { - let is_reserved = |name: &str| { - name == PKEY_VECTOR_POSITION_COLUMN - || name == SEARCH_SCORE_COLUMN - || name == ROW_ID_FIELD_NAME - || name == "_PKEY_VECTOR_SCORE" - }; - let reserved_err = |name: &str| crate::Error::DataInvalid { - message: format!( - "vector search read projection must not request reserved column '{name}'" - ), - source: None, - }; let fields = match &self.projection { None => self.table.schema().fields().to_vec(), Some(names) => { - // Reject a requested reserved name on the raw list first: these - // names are not real table columns, so `resolve_projected_fields` - // would otherwise fail with a confusing "not found" error. for name in names { - if is_reserved(name) { - return Err(reserved_err(name)); + if is_reserved_read_column(name) { + return Err(crate::Error::DataInvalid { + message: format!( + "vector search read projection must not request reserved column '{name}'" + ), + source: None, + }); } } resolve_projected_fields( @@ -793,19 +545,363 @@ impl<'a> VectorSearchBuilder<'a> { )? } }; - // Reject reserved output-column names on the RESOLVED field list too — this - // is what catches the default (no `with_projection`) case where a user table - // column is literally named `__paimon_search_score` (or the legacy - // `_PKEY_VECTOR_SCORE` alias): it would otherwise survive and collide with - // the score column appended during materialization, producing two - // identically named output columns. - for field in &fields { - if is_reserved(field.name()) { - return Err(reserved_err(field.name())); + // The default projection returns every user column, so a user column + // whose name collides with an injected metadata column must be rejected + // on the resolved field list too — not only when explicitly requested. + ensure_no_reserved_read_columns(&fields)?; + Ok(fields) + } +} + +/// Names a read injects as metadata columns — `__paimon_search_score`, +/// `_PKEY_VECTOR_POSITION`, and `_ROW_ID` — that a materialized read type must +/// not reuse for a user column. +fn is_reserved_read_column(name: &str) -> bool { + name == PKEY_VECTOR_POSITION_COLUMN || name == SEARCH_SCORE_COLUMN || name == ROW_ID_FIELD_NAME +} + +/// Reject a materialized read type whose resolved fields contain a reserved +/// metadata column name. Applied to the RESOLVED field list so the default +/// (all user columns) projection is covered, not only an explicit one. +fn ensure_no_reserved_read_columns(fields: &[DataField]) -> crate::Result<()> { + for field in fields { + if is_reserved_read_column(field.name()) { + return Err(crate::Error::DataInvalid { + message: format!( + "vector search read projection must not include reserved column '{}'", + field.name() + ), + source: None, + }); + } + } + Ok(()) +} + +/// Batch PK-vector search core shared by the single and batch builders: plan ONE +/// per-bucket split set, segment preload, ANN scorer, exact-fallback search +/// closure, and residual allow-list (all query-independent), then run +/// `search_candidates_batch` ONCE so N queries share the opened readers. Per +/// query, the approximate candidates are exact-reranked (when a refine factor is +/// set) and merged with the exact fallback into one best-first list bounded to +/// `limit`. Returns one candidate list per query (outer index aligned to +/// `queries`), together with the shared plan and resolved metric so the caller can +/// materialize each query's rows. An empty plan yields one empty candidate list +/// per query. +/// +/// The residual allow-list depends only on `filter` and the plan, NOT the query +/// vector, so it is computed once and the SAME slice is shared across all queries. +/// Rerank stays per-query (each query reranks its own indexed list). +#[allow(clippy::too_many_arguments)] +async fn plan_and_search_pk_candidates_batch( + table: &Table, + query_options: &HashMap, + filter: Option<&Predicate>, + core: &CoreOptions<'_>, + pk_col: &str, + queries: &[&[f32]], + limit: usize, +) -> crate::Result<( + Vec>, + PkVectorScanPlan, + VectorSearchMetric, +)> { + // Residual pre-filter guard, mirroring Java `PrimaryKeyVectorScan`. A data + // predicate set via `with_filter` is applied post-recall by re-reading each + // candidate file's physical rows (see below). That physical-position filtering + // only agrees with the bucket search when the table exposes physical rows + // directly: deletion vectors enabled and merge-on-read disabled. Under + // merge-on-read (or without deletion vectors) a read merges multiple key + // versions, so a scalar filter could retain a stale version whose live version + // does not match — a silent wrong-read. Reject such queries rather than answer + // them incorrectly. No filter → nothing to guard, so the search-only and read + // paths are unaffected. + let physical_row_read = + core.deletion_vectors_enabled() && !core.deletion_vectors_merge_on_read(); + if filter.is_some() && !physical_row_read { + return Err(crate::Error::DataInvalid { + message: + "primary-key vector pre-filter requires deletion vectors without merge-on-read" + .to_string(), + source: None, + }); + } + // `primary_key_vector_distance_metric` returns a validated name; re-parse into + // the enum for the numeric semantics. + let metric = VectorSearchMetric::parse(&core.primary_key_vector_distance_metric(pk_col)?)?; + let index_type = core.primary_key_vector_index_type(pk_col)?; + let field_id = find_field_id_by_name(table.schema().fields(), pk_col).ok_or_else(|| { + crate::Error::DataInvalid { + message: format!("PK-vector column '{pk_col}' not found in schema"), + source: None, + } + })?; + let vector_field = table + .schema() + .fields() + .iter() + .find(|f| f.name() == pk_col) + .cloned() + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("PK-vector column '{pk_col}' not found in schema"), + source: None, + })?; + + let search_mode = core.global_index_search_mode()?; + let skip_exact_fallback = search_mode == GlobalIndexSearchMode::Fast; + + // A non-positive limit is invalid regardless of the plan; reject it before + // planning so an empty plan cannot mask it with empty results. + if limit == 0 { + return Err(crate::Error::DataInvalid { + message: "vector search limit must be positive".to_string(), + source: None, + }); + } + + // Resolve the refine factor from the query options first, then fall back to + // the table options; a positive factor over-fetches indexed (approximate) + // candidates so the exact rerank below has a wider pool to reorder. Factor 0 + // (unset) leaves `indexed_limit == limit`, byte-identical to the no-rerank + // path. The two option maps are kept distinct (query options passed + // separately from table options) so a broad query key cannot be overridden + // by a more specific table key: query options take precedence as a whole. + // Resolved before planning so an invalid factor (e.g. a non-numeric value) + // fails loud regardless of whether the table currently has searchable data. + let refine_factor = + configured_refine_factor(query_options, table.schema().options(), pk_col, &index_type)?; + let indexed_limit = indexed_search_limit(limit, refine_factor)?; + + // Validate every query against the vector column's dimension (and finiteness) + // before planning or any read, so a malformed query fails loud even when the + // plan turns out empty. VECTOR carries the dimension in its type; + // ARRAY gets the index dimension from the same vindex option resolver + // used by index reads. Both valid PK-vector column shapes must reject NaN/Inf + // up front, not only after a non-empty plan opens readers. + if let Some(dimension) = pk_vector_query_dimension( + table.schema().options(), + query_options, + &index_type, + &vector_field, + )? { + for query in queries { + validate_query(query, dimension)?; + } + } + + let plan = PkVectorScan::new(table, field_id, index_type.clone(), filter.cloned()) + .plan() + .await?; + if plan.splits.is_empty() { + return Ok((vec![Vec::new(); queries.len()], plan, metric)); + } + + // Resolve the vector index backend from the single configured index type. + // Java enforces one index type per PK table and Rust filters segments to it, + // so one backend serves every segment. Computed after the empty-plan return so + // an empty table never errors on an unrecognized type. + let backend = VectorIndexBackend::from_index_type(&index_type).ok_or_else(|| { + crate::Error::DataInvalid { + message: format!("unsupported PK vector index backend/type: '{index_type}'"), + source: None, + } + })?; + + // Production data-file reader, mirroring `table_read.rs::new_data_file_reader` + // but projecting only the vector column with no predicates. + let reader = DataFileReader::new( + table.file_io().clone(), + table.schema_manager().clone(), + table.schema().id(), + table.schema().fields().to_vec(), + vec![vector_field.clone()], + Vec::new(), + ); + + // Real ANN scorer: preload each segment's bytes (keyed by resolved, globally + // unique path) and drive the vindex reader from memory. The reader is opened + // once per segment and every query in the batch is searched against it, + // mirroring the shared-reader batch search. + let segment_bytes = preload_segment_bytes(table.file_io(), &plan.splits).await?; + // Fail loud on a config/segment metric mismatch before scoring, mirroring Java + // `PkVectorAnnSegmentSearcher.search`. + verify_pk_vector_segment_metrics(&plan.splits, &segment_bytes, metric, backend)?; + let options = { + let mut o = table.schema().options().clone(); + o.extend(query_options.clone()); + o + }; + let search_options = options.clone(); + let field_name = pk_col.to_string(); + let scorer: crate::vindex::pkvector::ann::BatchScorer = Box::new( + move |segment: &BucketAnnSegment, searches: &[VectorSearch]| { + let data = segment_bytes + .get(&segment.path) + .ok_or_else(|| crate::Error::DataInvalid { + message: "missing preloaded ANN bytes for segment".to_string(), + source: None, + })? + .clone(); + let io_meta = GlobalIndexIOMeta::new( + segment.path.clone(), + segment.file_size, + segment.index_meta.clone(), + ); + match backend { + VectorIndexBackend::Lumina => { + let mut reader = LuminaVectorGlobalIndexReader::new(io_meta, options.clone()); + reader.visit_batch_vector_search(searches, |_| Ok(Cursor::new(data))) + } + VectorIndexBackend::Vindex => { + let mut reader = VindexVectorGlobalIndexReader::new(io_meta, options.clone()); + reader.visit_batch_vector_search(searches, |_| Ok(Cursor::new(data))) + } } + }, + ); + let ann_searcher = VindexAnnSearcher::new(field_name, scorer); + + // Residual (post-recall) filtering: for each candidate file, re-read its + // physical rows and keep the positions whose rows satisfy the filter. The + // per-split allow-list is threaded into the bucket search so the residual folds + // into recall (best-first order and Top-K are preserved). Built only when a + // filter is set; otherwise `None` leaves the search unfiltered. The residual + // depends only on the filter and the plan, not the query vector, so it is + // computed once here and shared across every query in the batch. The residual + // reader projects only the predicate columns and carries no pushdown; + // `residual_positions_by_file` recovers each surviving row's file-local + // physical position from its ordinal in the unfiltered scan (no `_ROW_ID`, no + // `first_row_id`). A file the allow-list leaves empty is skipped by the bucket + // search without opening an exact reader. + let residual_by_split: Option>> = match filter { + Some(filter) => { + let file_predicates = FilePredicates { + predicates: vec![filter.clone()], + row_filter_factory: None, + file_fields: table.schema().fields().to_vec(), + }; + let residual_read_type = widen_scan_fields(&[], Some(&file_predicates)); + let residual_reader = DataFileReader::new( + table.file_io().clone(), + table.schema_manager().clone(), + table.schema().id(), + table.schema().fields().to_vec(), + residual_read_type, + Vec::new(), + ); + let mut per_split = Vec::with_capacity(plan.splits.len()); + for split in &plan.splits { + per_split.push( + residual_positions_by_file( + &residual_reader, + &split.data_split, + &split.active_files, + &file_predicates, + ) + .await?, + ); + } + Some(per_split) } - Ok(fields) + None => None, + }; + + // 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 + // search streams the file's vector column one Arrow batch at a time into + // per-query bounded heaps (all queries share one stream). + let reader_for_factory = reader.clone(); + let vector_field_for_factory = vector_field.clone(); + let factory = as_split_exact_file_search( + move |_split_index: usize, + split: &PkVectorSearchSplit, + 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(); + let active = BucketActiveFile { + 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)?; + 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 + }) + }, + ); + + // Resolve the refine factor from the query options first, then fall back to the + // table options; a positive factor over-fetches indexed (approximate) + // candidates so the exact rerank below has a wider pool to reorder. Factor 0 + // (unset) leaves `indexed_limit == limit`, byte-identical to the no-rerank + // path. The two option maps are kept distinct (query options passed separately + // from table options) so a broad query key cannot be overridden by a more + // specific table key: query options take precedence as a whole. `search_options` + // above is the merged view used only to drive the ANN read. + + let searches: Vec = PkVectorOrchestrator::new(reader) + .search_candidates_batch( + &plan.splits, + queries, + metric, + limit, + indexed_limit, + Some(&ann_searcher), + &factory, + &search_options, + skip_exact_fallback, + residual_by_split.as_deref(), + ) + .await?; + + // Per query: exact rerank of the approximate candidates when a refine factor is + // set (exact-fallback candidates are already exact and are not reranked), then + // merge the (possibly reranked) indexed list with the exact list into one + // best-first list bounded to the caller's limit. With no refine factor the + // rerank is a plain merge, byte-identical to the no-rerank path. Each query + // reranks its OWN indexed candidates. + let mut per_query_candidates = Vec::with_capacity(searches.len()); + for (query_index, search) in searches.into_iter().enumerate() { + let query_vector = queries[query_index]; + let indexed = if refine_factor > 0 && !search.indexed.is_empty() { + // Vector-only reader (project just the vector field); the position read + // appends _PKEY_VECTOR_POSITION itself and injects _ROW_ID internally. + let rerank_reader = DataFileReader::new( + table.file_io().clone(), + table.schema_manager().clone(), + table.schema().id(), + table.schema().fields().to_vec(), + vec![vector_field.clone()], + Vec::new(), + ); + rerank_indexed_positional( + &rerank_reader, + search.indexed, + &plan.splits, + query_vector, + metric, + limit, + &vector_field, + ) + .await? + } else { + search.indexed + }; + per_query_candidates.push(merge_candidates(indexed, search.exact, limit)); } + + Ok((per_query_candidates, plan, metric)) } impl<'a> BatchVectorSearchBuilder<'a> { @@ -816,6 +912,8 @@ impl<'a> BatchVectorSearchBuilder<'a> { query_vectors: None, limit: None, options: HashMap::new(), + projection: None, + filter: None, } } @@ -839,11 +937,33 @@ impl<'a> BatchVectorSearchBuilder<'a> { self } - pub async fn execute(&self) -> crate::Result> { - // Fail closed before validation and empty-table fast paths: batch search - // returns data-derived results outside `TableScan`/`TableRead`. - CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?; + /// Attach a residual scalar predicate applied *after* vector recall on the + /// primary-key vector path, shared across every query in the batch. Mirrors + /// the single [`VectorSearchBuilder::with_filter`]: only the primary-key + /// vector path (via [`execute_read`](Self::execute_read)) consumes it, and only + /// when the table exposes physical rows directly (deletion vectors without + /// merge-on-read); otherwise the query fails loud. + pub fn with_filter(&mut self, filter: Predicate) -> &mut Self { + self.filter = Some(filter); + self + } + /// Restrict the columns materialized by [`execute_read`](Self::execute_read) to + /// `cols` (plus the always-appended `__paimon_search_score`). Without this call + /// `execute_read` materializes every user table column. Only affects + /// `execute_read`; `execute` ignores it. + pub fn with_projection(&mut self, cols: &[&str]) -> &mut Self { + self.projection = Some(cols.iter().map(|c| c.to_string()).collect()); + self + } + + pub async fn execute(&self) -> crate::Result> { + // Fail closed: like `execute_read` and the single-query builder, this + // returns data-derived row ids/scores outside `TableScan`/`TableRead`, + // so it must refuse a `query-auth.enabled` table before any fast path + // (an empty snapshot would otherwise return empty results and bypass it). + let core = CoreOptions::new(self.table.schema().options()); + core.ensure_read_authorized()?; let vector_column = self.vector_column .as_deref() @@ -872,6 +992,38 @@ impl<'a> BatchVectorSearchBuilder<'a> { message: "Limit must be set via with_limit()".to_string(), })?; + // A primary-key vector table exposes no global row ids, so scored batch + // search is unsupported: `execute()` returns `SearchResult`s (global row + // ids). Fail loud and direct callers to the materialized batch + // `execute_read`, mirroring the single-query builder's PK guard. Membership + // is resolved via the non-erroring columns accessor so a malformed + // PK-vector config cannot abort an unrelated DE query. + if core.primary_key_vector_index_enabled() { + let targets_pk_column = core + .primary_key_vector_index_columns() + .ok() + .is_some_and(|cols| cols.iter().any(|c| c == vector_column)); + if targets_pk_column { + return Err(crate::Error::DataInvalid { + message: "primary-key vector search does not produce global row ids; use the materialized read (execute_read) instead".to_string(), + source: None, + }); + } + } + + // The data-evolution (global-index) fall-through path cannot honor a + // residual filter — it never reads physical rows. Rather than silently + // drop the predicate and return unfiltered results, fail loud when a + // filter is set on a batch that does not resolve to the primary-key + // vector path, mirroring the single-query builder. + if self.filter.is_some() { + return Err(crate::Error::DataInvalid { + message: "vector search filter is only supported on the primary-key vector path" + .to_string(), + source: None, + }); + } + let vector_searches = query_vectors .iter() .map(|vector| { @@ -909,6 +1061,147 @@ impl<'a> BatchVectorSearchBuilder<'a> { ) .await } + + /// Run a batch of vector searches and materialize each query's matching rows as + /// a best-first Arrow stream. Supported only for the primary-key vector path + /// (which alone can materialize physical rows). The returned `Vec` is aligned + /// strictly to the input query order and its length always equals the query + /// count — a query with no hits yields an empty stream, never a missing entry. + /// If ANY query errors (e.g. a malformed vector) the whole call fails loud with + /// no partial `Vec` of streams. Output columns are the projected user table + /// columns (all user columns by default, or those set via + /// [`with_projection`](Self::with_projection)) plus `__paimon_search_score`; + /// `_ROW_ID` and `_PKEY_VECTOR_POSITION` are always hidden. + /// + /// A data-evolution (global-index) table fails loud: its batch search returns + /// scored global row-ids, not materialized rows, so callers use + /// [`execute`](Self::execute) instead. + pub async fn execute_read(&self) -> crate::Result> { + // Fail closed: returns data outside `TableScan`/`TableRead`. + let core = CoreOptions::new(self.table.schema().options()); + core.ensure_read_authorized()?; + let vector_column = + self.vector_column + .as_deref() + .ok_or_else(|| crate::Error::ConfigInvalid { + message: "Vector column must be set via with_vector_column()".to_string(), + })?; + let query_vectors = + self.query_vectors + .as_ref() + .ok_or_else(|| crate::Error::ConfigInvalid { + message: "Query vectors must be set via with_query_vectors()".to_string(), + })?; + if query_vectors.is_empty() { + return Err(crate::Error::ConfigInvalid { + message: "Query vectors must be set via with_query_vectors()".to_string(), + }); + } + let limit = self.limit.ok_or_else(|| crate::Error::ConfigInvalid { + message: "Limit must be set via with_limit()".to_string(), + })?; + + // Only the primary-key vector path can materialize rows. The data-evolution + // (global-index) path returns data-derived row-ids, not table rows, so a + // batch read against it (or a non-PK-vector column) fails loud, directing + // callers to `execute()`. + let targets_pk_column = core.primary_key_vector_index_enabled() + && core + .primary_key_vector_index_columns() + .ok() + .is_some_and(|cols| cols.iter().any(|c| c == vector_column)); + if !targets_pk_column { + return Err(crate::Error::DataInvalid { + message: "batch vector read is only supported on the primary-key vector path; data-evolution batch search returns scored row ids, use execute() instead".to_string(), + source: None, + }); + } + + let pk_col = core.primary_key_vector_index_column()?; + let query_refs: Vec<&[f32]> = query_vectors.iter().map(|q| q.as_slice()).collect(); + + // Resolve the materialization read-type up front so an invalid projection + // (unknown column, or a reserved metadata / row-id name) fails loud + // unconditionally, before any read — a whole-call failure, not a partial + // Vec. + let read_type = self.resolve_materialize_read_type()?; + + // One shared plan / segment preload / residual across all N queries; the + // per-query candidate lists come back in strict input order. Any query + // error (or a shared-plan error) propagates here, so no partial Vec is + // returned. + let (per_query_candidates, plan, metric) = plan_and_search_pk_candidates_batch( + self.table, + &self.options, + self.filter.as_ref(), + &core, + &pk_col, + &query_refs, + limit, + ) + .await?; + + let materialize_reader = DataFileReader::new( + self.table.file_io().clone(), + self.table.schema_manager().clone(), + self.table.schema().id(), + self.table.schema().fields().to_vec(), + read_type, + Vec::new(), + ); + + // Materialize each query's candidates into its own stream, preserving + // arity: an empty candidate list yields an empty stream. Build every stream + // before returning so a materialization error fails the whole call with no + // partial Vec. + let mut streams = Vec::with_capacity(per_query_candidates.len()); + for candidates in per_query_candidates { + streams.push( + VectorSearchBuilder::materialize_candidates( + candidates, + &plan, + metric, + &materialize_reader, + ) + .await?, + ); + } + Ok(streams) + } + + /// Resolve the projected fields for the materialization read-type. Default + /// (no projection set) is all user table fields; otherwise the requested names + /// resolved via `resolve_projected_fields`. Rejects reserved metadata names and + /// `_ROW_ID` so a user cannot request a hidden column. Mirrors the single + /// builder's resolver. + fn resolve_materialize_read_type(&self) -> crate::Result> { + let fields = match &self.projection { + None => self.table.schema().fields().to_vec(), + Some(names) => { + for name in names { + if is_reserved_read_column(name) { + return Err(crate::Error::DataInvalid { + message: format!( + "vector search read projection must not request reserved column '{name}'" + ), + source: None, + }); + } + } + resolve_projected_fields( + self.table.identifier().full_name(), + self.table.schema().fields(), + names, + true, + )? + } + }; + // The default projection returns every user column, so a user column + // whose name collides with an injected metadata column must be rejected + // on the resolved field list too — not only when explicitly requested. + ensure_no_reserved_read_columns(&fields)?; + Ok(fields) + } } #[derive(Clone, Copy)] @@ -1343,6 +1636,47 @@ fn verify_pk_vector_segment_metrics( Ok(()) } +fn pk_vector_query_dimension( + table_options: &HashMap, + query_options: &HashMap, + index_type: &str, + vector_field: &DataField, +) -> crate::Result> { + match vector_field.data_type() { + DataType::Vector(vector_type) + if matches!(vector_type.element_type(), DataType::Float(_)) => + { + Ok(Some(vector_type.length() as usize)) + } + DataType::Array(array_type) if matches!(array_type.element_type(), DataType::Float(_)) => { + // Resolve the dimension per the configured backend. An `ARRAY` + // column carries no dimension in its type, so it comes from options — + // but the option shape differs by backend. Lumina is not a vindex + // index type, so routing it through `VindexVectorIndexOptions` would + // reject it as unsupported before planning (even on an empty table). + if is_lumina_index_type(index_type) { + // Lumina reads `lumina.index.dimension` (default 128) from the + // merged table+query options, matching `resolve_lumina_options`. + let mut merged = table_options.clone(); + merged.extend(query_options.clone()); + let dimension = LuminaVectorIndexOptions::new(&merged)?.dimension; + Ok(Some(dimension as usize)) + } else { + Ok(Some( + VindexVectorIndexOptions::new( + table_options, + query_options, + index_type, + vector_field, + )? + .dimension(), + )) + } + } + _ => Ok(None), + } +} + /// Rerank approximate (indexed) candidates by rereading ONLY their candidate /// positions and recomputing the exact distance, then keep the best `limit`. /// @@ -2864,55 +3198,6 @@ mod tests { ); } - #[tokio::test] - async fn test_batch_vector_search_auth_check_precedes_validation() { - let table = crate::table::query_auth_table(); - let err = table - .new_batch_vector_search_builder() - .execute() - .await - .unwrap_err(); - - assert!( - matches!(err, crate::Error::Unsupported { ref message } if message.contains("query-auth.enabled")), - "batch vector search must check query auth before validating parameters" - ); - } - - #[tokio::test] - async fn test_batch_vector_search_empty_table_fails_closed_when_query_auth_enabled() { - let table = crate::table::query_auth_table(); - let err = table - .new_batch_vector_search_builder() - .with_vector_column("id") - .with_query_vectors(vec![vec![1.0]]) - .with_limit(1) - .execute() - .await - .unwrap_err(); - - assert!( - matches!(err, crate::Error::Unsupported { ref message } if message.contains("query-auth.enabled")), - "batch vector search must fail closed before the empty-table fast path" - ); - } - - #[tokio::test] - async fn test_batch_vector_search_empty_table_unchanged_without_query_auth() { - let table = vector_test_table(); - let results = table - .new_batch_vector_search_builder() - .with_vector_column("embedding") - .with_query_vectors(vec![vec![1.0]]) - .with_limit(1) - .execute() - .await - .unwrap(); - - assert_eq!(results.len(), 1); - assert!(results[0].is_empty()); - } - #[tokio::test] async fn test_batch_evaluate_no_matching_field_returns_empty_per_query() { let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap(); @@ -3150,6 +3435,28 @@ mod tests { ); } + #[tokio::test] + async fn test_batch_execute_fails_closed_when_query_auth_enabled() { + // The batch scored entry returns data-derived row ids/scores outside + // `TableScan`/`TableRead`, so it must fail closed under + // `query-auth.enabled` exactly like the single-query builder. Its config + // is otherwise valid, so without the guard the empty-snapshot fast path + // would return empty results and silently bypass authorization. + let table = crate::table::query_auth_table(); + let err = table + .new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors(vec![vec![1.0, 2.0]]) + .with_limit(5) + .execute() + .await + .unwrap_err(); + assert!( + matches!(err, crate::Error::Unsupported { ref message } if message.contains("query-auth.enabled")), + "batch vector search must fail closed for a query-auth table, got: {err:?}" + ); + } + fn pk_data_file(name: &str, row_count: i64, first_row_id: Option) -> DataFileMeta { DataFileMeta { file_name: name.to_string(), @@ -4273,12 +4580,15 @@ mod tests { ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), ("fields.embedding.pk-vector.distance.metric", "l2"), ("deletion-vectors.enabled", "true"), + // Pin the index dimension so the query vector below matches it; the + // up-front dimension guard runs before this test's residual guard. + ("fields.embedding.dimension", "4"), ]); let filter = id_gt_filter(&table, 2); let mut stream = table .new_vector_search_builder() .with_vector_column("embedding") - .with_query_vector(vec![1.0]) + .with_query_vector(vec![1.0; 4]) .with_limit(5) .with_filter(filter) .execute_read() @@ -4689,93 +4999,6 @@ mod tests { ); } - #[tokio::test] - async fn execute_read_default_projection_rejects_reserved_score_column_name() { - // A user table column literally named `__paimon_search_score` must be - // rejected even under the default (no `with_projection`) projection — - // otherwise it survives and collides with the score column appended during - // materialization, producing two identically named output columns. - use crate::spec::{FloatType, IntType, Schema, TableSchema, VectorType}; - let schema = Schema::builder() - .column("id", DataType::Int(IntType::new())) - .column( - "embedding", - DataType::Vector( - VectorType::try_new(true, 2, DataType::Float(FloatType::new())).unwrap(), - ), - ) - .column(SEARCH_SCORE_COLUMN, DataType::Float(FloatType::new())) - .primary_key(["id"]) - .option("bucket", "1") - .option("pk-vector.index.columns", "embedding") - .option("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER) - .option("fields.embedding.pk-vector.distance.metric", "l2") - .build() - .unwrap(); - let table = Table::new( - FileIOBuilder::new("memory").build().unwrap(), - Identifier::new("default", "pk_vector_collision"), - "memory:/pk_vector_collision".to_string(), - TableSchema::new(0, &schema), - None, - ); - let err = match table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0, 0.0]) - .with_limit(5) - .execute_read() - .await - { - Ok(_) => panic!("default projection over a reserved-named column must fail loud"), - Err(e) => e, - }; - assert!( - matches!(&err, crate::Error::DataInvalid { message, .. } if message.contains("reserved column")), - "expected a reserved-column error, got: {err}" - ); - } - - #[tokio::test] - async fn execute_read_de_default_projection_rejects_reserved_score_column_name() { - // Same collision, on the data-evolution path: a table with no PK-vector - // index (so the query falls through to DE materialization) but a user - // column named `__paimon_search_score` must fail loud under the default - // projection rather than emit two identically named columns. - use crate::spec::{ArrayType, FloatType, IntType, Schema, TableSchema}; - let schema = Schema::builder() - .column("id", DataType::Int(IntType::new())) - .column( - "embedding", - DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), - ) - .column(SEARCH_SCORE_COLUMN, DataType::Float(FloatType::new())) - .build() - .unwrap(); - let table = Table::new( - FileIOBuilder::new("memory").build().unwrap(), - Identifier::new("default", "de_vector_collision"), - "memory:/de_vector_collision".to_string(), - TableSchema::new(0, &schema), - None, - ); - let err = match table - .new_vector_search_builder() - .with_vector_column("embedding") - .with_query_vector(vec![1.0]) - .with_limit(5) - .execute_read() - .await - { - Ok(_) => panic!("DE default projection over a reserved-named column must fail loud"), - Err(e) => e, - }; - assert!( - matches!(&err, crate::Error::DataInvalid { message, .. } if message.contains("reserved column")), - "expected a reserved-column error, got: {err}" - ); - } - #[tokio::test] async fn execute_read_empty_plan_reserved_projection_fails_loud() { // Empty plan (no snapshot) must still fail loud on a reserved-name @@ -4787,17 +5010,20 @@ mod tests { ("pk-vector.index.columns", "embedding"), ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), ("fields.embedding.pk-vector.distance.metric", "l2"), + // Pin the index dimension so the query vector below matches it; the + // up-front dimension guard runs before this test's reserved-projection + // guard, so a mismatched query would mask the error under test. + ("fields.embedding.dimension", "4"), ]); for reserved in [ ROW_ID_FIELD_NAME, PKEY_VECTOR_POSITION_COLUMN, SEARCH_SCORE_COLUMN, - "_PKEY_VECTOR_SCORE", ] { let mut builder = table.new_vector_search_builder(); builder .with_vector_column("embedding") - .with_query_vector(vec![1.0]) + .with_query_vector(vec![1.0; 4]) .with_limit(5) .with_projection(&["id", reserved]); let err = builder @@ -4813,6 +5039,37 @@ mod tests { } } + #[tokio::test] + async fn execute_read_empty_plan_lumina_array_float_is_admitted() { + // A Lumina PK-vector `ARRAY` column is a valid configuration, but + // batch query dimension validation routed every `ARRAY` column + // through the vindex resolver, which rejects `lumina` as an unsupported + // index type before planning — failing even an empty table. The + // dimension must be resolved per the configured backend, so a + // well-formed Lumina query is admitted and (with no snapshot) yields an + // empty stream rather than an "Unsupported vindex index type" error. + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ( + "fields.embedding.pk-vector.index.type", + crate::lumina::LUMINA_IDENTIFIER, + ), + ("fields.embedding.pk-vector.distance.metric", "l2"), + ("lumina.index.dimension", "4"), + ]); + let mut stream = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0; 4]) + .with_limit(5) + .execute_read() + .await + .expect( + "Lumina ARRAY query must be admitted, not rejected as unsupported vindex", + ); + assert!(stream.try_next().await.unwrap().is_none()); + } + #[tokio::test] async fn execute_read_projection_reserved_name_fails_loud() { // Projecting a reserved metadata / row-id column must fail loud. The guard @@ -4827,7 +5084,6 @@ mod tests { ROW_ID_FIELD_NAME, PKEY_VECTOR_POSITION_COLUMN, SEARCH_SCORE_COLUMN, - "_PKEY_VECTOR_SCORE", ] { let mut builder = table.new_vector_search_builder(); builder @@ -4860,6 +5116,60 @@ mod tests { assert_eq!(names, vec!["id", "embedding"]); } + /// A PK-vector table whose user schema carries an extra column named + /// `reserved`, used to prove reserved metadata names are rejected even when + /// they arrive via the default (all-columns) projection. + fn pk_vector_table_with_extra_column(reserved: &str) -> Table { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column( + "embedding", + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ) + .column(reserved, DataType::Int(IntType::new())) + .option("pk-vector.index.columns", "embedding") + .option("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER) + .option("fields.embedding.pk-vector.distance.metric", "l2") + .build() + .unwrap(); + Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "reserved_col_test"), + "memory:/reserved_col_test".to_string(), + TableSchema::new(0, &schema), + None, + ) + } + + #[test] + fn resolve_materialize_read_type_default_rejects_reserved_user_column() { + // The default (all-columns) projection must reject a user column whose + // name collides with an injected metadata column, not only columns named + // in an explicit projection. Otherwise it silently passes on an empty + // result and collides with the metadata columns the read attaches. + let table = pk_vector_table_with_extra_column(SEARCH_SCORE_COLUMN); + let builder = table.new_vector_search_builder(); + let err = builder.resolve_materialize_read_type().unwrap_err(); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } + if message.contains("reserved column")), + "single-query default projection must reject reserved user column, got: {err:?}" + ); + } + + #[test] + fn batch_resolve_materialize_read_type_default_rejects_reserved_user_column() { + // Same guard on the batch resolver. + let table = pk_vector_table_with_extra_column(PKEY_VECTOR_POSITION_COLUMN); + let builder = table.new_batch_vector_search_builder(); + let err = builder.resolve_materialize_read_type().unwrap_err(); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } + if message.contains("reserved column")), + "batch default projection must reject reserved user column, got: {err:?}" + ); + } + #[test] fn resolve_materialize_read_type_projection_selects_named_columns() { let table = pk_vector_table(&[ diff --git a/crates/paimon/src/vindex/pkvector/ann.rs b/crates/paimon/src/vindex/pkvector/ann.rs index 627c052e..bbd7088d 100644 --- a/crates/paimon/src/vindex/pkvector/ann.rs +++ b/crates/paimon/src/vindex/pkvector/ann.rs @@ -176,6 +176,25 @@ pub(crate) fn map_ann_results( /// points of the async search path (the returned future is spawned on a `Send` /// runtime by callers such as the DataFusion integration). pub(crate) trait PkVectorAnnSearcher: Send + Sync { + /// Search one ANN segment for a batch of query vectors, returning one + /// BEST_FIRST result list per query (outer index aligned to `queries`). The + /// live-row mask (residual ∩ DV) is query-independent, so it is built once and + /// shared across all queries; only the per-query scores differ. + #[allow(clippy::too_many_arguments)] + fn search_batch( + &self, + segment: &BucketAnnSegment, + queries: &[&[f32]], + metric: VectorSearchMetric, + limit: usize, + active_source_files: &HashSet, + deletion_vectors: &HashMap>, + search_options: &HashMap, + residual_ranges: Option<&HashMap>, + ) -> crate::Result>>; + + /// Single-query wrapper over `search_batch`: searches the one query and + /// returns its result list. Asserts the batch produced exactly one list. #[allow(clippy::too_many_arguments)] fn search( &self, @@ -187,77 +206,119 @@ pub(crate) trait PkVectorAnnSearcher: Send + Sync { deletion_vectors: &HashMap>, search_options: &HashMap, residual_ranges: Option<&HashMap>, - ) -> crate::Result>; + ) -> crate::Result> { + let mut results = self.search_batch( + segment, + &[query], + metric, + limit, + active_source_files, + deletion_vectors, + search_options, + residual_ranges, + )?; + if results.len() != 1 { + return Err(data_invalid(format!( + "ANN batch search returned {} result lists for a single query", + results.len() + ))); + } + Ok(results.pop().expect("length checked to be 1")) + } } -/// Scorer seam: drives the underlying vindex ANN reader. Returns `ordinal -> -/// score` (higher-is-better). Any negative labels are skipped by the existing -/// `vindex` reader (`collect_results` drops `row_id < 0`), so this seam only -/// ever yields non-negative `u64` ordinals — no signed-label handling is needed -/// downstream. +/// Batch scorer seam: drives the underlying vindex ANN reader for a batch of +/// searches over ONE segment, opening the reader once and searching each query +/// against it (mirroring Java's shared-reader `visitBatchVectorSearch`). Returns +/// one `ordinal -> score` map (higher-is-better) per input search, aligned to the +/// `searches` slice. Any negative labels are skipped by the existing `vindex` +/// reader (`collect_results` drops `row_id < 0`), so this seam only ever yields +/// non-negative `u64` ordinals — no signed-label handling is needed downstream. /// -/// The production scorer drives `VindexVectorGlobalIndexReader::visit_vector_search` +/// The production scorer drives `VindexVectorGlobalIndexReader::visit_batch_vector_search` /// with a segment's index bytes; tests inject a synthetic scorer. The adapter's /// own logic (live-row masking, ordinal mapping, deletion checks, ordering) is /// exercised independently of the scorer. -pub(crate) type Scorer = Box< - dyn Fn(&BucketAnnSegment, &VectorSearch) -> crate::Result>> +pub(crate) type BatchScorer = Box< + dyn Fn(&BucketAnnSegment, &[VectorSearch]) -> crate::Result>>> + Send + Sync, >; /// Structural vindex-backed `PkVectorAnnSearcher`. Composes the pure helpers -/// (`build_live_row_ids`, `map_ann_results`) around the scorer seam. +/// (`build_live_row_ids`, `map_ann_results`) around the batch scorer seam. pub(crate) struct VindexAnnSearcher { field_name: String, - scorer: Scorer, + scorer: BatchScorer, } impl VindexAnnSearcher { - pub(crate) fn new(field_name: String, scorer: Scorer) -> Self { + pub(crate) fn new(field_name: String, scorer: BatchScorer) -> Self { Self { field_name, scorer } } } impl PkVectorAnnSearcher for VindexAnnSearcher { - fn search( + fn search_batch( &self, segment: &BucketAnnSegment, - query: &[f32], + queries: &[&[f32]], metric: VectorSearchMetric, limit: usize, active_source_files: &HashSet, deletion_vectors: &HashMap>, search_options: &HashMap, residual_ranges: Option<&HashMap>, - ) -> crate::Result> { + ) -> crate::Result>> { if limit == 0 { return Err(data_invalid("vector search limit must be positive")); } let source_files = segment.source_meta.source_files(); - let mut search = VectorSearch::new(query.to_vec(), limit, self.field_name.clone())? - .with_options(search_options.clone()); - if let Some(live) = build_live_row_ids( + // The live-row mask depends only on the segment's sources, the active set, + // the deletion vectors, and the residual — none of which vary by query — + // so it is built once and shared across every query's search. + let live = build_live_row_ids( source_files, active_source_files, deletion_vectors, residual_ranges, - )? { - search = search.with_include_row_ids(live); + )?; + let mut searches = Vec::with_capacity(queries.len()); + for query in queries { + let mut search = VectorSearch::new(query.to_vec(), limit, self.field_name.clone())? + .with_options(search_options.clone()); + if let Some(live) = &live { + search = search.with_include_row_ids(live.clone()); + } + searches.push(search); } - let scored = match (self.scorer)(segment, &search)? { - Some(map) => map, - None => return Ok(Vec::new()), - }; - let scored: Vec<(u64, f32)> = scored.into_iter().collect(); - map_ann_results( - &scored, - &segment.source_meta, - active_source_files, - deletion_vectors, - residual_ranges, - metric, - ) + let scored_batch = (self.scorer)(segment, &searches)?; + if scored_batch.len() != queries.len() { + return Err(data_invalid(format!( + "ANN batch scorer returned {} result maps for {} queries", + scored_batch.len(), + queries.len() + ))); + } + let mut out = Vec::with_capacity(queries.len()); + for scored in scored_batch { + let results = match scored { + Some(map) => { + let scored: Vec<(u64, f32)> = map.into_iter().collect(); + map_ann_results( + &scored, + &segment.source_meta, + active_source_files, + deletion_vectors, + residual_ranges, + metric, + )? + } + None => Vec::new(), + }; + out.push(results); + } + Ok(out) } } @@ -424,14 +485,17 @@ mod tests { let scorer_has_filter = Arc::clone(&seen_has_filter); let searcher = VindexAnnSearcher::new( "embedding".to_string(), - Box::new(move |_segment: &BucketAnnSegment, search: &VectorSearch| { - *scorer_limit.lock().unwrap() = search.limit; - *scorer_has_filter.lock().unwrap() = search.include_row_ids.is_some(); - let mut scores = HashMap::new(); - scores.insert(3u64, 0.5f32); // -> (f1, 0) - scores.insert(0u64, 0.25f32); // -> (f0, 0), l2 dist 3.0 - Ok(Some(scores)) - }), + Box::new( + move |_segment: &BucketAnnSegment, searches: &[VectorSearch]| { + let search = &searches[0]; + *scorer_limit.lock().unwrap() = search.limit; + *scorer_has_filter.lock().unwrap() = search.include_row_ids.is_some(); + let mut scores = HashMap::new(); + scores.insert(3u64, 0.5f32); // -> (f1, 0) + scores.insert(0u64, 0.25f32); // -> (f0, 0), l2 dist 3.0 + Ok(vec![Some(scores)]) + }, + ), ); let segment = BucketAnnSegment::for_test({ use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta}; @@ -472,7 +536,9 @@ mod tests { fn test_vindex_adapter_rejects_non_positive_limit() { let searcher = VindexAnnSearcher::new( "embedding".to_string(), - Box::new(|_: &BucketAnnSegment, _: &VectorSearch| Ok(None)), + Box::new(|_: &BucketAnnSegment, searches: &[VectorSearch]| { + Ok(vec![None; searches.len()]) + }), ); let segment = BucketAnnSegment::for_test({ use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta}; @@ -498,7 +564,9 @@ mod tests { fn test_vindex_adapter_empty_scorer_result_is_empty() { let searcher = VindexAnnSearcher::new( "embedding".to_string(), - Box::new(|_: &BucketAnnSegment, _: &VectorSearch| Ok(None)), + Box::new(|_: &BucketAnnSegment, searches: &[VectorSearch]| { + Ok(vec![None; searches.len()]) + }), ); let segment = BucketAnnSegment::for_test({ use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta}; @@ -633,13 +701,15 @@ mod tests { let scorer_rows = Arc::clone(&seen_rows); let searcher = VindexAnnSearcher::new( "embedding".to_string(), - Box::new(move |_segment: &BucketAnnSegment, search: &VectorSearch| { - *scorer_rows.lock().unwrap() = search - .include_row_ids - .as_ref() - .map(|t| t.iter().collect::>()); - Ok(None) - }), + Box::new( + move |_segment: &BucketAnnSegment, searches: &[VectorSearch]| { + *scorer_rows.lock().unwrap() = searches[0] + .include_row_ids + .as_ref() + .map(|t| t.iter().collect::>()); + Ok(vec![None; searches.len()]) + }, + ), ); let segment = BucketAnnSegment::for_test(source_meta(&[("f0", 3)])); let mut residual = HashMap::new(); @@ -658,4 +728,124 @@ mod tests { .unwrap(); assert_eq!(seen_rows.lock().unwrap().clone(), Some(vec![0, 2])); } + + #[test] + fn test_search_batch_of_one_equals_single_query() { + // The single-query `search` wrapper must return exactly what + // `search_batch(&[q])[0]` returns for the same inputs. + let make = || { + VindexAnnSearcher::new( + "embedding".to_string(), + Box::new(|_: &BucketAnnSegment, searches: &[VectorSearch]| { + let mut out = Vec::with_capacity(searches.len()); + for _ in searches { + let mut scores = HashMap::new(); + scores.insert(3u64, 0.5f32); // -> (f1, 0) + scores.insert(0u64, 0.25f32); // -> (f0, 0) + out.push(Some(scores)); + } + Ok(out) + }), + ) + }; + let meta = source_meta(&[("f0", 3), ("f1", 5)]); + let single = make() + .search( + &BucketAnnSegment::for_test(meta.clone()), + &[0.0, 0.0], + VectorSearchMetric::L2, + 2, + &active_set(&["f0", "f1"]), + &HashMap::new(), + &HashMap::new(), + None, + ) + .unwrap(); + let query: &[f32] = &[0.0, 0.0]; + let batch = make() + .search_batch( + &BucketAnnSegment::for_test(meta), + &[query], + VectorSearchMetric::L2, + 2, + &active_set(&["f0", "f1"]), + &HashMap::new(), + &HashMap::new(), + None, + ) + .unwrap(); + assert_eq!(batch.len(), 1); + assert_eq!(batch[0], single); + } + + #[test] + fn test_search_batch_returns_independent_per_query_results() { + // Two queries route to different synthetic scores; each result list is + // mapped from that query's own scores, with a shared live-row mask. + let searcher = VindexAnnSearcher::new( + "embedding".to_string(), + Box::new(|_: &BucketAnnSegment, searches: &[VectorSearch]| { + let mut out = Vec::with_capacity(searches.len()); + for (i, _) in searches.iter().enumerate() { + let mut scores = HashMap::new(); + // Query 0 -> ordinal 0 (f0,0); query 1 -> ordinal 3 (f1,0). + if i == 0 { + scores.insert(0u64, 0.5f32); + } else { + scores.insert(3u64, 0.5f32); + } + out.push(Some(scores)); + } + Ok(out) + }), + ); + let q0: &[f32] = &[0.0, 0.0]; + let q1: &[f32] = &[1.0, 1.0]; + let results = searcher + .search_batch( + &BucketAnnSegment::for_test(source_meta(&[("f0", 3), ("f1", 5)])), + &[q0, q1], + VectorSearchMetric::L2, + 2, + &active_set(&["f0", "f1"]), + &HashMap::new(), + &HashMap::new(), + None, + ) + .unwrap(); + assert_eq!(results.len(), 2); + assert_eq!(results[0][0].data_file_name, "f0"); + assert_eq!(results[1][0].data_file_name, "f1"); + } + + #[test] + fn test_search_batch_fails_loud_on_result_count_mismatch() { + // A batch scorer that returns the wrong number of result maps is corruption + // and must fail loud, not be silently padded/truncated. + let searcher = VindexAnnSearcher::new( + "embedding".to_string(), + Box::new(|_: &BucketAnnSegment, _: &[VectorSearch]| { + // Only one map returned regardless of query count. + Ok(vec![None]) + }), + ); + let q0: &[f32] = &[0.0, 0.0]; + let q1: &[f32] = &[1.0, 1.0]; + let err = searcher + .search_batch( + &BucketAnnSegment::for_test(source_meta(&[("f0", 3)])), + &[q0, q1], + VectorSearchMetric::L2, + 2, + &active_set(&["f0"]), + &HashMap::new(), + &HashMap::new(), + None, + ) + .unwrap_err(); + assert!( + err.to_string().contains("2 queries") || err.to_string().contains("result maps"), + "got: {err}" + ); + } } diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs b/crates/paimon/src/vindex/pkvector/bucket.rs index 3b2f7fb1..56900e07 100644 --- a/crates/paimon/src/vindex/pkvector/bucket.rs +++ b/crates/paimon/src/vindex/pkvector/bucket.rs @@ -355,6 +355,245 @@ pub(crate) async fn bucket_search( Ok(BucketSearchResult { indexed, exact }) } +/// Multi-query ANN + exact fallback search for one snapshot bucket, returning one +/// [`BucketSearchResult`] per query (outer index aligned to `queries`). N queries +/// share one pass over the ANN segments (the reader is opened once per segment and +/// every query is searched against it) and one stream per uncovered exact file (the +/// stream is opened once and every query is scored from it), with independent +/// per-query bounded heaps so no query's candidates bleed into another's. +/// +/// `queries.len() == 1` short-circuits to the single-query [`bucket_search`] and +/// wraps its result in a one-element `Vec`, so a batch-of-one is byte-identical to +/// the single-query path (including tie ordering). The remaining arguments match +/// [`bucket_search`]. +#[allow(clippy::too_many_arguments)] +#[allow(clippy::type_complexity)] +pub(crate) async fn bucket_search_batch( + ann_searcher: Option<&dyn PkVectorAnnSearcher>, + ann_segments: &[BucketAnnSegment], + active_files: &[BucketActiveFile], + deletion_vectors: &HashMap>, + exact_file_search: &(dyn for<'a> Fn( + &'a BucketActiveFile, + &'a [&'a [f32]], + VectorSearchMetric, + usize, + &'a (dyn Fn(i64) -> bool + Sync), + ) -> ExactFileSearchFuture<'a> + + Send + + Sync), + queries: &[&[f32]], + metric: VectorSearchMetric, + indexed_limit: usize, + exact_limit: usize, + search_options: &HashMap, + skip_exact_fallback: bool, + residual_ranges: Option<&HashMap>, +) -> crate::Result> { + if queries.is_empty() { + return Err(data_invalid("vector search requires at least one query")); + } + // A batch-of-one routes to the single-query body so its tie ordering is + // byte-identical to the pre-batch path (the multi-query heaps could in + // principle differ). + if queries.len() == 1 { + let single = bucket_search( + ann_searcher, + ann_segments, + active_files, + deletion_vectors, + exact_file_search, + queries[0], + metric, + indexed_limit, + exact_limit, + search_options, + skip_exact_fallback, + residual_ranges, + ) + .await?; + return Ok(vec![single]); + } + + if indexed_limit == 0 { + return Err(data_invalid("vector search limit must be positive")); + } + if exact_limit == 0 { + return Err(data_invalid("vector search limit must be positive")); + } + // Validate-before-OPEN: any 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). + for query in queries { + 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 { + if file.row_count < 0 { + return Err(data_invalid(format!( + "active data file {} row count must not be negative: {}", + file.file_name, file.row_count + ))); + } + if files_by_name + .insert(file.file_name.as_str(), file) + .is_some() + { + return Err(data_invalid(format!( + "duplicate data file: {}", + file.file_name + ))); + } + } + + // Validate ANN segments mirror Java PkVectorBucketIndexState constructor checks: + // (1) payload file uniqueness, (2) no source file covered by multiple segments. + let mut segments_by_path: HashMap<&str, usize> = HashMap::new(); + let mut source_to_segment: HashMap<&str, &str> = HashMap::new(); + for (idx, segment) in ann_segments.iter().enumerate() { + if segments_by_path + .insert(segment.path.as_str(), idx) + .is_some() + { + return Err(data_invalid(format!( + "ANN segment payload {} appears more than once", + segment.path + ))); + } + for source in segment.source_meta.source_files() { + if let Some(&prior_segment_path) = source_to_segment.get(source.file_name()) { + return Err(data_invalid(format!( + "source data file {} is covered by both ANN segments {} and {}", + source.file_name(), + prior_segment_path, + segment.path + ))); + } + source_to_segment.insert(source.file_name(), segment.path.as_str()); + } + } + + let mut indexed_heaps: Vec> = (0..queries.len()) + .map(|_| BinaryHeap::with_capacity(indexed_limit + 1)) + .collect(); + let mut exact_heaps: Vec> = (0..queries.len()) + .map(|_| BinaryHeap::with_capacity(exact_limit + 1)) + .collect(); + let active_source_files: HashSet = + files_by_name.keys().map(|name| name.to_string()).collect(); + let covered = covered_source_files(ann_segments, active_files); + + for segment in ann_segments { + for source in segment.source_meta.source_files() { + if let Some(active) = files_by_name.get(source.file_name()) { + if active.row_count != source.row_count() { + return Err(data_invalid(format!( + "ANN source {} does not match the active data file", + source.file_name() + ))); + } + } + } + let searcher = ann_searcher.ok_or_else(|| data_invalid("ANN search is not configured"))?; + // One shared reader per segment searches all queries; fan the per-query + // hits into their own bounded heaps. + let per_query = searcher.search_batch( + segment, + queries, + metric, + indexed_limit, + &active_source_files, + deletion_vectors, + search_options, + residual_ranges, + )?; + if per_query.len() != queries.len() { + return Err(data_invalid(format!( + "ANN batch search returned {} result lists for {} queries", + per_query.len(), + queries.len() + ))); + } + for (results, heap) in per_query.into_iter().zip(indexed_heaps.iter_mut()) { + for result in results { + add_candidate(heap, result, indexed_limit); + } + } + } + + if !skip_exact_fallback { + for file in active_files { + if covered.contains(&file.file_name) { + continue; + } + let residual_allowed: Option<&roaring::RoaringTreemap> = match residual_ranges { + Some(ranges) => match ranges.get(&file.file_name) { + Some(allowed) if !allowed.is_empty() => Some(allowed), + _ => continue, + }, + None => None, + }; + let dv = deletion_vectors.get(&file.file_name).cloned(); + let is_excluded = move |position: i64| -> bool { + let dv_deleted = match &dv { + Some(dv) => u64::try_from(position) + .map(|p| dv.is_deleted(p)) + .unwrap_or(false), + None => false, + }; + if dv_deleted { + return true; + } + match residual_allowed { + None => false, + Some(allowed) => match u64::try_from(position) { + Ok(p) => !allowed.contains(p), + Err(_) => true, + }, + } + }; + // One shared stream per file scores every query into its own heap. + let per_query = exact_file_search( + file, + queries, + metric, + exact_limit, + &is_excluded as &(dyn Fn(i64) -> bool + Sync), + ) + .await?; + if per_query.len() != queries.len() { + return Err(data_invalid(format!( + "exact file search returned {} result lists for {} queries", + per_query.len(), + queries.len() + ))); + } + for (results, heap) in per_query.into_iter().zip(exact_heaps.iter_mut()) { + for result in results { + add_candidate(heap, result, exact_limit); + } + } + } + } + + let mut out = Vec::with_capacity(queries.len()); + for (indexed_heap, exact_heap) in indexed_heaps.into_iter().zip(exact_heaps) { + let mut indexed: Vec = + indexed_heap.into_iter().map(|w| w.0).collect(); + indexed.sort_by(best_first); + let mut exact: Vec = exact_heap.into_iter().map(|w| w.0).collect(); + exact.sort_by(best_first); + out.push(BucketSearchResult { indexed, exact }); + } + Ok(out) +} + #[cfg(test)] mod tests { use super::*; @@ -474,18 +713,18 @@ mod tests { result: Vec, } impl PkVectorAnnSearcher for FakeAnnSearcher { - fn search( + fn search_batch( &self, _segment: &BucketAnnSegment, - _query: &[f32], + queries: &[&[f32]], _metric: VectorSearchMetric, _limit: usize, _active_source_files: &HashSet, _dvs: &HashMap>, _opts: &HashMap, _residual_ranges: Option<&HashMap>, - ) -> crate::Result> { - Ok(self.result.clone()) + ) -> crate::Result>> { + Ok(queries.iter().map(|_| self.result.clone()).collect()) } } @@ -1352,4 +1591,264 @@ mod tests { "the exact-file search closure must not be invoked for a malformed query" ); } + + /// Build a per-file exact-search closure that runs the reference `exact_search` + /// over an in-memory `ArrayReader` of `vectors`, looping ALL queries into one + /// list per query (the outer `Vec` is aligned to `queries`). Used by the + /// multi-query `bucket_search_batch` tests. + #[allow(clippy::type_complexity)] + fn array_search_batch( + 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 file_name = file.file_name.clone(); + let owned_queries: Vec> = queries.iter().map(|q| q.to_vec()).collect(); + let vectors = Arc::clone(&vectors); + Box::pin(async move { + let mut out = Vec::with_capacity(owned_queries.len()); + for query in &owned_queries { + let mut reader = ArrayReader::new(dimension, (*vectors).clone()); + out.push(exact_search( + &file_name, + &mut reader, + query, + metric, + exact_limit, + is_excluded, + )?); + } + Ok(out) + }) + }, + ) + } + + #[tokio::test] + async fn test_bucket_search_batch_of_one_equals_single_query() { + // A batch-of-one must equal the single-query `bucket_search` exactly. + let segment = BucketAnnSegment::for_test(meta(&[("ann.mosaic", 3)])); + let ann = FakeAnnSearcher { + result: vec![ + PkVectorSearchResult { + data_file_name: "ann.mosaic".into(), + row_position: 0, + distance: 0.1, + }, + PkVectorSearchResult { + data_file_name: "ann.mosaic".into(), + row_position: 1, + distance: 0.2, + }, + ], + }; + let active_files = vec![active("ann.mosaic", 3), active("exact.mosaic", 3)]; + let dvs: HashMap> = HashMap::new(); + let opts = HashMap::new(); + let query = [1.0f32, 0.0]; + + let single = { + let factory = array_search(vec![ + Some(vec![1.0, 0.0]), + Some(vec![9.0, 0.0]), + Some(vec![8.0, 0.0]), + ]); + bucket_search( + Some(&ann), + &[BucketAnnSegment::for_test(meta(&[("ann.mosaic", 3)]))], + &active_files, + &dvs, + &factory, + &query, + VectorSearchMetric::L2, + 3, + 2, + &opts, + false, + None, + ) + .await + .unwrap() + }; + + let factory = array_search_batch(vec![ + Some(vec![1.0, 0.0]), + Some(vec![9.0, 0.0]), + Some(vec![8.0, 0.0]), + ]); + let query_ref: &[f32] = &query; + let batch = bucket_search_batch( + Some(&ann), + &[segment], + &active_files, + &dvs, + &factory, + &[query_ref], + VectorSearchMetric::L2, + 3, + 2, + &opts, + false, + None, + ) + .await + .unwrap(); + assert_eq!(batch.len(), 1); + assert_eq!(batch[0].indexed, single.indexed); + assert_eq!(batch[0].exact, single.exact); + } + + #[tokio::test] + async fn test_bucket_search_batch_per_query_heaps_independent() { + // Two queries over one exact file with a shared residual/DV; each query's + // Top-K comes from its own heap (no cross-query bleed). Query 0 is nearest + // pos 0; query 1 is nearest pos 3. + let factory = array_search_batch(vec![ + Some(vec![0.0, 0.0]), + Some(vec![1.0, 0.0]), + Some(vec![2.0, 0.0]), + Some(vec![3.0, 0.0]), + ]); + let q0: &[f32] = &[0.0, 0.0]; + let q1: &[f32] = &[3.0, 0.0]; + let out = bucket_search_batch( + None, + &[], + &[active("data-1", 4)], + &HashMap::new(), + &factory, + &[q0, q1], + VectorSearchMetric::L2, + 1, + 1, + &HashMap::new(), + false, + None, + ) + .await + .unwrap(); + assert_eq!(out.len(), 2); + assert_eq!(out[0].exact.len(), 1); + assert_eq!(out[0].exact[0].row_position, 0); + assert_eq!(out[1].exact.len(), 1); + assert_eq!(out[1].exact[0].row_position, 3); + } + + #[tokio::test] + async fn test_bucket_search_batch_shares_one_read_per_file_across_queries() { + // The exact-file closure is invoked ONCE per uncovered file regardless of + // the query count (the shared stream scores all queries). + let calls = std::sync::Mutex::new(Vec::::new()); + 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 owned: Vec> = queries.iter().map(|q| q.to_vec()).collect(); + Box::pin(async move { + let mut out = Vec::with_capacity(owned.len()); + for query in &owned { + let mut reader = + ArrayReader::new(2, vec![Some(vec![1.0, 0.0]), Some(vec![3.0, 0.0])]); + out.push(exact_search( + &file_name, + &mut reader, + query, + metric, + exact_limit, + is_excluded, + )?); + } + Ok(out) + }) + }, + ); + let q0: &[f32] = &[0.0, 0.0]; + let q1: &[f32] = &[4.0, 0.0]; + let out = bucket_search_batch( + None, + &[], + &[active("data-1", 2)], + &HashMap::new(), + &factory, + &[q0, q1], + VectorSearchMetric::L2, + 2, + 2, + &HashMap::new(), + false, + None, + ) + .await + .unwrap(); + assert_eq!(out.len(), 2); + assert_eq!( + calls.lock().unwrap().as_slice(), + &["data-1".to_string()], + "one shared read per file for the whole batch" + ); + } + + #[tokio::test] + async fn test_bucket_search_batch_malformed_query_fails_before_opening_any_file() { + // A non-finite element in ANY query fails loud before the exact-file search + // closure is ever invoked (validate-before-OPEN). + 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 good: &[f32] = &[0.0, 0.0]; + let bad: &[f32] = &[f32::NAN, 0.0]; + let err = bucket_search_batch( + None, + &[], + &[active("data-1", 2)], + &HashMap::new(), + &factory, + &[good, bad], + 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 batch" + ); + } } diff --git a/crates/paimon/tests/pk_vector_batch_test.rs b/crates/paimon/tests/pk_vector_batch_test.rs new file mode 100644 index 00000000..84790e77 --- /dev/null +++ b/crates/paimon/tests/pk_vector_batch_test.rs @@ -0,0 +1,787 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! End-to-end acceptance gate for BATCH primary-key vector search. +//! +//! Builds a complete, self-contained primary-key vector table entirely from Rust +//! (mirroring `pk_vector_baseline_test`), then reads it back through the public +//! batch surface `new_batch_vector_search_builder().execute_read()` and asserts: +//! - batch-of-one == the single-query `execute_read`; +//! - an N-query batch yields one stream per query, each matching the +//! corresponding independent single-query read (arity + order + independence); +//! - an empty snapshot yields N empty streams (arity preserved); +//! - the PK batch `execute()` (scored) fails loud, directing to `execute_read`; +//! - a shared residual filter reshapes every query's rows with no cross-query +//! bleed. + +use std::collections::HashMap; + +use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; +use arrow_array::{Array, ArrayRef, Float32Array, Int32Array, RecordBatch}; +use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; +use bytes::Bytes; +use futures::TryStreamExt; +use paimon::catalog::Identifier; +use paimon::io::{FileIO, FileIOBuilder}; +use paimon::spec::{ + ArrayType, DataFileMeta, DataType, Datum, FloatType, GlobalIndexMeta, IndexFileMeta, IntType, + Predicate, PredicateBuilder, Schema, TableSchema, VectorType, +}; +use paimon::table::{ArrowRecordBatchStream, CommitMessage, SchemaManager, Table, TableCommit}; +use paimon_vindex_core::index::{VectorIndexConfig, VectorIndexTrainer, VectorIndexWriter}; +use paimon_vindex_core::io::PosWriter; +use std::sync::Arc; + +const DIM: usize = 4; +const VECTOR_COLUMN: &str = "embedding"; +const INDEX_TYPE: &str = "ivf-flat"; + +fn analytic_topk(query: &[f32], vectors: &[[f32; DIM]], k: usize) -> Vec<(u64, f32)> { + let mut scored: Vec<(u64, f32)> = vectors + .iter() + .enumerate() + .map(|(pos, v)| { + let dist: f32 = v + .iter() + .zip(query.iter()) + .map(|(a, b)| (a - b) * (a - b)) + .sum(); + (pos as u64, dist) + }) + .collect(); + scored.sort_by(|a, b| a.1.total_cmp(&b.1)); + scored.truncate(k); + scored +} + +fn table_options() -> Vec<(String, String)> { + vec![ + ("bucket".to_string(), "1".to_string()), + ("deletion-vectors.enabled".to_string(), "true".to_string()), + ( + "pk-vector.index.columns".to_string(), + VECTOR_COLUMN.to_string(), + ), + ( + format!("fields.{VECTOR_COLUMN}.pk-vector.index.type"), + INDEX_TYPE.to_string(), + ), + ( + format!("fields.{VECTOR_COLUMN}.pk-vector.distance.metric"), + "l2".to_string(), + ), + ] +} + +fn pk_vector_schema() -> TableSchema { + let mut builder = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column( + VECTOR_COLUMN, + DataType::Vector( + VectorType::try_new(true, DIM as u32, DataType::Float(FloatType::new())).unwrap(), + ), + ) + .primary_key(["id"]); + for (k, v) in table_options() { + builder = builder.option(k, v); + } + TableSchema::new(0, &builder.build().unwrap()) +} + +fn pk_vector_array_schema() -> TableSchema { + let mut builder = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column( + VECTOR_COLUMN, + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ) + .primary_key(["id"]); + for (k, v) in table_options() { + builder = builder.option(k, v); + } + builder = builder.option(format!("{INDEX_TYPE}.dimension"), DIM.to_string()); + TableSchema::new(0, &builder.build().unwrap()) +} + +fn data_batch(vectors: &[[f32; DIM]]) -> RecordBatch { + let ids: Vec = (0..vectors.len() as i32).collect(); + let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); + let mut vector_builder = FixedSizeListBuilder::new(Float32Builder::new(), DIM as i32) + .with_field(element_field.clone()); + for vector in vectors { + for &value in vector { + vector_builder.values().append_value(value); + } + vector_builder.append(true); + } + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new( + VECTOR_COLUMN, + ArrowDataType::FixedSizeList(element_field, DIM as i32), + true, + ), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(ids)) as ArrayRef, + Arc::new(vector_builder.finish()) as ArrayRef, + ], + ) + .unwrap() +} + +fn java_write_utf(s: &str) -> Vec { + let mut body = Vec::new(); + for c in s.encode_utf16() { + if (0x0001..=0x007F).contains(&c) { + body.push(c as u8); + } else if c > 0x07FF { + body.push(0xE0 | (c >> 12) as u8); + body.push(0x80 | ((c >> 6) & 0x3F) as u8); + body.push(0x80 | (c & 0x3F) as u8); + } else { + body.push(0xC0 | (c >> 6) as u8); + body.push(0x80 | (c & 0x3F) as u8); + } + } + let mut out = (body.len() as u16).to_be_bytes().to_vec(); + out.extend_from_slice(&body); + out +} + +fn source_meta_bytes(data_level: i32, files: &[(&str, i64)]) -> Vec { + let mut out = Vec::new(); + out.extend_from_slice(&1i32.to_be_bytes()); + out.extend_from_slice(&data_level.to_be_bytes()); + out.extend_from_slice(&(files.len() as i32).to_be_bytes()); + for (name, rows) in files { + out.extend_from_slice(&java_write_utf(name)); + out.extend_from_slice(&rows.to_be_bytes()); + } + out +} + +async fn write_ann_segment( + file_io: &FileIO, + table_location: &str, + file_name: &str, + vectors: &[[f32; DIM]], +) -> u64 { + let n = vectors.len(); + let flat: Vec = vectors.iter().flat_map(|v| v.iter().copied()).collect(); + let ids: Vec = (0..n as i64).collect(); + let native_options = HashMap::from([ + ("index.type".to_string(), "ivf_flat".to_string()), + ("dimension".to_string(), DIM.to_string()), + ("nlist".to_string(), "1".to_string()), + ("metric".to_string(), "l2".to_string()), + ]); + let config = VectorIndexConfig::from_options(&native_options).unwrap(); + let training = VectorIndexTrainer::train(config, &flat, n).unwrap(); + let mut writer = VectorIndexWriter::new(training); + writer.add_vectors(&ids, &flat, n).unwrap(); + let mut bytes = Vec::new(); + { + let mut output = PosWriter::new(&mut bytes); + writer.write(&mut output).unwrap(); + } + let index_dir = format!("{}/index", table_location.trim_end_matches('/')); + file_io.mkdirs(&index_dir).await.unwrap(); + let index_path = format!("{index_dir}/{file_name}"); + let file_size = bytes.len() as u64; + file_io + .new_output(&index_path) + .unwrap() + .write(Bytes::from(bytes)) + .await + .unwrap(); + file_size +} + +async fn open_table(file_io: &FileIO, location: &str) -> Table { + let schema = SchemaManager::new(file_io.clone(), location.to_string()) + .latest() + .await + .expect("failed to list schemas") + .expect("table has no schema"); + Table::new( + file_io.clone(), + Identifier::new("default", "pkvector_batch"), + location.to_string(), + (*schema).clone(), + None, + ) +} + +/// Build a complete self-contained primary-key vector table over `vectors` with a +/// real vindex ANN segment committed in one snapshot. Returns the temp dir (kept +/// alive by the caller) and the opened table. +async fn build_table(vectors: &[[f32; DIM]]) -> (tempfile::TempDir, Table) { + let tmp = tempfile::tempdir().expect("create temp dir"); + let location = format!("file://{}", tmp.path().display()); + let file_io = FileIOBuilder::new("file").build().unwrap(); + + for dir in ["schema", "snapshot", "manifest", "index"] { + file_io.mkdirs(&format!("{location}/{dir}")).await.unwrap(); + } + let schema = pk_vector_schema(); + file_io + .new_output(&format!("{location}/schema/schema-{}", schema.id())) + .unwrap() + .write(Bytes::from(serde_json::to_vec(&schema).unwrap())) + .await + .unwrap(); + + let table = open_table(&file_io, &location).await; + + let write_builder = table.new_write_builder(); + let mut writer = write_builder.new_write().unwrap(); + writer + .write_arrow_batch(&data_batch(vectors)) + .await + .unwrap(); + let write_messages = writer.prepare_commit().await.unwrap(); + let written = &write_messages[0]; + let base_meta = written.new_files[0].clone(); + let bucket = written.bucket; + let partition = written.partition.clone(); + let data_file_name = base_meta.file_name.clone(); + let row_count = base_meta.row_count; + + let indexed_meta = DataFileMeta { + level: 1, + file_source: Some(1), + first_row_id: Some(0), + ..base_meta + }; + + let index_file_name = "vector-ivf-flat-pkvector-batch.index".to_string(); + let index_file_size = write_ann_segment(&file_io, &location, &index_file_name, vectors).await; + + let vector_field_id = schema + .fields() + .iter() + .find(|f| f.name() == VECTOR_COLUMN) + .expect("vector field present") + .id(); + let index_file = IndexFileMeta { + index_type: INDEX_TYPE.to_string(), + file_name: index_file_name, + file_size: i32::try_from(index_file_size).unwrap(), + row_count: i32::try_from(row_count).unwrap(), + deletion_vectors_ranges: None, + global_index_meta: Some(GlobalIndexMeta { + row_range_start: 0, + row_range_end: row_count - 1, + index_field_id: vector_field_id, + extra_field_ids: None, + source_meta: Some(source_meta_bytes( + indexed_meta.level, + &[(&data_file_name, row_count)], + )), + index_meta: None, + }), + }; + + let mut message = CommitMessage::new(partition, bucket, vec![indexed_meta]); + message.new_index_files = vec![index_file]; + TableCommit::new(table.clone(), "pkvector-batch".to_string()) + .commit(vec![message]) + .await + .unwrap(); + + (tmp, table) +} + +/// Build a bare table with a schema but NO snapshot (empty). Used for the +/// empty-snapshot arity test. +async fn build_empty_table() -> (tempfile::TempDir, Table) { + let tmp = tempfile::tempdir().expect("create temp dir"); + let location = format!("file://{}", tmp.path().display()); + let file_io = FileIOBuilder::new("file").build().unwrap(); + for dir in ["schema", "snapshot", "manifest", "index"] { + file_io.mkdirs(&format!("{location}/{dir}")).await.unwrap(); + } + let schema = pk_vector_schema(); + file_io + .new_output(&format!("{location}/schema/schema-{}", schema.id())) + .unwrap() + .write(Bytes::from(serde_json::to_vec(&schema).unwrap())) + .await + .unwrap(); + let table = open_table(&file_io, &location).await; + (tmp, table) +} + +async fn build_empty_array_table() -> (tempfile::TempDir, Table) { + let tmp = tempfile::tempdir().expect("create temp dir"); + let location = format!("file://{}", tmp.path().display()); + let file_io = FileIOBuilder::new("file").build().unwrap(); + for dir in ["schema", "snapshot", "manifest", "index"] { + file_io.mkdirs(&format!("{location}/{dir}")).await.unwrap(); + } + let schema = pk_vector_array_schema(); + file_io + .new_output(&format!("{location}/schema/schema-{}", schema.id())) + .unwrap() + .write(Bytes::from(serde_json::to_vec(&schema).unwrap())) + .await + .unwrap(); + let table = open_table(&file_io, &location).await; + (tmp, table) +} + +/// Flatten a materialized best-first stream into per-row `(id, score)` tuples. +async fn drain_ids_and_scores(stream: ArrowRecordBatchStream) -> (Vec, Vec) { + let batches = stream + .try_collect::>() + .await + .expect("collecting read batches failed"); + let ids: Vec = batches + .iter() + .flat_map(|b| { + let idx = b.schema().index_of("id").unwrap(); + b.column(idx) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + let scores: Vec = batches + .iter() + .flat_map(|b| { + let idx = b.schema().index_of("__paimon_search_score").unwrap(); + b.column(idx) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + (ids, scores) +} + +/// Single-query `execute_read` into `(id, score)` tuples. +async fn single_read(table: &Table, query: Vec, limit: usize) -> (Vec, Vec) { + let mut builder = table.new_vector_search_builder(); + builder + .with_vector_column(VECTOR_COLUMN) + .with_query_vector(query) + .with_limit(limit); + let stream = builder + .execute_read() + .await + .expect("single-query read failed"); + drain_ids_and_scores(stream).await +} + +fn fixture() -> Vec<[f32; DIM]> { + // Distances make each query's top-k order unique. + vec![ + [0.0, 4.0, 0.0, 0.0], // pos 0 + [8.0, 0.0, 0.0, 0.0], // pos 1 + [0.0, 0.0, 5.0, 0.0], // pos 2 + [7.0, 0.0, 0.0, 0.0], // pos 3 + [0.0, 0.0, 0.0, 6.0], // pos 4 + [9.0, 0.0, 0.0, 0.0], // pos 5 + ] +} + +// Gated off Windows for the same `file://` tempdir reason as `pk_vector_baseline_test`. +#[cfg(not(windows))] +#[tokio::test] +async fn batch_of_one_equals_single_read() { + let vectors = fixture(); + let (_tmp, table) = build_table(&vectors).await; + let query = vec![10.0, 0.0, 0.0, 0.0]; + + let (single_ids, single_scores) = single_read(&table, query.clone(), 3).await; + + let mut builder = table.new_batch_vector_search_builder(); + let mut streams = builder + .with_vector_column(VECTOR_COLUMN) + .with_query_vectors(vec![query]) + .with_limit(3) + .execute_read() + .await + .expect("batch-of-one read failed"); + assert_eq!(streams.len(), 1, "batch-of-one yields exactly one stream"); + let (batch_ids, batch_scores) = drain_ids_and_scores(streams.remove(0)).await; + + assert_eq!(batch_ids, single_ids); + assert_eq!(batch_scores.len(), single_scores.len()); + for (got, want) in batch_scores.iter().zip(&single_scores) { + assert!((got - want).abs() < 1e-6, "score diverges: {got} vs {want}"); + } +} + +#[cfg(not(windows))] +#[tokio::test] +async fn n_query_batch_matches_n_independent_single_reads() { + let vectors = fixture(); + let (_tmp, table) = build_table(&vectors).await; + let queries = vec![ + vec![10.0, 0.0, 0.0, 0.0], // nearest the x-axis vectors + vec![0.0, 0.0, 6.0, 0.0], // nearest pos 2 / pos 4 + vec![0.0, 5.0, 0.0, 0.0], // nearest pos 0 + ]; + + // Independent single reads, one per query. + let mut expected = Vec::new(); + for q in &queries { + expected.push(single_read(&table, q.clone(), 3).await); + } + + let mut builder = table.new_batch_vector_search_builder(); + let streams = builder + .with_vector_column(VECTOR_COLUMN) + .with_query_vectors(queries.clone()) + .with_limit(3) + .execute_read() + .await + .expect("batch read failed"); + assert_eq!( + streams.len(), + queries.len(), + "one stream per query, in input order" + ); + + for (i, stream) in streams.into_iter().enumerate() { + let (ids, scores) = drain_ids_and_scores(stream).await; + let (want_ids, want_scores) = &expected[i]; + assert_eq!(&ids, want_ids, "query {i} rows must match its single read"); + assert_eq!(scores.len(), want_scores.len()); + for (got, want) in scores.iter().zip(want_scores) { + assert!( + (got - want).abs() < 1e-6, + "query {i} score diverges: {got} vs {want}" + ); + } + } +} + +#[cfg(not(windows))] +#[tokio::test] +async fn empty_snapshot_yields_n_empty_streams() { + let (_tmp, table) = build_empty_table().await; + let queries = vec![ + vec![1.0, 0.0, 0.0, 0.0], + vec![0.0, 1.0, 0.0, 0.0], + vec![0.0, 0.0, 1.0, 0.0], + ]; + + let mut builder = table.new_batch_vector_search_builder(); + let streams = builder + .with_vector_column(VECTOR_COLUMN) + .with_query_vectors(queries.clone()) + .with_limit(3) + .execute_read() + .await + .expect("empty-snapshot batch read failed"); + assert_eq!( + streams.len(), + queries.len(), + "arity preserved: one empty stream per query" + ); + for stream in streams { + let (ids, _scores) = drain_ids_and_scores(stream).await; + assert!(ids.is_empty(), "no-hit query must yield an empty stream"); + } +} + +#[cfg(not(windows))] +#[tokio::test] +async fn pk_batch_execute_scored_fails_loud() { + let vectors = fixture(); + let (_tmp, table) = build_table(&vectors).await; + let err = table + .new_batch_vector_search_builder() + .with_vector_column(VECTOR_COLUMN) + .with_query_vectors(vec![vec![10.0, 0.0, 0.0, 0.0]]) + .with_limit(3) + .execute() + .await + .expect_err("PK batch execute() must fail loud"); + assert!( + format!("{err:?}").contains("execute_read"), + "PK batch execute() should point at execute_read, got: {err:?}" + ); +} + +/// A shared residual filter reshapes every query's rows (excluding low ids) with +/// no cross-query bleed: each query still ranks only the residual-allowed rows by +/// its own distance. +#[cfg(not(windows))] +#[tokio::test] +async fn shared_residual_filter_applies_per_query_without_bleed() { + // Fixture where filtering id >= 3 changes each query's result set. + let vectors = vec![ + [10.0, 0.0, 0.0, 0.0], // pos 0 + [9.0, 0.0, 0.0, 0.0], // pos 1 + [8.0, 0.0, 0.0, 0.0], // pos 2 + [5.0, 0.0, 0.0, 0.0], // pos 3 + [7.0, 0.0, 0.0, 0.0], // pos 4 + [6.0, 0.0, 0.0, 0.0], // pos 5 + ]; + let (_tmp, table) = build_table(&vectors).await; + let threshold = 3; + let build_filter = || -> Predicate { + PredicateBuilder::new(table.schema().fields()) + .greater_or_equal("id", Datum::Int(threshold)) + .expect("build residual predicate on id") + }; + + let queries = vec![vec![10.0, 0.0, 0.0, 0.0], vec![5.0, 0.0, 0.0, 0.0]]; + + // Expected per-query truth: rank residual-allowed rows (id >= 3) by distance. + let expected: Vec> = queries + .iter() + .map(|q| { + let mut ranked: Vec<(u64, f32)> = analytic_topk(q, &vectors, vectors.len()) + .into_iter() + .filter(|(id, _)| *id >= threshold as u64) + .collect(); + ranked.truncate(3); + ranked.iter().map(|(id, _)| *id as i32).collect() + }) + .collect(); + // Guard: the two queries produce different residual orders (no accidental + // bleed-agnostic fixture). + assert_ne!( + expected[0], expected[1], + "queries must differ post-residual" + ); + + let mut builder = table.new_batch_vector_search_builder(); + let streams = builder + .with_vector_column(VECTOR_COLUMN) + .with_query_vectors(queries.clone()) + .with_limit(3) + .with_filter(build_filter()) + .execute_read() + .await + .expect("residual batch read failed"); + assert_eq!(streams.len(), queries.len()); + + for (i, stream) in streams.into_iter().enumerate() { + let (ids, _scores) = drain_ids_and_scores(stream).await; + for &id in &ids { + assert!( + id >= threshold, + "query {i} returned id {id} failing residual id >= {threshold}" + ); + } + assert_eq!(ids, expected[i], "query {i} residual order must be its own"); + } +} + +/// A table with no primary-key vector index is not a valid target for the batch +/// materialized read: it produces scored global row ids, not physical rows. The +/// batch `execute_read` must reject it up front rather than silently route it +/// through the primary-key materialization path. +// Gated off Windows for the same `file://` tempdir reason as `pk_vector_baseline_test`. +#[cfg(not(windows))] +#[tokio::test] +async fn batch_execute_read_on_non_pk_vector_table_fails_loud() { + let tmp = tempfile::tempdir().expect("create temp dir"); + let location = format!("file://{}", tmp.path().display()); + let file_io = FileIOBuilder::new("file").build().unwrap(); + for dir in ["schema", "snapshot", "manifest", "index"] { + file_io.mkdirs(&format!("{location}/{dir}")).await.unwrap(); + } + + // A plain primary-key table over a vector column but WITHOUT the + // `pk-vector.index.columns` option, so no PK-vector index is configured. + let mut builder = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column( + VECTOR_COLUMN, + DataType::Vector( + VectorType::try_new(true, DIM as u32, DataType::Float(FloatType::new())).unwrap(), + ), + ) + .primary_key(["id"]); + builder = builder.option("bucket".to_string(), "1".to_string()); + let schema = TableSchema::new(0, &builder.build().unwrap()); + file_io + .new_output(&format!("{location}/schema/schema-{}", schema.id())) + .unwrap() + .write(Bytes::from(serde_json::to_vec(&schema).unwrap())) + .await + .unwrap(); + + let table = open_table(&file_io, &location).await; + + let mut batch = table.new_batch_vector_search_builder(); + let result = batch + .with_vector_column(VECTOR_COLUMN) + .with_query_vectors(vec![vec![1.0, 0.0, 0.0, 0.0]]) + .with_limit(3) + .execute_read() + .await; + let err = match result { + Ok(_) => panic!("batch execute_read on a non-PK-vector table must fail loud"), + Err(e) => e, + }; + assert!( + err.to_string().contains("primary-key vector path"), + "expected a message directing to the primary-key vector path, got: {err}" + ); +} + +/// A malformed query (wrong dimension) must fail loud even when the plan is +/// empty. Query validation runs before planning, so an empty snapshot cannot +/// mask a bad query by returning empty streams. +// Gated off Windows for the same `file://` tempdir reason as `pk_vector_baseline_test`. +#[cfg(not(windows))] +#[tokio::test] +async fn empty_snapshot_still_rejects_malformed_query() { + let (_tmp, table) = build_empty_table().await; + // DIM is 4; a 3-element query is wrong-dimension. + let queries = vec![vec![1.0, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]]; + + let mut builder = table.new_batch_vector_search_builder(); + let result = builder + .with_vector_column(VECTOR_COLUMN) + .with_query_vectors(queries) + .with_limit(3) + .execute_read() + .await; + let err = match result { + Ok(_) => panic!("a wrong-dimension query must fail loud even on an empty snapshot"), + Err(e) => e, + }; + assert!( + err.to_string().contains("dimension does not match"), + "expected a dimension-mismatch error, got: {err}" + ); +} + +/// ARRAY is a valid PK-vector column. Batch query validation must still +/// reject NaN/Inf before planning, otherwise an empty snapshot would mask the bad +/// input by returning normal empty streams. +// Gated off Windows for the same `file://` tempdir reason as `pk_vector_baseline_test`. +#[cfg(not(windows))] +#[tokio::test] +async fn empty_array_snapshot_still_rejects_non_finite_query() { + let (_tmp, table) = build_empty_array_table().await; + let queries = vec![vec![1.0, f32::NAN, 0.0, 0.0], vec![0.0, 1.0, 0.0, 0.0]]; + + let mut builder = table.new_batch_vector_search_builder(); + let result = builder + .with_vector_column(VECTOR_COLUMN) + .with_query_vectors(queries) + .with_limit(3) + .execute_read() + .await; + let err = match result { + Ok(_) => panic!("a NaN query must fail loud for ARRAY even on an empty snapshot"), + Err(e) => e, + }; + assert!( + err.to_string().contains("must be finite"), + "expected a finite-query error, got: {err}" + ); +} + +/// A non-positive limit must fail loud even when the plan is empty. Limit +/// validation runs before planning, so an empty snapshot cannot mask a zero +/// limit by returning empty streams. +// Gated off Windows for the same `file://` tempdir reason as `pk_vector_baseline_test`. +#[cfg(not(windows))] +#[tokio::test] +async fn empty_snapshot_still_rejects_zero_limit() { + let (_tmp, table) = build_empty_table().await; + + let mut builder = table.new_batch_vector_search_builder(); + let result = builder + .with_vector_column(VECTOR_COLUMN) + .with_query_vectors(vec![vec![1.0, 0.0, 0.0, 0.0]]) + .with_limit(0) + .execute_read() + .await; + let err = match result { + Ok(_) => panic!("a zero limit must fail loud even on an empty snapshot"), + Err(e) => e, + }; + assert!( + err.to_string().contains("limit must be positive"), + "expected a positive-limit error, got: {err}" + ); +} + +/// A filter set on a batch `execute()` (the scored / data-evolution path) must +/// fail loud rather than silently drop the predicate: that path never reads +/// physical rows, so it cannot honor a residual filter. Mirrors the single-query +/// `execute_scored` guard. +// Gated off Windows for the same `file://` tempdir reason as `pk_vector_baseline_test`. +#[cfg(not(windows))] +#[tokio::test] +async fn batch_execute_with_filter_on_non_pk_vector_table_fails_loud() { + let tmp = tempfile::tempdir().expect("create temp dir"); + let location = format!("file://{}", tmp.path().display()); + let file_io = FileIOBuilder::new("file").build().unwrap(); + for dir in ["schema", "snapshot", "manifest", "index"] { + file_io.mkdirs(&format!("{location}/{dir}")).await.unwrap(); + } + + // A plain table over a vector column WITHOUT a PK-vector index, so a scored + // batch `execute()` takes the data-evolution fall-through path. + let mut builder = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column( + VECTOR_COLUMN, + DataType::Vector( + VectorType::try_new(true, DIM as u32, DataType::Float(FloatType::new())).unwrap(), + ), + ) + .primary_key(["id"]); + builder = builder.option("bucket".to_string(), "1".to_string()); + let schema = TableSchema::new(0, &builder.build().unwrap()); + file_io + .new_output(&format!("{location}/schema/schema-{}", schema.id())) + .unwrap() + .write(Bytes::from(serde_json::to_vec(&schema).unwrap())) + .await + .unwrap(); + + let table = open_table(&file_io, &location).await; + let filter = PredicateBuilder::new(table.schema().fields()) + .greater_or_equal("id", Datum::Int(1)) + .expect("build filter on id"); + + let mut batch = table.new_batch_vector_search_builder(); + let err = batch + .with_vector_column(VECTOR_COLUMN) + .with_query_vectors(vec![vec![1.0, 0.0, 0.0, 0.0]]) + .with_limit(3) + .with_filter(filter) + .execute() + .await + .expect_err("a filter on the data-evolution batch path must fail loud"); + assert!( + err.to_string() + .contains("only supported on the primary-key vector path"), + "expected a filter-unsupported error, got: {err}" + ); +}