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 f2c3a6dd..33067b6a 100644 --- a/crates/paimon/src/table/pk_vector_data_file_reader.rs +++ b/crates/paimon/src/table/pk_vector_data_file_reader.rs @@ -145,7 +145,7 @@ impl DataFilePkVectorReaderFactory { /// `None`). The column must be a `FixedSizeList`/`List` of `Float32`; every /// non-null row's child slice must have exactly `dimension` elements. Mirrors /// the layout handling in `vector_search_builder`. -fn append_batch_vectors( +pub(crate) fn append_batch_vectors( batch: &arrow_array::RecordBatch, field_name: &str, dimension: usize, @@ -206,6 +206,12 @@ fn append_batch_vectors( } let mut vector = Vec::with_capacity(dimension); for i in start..end { + if values.is_null(i) { + return Err(data_invalid(format!( + "vector row {row} has a null element at index {}", + i - start + ))); + } vector.push(values.value(i)); } out.push(Some(vector)); @@ -349,6 +355,35 @@ mod integration_tests { } } + /// A present (non-null) vector row whose child slice contains a NULL + /// element must fail loud rather than silently defaulting the element to + /// `0.0` and corrupting the distance. + #[test] + fn append_batch_vectors_fails_loud_on_null_element() { + let field = vector_field(); + let read_fields = vec![field.clone()]; + let arrow_schema = build_target_arrow_schema(&read_fields).unwrap(); + + // Row is present, but element index 1 in its child slice is NULL: [1.0, null]. + let mut builder = FixedSizeListBuilder::new(Float32Builder::new(), 2).with_field(Arc::new( + ArrowField::new("element", ArrowDataType::Float32, true), + )); + builder.values().append_value(1.0); + builder.values().append_null(); + builder.append(true); + let vec_array = builder.finish(); + let batch = RecordBatch::try_new(arrow_schema, vec![Arc::new(vec_array)]).unwrap(); + + let mut out: Vec>> = Vec::new(); + let err = append_batch_vectors(&batch, field.name(), 2, &mut out) + .expect_err("null child element must fail loud"); + let msg = err.to_string(); + assert!( + msg.contains("null") && msg.contains("element"), + "got: {msg}" + ); + } + /// Write a FixedSizeList vector column /// (`[1,2]`, NULL, `[3,4]`) as a parquet data file, build the factory over /// its split, preload via `create`, and assert the whole-file sequential 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 2dcb6e7f..b6383c99 100644 --- a/crates/paimon/src/table/pk_vector_indexed_split_read.rs +++ b/crates/paimon/src/table/pk_vector_indexed_split_read.rs @@ -61,7 +61,7 @@ pub(crate) struct PkVectorIndexedSplit { /// positions, validating bounds and ordering. Ranges must be non-empty, each /// within `[0, row_count)`, strictly ascending and non-overlapping (touching /// ranges allowed). Expansion is inclusive `from..=to`. -fn expand_ranges(ranges: &[RowRange], row_count: i64) -> crate::Result> { +pub(crate) fn expand_ranges(ranges: &[RowRange], row_count: i64) -> crate::Result> { if ranges.is_empty() { return Err(data_invalid("indexed split must select at least one row")); } diff --git a/crates/paimon/src/table/pk_vector_orchestrator.rs b/crates/paimon/src/table/pk_vector_orchestrator.rs index 0479781f..358211a5 100644 --- a/crates/paimon/src/table/pk_vector_orchestrator.rs +++ b/crates/paimon/src/table/pk_vector_orchestrator.rs @@ -112,6 +112,7 @@ pub(crate) struct PkVectorSearchSplit { /// the cross-bucket merge dimensions a lone `PkVectorSearchResult` lacks; /// `split_index` is the re-association handle back to /// `splits[split_index].data_split`. +#[derive(Clone)] pub(crate) struct PkVectorCandidate { pub split_index: usize, pub partition: BinaryRow, @@ -121,6 +122,14 @@ pub(crate) struct PkVectorCandidate { pub distance: f32, } +/// Best-first indexed (approximate) and exact-fallback candidate lists, each +/// already globally bounded. The indexed list may be over-fetched for a later +/// exact rerank; the exact list is bounded to the caller's final limit. +pub(crate) struct OrchestratorSearchResult { + pub(crate) indexed: Vec, + pub(crate) exact: Vec, +} + /// 5-level BEST_FIRST (smallest = best) key. Level 1 orders distance with /// `java_float_compare` so a NaN distance (e.g. from a non-finite stored vector /// under inner product) sorts last rather than winning Top-1. Level 2 uses the @@ -145,6 +154,20 @@ fn global_top_k(mut candidates: Vec, limit: usize) -> Vec, + exact: Vec, + limit: usize, +) -> Vec { + let mut all = indexed; + all.extend(exact); + global_top_k(all, limit) +} + /// Group Top-K survivors by `(partition, bucket, data_file_name)`, re-associate /// each group's file to its real `DataFileMeta` + aligned deletion file in the /// source bucket split, and build one `PkVectorIndexedSplit` per file. Groups are @@ -197,7 +220,15 @@ pub(crate) fn build_indexed_splits( } // Re-associate the file to its DataFileMeta + aligned deletion file. - let source = &splits[split_index].data_split; + let source = &splits + .get(split_index) + .ok_or_else(|| { + data_invalid(format!( + "vector search hit references split index {split_index} out of range (splits: {})", + splits.len() + )) + })? + .data_split; let file_idx = source .data_files() .iter() @@ -301,10 +332,13 @@ impl PkVectorOrchestrator { } /// Run the eager per-bucket search + cross-bucket global Top-K and return the - /// best-first survivors (through the full 5-level tie-break, raw distance - /// preserved). The exact-reader factory is split-scoped: it receives the - /// current split index and split so a caller can build a reader keyed to the - /// specific split/file. `skip_exact_fallback` forwards to `bucket_search`. + /// indexed (approximate) and exact-fallback survivors as two separate + /// best-first lists (each through the full 5-level tie-break, raw distance + /// preserved). The indexed list is bounded to `indexed_limit` (over-fetched + /// for a later exact rerank); the exact list is bounded to `limit`. The + /// exact-reader factory is split-scoped: it receives the current split index + /// and split so a caller can build a reader keyed to the specific split/file. + /// `skip_exact_fallback` forwards to `bucket_search`. /// /// `residual_by_split`, when present, carries one per-file allow-list of /// physical row positions per split (indexed parallel to `splits`): only @@ -320,6 +354,7 @@ impl PkVectorOrchestrator { query: &[f32], metric: VectorSearchMetric, limit: usize, + indexed_limit: usize, ann_searcher: Option<&dyn PkVectorAnnSearcher>, exact_reader_factory: &mut (dyn for<'s, 'f> FnMut( usize, @@ -330,11 +365,14 @@ impl PkVectorOrchestrator { search_options: &HashMap, skip_exact_fallback: bool, residual_by_split: Option<&[HashMap]>, - ) -> crate::Result> { + ) -> crate::Result { // Eager input-shape validation (Java checkArgument parity). 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")); } @@ -346,8 +384,9 @@ impl PkVectorOrchestrator { } } - // Eager per-bucket search -> tagged candidates. - let mut candidates: Vec = Vec::new(); + // Eager per-bucket search -> tagged candidates, kept split by path. + let mut indexed_candidates: Vec = Vec::new(); + let mut exact_candidates: Vec = Vec::new(); for (split_index, split) in splits.iter().enumerate() { let dvs = build_bucket_dv_map(&self.reader, split).await?; // Wrap the split-scoped factory into bucket_search's per-file signature. @@ -357,7 +396,7 @@ impl PkVectorOrchestrator { exact_reader_factory(split_index, split, file) }); let residual_ranges = residual_by_split.map(|per_split| &per_split[split_index]); - let results = bucket_search( + let result = bucket_search( ann_searcher, &split.ann_segments, &split.active_files, @@ -365,30 +404,33 @@ impl PkVectorOrchestrator { &mut bucket_factory, query, metric, + indexed_limit, limit, search_options, skip_exact_fallback, residual_ranges, ) .await?; - for PkVectorSearchResult { + let tag = |PkVectorSearchResult { + data_file_name, + row_position, + distance, + }: PkVectorSearchResult| PkVectorCandidate { + split_index, + partition: split.data_split.partition().clone(), + bucket: split.data_split.bucket(), data_file_name, row_position, distance, - } in results - { - candidates.push(PkVectorCandidate { - split_index, - partition: split.data_split.partition().clone(), - bucket: split.data_split.bucket(), - data_file_name, - row_position, - distance, - }); - } + }; + indexed_candidates.extend(result.indexed.into_iter().map(&tag)); + exact_candidates.extend(result.exact.into_iter().map(&tag)); } - Ok(global_top_k(candidates, limit)) + Ok(OrchestratorSearchResult { + indexed: global_top_k(indexed_candidates, indexed_limit), + exact: global_top_k(exact_candidates, limit), + }) } } @@ -547,6 +589,18 @@ mod tests { assert!(survivors.is_empty()); } + #[test] + fn merge_candidates_takes_global_top_k_over_union() { + // indexed and exact each already bounded; the union's global Top-K must be + // independent of how each side was bounded. + let indexed = vec![cand(0, 0, "f", 0, 0.1), cand(0, 0, "f", 1, 0.4)]; + let exact = vec![cand(0, 0, "g", 0, 0.2)]; + let out = merge_candidates(indexed, exact, 2); + assert_eq!(out.len(), 2); + assert_eq!(out[0].distance, 0.1); + assert_eq!(out[1].distance, 0.2); // 0.2 (exact) beats 0.4 (indexed) + } + #[test] fn builds_two_splits_with_ascending_position_ordered_scores() { // One bucket, two files. file-a hits at global order [pos=10, pos=2]; @@ -628,6 +682,18 @@ mod tests { assert!(format!("{err:?}").contains("duplicate"), "got: {err:?}"); } + #[test] + fn build_indexed_splits_fails_loud_on_out_of_range_split_index() { + // A malformed candidate whose split_index is beyond the splits slice must + // fail loud, not panic on a bare slice index. + let splits = vec![search_split(0, vec![data_file("f", 10)])]; + let survivors = vec![cand(5, 0, "f", 0, 1.0)]; + let err = build_indexed_splits(survivors, &splits, VectorSearchMetric::L2) + .map(|_| ()) + .expect_err("out-of-range split index must error"); + assert!(format!("{err:?}").contains("out of range"), "got: {err:?}"); + } + #[test] fn rejects_same_group_key_from_different_splits() { // Two buckets share (partition, bucket, file_name) but sit at different @@ -949,12 +1015,13 @@ mod e2e_tests { // expects; the split index/split are unused here. let mut wrapped = as_split_factory(|_: usize, _: &PkVectorSearchSplit, f: &BucketActiveFile| factory(f)); - let survivors = orch + let result = orch .search_candidates( splits, query, metric, limit, + limit, ann, &mut wrapped, opts, @@ -962,6 +1029,9 @@ mod e2e_tests { None, ) .await?; + // Merge the two bounded lists into the best-first survivors the + // materialization path expects. + let survivors = merge_candidates(result.indexed, result.exact, limit); let indexed_splits = build_indexed_splits(survivors, splits, metric)?; let mut out = Vec::new(); for indexed in indexed_splits { @@ -991,6 +1061,7 @@ mod e2e_tests { &[0.0, 0.0], VectorSearchMetric::L2, 0, + 0, None, &mut factory, &opts, @@ -1020,6 +1091,7 @@ mod e2e_tests { &[], VectorSearchMetric::L2, 5, + 5, None, &mut factory, &opts, @@ -1360,12 +1432,13 @@ mod e2e_tests { }, ); let opts = HashMap::new(); - let cands = PkVectorOrchestrator::new(make_reader(file_io, table_path)) + let result = PkVectorOrchestrator::new(make_reader(file_io, table_path)) .search_candidates( &[split], &[0.0, 0.0], VectorSearchMetric::L2, 2, + 2, None, &mut factory, &opts, @@ -1374,6 +1447,8 @@ mod e2e_tests { ) .await .unwrap(); + // Merge the two bounded lists into the best-first survivors. + let cands = merge_candidates(result.indexed, result.exact, 2); // Best-first: pos1 (d=1), pos2 (d=4). assert_eq!( cands @@ -1432,12 +1507,13 @@ mod e2e_tests { allowed.insert(2); let residual_by_split = vec![HashMap::from([("r.mosaic".to_string(), allowed)])]; let opts = HashMap::new(); - let cands = PkVectorOrchestrator::new(make_reader(file_io, table_path)) + let result = PkVectorOrchestrator::new(make_reader(file_io, table_path)) .search_candidates( &[split], &[0.0, 0.0], VectorSearchMetric::L2, 3, + 3, None, &mut factory, &opts, @@ -1446,6 +1522,8 @@ mod e2e_tests { ) .await .unwrap(); + // Merge the two bounded lists into the best-first survivors. + let cands = merge_candidates(result.indexed, result.exact, 3); // Best-first among allowed positions: pos2 (d=4) then pos0 (d=9). assert_eq!( cands @@ -1493,6 +1571,7 @@ mod e2e_tests { &[0.0, 0.0], VectorSearchMetric::L2, 3, + 3, None, &mut factory, &opts, @@ -1533,12 +1612,13 @@ mod e2e_tests { }, ); let opts = HashMap::new(); - let cands = PkVectorOrchestrator::new(make_reader(file_io, table_path)) + let result = PkVectorOrchestrator::new(make_reader(file_io, table_path)) .search_candidates( &[split], &[0.0, 0.0], VectorSearchMetric::L2, 2, + 2, None, &mut factory, &opts, @@ -1547,7 +1627,8 @@ mod e2e_tests { ) .await .unwrap(); - assert!(cands.is_empty()); + assert!(result.indexed.is_empty()); + assert!(result.exact.is_empty()); } #[test] diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index 992c4cd4..9f50ee46 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -29,13 +29,17 @@ use crate::table::global_index_scanner::{ deleted_row_ranges_for_data_evolution_dvs, search_limit_with_deleted_rows, unindexed_ranges_for_global_index_entries, RowRangeIndex, }; -use crate::table::pk_vector_data_file_reader::DataFilePkVectorReaderFactory; -use crate::table::pk_vector_indexed_split_read::PkVectorIndexedSplitRead; +use crate::table::pk_vector_data_file_reader::{ + append_batch_vectors, DataFilePkVectorReaderFactory, +}; +use crate::table::pk_vector_indexed_split_read::{expand_ranges, PkVectorIndexedSplitRead}; use crate::table::pk_vector_orchestrator::{ - as_split_exact_reader_factory, build_indexed_splits, PkVectorCandidate, PkVectorOrchestrator, - PkVectorSearchSplit, + as_split_exact_reader_factory, build_indexed_splits, merge_candidates, + OrchestratorSearchResult, PkVectorCandidate, PkVectorOrchestrator, PkVectorSearchSplit, +}; +use crate::table::pk_vector_position_read::{ + PkVectorPositionRead, PKEY_VECTOR_POSITION_COLUMN, SEARCH_SCORE_COLUMN, }; -use crate::table::pk_vector_position_read::{PKEY_VECTOR_POSITION_COLUMN, SEARCH_SCORE_COLUMN}; use crate::table::pk_vector_scan::{PkVectorScan, PkVectorScanPlan}; use crate::table::read_builder::resolve_projected_fields; use crate::table::source::DataSplit; @@ -437,6 +441,23 @@ impl<'a> VectorSearchBuilder<'a> { 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( + &self.options, + self.table.schema().options(), + pk_col, + &index_type, + )?; + let indexed_limit = indexed_search_limit(limit, refine_factor)?; + let plan = PkVectorScan::new( self.table, field_id, @@ -583,12 +604,13 @@ impl<'a> VectorSearchBuilder<'a> { }, ); - let candidates = PkVectorOrchestrator::new(reader) + let search: OrchestratorSearchResult = PkVectorOrchestrator::new(reader) .search_candidates( &plan.splits, query_vector, metric, limit, + indexed_limit, Some(&ann_searcher), &mut factory, &search_options, @@ -597,6 +619,38 @@ impl<'a> VectorSearchBuilder<'a> { ) .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)) } @@ -1279,6 +1333,168 @@ fn verify_pk_vector_segment_metrics( Ok(()) } +/// Rerank approximate (indexed) candidates by rereading ONLY their candidate +/// positions and recomputing the exact distance, then keep the best `limit`. +/// +/// Unlike a whole-column preload, this reuses [`PkVectorPositionRead`] to read +/// just the selected physical rows of each hit file (positions -> row ranges -> +/// local ranges), so a rerank over a large ANN-covered file touches only the +/// candidate rows. Mirrors Java's IndexedSplit rerank. +/// +/// Each returned row is matched back to its candidate by the +/// `_PKEY_VECTOR_POSITION` column VALUE (never batch order). The recomputed +/// distance is written into the ORIGINAL candidate so `split_index` / +/// partition / bucket survive (`build_indexed_splits` does not carry +/// `split_index`). A DV loaded exactly as [`PkVectorIndexedSplitRead::read`] +/// does drops deleted positions, so a candidate at a deleted position returns no +/// row and trips the leftover guard — a deleted candidate reaching rerank is a +/// real inconsistency (the search path already DV-filters), so fail loud. +#[allow(clippy::too_many_arguments)] +async fn rerank_indexed_positional( + rerank_reader: &DataFileReader, + indexed: Vec, + plan_splits: &[PkVectorSearchSplit], + query_vector: &[f32], + metric: VectorSearchMetric, + limit: usize, + vector_field: &DataField, +) -> crate::Result> { + // Original per-position candidates keyed by (split_index, file, position); + // the recomputed distance is written back into these so split_index and + // partition/bucket survive (build_indexed_splits does not carry split_index). + let mut by_key: HashMap<(usize, String, i64), PkVectorCandidate> = HashMap::new(); + for c in &indexed { + if by_key + .insert( + (c.split_index, c.data_file_name.clone(), c.row_position), + c.clone(), + ) + .is_some() + { + return Err(crate::Error::DataInvalid { + message: "duplicate primary-key vector candidate for reranking".to_string(), + source: None, + }); + } + } + + // Rebuild the split_index lookup by (partition bytes, bucket, file): the + // indexed split exposes partition/bucket/file but not split_index. + let mut split_index_of: HashMap<(Vec, i32, String), usize> = HashMap::new(); + for (i, s) in plan_splits.iter().enumerate() { + let p = s.data_split.partition().to_serialized_bytes(); + let b = s.data_split.bucket(); + for f in s.data_split.data_files() { + split_index_of.insert((p.clone(), b, f.file_name.clone()), i); + } + } + + // Every candidate must reference a (partition, bucket, file) that the plan + // actually carries. Checking up front — before build_indexed_splits, which + // indexes plan_splits by split_index — turns an absent file into a fail-loud + // error rather than an out-of-range panic, and keeps the per-split lookup + // below a self-consistent backstop. + for c in &indexed { + let key = ( + c.partition.to_serialized_bytes(), + c.bucket, + c.data_file_name.clone(), + ); + if !split_index_of.contains_key(&key) { + return Err(crate::Error::DataInvalid { + message: format!("rerank split for {} not found in plan", c.data_file_name), + source: None, + }); + } + } + + // Group the candidates into per-file indexed splits (position ranges + file + // meta), reusing the exact grouping/validation the materialization path uses. + let indexed_splits = build_indexed_splits(indexed, plan_splits, metric)?; + + let dimension = query_vector.len(); + let mut reranked: Vec = Vec::new(); + for split in indexed_splits { + let data_split = split.split.clone(); + let file_meta = data_split.data_files()[0].clone(); + let file_name = file_meta.file_name.clone(); + let partition_bytes = data_split.partition().to_serialized_bytes(); + let bucket = data_split.bucket(); + let split_index = *split_index_of + .get(&(partition_bytes, bucket, file_name.clone())) + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("rerank split for {file_name} not found in plan"), + source: None, + })?; + + // DV loaded exactly as PkVectorIndexedSplitRead::read does; skipping it + // would score deleted rows. + let dv_factory = rerank_reader.build_split_dv_factory(&data_split).await?; + let dv = DataFileReader::deletion_vector_for_file(dv_factory.as_ref(), &file_name); + let data_fields = rerank_reader.derive_data_fields(&file_meta).await?; + + // Positions from the split's row_ranges (ascending); read only those. + let positions = expand_ranges(&split.row_ranges, file_meta.row_count)?; + let mut stream = PkVectorPositionRead::new(rerank_reader).read( + &data_split, + file_meta, + data_fields, + dv, + positions, + None, // no scores; rerank recomputes distance + )?; + + while let Some(batch) = stream.try_next().await? { + let pos_idx = batch + .schema() + .index_of(PKEY_VECTOR_POSITION_COLUMN) + .map_err(|_| crate::Error::DataInvalid { + message: format!("rerank batch missing {PKEY_VECTOR_POSITION_COLUMN} column"), + source: None, + })?; + let pos_col = batch + .column(pos_idx) + .as_any() + .downcast_ref::() + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("{PKEY_VECTOR_POSITION_COLUMN} column is not Int64"), + source: None, + })?; + let mut vectors: Vec>> = Vec::new(); + append_batch_vectors(&batch, vector_field.name(), dimension, &mut vectors)?; + for (row, vector) in vectors.iter().enumerate() { + let position = pos_col.value(row); + let mut candidate = by_key + .remove(&(split_index, file_name.clone(), position)) + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("rerank read unexpected position {file_name}@{position}"), + source: None, + })?; + let vector = vector.as_ref().ok_or_else(|| crate::Error::DataInvalid { + message: format!( + "primary-key vector candidate {file_name}@{position} contains a null vector" + ), + source: None, + })?; + candidate.distance = metric.compute_distance(query_vector, vector); + reranked.push(candidate); + } + } + } + + if !by_key.is_empty() { + return Err(crate::Error::DataInvalid { + message: format!( + "failed to read {} primary-key vector candidate(s) for reranking", + by_key.len() + ), + source: None, + }); + } + + Ok(merge_candidates(reranked, Vec::new(), limit)) +} + /// One materialized row tagged with its best-first `rank` and its `(batch_index, /// row_index)` location in the retained materialization batches. struct RankedRow { @@ -2965,6 +3181,633 @@ mod tests { } } + fn pk_candidate( + split_index: usize, + bucket: i32, + file: &str, + pos: i64, + distance: f32, + ) -> PkVectorCandidate { + PkVectorCandidate { + split_index, + partition: BinaryRow::new(0), + bucket, + data_file_name: file.to_string(), + row_position: pos, + distance, + } + } + + // Candidate with a fixed empty (arity-0) partition and bucket 0, keyed only by + // (split_index, file, position) — the dimensions the rerank core groups on. + fn cand_at(split_index: usize, file: &str, pos: i64, dist: f32) -> PkVectorCandidate { + pk_candidate(split_index, 0, file, pos, dist) + } + + /// The single data-file name every rerank fixture writes. + const RERANK_FILE: &str = "part-0.parquet"; + + /// Serialize a Paimon deletion-vector blob covering `deleted_rows` and write it + /// at `path`, returning the matching `DeletionFile`. Byte layout mirrors the + /// position-read tests: `[length][magic][roaring bitmap][0]`. + async fn write_deletion_blob( + file_io: &FileIO, + path: &str, + deleted_rows: &[u32], + ) -> crate::table::source::DeletionFile { + use roaring::RoaringBitmap; + + const MAGIC_NUMBER: i32 = 1581511376; + let mut bitmap = RoaringBitmap::new(); + for row in deleted_rows { + bitmap.insert(*row); + } + let mut bitmap_bytes = Vec::new(); + bitmap.serialize_into(&mut bitmap_bytes).unwrap(); + let bitmap_length = 4 + bitmap_bytes.len() as i32; + let mut blob = Vec::new(); + blob.extend_from_slice(&bitmap_length.to_be_bytes()); + blob.extend_from_slice(&MAGIC_NUMBER.to_be_bytes()); + blob.extend_from_slice(&bitmap_bytes); + blob.extend_from_slice(&0i32.to_be_bytes()); + file_io + .new_output(path) + .unwrap() + .write(bytes::Bytes::from(blob)) + .await + .unwrap(); + crate::table::source::DeletionFile::new( + path.to_string(), + 0, + bitmap_length as i64, + Some(deleted_rows.len() as i64), + ) + } + + /// Write a single-file vector data file (`FixedSizeList` of width + /// `dim`) holding `rows` (a `None` entry is a NULL vector row) as Parquet, and + /// return a vector-only `DataFileReader`, the enclosing `PkVectorSearchSplit`, + /// and the vector `DataField`. When `deleted_rows` is non-empty a deletion + /// vector covering those physical positions is attached to the split, so the + /// position read drops them exactly as `PkVectorIndexedSplitRead::read` does. + /// + /// This is the position-only analogue of the old `ArrayReader`: rerank now + /// re-reads real stored rows through `PkVectorPositionRead`, so the fixtures + /// exercise that path rather than an in-memory preloaded column. + async fn vector_rerank_fixture( + table_path: &str, + dim: u32, + rows: &[Option>], + deleted_rows: &[u32], + ) -> (DataFileReader, PkVectorSearchSplit, DataField) { + use crate::arrow::build_target_arrow_schema; + use crate::arrow::format::{FormatFileWriter, ParquetFormatWriter}; + use crate::spec::VectorType; + use crate::table::schema_manager::SchemaManager; + + let vector_type = + VectorType::try_new(true, dim, DataType::Float(FloatType::new())).unwrap(); + let vector_field = + DataField::new(0, "embedding".to_string(), DataType::Vector(vector_type)); + let read_fields = vec![vector_field.clone()]; + let arrow_schema = build_target_arrow_schema(&read_fields).unwrap(); + + let mut builder = FixedSizeListBuilder::new(Float32Builder::new(), dim as i32).with_field( + Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)), + ); + for row in rows { + match row { + Some(values) => { + for v in values { + builder.values().append_value(*v); + } + builder.append(true); + } + None => { + for _ in 0..dim { + builder.values().append_value(0.0); + } + builder.append(false); + } + } + } + let vec_array = builder.finish(); + let batch = + arrow_array::RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(vec_array)]) + .unwrap(); + + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let bucket_path = format!("{table_path}/bucket-0"); + let output = file_io + .new_output(&format!("{bucket_path}/{RERANK_FILE}")) + .unwrap(); + let mut writer: Box = Box::new( + ParquetFormatWriter::new( + &output, + arrow_schema.clone(), + "zstd", + 1, + None, + &HashMap::new(), + ) + .await + .unwrap(), + ); + writer.write(&batch).await.unwrap(); + let file_size = writer.close().await.unwrap().file_size; + + let schema_id = 1; + let file_meta = pk_data_file(RERANK_FILE, rows.len() as i64, Some(0)); + let file_meta = DataFileMeta { + file_size: file_size as i64, + schema_id, + ..file_meta + }; + + let mut split_builder = 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![file_meta]); + if !deleted_rows.is_empty() { + let df = + write_deletion_blob(&file_io, &format!("{table_path}/index/dv-0"), deleted_rows) + .await; + split_builder = split_builder.with_data_deletion_files(vec![Some(df)]); + } + let data_split = split_builder.build().unwrap(); + let split = PkVectorSearchSplit { + data_split, + ann_segments: Vec::new(), + active_files: Vec::new(), + }; + + let schema_manager = SchemaManager::new(file_io.clone(), table_path.to_string()); + let reader = DataFileReader::new( + file_io, + schema_manager, + schema_id, + read_fields.clone(), + read_fields, + Vec::new(), + ); + (reader, split, vector_field) + } + + #[tokio::test] + async fn rerank_aligns_recomputed_distance_by_position_column() { + use crate::arrow::build_target_arrow_schema; + use crate::arrow::format::{FormatFileWriter, ParquetFormatWriter}; + use crate::spec::VectorType; + use crate::table::schema_manager::SchemaManager; + + // A vector data file with 4 physical rows: positions 0,1,3 hold vectors + // and position 2 (a NON-candidate) holds a NULL vector. Candidates sit at + // non-contiguous positions {1, 3}. The ANN-reported distances are + // deliberately reversed relative to the true stored vectors; after rerank + // each candidate must carry compute_distance(query, vec_at_its_position), + // proving alignment is by the _PKEY_VECTOR_POSITION column value, not batch + // order. Position 2's NULL is never read (it is not a candidate), so it + // cannot trip the null-vector guard. + let vector_type = VectorType::try_new(true, 2, DataType::Float(FloatType::new())).unwrap(); + let vector_field = + DataField::new(0, "embedding".to_string(), DataType::Vector(vector_type)); + let read_fields = vec![vector_field.clone()]; + let arrow_schema = build_target_arrow_schema(&read_fields).unwrap(); + + // pos0=[7,0], pos1=[1,0], pos2=NULL, pos3=[4,0]. + let mut builder = FixedSizeListBuilder::new(Float32Builder::new(), 2).with_field(Arc::new( + ArrowField::new("element", ArrowDataType::Float32, true), + )); + for row in [ + Some([7.0f32, 0.0]), + Some([1.0, 0.0]), + None, + Some([4.0, 0.0]), + ] { + match row { + Some([a, b]) => { + builder.values().append_value(a); + builder.values().append_value(b); + builder.append(true); + } + None => { + builder.values().append_value(0.0); + builder.values().append_value(0.0); + builder.append(false); + } + } + } + let vec_array = builder.finish(); + let batch = + arrow_array::RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(vec_array)]) + .unwrap(); + + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let table_path = "memory:/rerank_positional"; + let bucket_path = format!("{table_path}/bucket-0"); + let file_name = "part-0.parquet"; + let output = file_io + .new_output(&format!("{bucket_path}/{file_name}")) + .unwrap(); + let mut writer: Box = Box::new( + ParquetFormatWriter::new( + &output, + arrow_schema.clone(), + "zstd", + 1, + None, + &HashMap::new(), + ) + .await + .unwrap(), + ); + writer.write(&batch).await.unwrap(); + let file_size = writer.close().await.unwrap().file_size; + + let schema_id = 1; + let file_meta = pk_data_file(file_name, 4, Some(0)); + let file_meta = DataFileMeta { + file_size: file_size as i64, + schema_id, + ..file_meta + }; + let 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![file_meta]) + .build() + .unwrap(); + let split = PkVectorSearchSplit { + data_split, + ann_segments: Vec::new(), + active_files: Vec::new(), + }; + + let schema_manager = SchemaManager::new(file_io.clone(), table_path.to_string()); + let reader = DataFileReader::new( + file_io, + schema_manager, + schema_id, + read_fields.clone(), + read_fields.clone(), + Vec::new(), + ); + + let query = vec![1.0f32, 0.0]; + // ANN-reported distances reversed vs. truth: pos1 reported worse (0.9) than + // pos3 (0.1), but the true L2 distances are pos1=0 and pos3=9. + let indexed = vec![cand_at(0, file_name, 1, 0.9), cand_at(0, file_name, 3, 0.1)]; + + let out = rerank_indexed_positional( + &reader, + indexed, + &[split], + &query, + VectorSearchMetric::L2, + 2, + &vector_field, + ) + .await + .unwrap(); + + // Best-first after exact recompute: pos1 (d=0) then pos3 (d=9), each + // carrying the distance computed from its OWN position's stored vector. + assert_eq!(out.len(), 2); + assert_eq!(out[0].row_position, 1); + assert_eq!(out[0].distance, 0.0); + assert_eq!(out[1].row_position, 3); + assert_eq!(out[1].distance, 9.0); + } + + #[tokio::test] + async fn rerank_recomputes_distance_and_reorders() { + // pos0=[9,0], pos1=[1,0]; query=[1,0]. The ANN-reported distances are + // reversed relative to the truth (pos0 reported best at 0.1, pos1 worst at + // 0.9), so an implementation that trusted the ANN order would emit pos0 + // first. Exact L2 recompute yields pos0=64, pos1=0, so the output must + // reorder to pos1-then-pos0 with the recomputed distances. + let (reader, split, vector_field) = vector_rerank_fixture( + "memory:/rerank_reorder", + 2, + &[Some(vec![9.0, 0.0]), Some(vec![1.0, 0.0])], + &[], + ) + .await; + let query = vec![1.0f32, 0.0]; + let indexed = vec![ + cand_at(0, RERANK_FILE, 0, 0.1), + cand_at(0, RERANK_FILE, 1, 0.9), + ]; + + let out = rerank_indexed_positional( + &reader, + indexed, + &[split], + &query, + VectorSearchMetric::L2, + 2, + &vector_field, + ) + .await + .unwrap(); + + assert_eq!(out.len(), 2); + assert_eq!(out[0].row_position, 1); + assert_eq!(out[0].distance, 0.0); + assert_eq!(out[1].row_position, 0); + assert_eq!(out[1].distance, 64.0); + // Order genuinely changed vs. the ANN-reported best-first (which was pos0). + assert!(out[0].distance < out[1].distance); + } + + #[tokio::test] + async fn rerank_is_independent_of_fast_mode_reranks_indexed() { + // The rerank core takes only the indexed (fast-path) candidates and always + // recomputes their true distance; there is no fast/exact switch that can + // skip it. The single candidate carries a bogus ANN distance (0.42) but its + // stored vector equals the query, so the recomputed L2 distance is exactly + // 0.0 — proving the indexed candidate WAS reranked rather than passed + // through with its ANN distance. + let (reader, split, vector_field) = + vector_rerank_fixture("memory:/rerank_indexed", 2, &[Some(vec![1.0, 0.0])], &[]).await; + let query = vec![1.0f32, 0.0]; + let indexed = vec![cand_at(0, RERANK_FILE, 0, 0.42)]; + + let out = rerank_indexed_positional( + &reader, + indexed, + &[split], + &query, + VectorSearchMetric::L2, + 1, + &vector_field, + ) + .await + .unwrap(); + + assert_eq!(out.len(), 1); + assert_eq!(out[0].row_position, 0); + assert_ne!(out[0].distance, 0.42); + assert_eq!(out[0].distance, 0.0); + } + + #[tokio::test] + async fn rerank_fails_loud_on_null_vector() { + // A NULL vector stored AT a candidate position must fail loud rather than + // silently scoring it: the candidate genuinely has no vector to rerank on. + let (reader, split, vector_field) = + vector_rerank_fixture("memory:/rerank_null", 2, &[None], &[]).await; + let query = vec![1.0f32, 0.0]; + let indexed = vec![cand_at(0, RERANK_FILE, 0, 0.1)]; + + let err = rerank_indexed_positional( + &reader, + indexed, + &[split], + &query, + VectorSearchMetric::L2, + 1, + &vector_field, + ) + .await + .err() + .expect("null vector at a candidate position must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("null vector")), + "unexpected error: {err:?}" + ); + } + + #[tokio::test] + async fn rerank_fails_loud_on_leftover_candidate() { + // pos1 is deleted by the deletion vector, so the position read returns no + // row for it. The search path already DV-filters, so a deleted candidate + // reaching rerank is a real inconsistency: the leftover guard must fail + // loud rather than silently dropping the candidate. + let (reader, split, vector_field) = vector_rerank_fixture( + "memory:/rerank_leftover", + 2, + &[Some(vec![1.0, 0.0]), Some(vec![2.0, 0.0])], + &[1], + ) + .await; + let query = vec![1.0f32, 0.0]; + let indexed = vec![ + cand_at(0, RERANK_FILE, 0, 0.1), + cand_at(0, RERANK_FILE, 1, 0.9), + ]; + + let err = rerank_indexed_positional( + &reader, + indexed, + &[split], + &query, + VectorSearchMetric::L2, + 2, + &vector_field, + ) + .await + .err() + .expect("a candidate returning no row must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("failed to read")), + "unexpected error: {err:?}" + ); + } + + #[tokio::test] + async fn rerank_fails_loud_on_dimension_mismatch() { + // Stored vectors are 3-dimensional but the query is 2-dimensional. The + // vector extraction validates each stored row against the query dimension + // and fails loud, so the recompute never runs against mismatched vectors. + let (reader, split, vector_field) = + vector_rerank_fixture("memory:/rerank_dim", 3, &[Some(vec![1.0, 0.0, 0.0])], &[]).await; + let query = vec![1.0f32, 0.0]; + let indexed = vec![cand_at(0, RERANK_FILE, 0, 0.1)]; + + let err = rerank_indexed_positional( + &reader, + indexed, + &[split], + &query, + VectorSearchMetric::L2, + 1, + &vector_field, + ) + .await + .err() + .expect("dimension mismatch must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("dimension")), + "unexpected error: {err:?}" + ); + } + + #[tokio::test] + async fn rerank_fails_loud_on_duplicate_candidate_position() { + // Two candidates addressing the same (split_index, file, position) is a + // programming error upstream: the dedup guard fires before any read. + let (reader, split, vector_field) = + vector_rerank_fixture("memory:/rerank_dup", 2, &[Some(vec![1.0, 0.0])], &[]).await; + let query = vec![1.0f32, 0.0]; + let indexed = vec![ + cand_at(0, RERANK_FILE, 0, 0.1), + cand_at(0, RERANK_FILE, 0, 0.9), + ]; + + let err = rerank_indexed_positional( + &reader, + indexed, + &[split], + &query, + VectorSearchMetric::L2, + 2, + &vector_field, + ) + .await + .err() + .expect("duplicate candidate position must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("duplicate")), + "unexpected error: {err:?}" + ); + } + + #[tokio::test] + async fn rerank_fails_loud_on_unexpected_position() { + // Every position the read surfaces must resolve to a candidate keyed by + // (split_index, file, position). Here the plan carries two splits for the + // SAME (partition, bucket, file), so `split_index_of` resolves the file to + // the LAST plan index (1). The single candidate is tagged with split_index + // 0, so its by_key entry is (0, file, 0) while the read looks up + // (1, file, 0). The lookup misses and the unexpected-position guard fires + // rather than silently dropping the surfaced row. + let (reader, split, vector_field) = + vector_rerank_fixture("memory:/rerank_unexpected", 2, &[Some(vec![1.0, 0.0])], &[]) + .await; + let query = vec![1.0f32, 0.0]; + let indexed = vec![cand_at(0, RERANK_FILE, 0, 0.1)]; + + // Two plan entries for the same file: split_index_of ends up mapping the + // file to plan index 1, not the candidate's split_index 0. + let dup = PkVectorSearchSplit { + data_split: split.data_split.clone(), + ann_segments: Vec::new(), + active_files: Vec::new(), + }; + let plan = vec![dup, split]; + + let err = rerank_indexed_positional( + &reader, + indexed, + &plan, + &query, + VectorSearchMetric::L2, + 1, + &vector_field, + ) + .await + .err() + .expect("a read position absent from the candidate map must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("unexpected position")), + "unexpected error: {err:?}" + ); + } + + #[tokio::test] + async fn rerank_fails_loud_on_file_not_in_plan() { + // A candidate references a (partition, bucket, file) that is absent from + // plan_splits. build_indexed_splits groups it into an indexed split, but the + // split_index_of lookup — built only from plan_splits — has no entry, so the + // kernel fails loud rather than reading an unplanned file. + let (reader, _split, vector_field) = + vector_rerank_fixture("memory:/rerank_noplan", 2, &[Some(vec![1.0, 0.0])], &[]).await; + let query = vec![1.0f32, 0.0]; + let indexed = vec![cand_at(0, RERANK_FILE, 0, 0.1)]; + + // Empty plan: the candidate's file resolves in no plan split. + let err = rerank_indexed_positional( + &reader, + indexed, + &[], + &query, + VectorSearchMetric::L2, + 1, + &vector_field, + ) + .await + .err() + .expect("a candidate file absent from the plan must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("not found in plan")), + "unexpected error: {err:?}" + ); + } + + #[tokio::test] + async fn rerank_reads_only_candidate_positions_not_whole_column() { + // A 6-row file where every NON-candidate position (0, 2, 4, 5) holds a NULL + // vector "poison" and only the two candidate positions (1, 3) hold real + // vectors. The rerank read is told to fetch only positions {1, 3}; every + // row it surfaces is looked up in the candidate map, and any position not in + // the map trips the "unexpected position" guard (a surfaced NULL row would + // additionally trip the null-vector guard). So if the read had surfaced any + // of the poison rows, rerank would fail. It succeeds and returns exactly the + // two candidates at positions {1, 3}, which proves the position selection + // reaching the read contained only the candidate positions (not the whole + // column). + let rows = &[ + None, // pos0 poison (non-candidate) + Some(vec![1.0, 0.0]), // pos1 candidate + None, // pos2 poison (non-candidate) + Some(vec![3.0, 0.0]), // pos3 candidate + None, // pos4 poison (non-candidate) + None, // pos5 poison (non-candidate) + ]; + let (reader, split, vector_field) = + vector_rerank_fixture("memory:/rerank_spy", 2, rows, &[]).await; + let query = vec![1.0f32, 0.0]; + let indexed = vec![ + cand_at(0, RERANK_FILE, 1, 0.9), + cand_at(0, RERANK_FILE, 3, 0.1), + ]; + + let out = rerank_indexed_positional( + &reader, + indexed, + &[split], + &query, + VectorSearchMetric::L2, + 2, + &vector_field, + ) + .await + .unwrap_or_else(|e| { + panic!("only candidate positions are read, so the poison NULLs never decode: {e:?}") + }); + + assert_eq!(out.len(), 2, "exactly the candidate count of rows was read"); + let mut positions: Vec = out.iter().map(|c| c.row_position).collect(); + positions.sort_unstable(); + assert_eq!( + positions, + vec![1, 3], + "only candidate positions reached the read" + ); + // Recomputed distances confirm each surviving row is its own candidate's vector. + assert_eq!(out[0].row_position, 1); + assert_eq!(out[0].distance, 0.0); + assert_eq!(out[1].row_position, 3); + assert_eq!(out[1].distance, 4.0); + } + /// Build a real vindex IVF-flat segment trained with `metric`, returning the /// serialized bytes. `nlist = 1` keeps training trivial and deterministic; the /// only thing the metric check cares about is the persisted metadata metric. diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs b/crates/paimon/src/vindex/pkvector/bucket.rs index 137d7fc0..f3cfc35f 100644 --- a/crates/paimon/src/vindex/pkvector/bucket.rs +++ b/crates/paimon/src/vindex/pkvector/bucket.rs @@ -142,6 +142,15 @@ pub(crate) fn covered_source_files( covered } +/// Separately bounded approximate-index and exact-fallback candidates for one +/// bucket. The approximate list may be over-fetched (for later exact reranking) +/// while the exact-fallback list stays bounded to the caller's final limit. +#[derive(Debug)] +pub(crate) struct BucketSearchResult { + pub(crate) indexed: Vec, + pub(crate) exact: Vec, +} + /// ANN + exact data-file fallback search for one snapshot bucket. Mirrors Java /// `org.apache.paimon.index.pkvector.PrimaryKeyVectorBucketSearch.search`. /// @@ -164,12 +173,16 @@ pub(crate) async fn bucket_search( + Send), query: &[f32], metric: VectorSearchMetric, - limit: usize, + indexed_limit: usize, + exact_limit: usize, search_options: &HashMap, skip_exact_fallback: bool, residual_ranges: Option<&HashMap>, -) -> crate::Result> { - if limit == 0 { +) -> crate::Result { + 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")); } @@ -219,7 +232,8 @@ pub(crate) async fn bucket_search( } } - let mut heap: BinaryHeap = BinaryHeap::with_capacity(limit + 1); + let mut indexed_heap: BinaryHeap = BinaryHeap::with_capacity(indexed_limit + 1); + let mut exact_heap: BinaryHeap = BinaryHeap::with_capacity(exact_limit + 1); let active_source_files: HashSet = files_by_name.keys().map(|name| name.to_string()).collect(); // Active files whose rows an ANN segment already covers; the exact fallback @@ -247,13 +261,13 @@ pub(crate) async fn bucket_search( segment, query, metric, - limit, + indexed_limit, &active_source_files, deletion_vectors, search_options, residual_ranges, )? { - add_candidate(&mut heap, result, limit); + add_candidate(&mut indexed_heap, result, indexed_limit); } } @@ -299,17 +313,19 @@ pub(crate) async fn bucket_search( reader.as_mut(), query, metric, - limit, + exact_limit, &is_excluded, )? { - add_candidate(&mut heap, result, limit); + add_candidate(&mut exact_heap, result, exact_limit); } } } - let mut results: Vec = heap.into_iter().map(|w| w.0).collect(); - results.sort_by(best_first); - Ok(results) + 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); + Ok(BucketSearchResult { indexed, exact }) } #[cfg(test)] @@ -369,6 +385,75 @@ mod tests { } } + #[tokio::test] + async fn dual_limit_bounds_indexed_and_exact_independently() { + // One ANN segment covering "ann.mosaic" (3 rows) and one uncovered exact + // file "exact.mosaic" (3 rows). indexed_limit = 3 lets all ANN hits through; + // exact_limit = 1 keeps only the single closest exact hit. + let segment = BucketAnnSegment::for_test(meta(&[("ann.mosaic", 3)])); + // Fake ANN searcher returns three hits at increasing distance. + 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, + }, + PkVectorSearchResult { + data_file_name: "ann.mosaic".into(), + row_position: 2, + distance: 0.3, + }, + ], + }; + // Exact file has three rows; query nearest is position 0. + let mut factory = as_factory(|_f: &BucketActiveFile| -> ExactReaderFuture<'_> { + let reader = ArrayReader::new( + 2, + vec![ + Some(vec![1.0, 0.0]), + Some(vec![9.0, 0.0]), + Some(vec![8.0, 0.0]), + ], + ); + Box::pin(async move { Ok(Box::new(reader) as Box) }) + }); + let active_files = vec![active("ann.mosaic", 3), active("exact.mosaic", 3)]; + let dvs: HashMap> = HashMap::new(); + let opts = HashMap::new(); + + let out = bucket_search( + Some(&ann), + &[segment], + &active_files, + &dvs, + &mut factory, + &[1.0, 0.0], + VectorSearchMetric::L2, + 3, // indexed_limit + 1, // exact_limit + &opts, + false, + None, + ) + .await + .unwrap(); + + assert_eq!( + out.indexed.len(), + 3, + "indexed heap keeps indexed_limit hits" + ); + assert_eq!(out.exact.len(), 1, "exact heap bounded to exact_limit"); + assert_eq!(out.exact[0].data_file_name, "exact.mosaic"); + assert_eq!(out.exact[0].row_position, 0); + } + #[tokio::test] async fn test_rejects_non_positive_limit() { let mut factory = as_factory(|_: &BucketActiveFile| -> ExactReaderFuture<'_> { @@ -383,6 +468,7 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 0, + 0, &HashMap::new(), false, None, @@ -417,7 +503,7 @@ mod tests { let mut factory = as_factory(|_: &BucketActiveFile| -> ExactReaderFuture<'_> { Box::pin(async { unreachable!() }) }); - let results = bucket_search( + let out = bucket_search( Some(&ann), &[segment], &[active("data-1", 3)], @@ -426,12 +512,17 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 3, + 3, &HashMap::new(), false, None, ) .await .unwrap(); + let mut results = out.indexed.clone(); + results.extend(out.exact.clone()); + results.sort_by(best_first); + results.truncate(3); // Top-3 BEST_FIRST: (data-1,0), (data-1,1), (data-1,2) — the larger // data_file_name "data-2" entries are evicted despite equal distance. assert_eq!( @@ -469,7 +560,7 @@ mod tests { let mut factory = as_factory(|_: &BucketActiveFile| -> ExactReaderFuture<'_> { Box::pin(async { unreachable!() }) }); - let results = bucket_search( + let out = bucket_search( Some(&ann), &[segment], &[active("data-1", 2)], @@ -478,12 +569,17 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 1, + 1, &HashMap::new(), false, None, ) .await .unwrap(); + let mut results = out.indexed.clone(); + results.extend(out.exact.clone()); + results.sort_by(best_first); + results.truncate(1); assert_eq!(results.len(), 1); assert_eq!(results[0].row_position, 1); assert_eq!(results[0].distance, -1.0); @@ -512,7 +608,7 @@ mod tests { )) as Box) }) }); - let results = bucket_search( + let out = bucket_search( Some(&ann), &[segment], &[active("data-1", 2), active("data-2", 2)], @@ -521,12 +617,17 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 2, + 2, &HashMap::new(), false, None, ) .await .unwrap(); + let mut results = out.indexed.clone(); + results.extend(out.exact.clone()); + results.sort_by(best_first); + results.truncate(2); assert_eq!( results, vec![ @@ -565,7 +666,7 @@ mod tests { bm.insert(0); // data-1 position 0 deleted dvs.insert("data-1".into(), Arc::new(DeletionVector::from_bitmap(bm))); - let results = bucket_search( + let out = bucket_search( None, &[], &[active("data-1", 2), active("data-2", 2)], @@ -574,6 +675,7 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 2, + 2, &HashMap::new(), false, None, @@ -582,6 +684,10 @@ mod tests { .unwrap(); // Candidates: data-2 pos0 {1,0} dist 1.0; data-1 pos1 {2,0} dist 4.0. // (data-1 pos0 deleted, data-2 pos1 null.) + let mut results = out.indexed.clone(); + results.extend(out.exact.clone()); + results.sort_by(best_first); + results.truncate(2); assert_eq!( results, vec![ @@ -613,6 +719,7 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 1, + 1, &HashMap::new(), false, None, @@ -640,6 +747,7 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 1, + 1, &HashMap::new(), false, None, @@ -671,7 +779,7 @@ mod tests { calls.lock().unwrap().push(f.file_name.clone()); Box::pin(async { unreachable!("only data-1 is active and it is ANN-covered") }) }); - let results = bucket_search( + let out = bucket_search( Some(&ann), &[segment], &[active("data-1", 2)], @@ -680,12 +788,17 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 2, + 2, &HashMap::new(), false, None, ) .await .unwrap(); + let mut results = out.indexed.clone(); + results.extend(out.exact.clone()); + results.sort_by(best_first); + results.truncate(2); assert_eq!( results, vec![PkVectorSearchResult { @@ -713,6 +826,7 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 1, + 1, &HashMap::new(), false, None, @@ -732,7 +846,7 @@ mod tests { let mut factory = as_factory(|_: &BucketActiveFile| -> ExactReaderFuture<'_> { Box::pin(async { unreachable!() }) }); - let results = bucket_search( + let out = bucket_search( None, &[], &[active("data-1", 2), active("data-2", 2)], @@ -741,13 +855,15 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 2, + 2, &HashMap::new(), true, // skip_exact_fallback None, ) .await .unwrap(); - assert!(results.is_empty()); + assert!(out.indexed.is_empty()); + assert!(out.exact.is_empty()); } #[tokio::test] @@ -777,6 +893,7 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 1, + 1, &HashMap::new(), false, None, @@ -816,6 +933,7 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 1, + 1, &HashMap::new(), false, None, @@ -844,6 +962,7 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 1, + 1, &HashMap::new(), false, None, @@ -909,7 +1028,7 @@ mod tests { }); let mut residual: HashMap = HashMap::new(); residual.insert("data-1".into(), treemap(&[0, 2])); - let results = bucket_search( + let out = bucket_search( None, &[], &[active("data-1", 3)], @@ -918,12 +1037,17 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 5, + 5, &HashMap::new(), false, Some(&residual), ) .await .unwrap(); + let mut results = out.indexed.clone(); + results.extend(out.exact.clone()); + results.sort_by(best_first); + results.truncate(5); assert_eq!( results, vec![ @@ -957,7 +1081,7 @@ mod tests { }); let mut residual: HashMap = HashMap::new(); residual.insert("data-1".into(), treemap(&[0, 1])); - let results = bucket_search( + let out = bucket_search( None, &[], &[active("data-1", 2), active("data-2", 2)], @@ -966,12 +1090,17 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 5, + 5, &HashMap::new(), false, Some(&residual), ) .await .unwrap(); + let mut results = out.indexed.clone(); + results.extend(out.exact.clone()); + results.sort_by(best_first); + results.truncate(5); // Only data-1 rows appear; data-2 was never read. assert!(results.iter().all(|r| r.data_file_name == "data-1")); assert_eq!(calls.lock().unwrap().as_slice(), &["data-1".to_string()]); @@ -988,7 +1117,7 @@ mod tests { }); let mut residual: HashMap = HashMap::new(); residual.insert("data-1".into(), treemap(&[])); - let results = bucket_search( + let out = bucket_search( None, &[], &[active("data-1", 3)], @@ -997,13 +1126,15 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 5, + 5, &HashMap::new(), false, Some(&residual), ) .await .unwrap(); - assert!(results.is_empty()); + assert!(out.indexed.is_empty()); + assert!(out.exact.is_empty()); assert_eq!(*calls.lock().unwrap(), 0); } @@ -1029,7 +1160,7 @@ mod tests { dvs.insert("data-1".into(), Arc::new(DeletionVector::from_bitmap(bm))); let mut residual: HashMap = HashMap::new(); residual.insert("data-1".into(), treemap(&[0, 1, 2])); - let results = bucket_search( + let out = bucket_search( None, &[], &[active("data-1", 3)], @@ -1038,12 +1169,17 @@ mod tests { &[0.0, 0.0], VectorSearchMetric::L2, 5, + 5, &HashMap::new(), false, Some(&residual), ) .await .unwrap(); + let mut results = out.indexed.clone(); + results.extend(out.exact.clone()); + results.sort_by(best_first); + results.truncate(5); assert_eq!( results.iter().map(|r| r.row_position).collect::>(), vec![1, 2] diff --git a/crates/paimon/tests/pk_vector_baseline_test.rs b/crates/paimon/tests/pk_vector_baseline_test.rs index 8ec2fa62..97d9e18c 100644 --- a/crates/paimon/tests/pk_vector_baseline_test.rs +++ b/crates/paimon/tests/pk_vector_baseline_test.rs @@ -1006,3 +1006,487 @@ async fn pk_vector_residual_filter_excludes_non_matching_rows() { ); } } + +// --- Exact-rerank coverage -------------------------------------------------- +// +// The refine factor over-fetches approximate (ANN) candidates and reorders them +// by exact distance re-read from the data file. The two tests below drive that +// end-to-end through the public builder: one where the approximate order and the +// exact order genuinely differ (so the reorder is observable), and one where the +// search yields no indexed candidates (so the rerank must be skipped cleanly). + +/// The query-side option key that requests exact rerank for the `embedding` +/// PK-vector column. `fields..ivf.refine-factor` is one of the aliases the +/// builder accepts for an `ivf-flat` index (the `ivf.` prefix, `refine-factor` +/// suffix); it is set as a query option so it takes precedence over table options. +fn refine_factor_option(factor: usize) -> HashMap { + HashMap::from([( + format!("fields.{VECTOR_COLUMN}.ivf.refine-factor"), + factor.to_string(), + )]) +} + +/// Primary-key schema `(id INT PRIMARY KEY, embedding VECTOR)` with extra +/// table options merged on top of [`table_options`]. Used to pin table-level +/// options (e.g. `global-index.search-mode`) that the query-side `with_options` +/// cannot influence. +fn pk_vector_schema_with(extra: &[(&str, &str)]) -> 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); + } + for (k, v) in extra { + builder = builder.option(k.to_string(), v.to_string()); + } + TableSchema::new(0, &builder.build().unwrap()) +} + +/// Persist `schema` and write one real data file over `data_vectors` in a fresh +/// temp dir, returning the temp dir, the opened table, the `FileIO`, the table +/// location, the written file's meta rewritten to satisfy the two PK-vector +/// constraints (compacted, non-level-0, `first_row_id = 0`), and its +/// bucket / partition. +async fn write_schema_and_data( + schema: &TableSchema, + data_vectors: &[[f32; DIM]], +) -> ( + tempfile::TempDir, + Table, + FileIO, + String, + DataFileMeta, + Vec, + i32, +) { + 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(); + } + 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(data_vectors)) + .await + .unwrap(); + let write_messages = writer.prepare_commit().await.unwrap(); + assert_eq!( + write_messages.len(), + 1, + "single bucket -> one write message" + ); + let written = &write_messages[0]; + assert_eq!(written.new_files.len(), 1, "single data file expected"); + let base_meta = written.new_files[0].clone(); + + let bucket = written.bucket; + let partition = written.partition.clone(); + let indexed_meta = DataFileMeta { + level: 1, + file_source: Some(1), + first_row_id: Some(0), + ..base_meta + }; + ( + tmp, + table, + file_io, + location, + indexed_meta, + partition, + bucket, + ) +} + +/// Read `execute_read()` into `(id, score)` tuples (best-first) plus the batches, +/// mirroring [`read_id_and_scores`] but with caller-supplied query options so the +/// refine factor can be requested. +async fn read_id_and_scores_with_options( + table: &Table, + query: Vec, + limit: usize, + options: HashMap, +) -> (Vec, Vec, Vec) { + let mut builder = table.new_vector_search_builder(); + builder + .with_vector_column(VECTOR_COLUMN) + .with_query_vector(query) + .with_limit(limit) + .with_options(options); + let batches = builder + .execute_read() + .await + .expect("primary-key vector rerank read failed") + .try_collect::>() + .await + .expect("collecting rerank 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, batches) +} + +/// Fixture where the ANN-approximate order and the exact-distance order are +/// disjoint. The data file holds `data_vectors` (the truth the rerank re-reads), +/// while the ANN segment is built over `index_vectors` (a perturbed copy), so the +/// approximate recall order differs from the exact order over the true data. +/// +/// query [10,0,0,0] +/// position | index_vectors (ANN sees) | data_vectors (truth) +/// ---------+----------------------------------+--------------------------- +/// 0 | [0,0,0,0] -> 100 | [9,0,0,0] -> 1 (exact 1st) +/// 1 | [0,0,1,0] -> 101 | [8,0,0,0] -> 4 (exact 2nd) +/// 2 | [0,0,2,0] -> 104 | [7,0,0,0] -> 9 (exact 3rd) +/// 3 | [9,0,0,0] -> 1 (ANN 3rd) | [0,0,0,0] -> 100 +/// 4 | [9.5,0,0,0]-> 0.25 (ANN 2nd) | [0,5,0,0] -> 125 +/// 5 | [10,0,0,0] -> 0 (ANN 1st) | [0,0,6,0] -> 136 +/// +/// ANN top-3 = [5, 4, 3]; exact top-3 = [0, 1, 2]. The two are disjoint, so if the +/// rerank never ran (gate broken or factor unset) the output would be [5, 4, 3]. +fn fixture_rerank_divergent() -> ([f32; DIM], Vec<[f32; DIM]>, Vec<[f32; DIM]>) { + let query = [10.0, 0.0, 0.0, 0.0]; + let index_vectors = vec![ + [0.0, 0.0, 0.0, 0.0], // pos 0 + [0.0, 0.0, 1.0, 0.0], // pos 1 + [0.0, 0.0, 2.0, 0.0], // pos 2 + [9.0, 0.0, 0.0, 0.0], // pos 3 + [9.5, 0.0, 0.0, 0.0], // pos 4 + [10.0, 0.0, 0.0, 0.0], // pos 5 + ]; + let data_vectors = vec![ + [9.0, 0.0, 0.0, 0.0], // pos 0 + [8.0, 0.0, 0.0, 0.0], // pos 1 + [7.0, 0.0, 0.0, 0.0], // pos 2 + [0.0, 0.0, 0.0, 0.0], // pos 3 + [0.0, 5.0, 0.0, 0.0], // pos 4 + [0.0, 0.0, 6.0, 0.0], // pos 5 + ]; + (query, index_vectors, data_vectors) +} + +/// refine_factor > 0 reorders the emitted rows to the exact-distance ground truth: +/// the ANN-approximate order (built over a perturbed copy of the vectors) differs +/// from the exact order over the true data, and the rerank corrects it. +/// +/// The fixture builds the ANN segment over `index_vectors` while the data file +/// holds `data_vectors`, so the approximate recall order [5, 4, 3] is disjoint +/// from the exact order [0, 1, 2]. With `refine-factor = 2` the search over-fetches +/// `limit * 2 = 6` candidates (the whole bucket), then re-reads each candidate's +/// true vector and reorders by exact distance. The emitted rows must be the exact +/// top-3 [0, 1, 2] with scores computed from the true distances. +// Gated off Windows for the same `file://` tempdir reason as the tests above. +#[cfg(not(windows))] +#[tokio::test] +async fn pk_vector_refine_factor_matches_exact_ground_truth() { + let (query, index_vectors, data_vectors) = fixture_rerank_divergent(); + let k = 3; + + // Ground truths: the ANN recall order (exact over the perturbed index vectors, + // since nlist = 1 is exhaustive) and the exact order over the true data. + let ann_order: Vec = analytic_topk(&query, &index_vectors, k) + .iter() + .map(|(id, _)| *id as i32) + .collect(); + let exact = analytic_topk(&query, &data_vectors, k); + let exact_ids: Vec = exact.iter().map(|(id, _)| *id as i32).collect(); + let exact_scores: Vec = exact.iter().map(|(_, d)| l2_score(*d)).collect(); + assert_eq!(ann_order, vec![5, 4, 3], "fixture guard: ANN order"); + assert_eq!(exact_ids, vec![0, 1, 2], "fixture guard: exact order"); + assert_ne!( + ann_order, exact_ids, + "fixture is only discriminating if ANN order differs from exact order" + ); + + // Build a table whose ANN segment encodes the perturbed order while the data + // file holds the true vectors, and confirm the segment really recalls the + // perturbed order before the read path runs. + let schema = pk_vector_schema(); + let (_tmp, table, file_io, location, indexed_meta, partition, bucket) = + write_schema_and_data(&schema, &data_vectors).await; + let data_file_name = indexed_meta.file_name.clone(); + let row_count = indexed_meta.row_count; + + let index_file_name = "vector-ivf-flat-pkvector-rerank.index".to_string(); + let index_file_size = + write_ann_segment(&file_io, &location, &index_file_name, &index_vectors).await; + { + let bytes = file_io + .new_input(&format!("{location}/index/{index_file_name}")) + .unwrap() + .read() + .await + .unwrap(); + // The segment recalls the perturbed order, not the true-data order. + assert_segment_reads_back(&bytes, &query, &analytic_topk(&query, &index_vectors, k)); + } + + 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-rerank".to_string()) + .commit(vec![message]) + .await + .unwrap(); + + // Sanity anchor: with NO refine factor the same fixture emits the raw ANN + // recall order [5, 4, 3] (the rerank never runs), so any reordering below is + // attributable to the refine factor and not to the fixture happening to agree. + let (unset_ids, _unset_scores, _unset_batches) = + read_id_and_scores_with_options(&table, query.to_vec(), k, HashMap::new()).await; + assert_eq!( + unset_ids, ann_order, + "without a refine factor the rows must stay in raw ANN order [5, 4, 3]" + ); + + // With the refine factor set, the emitted rows are the EXACT top-3 [0, 1, 2], + // not the approximate order [5, 4, 3]: the rerank reordered them. + let (ids, scores, batches) = + read_id_and_scores_with_options(&table, query.to_vec(), k, refine_factor_option(2)).await; + assert_eq!( + ids, exact_ids, + "refine factor must reorder rows to exact top-3 [0, 1, 2], not ANN order [5, 4, 3]" + ); + + // Row content: each emitted row's vector equals the TRUE data vector at that + // physical position (the rerank read the data file, not the ANN segment). + let got_vectors = collect_vectors(&batches); + assert_eq!(got_vectors.len(), k, "three rows expected"); + for (row_idx, (id, _)) in exact.iter().enumerate() { + assert_eq!( + got_vectors[row_idx], + data_vectors[*id as usize].to_vec(), + "materialized vector for row id {id} diverges from true data" + ); + } + + // Scores are computed from the recomputed exact distances over the true data. + assert_eq!(scores.len(), k); + for (got, want) in scores.iter().zip(&exact_scores) { + assert!( + (got - want).abs() < 1e-4, + "reranked score diverges: got {got}, want {want}" + ); + } +} + +/// Gating: with a positive refine factor set but a search that yields ZERO +/// indexed (ANN) candidates, the exact rerank must be skipped — no error, no +/// spurious position read — and the exact-fallback rows come through unchanged. +/// +/// The table carries a compacted, non-level-0 data file but NO ANN index segment, +/// and `global-index.search-mode = full` enables the exact data-file fallback. The +/// search therefore produces exact-fallback candidates while `search.indexed` is +/// empty, which is exactly the `refine_factor > 0 && !search.indexed.is_empty()` +/// gate's short-circuit path. The emitted rows must be the exact top-3 the fallback +/// found, identical to what a refine-unset run would emit. +// Gated off Windows for the same `file://` tempdir reason as the tests above. +#[cfg(not(windows))] +#[tokio::test] +async fn pk_vector_refine_factor_with_no_indexed_candidates_is_noop() { + // Reuse the discriminating fixture's true-data layout so best-first order + // [5, 1, 3] is distinct from ascending position, proving the exact fallback + // ranked correctly rather than emitting rows in file order. + let (query, data_vectors) = fixture_discriminating(); + let k = 3; + let exact = analytic_topk(&query, &data_vectors, k); + let exact_ids: Vec = exact.iter().map(|(id, _)| *id as i32).collect(); + let exact_scores: Vec = exact.iter().map(|(_, d)| l2_score(*d)).collect(); + assert_eq!(exact_ids, vec![5, 1, 3], "fixture guard: exact order"); + + // FULL search mode enables the exact data-file fallback; no ANN segment is + // committed, so the ANN candidate list stays empty. + let schema = pk_vector_schema_with(&[("global-index.search-mode", "full")]); + let (_tmp, table, _file_io, _location, indexed_meta, partition, bucket) = + write_schema_and_data(&schema, &data_vectors).await; + + // Commit the data file WITHOUT any index segment: the plan has a searchable + // bucket but zero ANN segments -> zero indexed candidates. + let message = CommitMessage::new(partition, bucket, vec![indexed_meta]); + TableCommit::new(table.clone(), "pkvector-rerank-noop".to_string()) + .commit(vec![message]) + .await + .unwrap(); + + // A positive refine factor is set, but with no indexed candidates the rerank is + // gated off: execute_read must not error, and returns the exact-fallback rows. + let (ids, scores, batches) = + read_id_and_scores_with_options(&table, query.to_vec(), k, refine_factor_option(2)).await; + assert_eq!( + ids, exact_ids, + "with no indexed candidates the exact-fallback rows must pass through unchanged" + ); + + let got_vectors = collect_vectors(&batches); + assert_eq!(got_vectors.len(), k, "three rows expected"); + for (row_idx, (id, _)) in exact.iter().enumerate() { + assert_eq!( + got_vectors[row_idx], + data_vectors[*id as usize].to_vec(), + "materialized vector for row id {id} diverges from true data" + ); + } + assert_eq!(scores.len(), k); + for (got, want) in scores.iter().zip(&exact_scores) { + assert!( + (got - want).abs() < 1e-4, + "exact-fallback score diverges: got {got}, want {want}" + ); + } +} + +/// An invalid refine factor must fail loud even when the table is empty: the +/// factor is resolved before planning, so config validity does not depend on +/// whether the table currently has searchable data. Regression for a resolution +/// that ran only after the empty-plan early return. Covers both invalid forms +/// the reviewer named (`0` and a non-integer) and both option sources (query and +/// table), since the fix moved the resolution of both sources ahead of planning. +#[cfg(not(target_os = "windows"))] +#[tokio::test] +async fn pk_vector_invalid_refine_factor_fails_loud_on_empty_table() { + // Persist only the schema (optionally with extra table options): no snapshot, + // so the PK-vector plan is empty. + async fn empty_table(file_io: &FileIO, location: &str, extra: &[(&str, &str)]) -> Table { + file_io.mkdirs(&format!("{location}/schema")).await.unwrap(); + let schema = pk_vector_schema_with(extra); + file_io + .new_output(&format!("{location}/schema/schema-{}", schema.id())) + .unwrap() + .write(Bytes::from(serde_json::to_vec(&schema).unwrap())) + .await + .unwrap(); + open_table(file_io, location).await + } + + let refine_key = format!("fields.{VECTOR_COLUMN}.ivf.refine-factor"); + + // Case 1: non-integer via a QUERY option -> "Invalid ... Must be an integer". + { + let tmp = tempfile::tempdir().expect("create temp dir"); + let location = format!("file://{}", tmp.path().display()); + let file_io = FileIOBuilder::new("file").build().unwrap(); + let table = empty_table(&file_io, &location, &[]).await; + let mut builder = table.new_vector_search_builder(); + builder + .with_vector_column(VECTOR_COLUMN) + .with_query_vector(vec![0.0f32; DIM]) + .with_limit(3) + .with_options(HashMap::from([(refine_key.clone(), "abc".to_string())])); + let err = match builder.execute_read().await { + Ok(_) => panic!("a non-integer refine factor must fail loud on an empty table"), + Err(e) => e, + }; + assert!( + err.to_string().contains("Invalid vector refine factor"), + "expected a non-integer refine-factor error, got: {err}" + ); + } + + // Case 2: explicit `0` via a QUERY option -> "must be positive" (0 means + // "omit"; it cannot be requested). + { + let tmp = tempfile::tempdir().expect("create temp dir"); + let location = format!("file://{}", tmp.path().display()); + let file_io = FileIOBuilder::new("file").build().unwrap(); + let table = empty_table(&file_io, &location, &[]).await; + let mut builder = table.new_vector_search_builder(); + builder + .with_vector_column(VECTOR_COLUMN) + .with_query_vector(vec![0.0f32; DIM]) + .with_limit(3) + .with_options(HashMap::from([(refine_key.clone(), "0".to_string())])); + let err = match builder.execute_read().await { + Ok(_) => panic!("a zero refine factor must fail loud on an empty table"), + Err(e) => e, + }; + assert!( + err.to_string().contains("must be positive"), + "expected a positive-refine-factor error, got: {err}" + ); + } + + // Case 3: non-integer via a TABLE option (the fix moved table-option + // resolution ahead of planning too). + { + let tmp = tempfile::tempdir().expect("create temp dir"); + let location = format!("file://{}", tmp.path().display()); + let file_io = FileIOBuilder::new("file").build().unwrap(); + let table = empty_table(&file_io, &location, &[(refine_key.as_str(), "abc")]).await; + let mut builder = table.new_vector_search_builder(); + builder + .with_vector_column(VECTOR_COLUMN) + .with_query_vector(vec![0.0f32; DIM]) + .with_limit(3); + let err = match builder.execute_read().await { + Ok(_) => panic!("a non-integer table refine factor must fail loud on an empty table"), + Err(e) => e, + }; + assert!( + err.to_string().contains("Invalid vector refine factor"), + "expected a non-integer refine-factor error, got: {err}" + ); + } +}