From d20be5686ded0bf0241741e8920e18e23a8c47f5 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Mon, 20 Jul 2026 00:54:14 +0800 Subject: [PATCH 1/2] feat(table): rerank primary-key vector candidates with exact distances MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reconcile the primary-key vector exact rerank feature onto the physical-coordinate read path. On a PK-vector table, execute_read() materializes rows in file-local physical coordinates. This adds: - Dual-limit bucket search: ANN/indexed candidates can over-fetch (indexed_limit = limit × refine_factor) while exact-fallback remains bounded by the final limit. - When a positive refine factor is configured, rerank the over-fetched ANN candidates by rereading their vectors from data files, recomputing exact distance, and re-running the global Top-K merge. - Fail-loud validation: null vector elements and out-of-range split indices return DataInvalid errors instead of panicking or silently corrupting results. Changes: - bucket_search returns separate indexed/exact lists, each bounded independently - PkVectorOrchestrator.search_candidates takes dual limits, returns OrchestratorSearchResult - merge_candidates performs global Top-K over the union of indexed and exact channels - rerank_indexed_positional reads vectors by file-local position, validates ownership order - append_batch_vectors rejects null child elements in vector arrays - build_indexed_splits validates split_index is in range before indexing Tests added: dual-limit bucket behavior, merge_candidates union Top-K, rerank kernel (ownership/dedup/NULL/dimension/file-not-in-plan), e2e refine integration, baseline regression (refine=0 preserves exact output), null element fail-loud, out-of-range split_index fail-loud. --- .../src/table/pk_vector_data_file_reader.rs | 37 +- .../src/table/pk_vector_indexed_split_read.rs | 2 +- .../src/table/pk_vector_orchestrator.rs | 135 ++- .../paimon/src/table/vector_search_builder.rs | 854 +++++++++++++++++- crates/paimon/src/vindex/pkvector/bucket.rs | 182 +++- .../paimon/tests/pk_vector_baseline_test.rs | 391 ++++++++ 6 files changed, 1543 insertions(+), 58 deletions(-) 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..2a8de144 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; @@ -583,12 +587,29 @@ impl<'a> VectorSearchBuilder<'a> { }, ); - let candidates = PkVectorOrchestrator::new(reader) + // 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 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 search: OrchestratorSearchResult = PkVectorOrchestrator::new(reader) .search_candidates( &plan.splits, query_vector, metric, limit, + indexed_limit, Some(&ann_searcher), &mut factory, &search_options, @@ -597,6 +618,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 +1332,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 +3180,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..4e6db812 100644 --- a/crates/paimon/tests/pk_vector_baseline_test.rs +++ b/crates/paimon/tests/pk_vector_baseline_test.rs @@ -1006,3 +1006,394 @@ 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}" + ); + } +} From 51c5c422c875ebf9fe57f69a4aeedca41e97aa51 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Mon, 20 Jul 2026 16:30:51 +0800 Subject: [PATCH 2/2] fix(table): resolve vector refine factor before planning The refine factor was resolved only after PkVectorScan::plan() and the empty-plan early return, so an empty primary-key vector table accepted an invalid or zero refine factor and returned an empty reader, while the same query started failing once the table had a searchable split. Configuration validity must not depend on whether data currently exists. Resolve the refine factor (and its search limit) before planning so an invalid factor fails loud regardless of table state. Add an empty-table regression test. --- .../paimon/src/table/vector_search_builder.rs | 33 +++---- .../paimon/tests/pk_vector_baseline_test.rs | 93 +++++++++++++++++++ 2 files changed, 110 insertions(+), 16 deletions(-) diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index 2a8de144..9f50ee46 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -441,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, @@ -587,22 +604,6 @@ impl<'a> VectorSearchBuilder<'a> { }, ); - // 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 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 search: OrchestratorSearchResult = PkVectorOrchestrator::new(reader) .search_candidates( &plan.splits, diff --git a/crates/paimon/tests/pk_vector_baseline_test.rs b/crates/paimon/tests/pk_vector_baseline_test.rs index 4e6db812..97d9e18c 100644 --- a/crates/paimon/tests/pk_vector_baseline_test.rs +++ b/crates/paimon/tests/pk_vector_baseline_test.rs @@ -1397,3 +1397,96 @@ async fn pk_vector_refine_factor_with_no_indexed_candidates_is_noop() { ); } } + +/// 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}" + ); + } +}