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 8e44700 [core] Add primary-key vector bucket search kernel (#516)
8e44700 is described below
commit 8e44700eb38f2267f6f0cf7746e622bfc7495099
Author: Junrui Lee <[email protected]>
AuthorDate: Tue Jul 14 11:24:04 2026 +0800
[core] Add primary-key vector bucket search kernel (#516)
---
crates/paimon/src/deletion_vector/core.rs | 20 +
crates/paimon/src/spec/pk_vector_source.rs | 4 +-
crates/paimon/src/vindex/mod.rs | 2 +
crates/paimon/src/vindex/pkvector/ann.rs | 459 +++++++++++++++++++++++
crates/paimon/src/vindex/pkvector/bucket.rs | 554 ++++++++++++++++++++++++++++
crates/paimon/src/vindex/pkvector/exact.rs | 289 +++++++++++++++
crates/paimon/src/vindex/pkvector/metric.rs | 211 +++++++++++
crates/paimon/src/vindex/pkvector/mod.rs | 42 +++
crates/paimon/src/vindex/pkvector/reader.rs | 80 ++++
crates/paimon/src/vindex/pkvector/result.rs | 27 ++
10 files changed, 1686 insertions(+), 2 deletions(-)
diff --git a/crates/paimon/src/deletion_vector/core.rs
b/crates/paimon/src/deletion_vector/core.rs
index 11c2d06..4164da6 100644
--- a/crates/paimon/src/deletion_vector/core.rs
+++ b/crates/paimon/src/deletion_vector/core.rs
@@ -59,6 +59,15 @@ impl DeletionVector {
self.bitmap.len()
}
+ /// Returns true if `position` is deleted. Positions above `u32::MAX`
cannot be
+ /// present in a roaring32 bitmap and are therefore reported as not
deleted.
+ /// Mirrors Java `BitmapDeletionVector#isDeleted` / the searchers'
`LongPredicate`.
+ pub fn is_deleted(&self, position: u64) -> bool {
+ u32::try_from(position)
+ .ok()
+ .is_some_and(|p| self.bitmap.contains(p))
+ }
+
/// Returns an iterator over deleted positions that supports
[DeletionVectorIterator::advance_to].
/// Required for efficient row selection building when skipping row groups
(avoid re-scanning
/// deletes in skipped ranges).
@@ -249,4 +258,15 @@ mod tests {
let expected_bitmap = RoaringBitmap::from_iter([1u32, 2u32]);
assert_eq!(dv.bitmap(), &expected_bitmap, "bitmap should be [1, 2]");
}
+
+ #[test]
+ fn test_is_deleted_reports_membership_and_guards_u32_overflow() {
+ let mut bitmap = RoaringBitmap::new();
+ bitmap.insert(2);
+ let dv = DeletionVector::from_bitmap(bitmap);
+ assert!(dv.is_deleted(2), "position 2 was deleted");
+ assert!(!dv.is_deleted(0), "position 0 was not deleted");
+ // Positions above u32::MAX cannot exist in a roaring32 bitmap -> not
deleted.
+ assert!(!dv.is_deleted(u64::from(u32::MAX) + 1));
+ }
}
diff --git a/crates/paimon/src/spec/pk_vector_source.rs
b/crates/paimon/src/spec/pk_vector_source.rs
index 2685f45..2c859e0 100644
--- a/crates/paimon/src/spec/pk_vector_source.rs
+++ b/crates/paimon/src/spec/pk_vector_source.rs
@@ -157,8 +157,8 @@ impl PkVectorSourceMeta {
}
/// Minimal big-endian reader mirroring the Java `DataInput` primitives the
-/// `_SOURCE_META` frame uses. Module-private by design (see the PR1 spec: no
-/// shared `common/` abstraction until a second consumer exists).
+/// `_SOURCE_META` frame uses. Module-private by design: no shared `common/`
+/// abstraction until a second consumer exists.
struct DataInputCursor<'a> {
bytes: &'a [u8],
position: usize,
diff --git a/crates/paimon/src/vindex/mod.rs b/crates/paimon/src/vindex/mod.rs
index c0d566a..77457ef 100644
--- a/crates/paimon/src/vindex/mod.rs
+++ b/crates/paimon/src/vindex/mod.rs
@@ -17,6 +17,8 @@
pub mod reader;
+pub mod pkvector;
+
use crate::spec::{DataField, DataType};
use paimon_vindex_core::index::VectorIndexConfig;
use std::collections::HashMap;
diff --git a/crates/paimon/src/vindex/pkvector/ann.rs
b/crates/paimon/src/vindex/pkvector/ann.rs
new file mode 100644
index 0000000..f87e7f9
--- /dev/null
+++ b/crates/paimon/src/vindex/pkvector/ann.rs
@@ -0,0 +1,459 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::collections::{HashMap, HashSet};
+use std::sync::Arc;
+
+use super::bucket::BucketAnnSegment;
+use super::data_invalid;
+use super::metric::VectorSearchMetric;
+use super::result::PkVectorSearchResult;
+use crate::deletion_vector::DeletionVector;
+use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta};
+use crate::vector_search::VectorSearch;
+
+/// Build the live-row-id mask for the ANN reader's `include_row_ids` filter,
in
+/// segment-ordinal space (source files concatenated in order). Mirrors Java
+/// `PkVectorAnnSegmentSearcher.liveRowPositions`.
+///
+/// Only source files present in `active_source_files` contribute live
ordinals;
+/// inactive sources' ordinal ranges are masked out entirely (their rows are no
+/// longer readable in this snapshot). Deletion vectors are applied only to
active
+/// sources.
+///
+/// Returns `None` only when every source file is active AND no deletion
vector is
+/// relevant — nothing to mask. Otherwise returns the masked live ids.
+pub(crate) fn build_live_row_ids(
+ source_files: &[PkVectorSourceFile],
+ active_source_files: &HashSet<String>,
+ deletion_vectors: &HashMap<String, Arc<DeletionVector>>,
+) -> crate::Result<Option<roaring::RoaringTreemap>> {
+ let all_active = source_files
+ .iter()
+ .all(|f| active_source_files.contains(f.file_name()));
+ let has_relevant_dv = source_files
+ .iter()
+ .any(|f| deletion_vectors.contains_key(f.file_name()));
+ if all_active && !has_relevant_dv {
+ return Ok(None);
+ }
+
+ let mut live = roaring::RoaringTreemap::new();
+ let mut deleted = roaring::RoaringTreemap::new();
+ let mut file_offset: u64 = 0;
+ for source_file in source_files {
+ let row_count = u64::try_from(source_file.row_count())
+ .map_err(|_| data_invalid("vector source row count must not be
negative"))?;
+ let end = file_offset
+ .checked_add(row_count)
+ .ok_or_else(|| data_invalid("vector source row counts overflow
u64"))?;
+ let active = active_source_files.contains(source_file.file_name());
+ if active && row_count > 0 {
+ live.insert_range(file_offset..end);
+ }
+ if active {
+ if let Some(dv) = deletion_vectors.get(source_file.file_name()) {
+ for position in dv.iter() {
+ let global =
file_offset.checked_add(position).ok_or_else(|| {
+ data_invalid("vector source deleted position overflows
u64")
+ })?;
+ deleted.insert(global);
+ }
+ }
+ }
+ file_offset = end;
+ }
+ live -= deleted;
+ Ok(Some(live))
+}
+
+/// Map ANN `(ordinal, score)` pairs to physical `(data file, position)`
results,
+/// validating ordinals against source metadata, rejecting hits that resolve
to an
+/// inactive source file, and rejecting hits on snapshot-deleted rows. Mirrors
the
+/// post-processing loop of Java `PkVectorAnnSegmentSearcher.search`. Results
are
+/// sorted BEST_FIRST.
+pub(crate) fn map_ann_results(
+ scored: &[(u64, f32)],
+ source_meta: &PkVectorSourceMeta,
+ active_source_files: &HashSet<String>,
+ deletion_vectors: &HashMap<String, Arc<DeletionVector>>,
+ metric: VectorSearchMetric,
+) -> crate::Result<Vec<PkVectorSearchResult>> {
+ let mut results = Vec::with_capacity(scored.len());
+ for &(ordinal, score) in scored {
+ let ordinal_i64 = i64::try_from(ordinal)
+ .map_err(|_| data_invalid(format!("ANN ordinal {ordinal} exceeds
i64::MAX")))?;
+ let (data_file_name, row_position) = source_meta.resolve(ordinal_i64)?;
+ if !active_source_files.contains(&data_file_name) {
+ return Err(data_invalid(format!(
+ "ANN segment returned inactive source {data_file_name}"
+ )));
+ }
+ if let Some(dv) = deletion_vectors.get(&data_file_name) {
+ let pos = u64::try_from(row_position)
+ .map_err(|_| data_invalid("resolved row position must not be
negative"))?;
+ if dv.is_deleted(pos) {
+ return Err(data_invalid(format!(
+ "ANN segment returned snapshot-deleted row position
{row_position} in {data_file_name}"
+ )));
+ }
+ }
+ results.push(PkVectorSearchResult {
+ data_file_name,
+ row_position,
+ distance: metric.score_to_distance(score),
+ });
+ }
+ results.sort_by(|a, b| {
+ a.distance
+ .total_cmp(&b.distance)
+ .then_with(|| a.data_file_name.cmp(&b.data_file_name))
+ .then_with(|| a.row_position.cmp(&b.row_position))
+ });
+ Ok(results)
+}
+
+/// One ANN segment's search dependency for the bucket kernel. Bucket tests
fake
+/// this (mirroring Java's mock of `PkVectorAnnSegmentSearcher`).
+pub(crate) trait PkVectorAnnSearcher {
+ #[allow(clippy::too_many_arguments)]
+ fn search(
+ &self,
+ segment: &BucketAnnSegment,
+ query: &[f32],
+ metric: VectorSearchMetric,
+ limit: usize,
+ active_source_files: &HashSet<String>,
+ deletion_vectors: &HashMap<String, Arc<DeletionVector>>,
+ search_options: &HashMap<String, String>,
+ ) -> crate::Result<Vec<PkVectorSearchResult>>;
+}
+
+/// Scorer seam: drives the underlying vindex ANN reader. Returns `ordinal ->
+/// score` (higher-is-better). Any negative labels are skipped by the existing
+/// `vindex` reader (`collect_results` drops `row_id < 0`), so this seam only
+/// ever yields non-negative `u64` ordinals — no signed-label handling is
needed
+/// downstream.
+///
+/// The production scorer drives
`VindexVectorGlobalIndexReader::visit_vector_search`
+/// with a segment's index bytes; tests inject a synthetic scorer. The
adapter's
+/// own logic (live-row masking, ordinal mapping, deletion checks, ordering) is
+/// exercised independently of the scorer.
+type Scorer = Box<dyn Fn(&VectorSearch) -> crate::Result<Option<HashMap<u64,
f32>>>>;
+
+/// Structural vindex-backed `PkVectorAnnSearcher`. Composes the pure helpers
+/// (`build_live_row_ids`, `map_ann_results`) around the scorer seam.
+pub(crate) struct VindexAnnSearcher {
+ field_name: String,
+ scorer: Scorer,
+}
+
+impl VindexAnnSearcher {
+ pub(crate) fn new(field_name: String, scorer: Scorer) -> Self {
+ Self { field_name, scorer }
+ }
+}
+
+impl PkVectorAnnSearcher for VindexAnnSearcher {
+ fn search(
+ &self,
+ segment: &BucketAnnSegment,
+ query: &[f32],
+ metric: VectorSearchMetric,
+ limit: usize,
+ active_source_files: &HashSet<String>,
+ deletion_vectors: &HashMap<String, Arc<DeletionVector>>,
+ search_options: &HashMap<String, String>,
+ ) -> crate::Result<Vec<PkVectorSearchResult>> {
+ if limit == 0 {
+ return Err(data_invalid("vector search limit must be positive"));
+ }
+ let source_files = segment.source_meta.source_files();
+ let mut search = VectorSearch::new(query.to_vec(), limit,
self.field_name.clone())?
+ .with_options(search_options.clone());
+ if let Some(live) = build_live_row_ids(source_files,
active_source_files, deletion_vectors)?
+ {
+ search = search.with_include_row_ids(live);
+ }
+ let scored = match (self.scorer)(&search)? {
+ Some(map) => map,
+ None => return Ok(Vec::new()),
+ };
+ let scored: Vec<(u64, f32)> = scored.into_iter().collect();
+ map_ann_results(
+ &scored,
+ &segment.source_meta,
+ active_source_files,
+ deletion_vectors,
+ metric,
+ )
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use roaring::RoaringBitmap;
+
+ fn source_meta(files: &[(&str, i64)]) -> PkVectorSourceMeta {
+ let files = files
+ .iter()
+ .map(|(name, rows)| PkVectorSourceFile::new((*name).to_string(),
*rows).unwrap())
+ .collect();
+ PkVectorSourceMeta::new(files).unwrap()
+ }
+
+ fn dv(deleted: &[u32]) -> Arc<DeletionVector> {
+ let mut bitmap = RoaringBitmap::new();
+ for &p in deleted {
+ bitmap.insert(p);
+ }
+ Arc::new(DeletionVector::from_bitmap(bitmap))
+ }
+
+ fn active_set(names: &[&str]) -> HashSet<String> {
+ names.iter().map(|n| (*n).to_string()).collect()
+ }
+
+ #[test]
+ fn test_build_live_row_ids_none_when_all_active_and_no_relevant_dv() {
+ let files = [PkVectorSourceFile::new("f0".into(), 3).unwrap()];
+ let active = active_set(&["f0"]);
+ // All active + empty map -> None.
+ assert!(build_live_row_ids(&files, &active, &HashMap::new())
+ .unwrap()
+ .is_none());
+ // All active + non-empty map but no matching file name -> None.
+ let mut dvs = HashMap::new();
+ dvs.insert("other".to_string(), dv(&[0]));
+ assert!(build_live_row_ids(&files, &active, &dvs).unwrap().is_none());
+ }
+
+ #[test]
+ fn test_build_live_row_ids_masks_inactive_source_ordinal_range() {
+ // f0 rows 0..3 (global 0,1,2), f1 rows 0..2 (global 3,4). f1 is
inactive,
+ // so its whole ordinal range is masked out; f0 stays fully live. No
DV.
+ let files = vec![
+ PkVectorSourceFile::new("f0".into(), 3).unwrap(),
+ PkVectorSourceFile::new("f1".into(), 2).unwrap(),
+ ];
+ let live = build_live_row_ids(&files, &active_set(&["f0"]),
&HashMap::new())
+ .unwrap()
+ .unwrap();
+ assert_eq!(live.iter().collect::<Vec<u64>>(), vec![0, 1, 2]);
+ }
+
+ #[test]
+ fn test_build_live_row_ids_masks_deleted_positions_with_file_offsets() {
+ // f0 rows 0..3 (global 0,1,2), f1 rows 0..2 (global 3,4).
+ let files = vec![
+ PkVectorSourceFile::new("f0".into(), 3).unwrap(),
+ PkVectorSourceFile::new("f1".into(), 2).unwrap(),
+ ];
+ let mut dvs = HashMap::new();
+ dvs.insert("f0".to_string(), dv(&[1])); // deletes global 1
+ dvs.insert("f1".to_string(), dv(&[0])); // deletes global 3
+ let live = build_live_row_ids(&files, &active_set(&["f0", "f1"]), &dvs)
+ .unwrap()
+ .unwrap();
+ assert_eq!(live.iter().collect::<Vec<u64>>(), vec![0, 2, 4]);
+ }
+
+ #[test]
+ fn test_map_ann_results_maps_ordinals_to_positions_and_scores() {
+ let meta = source_meta(&[("f0", 3), ("f1", 5)]);
+ // ordinal 3 -> (f1, 0); ordinal 0 -> (f0, 0). l2
score_to_distance(0.5)=1.0.
+ let scored = [(3u64, 0.5f32), (0u64, 0.5f32)];
+ let results = map_ann_results(
+ &scored,
+ &meta,
+ &active_set(&["f0", "f1"]),
+ &HashMap::new(),
+ VectorSearchMetric::L2,
+ )
+ .unwrap();
+ assert_eq!(
+ results,
+ vec![
+ PkVectorSearchResult {
+ data_file_name: "f0".into(),
+ row_position: 0,
+ distance: 1.0
+ },
+ PkVectorSearchResult {
+ data_file_name: "f1".into(),
+ row_position: 0,
+ distance: 1.0
+ },
+ ]
+ );
+ }
+
+ #[test]
+ fn test_map_ann_results_rejects_out_of_range_ordinal() {
+ let meta = source_meta(&[("f0", 3)]);
+ let err = map_ann_results(
+ &[(3u64, 0.5)],
+ &meta,
+ &active_set(&["f0"]),
+ &HashMap::new(),
+ VectorSearchMetric::L2,
+ )
+ .unwrap_err();
+ assert!(err.to_string().contains("out of range") ||
err.to_string().contains("ordinal"));
+ }
+
+ #[test]
+ fn test_map_ann_results_rejects_hit_resolving_to_inactive_source() {
+ // ordinal 3 resolves to f1, which is not in the active set -> error.
+ let meta = source_meta(&[("f0", 3), ("f1", 5)]);
+ let err = map_ann_results(
+ &[(3u64, 0.5)],
+ &meta,
+ &active_set(&["f0"]),
+ &HashMap::new(),
+ VectorSearchMetric::L2,
+ )
+ .unwrap_err();
+ assert!(err.to_string().contains("inactive"));
+ }
+
+ #[test]
+ fn test_map_ann_results_rejects_hit_on_deleted_position() {
+ let meta = source_meta(&[("f0", 3)]);
+ let mut dvs = HashMap::new();
+ dvs.insert("f0".to_string(), dv(&[1])); // position 1 deleted
+ let err = map_ann_results(
+ &[(1u64, 0.5)],
+ &meta,
+ &active_set(&["f0"]),
+ &dvs,
+ VectorSearchMetric::L2,
+ )
+ .unwrap_err();
+ assert!(err.to_string().contains("deleted"));
+ }
+
+ #[test]
+ fn test_vindex_adapter_composes_live_rows_and_maps_results() {
+ // Scorer records the VectorSearch it received and returns synthetic
ordinals.
+ // The scorer must be `'static`, so share the recording cells via `Rc`
moved
+ // into the closure rather than borrowing locals.
+ use std::cell::RefCell;
+ use std::rc::Rc;
+ let seen_limit = Rc::new(RefCell::new(0usize));
+ let seen_has_filter = Rc::new(RefCell::new(false));
+ let scorer_limit = Rc::clone(&seen_limit);
+ let scorer_has_filter = Rc::clone(&seen_has_filter);
+ let searcher = VindexAnnSearcher::new(
+ "embedding".to_string(),
+ Box::new(move |search: &VectorSearch| {
+ *scorer_limit.borrow_mut() = search.limit;
+ *scorer_has_filter.borrow_mut() =
search.include_row_ids.is_some();
+ let mut scores = HashMap::new();
+ scores.insert(3u64, 0.5f32); // -> (f1, 0)
+ scores.insert(0u64, 0.25f32); // -> (f0, 0), l2 dist 3.0
+ Ok(Some(scores))
+ }),
+ );
+ let segment = BucketAnnSegment {
+ source_meta: {
+ use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta};
+ PkVectorSourceMeta::new(vec![
+ PkVectorSourceFile::new("f0".into(), 3).unwrap(),
+ PkVectorSourceFile::new("f1".into(), 5).unwrap(),
+ ])
+ .unwrap()
+ },
+ };
+ let mut dvs = HashMap::new();
+ dvs.insert("f0".to_string(), dv(&[1]));
+ let results = searcher
+ .search(
+ &segment,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 2,
+ &active_set(&["f0", "f1"]),
+ &dvs,
+ &HashMap::new(),
+ )
+ .unwrap();
+ // Sorted BEST_FIRST by distance: (f1,0) dist 1.0 then (f0,0) dist 3.0.
+ assert_eq!(results[0].data_file_name, "f1");
+ assert_eq!(results[1].data_file_name, "f0");
+ assert_eq!(*seen_limit.borrow(), 2);
+ assert!(
+ *seen_has_filter.borrow(),
+ "DV present -> include_row_ids set"
+ );
+ }
+
+ #[test]
+ fn test_vindex_adapter_rejects_non_positive_limit() {
+ let searcher = VindexAnnSearcher::new(
+ "embedding".to_string(),
+ Box::new(|_: &VectorSearch| Ok(None)),
+ );
+ let segment = BucketAnnSegment {
+ source_meta: {
+ use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta};
+
PkVectorSourceMeta::new(vec![PkVectorSourceFile::new("f0".into(), 1).unwrap()])
+ .unwrap()
+ },
+ };
+ let err = searcher
+ .search(
+ &segment,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 0,
+ &active_set(&["f0"]),
+ &HashMap::new(),
+ &HashMap::new(),
+ )
+ .unwrap_err();
+ assert!(err.to_string().contains("positive"));
+ }
+
+ #[test]
+ fn test_vindex_adapter_empty_scorer_result_is_empty() {
+ let searcher = VindexAnnSearcher::new(
+ "embedding".to_string(),
+ Box::new(|_: &VectorSearch| Ok(None)),
+ );
+ let segment = BucketAnnSegment {
+ source_meta: {
+ use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta};
+
PkVectorSourceMeta::new(vec![PkVectorSourceFile::new("f0".into(), 1).unwrap()])
+ .unwrap()
+ },
+ };
+ let results = searcher
+ .search(
+ &segment,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 2,
+ &active_set(&["f0"]),
+ &HashMap::new(),
+ &HashMap::new(),
+ )
+ .unwrap();
+ assert!(results.is_empty());
+ }
+}
diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs
b/crates/paimon/src/vindex/pkvector/bucket.rs
new file mode 100644
index 0000000..23a2b1c
--- /dev/null
+++ b/crates/paimon/src/vindex/pkvector/bucket.rs
@@ -0,0 +1,554 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::cmp::Ordering;
+use std::collections::{BinaryHeap, HashMap, HashSet};
+use std::sync::Arc;
+
+use super::ann::PkVectorAnnSearcher;
+use super::data_invalid;
+use super::exact::exact_search;
+use super::metric::VectorSearchMetric;
+use super::reader::PkVectorReader;
+use super::result::PkVectorSearchResult;
+use crate::deletion_vector::DeletionVector;
+use crate::spec::PkVectorSourceMeta;
+
+/// One ANN segment to be searched by the bucket kernel: the source metadata
+/// resolving segment ordinals back to physical `(data file, position)`. Only
+/// `source_meta` is needed for ordinal mapping and live-row masking.
+pub(crate) struct BucketAnnSegment {
+ pub source_meta: PkVectorSourceMeta,
+}
+
+/// A data file participating in the bucket search, with its row count. Used by
+/// the bucket kernel to plan exact vs. ANN search over active files.
+pub(crate) struct BucketActiveFile {
+ pub file_name: String,
+ pub row_count: i64,
+}
+
+/// Total BEST_FIRST order over results: distance ASC, then data_file_name ASC,
+/// then row_position ASC. `total_cmp` keeps it NaN-safe and panic-free.
+fn best_first(a: &PkVectorSearchResult, b: &PkVectorSearchResult) -> Ordering {
+ a.distance
+ .total_cmp(&b.distance)
+ .then_with(|| a.data_file_name.cmp(&b.data_file_name))
+ .then_with(|| a.row_position.cmp(&b.row_position))
+}
+
+/// A candidate wrapped so a max-heap keeps the WORST (BEST_FIRST-largest)
+/// candidate on top; popping evicts the least-wanted one. Mirrors the
+/// `PriorityQueue<>(limit, BEST_FIRST.reversed())` in Java
+/// `PrimaryKeyVectorBucketSearch`.
+struct WorstFirst(PkVectorSearchResult);
+
+impl PartialEq for WorstFirst {
+ fn eq(&self, other: &Self) -> bool {
+ best_first(&self.0, &other.0) == Ordering::Equal
+ }
+}
+impl Eq for WorstFirst {}
+impl PartialOrd for WorstFirst {
+ fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+ Some(self.cmp(other))
+ }
+}
+impl Ord for WorstFirst {
+ fn cmp(&self, other: &Self) -> Ordering {
+ best_first(&self.0, &other.0)
+ }
+}
+
+/// Add `candidate` to a bounded (size `limit`) BEST_FIRST Top-K max-heap:
push if
+/// under capacity, else replace the current worst iff the candidate beats it.
+/// `O(log limit)` per call. Mirrors Java `PrimaryKeyVectorBucketSearch.add`.
+fn add_candidate(heap: &mut BinaryHeap<WorstFirst>, candidate:
PkVectorSearchResult, limit: usize) {
+ if heap.len() < limit {
+ heap.push(WorstFirst(candidate));
+ } else if heap
+ .peek()
+ .is_some_and(|worst| best_first(&candidate, &worst.0) ==
Ordering::Less)
+ {
+ heap.pop();
+ heap.push(WorstFirst(candidate));
+ }
+}
+
+/// ANN + exact data-file fallback search for one snapshot bucket. Mirrors Java
+/// `org.apache.paimon.index.pkvector.PrimaryKeyVectorBucketSearch.search`.
+///
+/// `ann_searcher` may be `None` only when there are no ANN segments; segments
+/// present with `None` is an error.
+#[allow(clippy::too_many_arguments)]
+pub(crate) fn bucket_search(
+ ann_searcher: Option<&dyn PkVectorAnnSearcher>,
+ ann_segments: &[BucketAnnSegment],
+ active_files: &[BucketActiveFile],
+ deletion_vectors: &HashMap<String, Arc<DeletionVector>>,
+ exact_reader_factory: &mut dyn FnMut(
+ &BucketActiveFile,
+ ) -> crate::Result<Box<dyn PkVectorReader>>,
+ query: &[f32],
+ metric: VectorSearchMetric,
+ limit: usize,
+ search_options: &HashMap<String, String>,
+) -> crate::Result<Vec<PkVectorSearchResult>> {
+ if limit == 0 {
+ return Err(data_invalid("vector search limit must be positive"));
+ }
+
+ let mut files_by_name: HashMap<&str, &BucketActiveFile> = HashMap::new();
+ for file in active_files {
+ if file.row_count < 0 {
+ return Err(data_invalid(format!(
+ "active data file {} row count must not be negative: {}",
+ file.file_name, file.row_count
+ )));
+ }
+ if files_by_name
+ .insert(file.file_name.as_str(), file)
+ .is_some()
+ {
+ return Err(data_invalid(format!(
+ "duplicate data file: {}",
+ file.file_name
+ )));
+ }
+ }
+
+ let mut heap: BinaryHeap<WorstFirst> = BinaryHeap::with_capacity(limit +
1);
+ let active_source_files: HashSet<String> =
+ files_by_name.keys().map(|name| name.to_string()).collect();
+ let mut covered: HashSet<String> = HashSet::new();
+
+ for segment in ann_segments {
+ for source in segment.source_meta.source_files() {
+ // An ANN source that is no longer an active file (e.g. compacted
away)
+ // is skipped, not rejected: its ordinal range is masked out of
the ANN
+ // live-row bitmap and the remaining active sources are still
searched.
+ // Mirrors Java master `PrimaryKeyVectorBucketSearch` (`file ==
null`
+ // -> continue). Active sources still require a row-count match.
+ match files_by_name.get(source.file_name()) {
+ Some(active) if active.row_count == source.row_count() => {
+ covered.insert(source.file_name().to_string());
+ }
+ Some(_) => {
+ return Err(data_invalid(format!(
+ "ANN source {} does not match the active data file",
+ source.file_name()
+ )));
+ }
+ None => continue,
+ }
+ }
+ let searcher = ann_searcher.ok_or_else(|| data_invalid("ANN search is
not configured"))?;
+ for result in searcher.search(
+ segment,
+ query,
+ metric,
+ limit,
+ &active_source_files,
+ deletion_vectors,
+ search_options,
+ )? {
+ add_candidate(&mut heap, result, limit);
+ }
+ }
+
+ for file in active_files {
+ if covered.contains(&file.file_name) {
+ continue;
+ }
+ let dv = deletion_vectors.get(&file.file_name).cloned();
+ let is_excluded = move |position: i64| -> bool {
+ match &dv {
+ Some(dv) => u64::try_from(position)
+ .map(|p| dv.is_deleted(p))
+ .unwrap_or(false),
+ None => false,
+ }
+ };
+ let mut reader = exact_reader_factory(file)?;
+ for result in exact_search(
+ &file.file_name,
+ reader.as_mut(),
+ query,
+ metric,
+ limit,
+ &is_excluded,
+ )? {
+ add_candidate(&mut heap, result, limit);
+ }
+ }
+
+ let mut results: Vec<PkVectorSearchResult> = heap.into_iter().map(|w|
w.0).collect();
+ results.sort_by(best_first);
+ Ok(results)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::spec::PkVectorSourceFile;
+ use crate::vindex::pkvector::ann::PkVectorAnnSearcher;
+ use crate::vindex::pkvector::reader::test_support::ArrayReader;
+ use roaring::RoaringBitmap;
+ use std::cell::RefCell;
+
+ fn meta(files: &[(&str, i64)]) -> PkVectorSourceMeta {
+ PkVectorSourceMeta::new(
+ files
+ .iter()
+ .map(|(n, r)| PkVectorSourceFile::new((*n).into(),
*r).unwrap())
+ .collect(),
+ )
+ .unwrap()
+ }
+
+ fn active(name: &str, rows: i64) -> BucketActiveFile {
+ BucketActiveFile {
+ file_name: name.into(),
+ row_count: rows,
+ }
+ }
+
+ /// Fake ANN searcher returning preset results and recording calls.
+ struct FakeAnnSearcher {
+ result: Vec<PkVectorSearchResult>,
+ }
+ impl PkVectorAnnSearcher for FakeAnnSearcher {
+ fn search(
+ &self,
+ _segment: &BucketAnnSegment,
+ _query: &[f32],
+ _metric: VectorSearchMetric,
+ _limit: usize,
+ _active_source_files: &HashSet<String>,
+ _dvs: &HashMap<String, Arc<DeletionVector>>,
+ _opts: &HashMap<String, String>,
+ ) -> crate::Result<Vec<PkVectorSearchResult>> {
+ Ok(self.result.clone())
+ }
+ }
+
+ #[test]
+ fn test_rejects_non_positive_limit() {
+ let mut factory =
+ |_: &BucketActiveFile| -> crate::Result<Box<dyn PkVectorReader>> {
unreachable!() };
+ let err = bucket_search(
+ None,
+ &[],
+ &[],
+ &HashMap::new(),
+ &mut factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 0,
+ &HashMap::new(),
+ )
+ .unwrap_err();
+ assert!(err.to_string().contains("positive"));
+ }
+
+ #[test]
+ fn test_bounded_heap_evicts_by_best_first_tiebreak_over_limit() {
+ // All candidates share distance 1.0, so eviction is decided purely by
the
+ // BEST_FIRST tie-break (data_file_name ASC, then row_position ASC).
Feed
+ // more than `limit` ANN hits and assert the kept set is the smallest
+ // (file, position) pairs in that order. Locks the bounded-heap merge.
+ let segment = BucketAnnSegment {
+ source_meta: meta(&[("data-1", 3)]),
+ };
+ let hit = |file: &str, pos: i64| PkVectorSearchResult {
+ data_file_name: file.into(),
+ row_position: pos,
+ distance: 1.0,
+ };
+ // Deliberately unsorted input across two files at the same distance.
+ let ann = FakeAnnSearcher {
+ result: vec![
+ hit("data-2", 0),
+ hit("data-1", 2),
+ hit("data-1", 0),
+ hit("data-2", 1),
+ hit("data-1", 1),
+ ],
+ };
+ let mut factory =
+ |_: &BucketActiveFile| -> crate::Result<Box<dyn PkVectorReader>> {
unreachable!() };
+ let results = bucket_search(
+ Some(&ann),
+ &[segment],
+ &[active("data-1", 3)],
+ &HashMap::new(),
+ &mut factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 3,
+ &HashMap::new(),
+ )
+ .unwrap();
+ // 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!(
+ results
+ .iter()
+ .map(|r| (r.data_file_name.as_str(), r.row_position))
+ .collect::<Vec<_>>(),
+ vec![("data-1", 0), ("data-1", 1), ("data-1", 2)]
+ );
+ }
+
+ #[test]
+ fn test_merges_ann_and_exact_without_rescanning_covered_files() {
+ // data-1 is ANN-covered; data-2 is exact fallback. Factory must never
be
+ // called for data-1.
+ let segment = BucketAnnSegment {
+ source_meta: meta(&[("data-1", 2)]),
+ };
+ let ann = FakeAnnSearcher {
+ result: vec![PkVectorSearchResult {
+ data_file_name: "data-1".into(),
+ row_position: 1,
+ distance: 0.5,
+ }],
+ };
+ let calls = RefCell::new(Vec::<String>::new());
+ let mut factory = |f: &BucketActiveFile| -> crate::Result<Box<dyn
PkVectorReader>> {
+ calls.borrow_mut().push(f.file_name.clone());
+ // data-2 vectors: pos0 {1,0} dist 1.0, pos1 {3,0} dist 9.0
+ Ok(Box::new(ArrayReader::new(
+ 2,
+ vec![Some(vec![1.0, 0.0]), Some(vec![3.0, 0.0])],
+ )))
+ };
+ let results = bucket_search(
+ Some(&ann),
+ &[segment],
+ &[active("data-1", 2), active("data-2", 2)],
+ &HashMap::new(),
+ &mut factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 2,
+ &HashMap::new(),
+ )
+ .unwrap();
+ assert_eq!(
+ results,
+ vec![
+ PkVectorSearchResult {
+ data_file_name: "data-1".into(),
+ row_position: 1,
+ distance: 0.5
+ },
+ PkVectorSearchResult {
+ data_file_name: "data-2".into(),
+ row_position: 0,
+ distance: 1.0
+ },
+ ]
+ );
+ assert_eq!(calls.borrow().as_slice(), &["data-2".to_string()]);
+ }
+
+ #[test]
+ fn test_exact_fallback_merges_files_and_applies_deletion_vectors() {
+ // No ANN. data-1 pos0 {0,0} deleted; remaining candidates merge
across files.
+ let calls = RefCell::new(0);
+ let mut factory = |f: &BucketActiveFile| -> crate::Result<Box<dyn
PkVectorReader>> {
+ *calls.borrow_mut() += 1;
+ let vectors = match f.file_name.as_str() {
+ "data-1" => vec![Some(vec![0.0, 0.0]), Some(vec![2.0, 0.0])],
+ "data-2" => vec![Some(vec![1.0, 0.0]), None],
+ _ => unreachable!(),
+ };
+ Ok(Box::new(ArrayReader::new(2, vectors)))
+ };
+ let mut dvs: HashMap<String, Arc<DeletionVector>> = HashMap::new();
+ let mut bm = RoaringBitmap::new();
+ bm.insert(0); // data-1 position 0 deleted
+ dvs.insert("data-1".into(), Arc::new(DeletionVector::from_bitmap(bm)));
+
+ let results = bucket_search(
+ None,
+ &[],
+ &[active("data-1", 2), active("data-2", 2)],
+ &dvs,
+ &mut factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 2,
+ &HashMap::new(),
+ )
+ .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.)
+ assert_eq!(
+ results,
+ vec![
+ PkVectorSearchResult {
+ data_file_name: "data-2".into(),
+ row_position: 0,
+ distance: 1.0
+ },
+ PkVectorSearchResult {
+ data_file_name: "data-1".into(),
+ row_position: 1,
+ distance: 4.0
+ },
+ ]
+ );
+ }
+
+ #[test]
+ fn test_rejects_duplicate_active_file_name() {
+ let mut factory =
+ |_: &BucketActiveFile| -> crate::Result<Box<dyn PkVectorReader>> {
unreachable!() };
+ let err = bucket_search(
+ None,
+ &[],
+ &[active("dup", 1), active("dup", 1)],
+ &HashMap::new(),
+ &mut factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 1,
+ &HashMap::new(),
+ )
+ .unwrap_err();
+ assert!(err.to_string().contains("duplicate") ||
err.to_string().contains("Duplicate"));
+ }
+
+ #[test]
+ fn test_rejects_ann_source_row_count_mismatch_for_active_file() {
+ let ann = FakeAnnSearcher { result: vec![] };
+ // Segment references data-1 with 2 rows, but the active file has 3
rows.
+ // An active source with a mismatched row count is still a hard error.
+ let segment = BucketAnnSegment {
+ source_meta: meta(&[("data-1", 2)]),
+ };
+ let mut factory =
+ |_: &BucketActiveFile| -> crate::Result<Box<dyn PkVectorReader>> {
unreachable!() };
+ let err = bucket_search(
+ Some(&ann),
+ &[segment],
+ &[active("data-1", 3)],
+ &HashMap::new(),
+ &mut factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 1,
+ &HashMap::new(),
+ )
+ .unwrap_err();
+ assert!(
+ err.to_string().contains("does not match") ||
err.to_string().contains("ANN source")
+ );
+ }
+
+ #[test]
+ fn test_skips_inactive_ann_source_and_searches_active_ones() {
+ // Segment covers [data-1, data-2] but only data-1 is still active
+ // (data-2 was compacted away). Java master skips the inactive source
+ // instead of failing the whole query; data-2 is neither covered (so it
+ // is not treated as ANN-covered) nor an active file (so it is not
exact
+ // scanned). The ANN searcher still runs for the segment.
+ let segment = BucketAnnSegment {
+ source_meta: meta(&[("data-1", 2), ("data-2", 2)]),
+ };
+ let ann = FakeAnnSearcher {
+ result: vec![PkVectorSearchResult {
+ data_file_name: "data-1".into(),
+ row_position: 0,
+ distance: 0.5,
+ }],
+ };
+ let calls = RefCell::new(Vec::<String>::new());
+ let mut factory = |f: &BucketActiveFile| -> crate::Result<Box<dyn
PkVectorReader>> {
+ calls.borrow_mut().push(f.file_name.clone());
+ unreachable!("only data-1 is active and it is ANN-covered")
+ };
+ let results = bucket_search(
+ Some(&ann),
+ &[segment],
+ &[active("data-1", 2)],
+ &HashMap::new(),
+ &mut factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 2,
+ &HashMap::new(),
+ )
+ .unwrap();
+ assert_eq!(
+ results,
+ vec![PkVectorSearchResult {
+ data_file_name: "data-1".into(),
+ row_position: 0,
+ distance: 0.5
+ }]
+ );
+ // No exact fallback ran: data-1 is ANN-covered, data-2 is not active.
+ assert!(calls.borrow().is_empty());
+ }
+
+ #[test]
+ fn test_rejects_segments_without_ann_searcher() {
+ let segment = BucketAnnSegment {
+ source_meta: meta(&[("data-1", 2)]),
+ };
+ let mut factory =
+ |_: &BucketActiveFile| -> crate::Result<Box<dyn PkVectorReader>> {
unreachable!() };
+ let err = bucket_search(
+ None,
+ &[segment],
+ &[active("data-1", 2)],
+ &HashMap::new(),
+ &mut factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 1,
+ &HashMap::new(),
+ )
+ .unwrap_err();
+ assert!(
+ err.to_string().contains("ANN search is not configured")
+ || err.to_string().contains("not configured")
+ );
+ }
+
+ #[test]
+ fn test_negative_active_row_count_rejected() {
+ let mut factory =
+ |_: &BucketActiveFile| -> crate::Result<Box<dyn PkVectorReader>> {
unreachable!() };
+ let err = bucket_search(
+ None,
+ &[],
+ &[active("data-1", -1)],
+ &HashMap::new(),
+ &mut factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 1,
+ &HashMap::new(),
+ )
+ .unwrap_err();
+ assert!(err.to_string().contains("row count") ||
err.to_string().contains("-1"));
+ }
+}
diff --git a/crates/paimon/src/vindex/pkvector/exact.rs
b/crates/paimon/src/vindex/pkvector/exact.rs
new file mode 100644
index 0000000..c9c0480
--- /dev/null
+++ b/crates/paimon/src/vindex/pkvector/exact.rs
@@ -0,0 +1,289 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::cmp::Ordering;
+use std::collections::BinaryHeap;
+
+use super::data_invalid;
+use super::metric::VectorSearchMetric;
+use super::reader::PkVectorReader;
+use super::result::PkVectorSearchResult;
+
+/// A candidate wrapped so a max-heap keeps the WORST candidate on top:
+/// worst = largest distance, ties broken by largest row_position. Popping the
+/// top therefore evicts the least-wanted candidate. Uses `total_cmp` for a
+/// deterministic total order over f32 (NaN-safe, no panic).
+struct WorstFirst(PkVectorSearchResult);
+
+impl PartialEq for WorstFirst {
+ fn eq(&self, other: &Self) -> bool {
+ self.cmp(other) == Ordering::Equal
+ }
+}
+impl Eq for WorstFirst {}
+impl PartialOrd for WorstFirst {
+ fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+ Some(self.cmp(other))
+ }
+}
+impl Ord for WorstFirst {
+ fn cmp(&self, other: &Self) -> Ordering {
+ self.0
+ .distance
+ .total_cmp(&other.0.distance)
+ .then_with(|| self.0.row_position.cmp(&other.0.row_position))
+ }
+}
+
+/// True if `candidate` ranks strictly better (BEST_FIRST) than the current
+/// worst-on-heap `weakest`: smaller distance, ties broken by smaller position.
+fn is_better_than(candidate: &PkVectorSearchResult, weakest:
&PkVectorSearchResult) -> bool {
+ candidate
+ .distance
+ .total_cmp(&weakest.distance)
+ .then_with(|| candidate.row_position.cmp(&weakest.row_position))
+ == Ordering::Less
+}
+
+/// Exact Top-K over one sequential physical-row vector source. Mirrors Java
+/// `PkVectorExactSearcher.search`. Results are sorted BEST_FIRST: distance
ASC,
+/// then row_position ASC (single file, so data_file_name is constant).
+pub(crate) fn exact_search(
+ data_file_name: &str,
+ reader: &mut dyn PkVectorReader,
+ query: &[f32],
+ metric: VectorSearchMetric,
+ limit: usize,
+ is_excluded: &dyn Fn(i64) -> bool,
+) -> crate::Result<Vec<PkVectorSearchResult>> {
+ if query.len() != reader.dimension() {
+ return Err(data_invalid(format!(
+ "query vector dimension does not match: index expects {}, got {}",
+ reader.dimension(),
+ query.len()
+ )));
+ }
+ if limit == 0 {
+ return Err(data_invalid("vector search limit must be positive"));
+ }
+ if let Some(i) = query.iter().position(|v| !v.is_finite()) {
+ return Err(data_invalid(format!(
+ "query vector element at position {i} must be finite"
+ )));
+ }
+ let row_count = reader.row_count();
+ if row_count < 0 {
+ return Err(data_invalid(format!(
+ "vector reader row count must not be negative: {row_count}"
+ )));
+ }
+
+ let mut reuse = vec![0.0f32; reader.dimension()];
+ let mut heap: BinaryHeap<WorstFirst> = BinaryHeap::with_capacity(limit +
1);
+ for position in 0..row_count {
+ let has_vector = reader.read_next_vector(&mut reuse)?;
+ if !has_vector || is_excluded(position) {
+ continue;
+ }
+ let candidate = PkVectorSearchResult {
+ data_file_name: data_file_name.to_string(),
+ row_position: position,
+ distance: metric.compute_distance(query, &reuse),
+ };
+ if heap.len() < limit {
+ heap.push(WorstFirst(candidate));
+ } else if heap
+ .peek()
+ .is_some_and(|worst| is_better_than(&candidate, &worst.0))
+ {
+ heap.pop();
+ heap.push(WorstFirst(candidate));
+ }
+ }
+
+ let mut results: Vec<PkVectorSearchResult> = heap.into_iter().map(|w|
w.0).collect();
+ results.sort_by(|a, b| {
+ a.distance
+ .total_cmp(&b.distance)
+ .then_with(|| a.row_position.cmp(&b.row_position))
+ });
+ Ok(results)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::vindex::pkvector::metric::VectorSearchMetric;
+ use crate::vindex::pkvector::reader::test_support::ArrayReader;
+
+ fn no_exclusion() -> impl Fn(i64) -> bool {
+ |_| false
+ }
+
+ #[test]
+ fn test_distances_for_supported_metrics() {
+ // Java testDistancesForSupportedMetrics: q=[2,0] over stored [1,0].
+ for (metric, expected) in [
+ (VectorSearchMetric::L2, 1.0f32),
+ (VectorSearchMetric::Cosine, 0.0),
+ (VectorSearchMetric::InnerProduct, -2.0),
+ ] {
+ let mut reader = ArrayReader::new(2, vec![Some(vec![1.0, 0.0])]);
+ let results = exact_search(
+ "data-file",
+ &mut reader,
+ &[2.0, 0.0],
+ metric,
+ 1,
+ &no_exclusion(),
+ )
+ .unwrap();
+ assert_eq!(results[0].distance, expected);
+ }
+ }
+
+ #[test]
+ fn test_rejects_dimension_mismatch() {
+ let mut reader = ArrayReader::new(2, vec![Some(vec![1.0, 0.0])]);
+ let err = exact_search(
+ "data-file",
+ &mut reader,
+ &[1.0],
+ VectorSearchMetric::L2,
+ 1,
+ &no_exclusion(),
+ )
+ .unwrap_err();
+ assert!(err.to_string().contains("dimension"));
+ }
+
+ #[test]
+ fn test_rejects_non_positive_limit() {
+ let mut reader = ArrayReader::new(2, vec![Some(vec![1.0, 0.0])]);
+ let err = exact_search(
+ "data-file",
+ &mut reader,
+ &[1.0, 0.0],
+ VectorSearchMetric::L2,
+ 0,
+ &no_exclusion(),
+ )
+ .unwrap_err();
+ assert!(err.to_string().contains("positive"));
+ }
+
+ #[test]
+ fn test_rejects_non_finite_query() {
+ let mut reader = ArrayReader::new(2, vec![Some(vec![1.0, 0.0])]);
+ let err = exact_search(
+ "data-file",
+ &mut reader,
+ &[f32::NAN, 0.0],
+ VectorSearchMetric::L2,
+ 1,
+ &no_exclusion(),
+ )
+ .unwrap_err();
+ assert!(err.to_string().contains("finite"));
+ }
+
+ #[test]
+ fn test_rejects_negative_row_count() {
+ struct NegativeReader;
+ impl PkVectorReader for NegativeReader {
+ fn dimension(&self) -> usize {
+ 2
+ }
+ fn row_count(&self) -> i64 {
+ -1
+ }
+ fn read_next_vector(&mut self, _reuse: &mut [f32]) ->
crate::Result<bool> {
+ unreachable!()
+ }
+ }
+ let mut reader = NegativeReader;
+ let err = exact_search(
+ "data-file",
+ &mut reader,
+ &[1.0, 0.0],
+ VectorSearchMetric::L2,
+ 1,
+ &no_exclusion(),
+ )
+ .unwrap_err();
+ assert!(err.to_string().contains("row count") ||
err.to_string().contains("-1"));
+ }
+
+ #[test]
+ fn test_preserves_null_and_excluded_physical_positions() {
+ // Java testPreservesNullAndDeletedPhysicalPositions:
+ // vectors [{3,0}, null, {1,0}, {2,0}], q=[0,0], l2, limit 2, exclude
pos==2.
+ let mut reader = ArrayReader::new(
+ 2,
+ vec![
+ Some(vec![3.0, 0.0]),
+ None,
+ Some(vec![1.0, 0.0]),
+ Some(vec![2.0, 0.0]),
+ ],
+ );
+ let results = exact_search(
+ "data-file",
+ &mut reader,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 2,
+ &|pos| pos == 2,
+ )
+ .unwrap();
+ assert_eq!(
+ results,
+ vec![
+ PkVectorSearchResult {
+ data_file_name: "data-file".into(),
+ row_position: 3,
+ distance: 4.0
+ },
+ PkVectorSearchResult {
+ data_file_name: "data-file".into(),
+ row_position: 0,
+ distance: 9.0
+ },
+ ]
+ );
+ }
+
+ #[test]
+ fn test_tie_break_prefers_smaller_row_position() {
+ // Equal distances -> smaller row_position ranks first.
+ let mut reader =
+ ArrayReader::new(1, vec![Some(vec![1.0]), Some(vec![1.0]),
Some(vec![1.0])]);
+ let results = exact_search(
+ "data-file",
+ &mut reader,
+ &[0.0],
+ VectorSearchMetric::L2,
+ 2,
+ &no_exclusion(),
+ )
+ .unwrap();
+ assert_eq!(
+ results.iter().map(|r| r.row_position).collect::<Vec<_>>(),
+ vec![0, 1]
+ );
+ }
+}
diff --git a/crates/paimon/src/vindex/pkvector/metric.rs
b/crates/paimon/src/vindex/pkvector/metric.rs
new file mode 100644
index 0000000..1dda56c
--- /dev/null
+++ b/crates/paimon/src/vindex/pkvector/metric.rs
@@ -0,0 +1,211 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use super::data_invalid;
+
+/// Normalize a metric name: lowercase and `-` → `_`. NO trim (deliberately
+/// stricter than the build-side `vindex::normalize_metric`, to match Java
+/// `VectorSearchMetric.normalize`).
+pub(crate) fn normalize_metric(metric: &str) -> String {
+ metric.to_ascii_lowercase().replace('-', "_")
+}
+
+/// True if the (normalized) metric is one of the three supported metrics.
+pub(crate) fn is_supported_metric(metric: &str) -> bool {
+ matches!(
+ normalize_metric(metric).as_str(),
+ "l2" | "cosine" | "inner_product"
+ )
+}
+
+/// Numeric semantics for a supported vector search metric. Mirrors Java
+/// `org.apache.paimon.globalindex.VectorSearchMetric`.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub(crate) enum VectorSearchMetric {
+ L2,
+ Cosine,
+ InnerProduct,
+}
+
+impl VectorSearchMetric {
+ /// Normalize, validate, and map to the enum. Errors on an unsupported
metric.
+ pub(crate) fn parse(metric: &str) -> crate::Result<Self> {
+ match normalize_metric(metric).as_str() {
+ "l2" => Ok(Self::L2),
+ "cosine" => Ok(Self::Cosine),
+ "inner_product" => Ok(Self::InnerProduct),
+ other => Err(data_invalid(format!(
+ "unsupported vector distance metric: {other}"
+ ))),
+ }
+ }
+
+ /// Higher-is-better score for exact vector search.
+ pub(crate) fn compute_score(&self, query: &[f32], stored: &[f32]) -> f32 {
+ match self {
+ Self::L2 => 1.0 / (1.0 + squared_l2(query, stored)),
+ Self::Cosine => cosine_similarity(query, stored),
+ Self::InnerProduct => inner_product(query, stored),
+ }
+ }
+
+ /// Lower-is-better distance for exact vector search.
+ pub(crate) fn compute_distance(&self, query: &[f32], stored: &[f32]) ->
f32 {
+ match self {
+ Self::L2 => squared_l2(query, stored),
+ Self::Cosine => cosine_distance(cosine_similarity(query, stored)),
+ Self::InnerProduct => -inner_product(query, stored),
+ }
+ }
+
+ /// Convert a higher-is-better standardized index score to a
lower-is-better
+ /// distance. For L2 with `score == 0.0` this yields `inf` (natural f32
+ /// behavior, matching Java — no clamp).
+ pub(crate) fn score_to_distance(&self, score: f32) -> f32 {
+ match self {
+ Self::L2 => 1.0 / score - 1.0,
+ Self::Cosine => cosine_distance(score),
+ Self::InnerProduct => -score,
+ }
+ }
+}
+
+fn squared_l2(query: &[f32], stored: &[f32]) -> f32 {
+ let mut squared = 0.0f32;
+ for i in 0..query.len() {
+ let delta = query[i] - stored[i];
+ squared += delta * delta;
+ }
+ squared
+}
+
+fn cosine_similarity(query: &[f32], stored: &[f32]) -> f32 {
+ let mut dot = 0.0f32;
+ let mut query_norm = 0.0f32;
+ let mut stored_norm = 0.0f32;
+ for i in 0..query.len() {
+ dot += query[i] * stored[i];
+ query_norm += query[i] * query[i];
+ stored_norm += stored[i] * stored[i];
+ }
+ let denominator = ((query_norm as f64).sqrt() * (stored_norm as
f64).sqrt()) as f32;
+ if denominator == 0.0 {
+ 0.0
+ } else {
+ dot / denominator
+ }
+}
+
+fn inner_product(query: &[f32], stored: &[f32]) -> f32 {
+ let mut dot = 0.0f32;
+ for i in 0..query.len() {
+ dot += query[i] * stored[i];
+ }
+ dot
+}
+
+fn cosine_distance(similarity: f32) -> f32 {
+ 1.0 - similarity.clamp(-1.0, 1.0)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn test_normalize_lowercases_and_replaces_hyphens_without_trimming() {
+ assert_eq!(normalize_metric("Inner-Product"), "inner_product");
+ assert_eq!(normalize_metric("L2"), "l2");
+ // No trim: surrounding whitespace is preserved (unlike the build-side
helper).
+ assert_eq!(normalize_metric(" l2 "), " l2 ");
+ }
+
+ #[test]
+ fn test_is_supported_only_for_three_metrics() {
+ assert!(is_supported_metric("L2"));
+ assert!(is_supported_metric("cosine"));
+ assert!(is_supported_metric("inner-product"));
+ assert!(!is_supported_metric("manhattan"));
+ assert!(!is_supported_metric(" l2 "));
+ }
+
+ #[test]
+ fn test_parse_rejects_unsupported_metric() {
+ assert!(VectorSearchMetric::parse("l2").is_ok());
+ assert!(VectorSearchMetric::parse("cosine").is_ok());
+ assert!(VectorSearchMetric::parse("inner_product").is_ok());
+ assert!(VectorSearchMetric::parse("manhattan").is_err());
+ }
+
+ #[test]
+ fn test_compute_distance_matches_java_anchor() {
+ // Java PkVectorExactSearcherTest.testDistancesForSupportedMetrics:
+ // q=[2,0], s=[1,0] -> l2=1.0, cosine=0.0, inner_product=-2.0
+ let q = [2.0f32, 0.0];
+ let s = [1.0f32, 0.0];
+ assert_eq!(VectorSearchMetric::L2.compute_distance(&q, &s), 1.0);
+ assert_eq!(VectorSearchMetric::Cosine.compute_distance(&q, &s), 0.0);
+ assert_eq!(
+ VectorSearchMetric::InnerProduct.compute_distance(&q, &s),
+ -2.0
+ );
+ }
+
+ #[test]
+ fn test_compute_score_higher_is_better() {
+ let q = [2.0f32, 0.0];
+ let s = [1.0f32, 0.0];
+ assert_eq!(VectorSearchMetric::L2.compute_score(&q, &s), 0.5); //
1/(1+1)
+ assert_eq!(VectorSearchMetric::Cosine.compute_score(&q, &s), 1.0); //
parallel
+ assert_eq!(VectorSearchMetric::InnerProduct.compute_score(&q, &s),
2.0);
+ }
+
+ #[test]
+ fn test_cosine_zero_norm_similarity_is_zero() {
+ let zero = [0.0f32, 0.0];
+ let s = [1.0f32, 0.0];
+ assert_eq!(VectorSearchMetric::Cosine.compute_score(&zero, &s), 0.0);
+ assert_eq!(VectorSearchMetric::Cosine.compute_distance(&zero, &s),
1.0);
+ }
+
+ #[test]
+ fn test_cosine_non_perfect_square_norm_uses_f64_sqrt() {
+ // query [0,3] -> norm 9, stored [1,2] -> norm 5; sqrt(5) is irrational
+ // so the f32-sqrt path (each sqrt taken in f32, then widened) and the
+ // f64-sqrt path (widen first, sqrt in f64) produce different f32 bits:
+ // buggy score 0.8944271 vs correct 0.8944272. Encode the f64 contract
in
+ // the expected value (dot / (sqrt(9.0f64) * sqrt(5.0f64)) as f32)
rather
+ // than a magic literal, so this pins the Java-matching f64 arithmetic.
+ let q = [0.0f32, 3.0];
+ let s = [1.0f32, 2.0];
+ let dot = 6.0f32;
+ let denominator = ((9.0f64).sqrt() * (5.0f64).sqrt()) as f32;
+ let expected = dot / denominator;
+ assert_eq!(VectorSearchMetric::Cosine.compute_score(&q, &s), expected);
+ }
+
+ #[test]
+ fn test_score_to_distance_l2_zero_score_is_infinite() {
+ assert!(VectorSearchMetric::L2.score_to_distance(0.0).is_infinite());
+ assert_eq!(VectorSearchMetric::L2.score_to_distance(0.5), 1.0); //
1/0.5 - 1
+ assert_eq!(
+ VectorSearchMetric::InnerProduct.score_to_distance(2.0),
+ -2.0
+ );
+ assert_eq!(VectorSearchMetric::Cosine.score_to_distance(1.0), 0.0);
+ }
+}
diff --git a/crates/paimon/src/vindex/pkvector/mod.rs
b/crates/paimon/src/vindex/pkvector/mod.rs
new file mode 100644
index 0000000..3a685ca
--- /dev/null
+++ b/crates/paimon/src/vindex/pkvector/mod.rs
@@ -0,0 +1,42 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Primary-key vector (bucket-local ANN) search kernel.
+//!
+//! Read-only bucket-local approximate-nearest-neighbour search over the
+//! primary-key vector index.
+
+// The kernel is crate-private and has no production caller yet, so its items
+// are unreachable outside their own tests. Suppress the resulting dead_code
+// lint at the module boundary until the read path wires it in.
+#![allow(dead_code)]
+
+pub(crate) mod ann;
+pub(crate) mod bucket;
+pub(crate) mod exact;
+pub(crate) mod metric;
+pub(crate) mod reader;
+pub(crate) mod result;
+
+/// Shared constructor for validation failures in this module (mirrors Java
+/// `checkArgument` / `IllegalArgumentException`).
+pub(crate) fn data_invalid(message: impl Into<String>) -> crate::Error {
+ crate::Error::DataInvalid {
+ message: message.into(),
+ source: None,
+ }
+}
diff --git a/crates/paimon/src/vindex/pkvector/reader.rs
b/crates/paimon/src/vindex/pkvector/reader.rs
new file mode 100644
index 0000000..3298013
--- /dev/null
+++ b/crates/paimon/src/vindex/pkvector/reader.rs
@@ -0,0 +1,80 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+/// Sequential exact-scan source of vectors for one data file. Mirrors Java
+/// `org.apache.paimon.index.pkvector.PkVectorReader`.
+pub(crate) trait PkVectorReader {
+ fn dimension(&self) -> usize;
+
+ fn row_count(&self) -> i64;
+
+ /// Read the next row's vector into `reuse` (`reuse.len() == dimension()`).
+ /// Returns `false` when the physical row is a NULL vector: it is not
scored,
+ /// but the physical position still advances by one. Each call advances
+ /// exactly one physical row.
+ fn read_next_vector(&mut self, reuse: &mut [f32]) -> crate::Result<bool>;
+}
+
+#[cfg(test)]
+pub(crate) mod test_support {
+ use super::PkVectorReader;
+
+ /// In-memory `PkVectorReader` for tests. `None` entries are NULL rows.
+ /// Mirrors Java's test `ArrayReader`.
+ pub(crate) struct ArrayReader {
+ dimension: usize,
+ vectors: Vec<Option<Vec<f32>>>,
+ position: usize,
+ }
+
+ impl ArrayReader {
+ pub(crate) fn new(dimension: usize, vectors: Vec<Option<Vec<f32>>>) ->
Self {
+ Self {
+ dimension,
+ vectors,
+ position: 0,
+ }
+ }
+ }
+
+ impl PkVectorReader for ArrayReader {
+ fn dimension(&self) -> usize {
+ self.dimension
+ }
+
+ fn row_count(&self) -> i64 {
+ self.vectors.len() as i64
+ }
+
+ fn read_next_vector(&mut self, reuse: &mut [f32]) ->
crate::Result<bool> {
+ assert_eq!(
+ reuse.len(),
+ self.dimension,
+ "reuse buffer must equal dimension"
+ );
+ let entry = &self.vectors[self.position];
+ self.position += 1;
+ match entry {
+ Some(vector) => {
+ reuse.copy_from_slice(vector);
+ Ok(true)
+ }
+ None => Ok(false),
+ }
+ }
+ }
+}
diff --git a/crates/paimon/src/vindex/pkvector/result.rs
b/crates/paimon/src/vindex/pkvector/result.rs
new file mode 100644
index 0000000..97bea5e
--- /dev/null
+++ b/crates/paimon/src/vindex/pkvector/result.rs
@@ -0,0 +1,27 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+/// One vector search result addressed by source data-file row position.
+/// Mirrors Java `org.apache.paimon.index.pkvector.PkVectorSearchResult`.
+#[derive(Clone, Debug, PartialEq)]
+pub(crate) struct PkVectorSearchResult {
+ pub data_file_name: String,
+ /// Zero-based physical row position within the source data file.
+ pub row_position: i64,
+ /// Lower-is-better distance.
+ pub distance: f32,
+}
