This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 4211ddad [core] Rerank primary-key vector candidates with exact
distances (#550)
4211ddad is described below
commit 4211ddad799665eede5933e9a39e7f12cc006a74
Author: Junrui Lee <[email protected]>
AuthorDate: Mon Jul 20 18:51:17 2026 +0800
[core] Rerank primary-key vector candidates with exact distances (#550)
---
.../paimon/src/table/pk_vector_data_file_reader.rs | 37 +-
.../src/table/pk_vector_indexed_split_read.rs | 2 +-
crates/paimon/src/table/pk_vector_orchestrator.rs | 135 +++-
crates/paimon/src/table/vector_search_builder.rs | 855 ++++++++++++++++++++-
crates/paimon/src/vindex/pkvector/bucket.rs | 182 ++++-
crates/paimon/tests/pk_vector_baseline_test.rs | 484 ++++++++++++
6 files changed, 1637 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<Option<Vec<f32>>> = 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<Float32, 2> 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<Vec<i64>> {
+pub(crate) fn expand_ranges(ranges: &[RowRange], row_count: i64) ->
crate::Result<Vec<i64>> {
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<PkVectorCandidate>,
+ pub(crate) exact: Vec<PkVectorCandidate>,
+}
+
/// 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<PkVectorCandidate>,
limit: usize) -> Vec<PkV
candidates
}
+/// Merge already-bounded indexed and exact candidate lists into a single
+/// best-first list truncated to `limit`. The final global Top-K over the union
+/// is what guarantees the merged result is independent of how the two inputs
+/// were individually bounded.
+pub(crate) fn merge_candidates(
+ indexed: Vec<PkVectorCandidate>,
+ exact: Vec<PkVectorCandidate>,
+ limit: usize,
+) -> Vec<PkVectorCandidate> {
+ 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<String, String>,
skip_exact_fallback: bool,
residual_by_split: Option<&[HashMap<String, RoaringTreemap>]>,
- ) -> crate::Result<Vec<PkVectorCandidate>> {
+ ) -> crate::Result<OrchestratorSearchResult> {
// 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<PkVectorCandidate> = Vec::new();
+ // Eager per-bucket search -> tagged candidates, kept split by path.
+ let mut indexed_candidates: Vec<PkVectorCandidate> = Vec::new();
+ let mut exact_candidates: Vec<PkVectorCandidate> = 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<PkVectorCandidate>,
+ plan_splits: &[PkVectorSearchSplit],
+ query_vector: &[f32],
+ metric: VectorSearchMetric,
+ limit: usize,
+ vector_field: &DataField,
+) -> crate::Result<Vec<PkVectorCandidate>> {
+ // 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<u8>, 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<PkVectorCandidate> = 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::<Int64Array>()
+ .ok_or_else(|| crate::Error::DataInvalid {
+ message: format!("{PKEY_VECTOR_POSITION_COLUMN} column is
not Int64"),
+ source: None,
+ })?;
+ let mut vectors: Vec<Option<Vec<f32>>> = 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<Float32>` 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<Vec<f32>>],
+ 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<dyn FormatFileWriter> = 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<dyn FormatFileWriter> = 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<i64> = 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<PkVectorSearchResult>,
+ pub(crate) exact: Vec<PkVectorSearchResult>,
+}
+
/// 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<String, String>,
skip_exact_fallback: bool,
residual_ranges: Option<&HashMap<String, roaring::RoaringTreemap>>,
-) -> crate::Result<Vec<PkVectorSearchResult>> {
- if limit == 0 {
+) -> crate::Result<BucketSearchResult> {
+ 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<WorstFirst> = BinaryHeap::with_capacity(limit +
1);
+ let mut indexed_heap: BinaryHeap<WorstFirst> =
BinaryHeap::with_capacity(indexed_limit + 1);
+ let mut exact_heap: BinaryHeap<WorstFirst> =
BinaryHeap::with_capacity(exact_limit + 1);
let active_source_files: HashSet<String> =
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<PkVectorSearchResult> = heap.into_iter().map(|w|
w.0).collect();
- results.sort_by(best_first);
- Ok(results)
+ let mut indexed: Vec<PkVectorSearchResult> =
indexed_heap.into_iter().map(|w| w.0).collect();
+ indexed.sort_by(best_first);
+ let mut exact: Vec<PkVectorSearchResult> = 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<dyn
PkVectorReader>) })
+ });
+ let active_files = vec![active("ann.mosaic", 3),
active("exact.mosaic", 3)];
+ let dvs: HashMap<String, Arc<DeletionVector>> = 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<dyn PkVectorReader>)
})
});
- 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<String, roaring::RoaringTreemap> =
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<String, roaring::RoaringTreemap> =
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<String, roaring::RoaringTreemap> =
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<String, roaring::RoaringTreemap> =
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<_>>(),
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.<col>.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<String, String> {
+ HashMap::from([(
+ format!("fields.{VECTOR_COLUMN}.ivf.refine-factor"),
+ factor.to_string(),
+ )])
+}
+
+/// Primary-key schema `(id INT PRIMARY KEY, embedding VECTOR<FLOAT>)` 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<u8>,
+ 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<f32>,
+ limit: usize,
+ options: HashMap<String, String>,
+) -> (Vec<i32>, Vec<f32>, Vec<RecordBatch>) {
+ 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::<Vec<_>>()
+ .await
+ .expect("collecting rerank read batches failed");
+
+ let ids: Vec<i32> = batches
+ .iter()
+ .flat_map(|b| {
+ let idx = b.schema().index_of("id").unwrap();
+ b.column(idx)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ })
+ .collect();
+ let scores: Vec<f32> = batches
+ .iter()
+ .flat_map(|b| {
+ let idx = b.schema().index_of("__paimon_search_score").unwrap();
+ b.column(idx)
+ .as_any()
+ .downcast_ref::<Float32Array>()
+ .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<i32> = 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<i32> = exact.iter().map(|(id, _)| *id as i32).collect();
+ let exact_scores: Vec<f32> = 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<i32> = exact.iter().map(|(id, _)| *id as i32).collect();
+ let exact_scores: Vec<f32> = 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}"
+ );
+ }
+}
