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 f55ad02 [core] Add primary-key vector read kernel (#517)
f55ad02 is described below
commit f55ad02fccab274e279c67925509311d3ddc3bf0
Author: Junrui Lee <[email protected]>
AuthorDate: Thu Jul 16 15:26:47 2026 +0800
[core] Add primary-key vector read kernel (#517)
---
crates/paimon/src/table/data_file_reader.rs | 161 ++-
crates/paimon/src/table/mod.rs | 3 +
.../src/table/pk_vector_indexed_split_read.rs | 778 +++++++++++++
crates/paimon/src/table/pk_vector_orchestrator.rs | 1181 ++++++++++++++++++++
crates/paimon/src/table/pk_vector_position_read.rs | 1079 ++++++++++++++++++
crates/paimon/src/vindex/pkvector/ann.rs | 5 +-
crates/paimon/src/vindex/pkvector/bucket.rs | 52 +-
crates/paimon/src/vindex/pkvector/exact.rs | 41 +-
crates/paimon/src/vindex/pkvector/metric.rs | 106 ++
9 files changed, 3362 insertions(+), 44 deletions(-)
diff --git a/crates/paimon/src/table/data_file_reader.rs
b/crates/paimon/src/table/data_file_reader.rs
index 05df8b3..3395f8d 100644
--- a/crates/paimon/src/table/data_file_reader.rs
+++ b/crates/paimon/src/table/data_file_reader.rs
@@ -71,6 +71,35 @@ impl DataFileReader {
self
}
+ // These three accessors exist for the sibling `pk_vector_position_read`,
+ // `pk_vector_indexed_split_read`, and `pk_vector_orchestrator` modules.
The
+ // read path that drives that chain lands in a later change, so under
clippy
+ // -D warnings they read as dead_code until then.
+ /// Return a copy with a replaced read-type. Used by
`pk_vector_position_read`
+ /// to inject the internal `_ROW_ID` column for physical-position recovery.
+ #[allow(dead_code)]
+ pub(super) fn with_read_type(mut self, read_type: Vec<DataField>) -> Self {
+ self.read_type = read_type;
+ self
+ }
+
+ /// The effective read-type (requested output fields) of this reader.
+ /// Exposed for the sibling `pk_vector_position_read` module.
+ #[allow(dead_code)]
+ pub(super) fn read_type(&self) -> &[DataField] {
+ &self.read_type
+ }
+
+ /// True if any configured predicate can actually drop rows. A lone
+ /// `Predicate::AlwaysTrue` keeps every row in order and is not
row-filtering,
+ /// matching `reject_row_id_with_predicate`'s notion.
+ #[allow(dead_code)]
+ pub(super) fn has_row_filtering_predicate(&self) -> bool {
+ self.predicates
+ .iter()
+ .any(|p| !matches!(p, Predicate::AlwaysTrue))
+ }
+
/// Reject projecting `_ROW_ID` alongside a data predicate. `_ROW_ID` is
/// assigned positionally from post-filter batch row counts, so a residual
/// filter that drops rows would desync it. (`_ROW_ID` predicates travel
via
@@ -111,35 +140,16 @@ impl DataFileReader {
Ok(try_stream! {
for split in splits {
// Create DV factory for this split only.
- let dv_factory = if split
- .data_deletion_files()
- .is_some_and(|files| files.iter().any(Option::is_some))
- {
- Some(
- DeletionVectorFactory::new(
- &reader.file_io,
- split.data_files(),
- split.data_deletion_files(),
- )
- .await?,
- )
- } else {
- None
- };
+ let dv_factory = reader.build_split_dv_factory(&split).await?;
for file_meta in split.data_files().to_vec() {
- let dv = dv_factory
- .as_ref()
- .and_then(|factory|
factory.get_deletion_vector(&file_meta.file_name))
- .cloned();
+ let dv = DataFileReader::deletion_vector_for_file(
+ dv_factory.as_ref(),
+ &file_meta.file_name,
+ );
// Load data file's schema if it differs from the table
schema.
- let data_fields: Option<Vec<DataField>> = if
file_meta.schema_id != reader.table_schema_id {
- let data_schema =
reader.schema_manager.schema(file_meta.schema_id).await?;
- Some(data_schema.fields().to_vec())
- } else {
- None
- };
+ let data_fields =
reader.derive_data_fields(&file_meta).await?;
let mut stream = reader.read_single_file_stream(
&split,
@@ -157,6 +167,55 @@ impl DataFileReader {
.boxed())
}
+ /// Build the deletion-vector factory for a split, or `None` when the split
+ /// carries no deletion files. One factory per split (not per file),
matching
+ /// `read`. Shared with `pk_vector_indexed_split_read` so that path reuses
the
+ /// exact production derivation instead of duplicating it.
+ pub(super) async fn build_split_dv_factory(
+ &self,
+ split: &DataSplit,
+ ) -> crate::Result<Option<DeletionVectorFactory>> {
+ if split
+ .data_deletion_files()
+ .is_some_and(|files| files.iter().any(Option::is_some))
+ {
+ Ok(Some(
+ DeletionVectorFactory::new(
+ &self.file_io,
+ split.data_files(),
+ split.data_deletion_files(),
+ )
+ .await?,
+ ))
+ } else {
+ Ok(None)
+ }
+ }
+
+ /// Look up the deletion vector for one file from a split-level factory.
+ pub(super) fn deletion_vector_for_file(
+ factory: Option<&DeletionVectorFactory>,
+ file_name: &str,
+ ) -> Option<Arc<DeletionVector>> {
+ factory
+ .and_then(|factory| factory.get_deletion_vector(file_name))
+ .cloned()
+ }
+
+ /// Load the data file's own schema fields when its `schema_id` differs
from
+ /// the table schema id (schema evolution); `None` when they match.
+ pub(super) async fn derive_data_fields(
+ &self,
+ file_meta: &DataFileMeta,
+ ) -> crate::Result<Option<Vec<DataField>>> {
+ if file_meta.schema_id != self.table_schema_id {
+ let data_schema =
self.schema_manager.schema(file_meta.schema_id).await?;
+ Ok(Some(data_schema.fields().to_vec()))
+ } else {
+ Ok(None)
+ }
+ }
+
/// Read a single parquet file from a split, returning a lazy stream of
batches.
/// Optionally applies a deletion vector.
///
@@ -994,6 +1053,58 @@ mod tests {
use roaring::RoaringBitmap;
use std::io;
+ #[test]
+ fn test_accessors_expose_read_type_and_row_filtering_predicate() {
+ use crate::spec::{DataField, DataType, IntType};
+ let fields = vec![DataField::new(
+ 0,
+ "id".to_string(),
+ DataType::Int(IntType::new()),
+ )];
+ let file_io = crate::io::FileIOBuilder::new("memory").build().unwrap();
+ let schema_manager = SchemaManager::new(file_io.clone(),
"memory:/acc".to_string());
+
+ let no_pred = DataFileReader::new(
+ file_io.clone(),
+ schema_manager.clone(),
+ 1,
+ fields.clone(),
+ fields.clone(),
+ vec![],
+ );
+ assert_eq!(no_pred.read_type().len(), 1);
+ assert!(!no_pred.has_row_filtering_predicate());
+
+ let always_true = DataFileReader::new(
+ file_io.clone(),
+ schema_manager.clone(),
+ 1,
+ fields.clone(),
+ fields.clone(),
+ vec![crate::spec::Predicate::AlwaysTrue],
+ );
+ assert!(
+ !always_true.has_row_filtering_predicate(),
+ "AlwaysTrue is not row-filtering"
+ );
+
+ let filtering = PredicateBuilder::new(&fields)
+ .equal("id", crate::spec::Datum::Int(10))
+ .unwrap();
+ let with_filter = DataFileReader::new(
+ file_io.clone(),
+ schema_manager.clone(),
+ 1,
+ fields.clone(),
+ fields.clone(),
+ vec![filtering],
+ );
+ assert!(
+ with_filter.has_row_filtering_predicate(),
+ "a real (non-AlwaysTrue) predicate is row-filtering"
+ );
+ }
+
struct MemOutputFile {
data: Vec<u8>,
}
diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs
index 617c1fb..2a239e9 100644
--- a/crates/paimon/src/table/mod.rs
+++ b/crates/paimon/src/table/mod.rs
@@ -55,6 +55,9 @@ mod lumina_index_build_builder;
pub(crate) mod merge_tree_split_generator;
mod partition_filter;
mod partition_stat;
+mod pk_vector_indexed_split_read;
+mod pk_vector_orchestrator;
+mod pk_vector_position_read;
mod postpone_file_writer;
mod prepared_files;
mod read_builder;
diff --git a/crates/paimon/src/table/pk_vector_indexed_split_read.rs
b/crates/paimon/src/table/pk_vector_indexed_split_read.rs
new file mode 100644
index 0000000..00bd1ae
--- /dev/null
+++ b/crates/paimon/src/table/pk_vector_indexed_split_read.rs
@@ -0,0 +1,778 @@
+// 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 indexed-split read-path contract (read-path subset of
+//! apache/paimon#8576).
+//!
+//! `PkVectorIndexedSplit` carries one data file + inclusive physical-position
+//! ranges + an optional aligned score array. `PkVectorIndexedSplitRead`
validates
+//! the split, expands the ranges into an ascending position set and a
+//! `position -> score` map, and delegates to the sibling
`PkVectorPositionRead`.
+//! It is a pure consumer: no bucket/ANN search, no cross-bucket
orchestration, no
+//! serialization.
+
+// This module wires the position reader into the indexed read contract; the
+// sibling `pk_vector_orchestrator` adds cross-bucket orchestration on top. The
+// read path that drives that chain lands in a later change, so under clippy
+// -D warnings these items read as dead_code until then. Suppress at the module
+// boundary.
+#![allow(dead_code)]
+
+use std::collections::BTreeMap;
+
+use futures::StreamExt;
+
+use crate::spec::DataFileMeta;
+use crate::table::data_file_reader::DataFileReader;
+use crate::table::pk_vector_position_read::PkVectorPositionRead;
+use crate::table::source::DataSplit;
+use crate::table::{ArrowRecordBatchStream, RowRange};
+
+fn data_invalid(message: impl Into<String>) -> crate::Error {
+ crate::Error::DataInvalid {
+ message: message.into(),
+ source: None,
+ }
+}
+
+/// A single-file indexed split for the PK-vector read path.
+///
+/// `row_ranges` are 0-based PHYSICAL positions within the one data file,
inclusive
+/// on both ends, required strictly ascending and non-overlapping. Adjacent
+/// (touching) ranges are allowed and need not be coalesced by the producer.
+/// `scores`, when present, is aligned to the expanded-range order (ascending
+/// position); `None` means no `_PKEY_VECTOR_SCORE` column in the output.
+///
+/// Deliberately NOT reusing `DataSplit.row_ranges`, whose ranges mean
stable/global
+/// row ids on the append/data-evolution path. Not serialized.
+pub(crate) struct PkVectorIndexedSplit {
+ pub split: DataSplit,
+ pub row_ranges: Vec<RowRange>,
+ pub scores: Option<Vec<f32>>,
+}
+
+/// Expand inclusive physical-position ranges into an ascending `Vec<i64>` of
+/// 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>> {
+ if ranges.is_empty() {
+ return Err(data_invalid("indexed split must select at least one row"));
+ }
+ let mut positions = Vec::new();
+ let mut prev_to: Option<i64> = None;
+ for range in ranges {
+ if range.from() < 0 || range.to() >= row_count {
+ return Err(data_invalid(format!(
+ "indexed-split range [{}, {}] is outside [0, {})",
+ range.from(),
+ range.to(),
+ row_count
+ )));
+ }
+ if let Some(prev) = prev_to {
+ // Rejects overlap AND descending order; touching ranges
+ // (range.from() == prev + 1) are accepted.
+ if range.from() <= prev {
+ return Err(data_invalid(
+ "indexed-split ranges must be ascending and
non-overlapping",
+ ));
+ }
+ }
+ for position in range.from()..=range.to() {
+ positions.push(position);
+ }
+ prev_to = Some(range.to());
+ }
+ Ok(positions)
+}
+
+/// Eagerly validate the split and expand it into `(file_meta, positions,
scores)`.
+/// Single-file check, range bounds/ordering, and score-length alignment happen
+/// here so their errors surface at the `read(...)` call site.
+#[allow(clippy::type_complexity)]
+fn validate_and_expand(
+ split: &PkVectorIndexedSplit,
+) -> crate::Result<(DataFileMeta, Vec<i64>, Option<BTreeMap<i64, f32>>)> {
+ let files = split.split.data_files();
+ if files.len() != 1 {
+ return Err(data_invalid(
+ "indexed split for a primary-key table must contain exactly one
data file",
+ ));
+ }
+ let file_meta = files[0].clone();
+ let positions = expand_ranges(&split.row_ranges, file_meta.row_count)?;
+
+ let score_map = match &split.scores {
+ Some(scores) => {
+ if scores.len() != positions.len() {
+ return Err(data_invalid(format!(
+ "indexed-split scores length {} does not match selected
positions {}",
+ scores.len(),
+ positions.len()
+ )));
+ }
+ // positions is ascending == expanded-range order; zip aligns each
+ // score to its position.
+ Some(
+ positions
+ .iter()
+ .copied()
+ .zip(scores.iter().copied())
+ .collect(),
+ )
+ }
+ None => None,
+ };
+
+ Ok((file_meta, positions, score_map))
+}
+
+/// Validates a `PkVectorIndexedSplit` and delegates its materialization to the
+/// sibling `PkVectorPositionRead`. Rust equivalent of Java
`PrimaryKeyIndexedSplitRead`.
+pub(crate) struct PkVectorIndexedSplitRead {
+ reader: DataFileReader,
+}
+
+impl PkVectorIndexedSplitRead {
+ pub(crate) fn new(reader: DataFileReader) -> Self {
+ Self { reader }
+ }
+
+ /// Validate the split eagerly, then return a lazy stream that derives
+ /// `data_fields`/deletion-vector and delegates to `PkVectorPositionRead`.
+ /// Eager errors (single-file, range bounds/ordering, score length) surface
+ /// here; lazy errors (schema/DV load, the position reader's own guards)
surface
+ /// on first poll.
+ pub(crate) fn read(
+ &self,
+ split: &PkVectorIndexedSplit,
+ ) -> crate::Result<ArrowRecordBatchStream> {
+ let (file_meta, positions, score_map) = validate_and_expand(split)?;
+ let reader = self.reader.clone();
+ let data_split = split.split.clone();
+
+ let stream = async_stream::try_stream! {
+ let dv_factory = reader.build_split_dv_factory(&data_split).await?;
+ let dv = DataFileReader::deletion_vector_for_file(
+ dv_factory.as_ref(),
+ &file_meta.file_name,
+ );
+ let data_fields = reader.derive_data_fields(&file_meta).await?;
+
+ let inner = PkVectorPositionRead::new(&reader).read(
+ &data_split,
+ file_meta,
+ data_fields,
+ dv,
+ positions,
+ score_map,
+ )?;
+ futures::pin_mut!(inner);
+ while let Some(batch) = inner.next().await {
+ yield batch?;
+ }
+ };
+ Ok(Box::pin(stream))
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::spec::stats::BinaryTableStats;
+ use crate::spec::DataFileMeta;
+ use crate::table::source::DataSplitBuilder;
+
+ fn data_file(row_count: i64) -> DataFileMeta {
+ DataFileMeta {
+ file_name: "part-0.mosaic".to_string(),
+ file_size: 1,
+ row_count,
+ min_key: Vec::new(),
+ max_key: Vec::new(),
+ key_stats: BinaryTableStats::empty(),
+ value_stats: BinaryTableStats::empty(),
+ min_sequence_number: 0,
+ max_sequence_number: 0,
+ schema_id: 1,
+ level: 0,
+ extra_files: Vec::new(),
+ creation_time: None,
+ delete_row_count: None,
+ embedded_index: None,
+ file_source: None,
+ value_stats_cols: None,
+ external_path: None,
+ first_row_id: Some(0),
+ write_cols: None,
+ }
+ }
+
+ fn split_with_files(files: Vec<DataFileMeta>) -> DataSplit {
+ DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(crate::spec::BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path("memory:/pkvisr/bucket-0".to_string())
+ .with_total_buckets(1)
+ .with_data_files(files)
+ .build()
+ .unwrap()
+ }
+
+ fn indexed_split(
+ row_count: i64,
+ row_ranges: Vec<RowRange>,
+ scores: Option<Vec<f32>>,
+ ) -> PkVectorIndexedSplit {
+ PkVectorIndexedSplit {
+ split: split_with_files(vec![data_file(row_count)]),
+ row_ranges,
+ scores,
+ }
+ }
+
+ #[test]
+ fn expand_single_contiguous_range() {
+ let positions = expand_ranges(&[RowRange::new(0, 3)], 10).unwrap();
+ assert_eq!(positions, vec![0, 1, 2, 3]);
+ }
+
+ #[test]
+ fn expand_multiple_ascending_ranges() {
+ // Non-overlapping with a gap, plus an adjacent (touching) pair.
+ let positions = expand_ranges(
+ &[
+ RowRange::new(0, 1),
+ RowRange::new(2, 2),
+ RowRange::new(5, 6),
+ ],
+ 10,
+ )
+ .unwrap();
+ assert_eq!(positions, vec![0, 1, 2, 5, 6]);
+ }
+
+ #[test]
+ fn expand_single_row_range() {
+ let positions = expand_ranges(&[RowRange::new(4, 4)], 5).unwrap();
+ assert_eq!(positions, vec![4]);
+ }
+
+ #[test]
+ fn expand_rejects_empty_ranges() {
+ let err = expand_ranges(&[], 5).expect_err("empty ranges must error");
+ assert!(
+ format!("{err:?}").contains("at least one row"),
+ "got: {err:?}"
+ );
+ }
+
+ #[test]
+ fn expand_rejects_range_at_or_past_row_count() {
+ let err = expand_ranges(&[RowRange::new(3, 5)], 5).expect_err("to >=
row_count must error");
+ assert!(
+ format!("{err:?}").contains("outside [0, 5)"),
+ "got: {err:?}"
+ );
+ }
+
+ #[test]
+ fn expand_rejects_negative_from() {
+ let err = expand_ranges(&[RowRange::new(-1, 2)],
5).expect_err("negative from must error");
+ assert!(format!("{err:?}").contains("outside"), "got: {err:?}");
+ }
+
+ #[test]
+ fn expand_rejects_overlapping_ranges() {
+ // range[1].from()==2 <= range[0].to()==3 -> overlap.
+ let err = expand_ranges(&[RowRange::new(0, 3), RowRange::new(2, 4)],
10)
+ .expect_err("overlapping ranges must error");
+ assert!(
+ format!("{err:?}").contains("ascending and non-overlapping"),
+ "got: {err:?}"
+ );
+ }
+
+ #[test]
+ fn expand_rejects_descending_ranges() {
+ let err = expand_ranges(&[RowRange::new(5, 6), RowRange::new(0, 1)],
10)
+ .expect_err("descending ranges must error");
+ assert!(
+ format!("{err:?}").contains("ascending and non-overlapping"),
+ "got: {err:?}"
+ );
+ }
+
+ #[test]
+ fn validate_rejects_multi_file_split() {
+ let split = PkVectorIndexedSplit {
+ split: split_with_files(vec![data_file(5), data_file(5)]),
+ row_ranges: vec![RowRange::new(0, 0)],
+ scores: None,
+ };
+ let err = validate_and_expand(&split).expect_err("multi-file split
must error");
+ assert!(
+ format!("{err:?}").contains("exactly one data file"),
+ "got: {err:?}"
+ );
+ }
+
+ #[test]
+ fn validate_builds_score_map_for_non_contiguous_ranges() {
+ let split = indexed_split(
+ 10,
+ vec![
+ RowRange::new(0, 0),
+ RowRange::new(2, 2),
+ RowRange::new(5, 5),
+ ],
+ Some(vec![0.9, 0.5, 0.1]),
+ );
+ let (_file, positions, score_map) =
validate_and_expand(&split).unwrap();
+ assert_eq!(positions, vec![0, 2, 5]);
+ let score_map = score_map.unwrap();
+ assert_eq!(score_map, BTreeMap::from([(0, 0.9f32), (2, 0.5), (5,
0.1)]));
+ }
+
+ #[test]
+ fn validate_no_scores_yields_none() {
+ let split = indexed_split(5, vec![RowRange::new(0, 1)], None);
+ let (_file, positions, score_map) =
validate_and_expand(&split).unwrap();
+ assert_eq!(positions, vec![0, 1]);
+ assert!(score_map.is_none());
+ }
+
+ #[test]
+ fn validate_rejects_score_length_mismatch() {
+ // select 2 positions but supply 1 score.
+ let split = indexed_split(5, vec![RowRange::new(0, 1)],
Some(vec![0.5]));
+ let err = validate_and_expand(&split).expect_err("score length
mismatch must error");
+ assert!(format!("{err:?}").contains("scores length"), "got: {err:?}");
+ }
+}
+
+#[cfg(test)]
+mod e2e_tests {
+ use super::*;
+ use crate::arrow::build_target_arrow_schema;
+ use crate::io::FileIOBuilder;
+ use crate::spec::stats::BinaryTableStats;
+ use crate::spec::{DataField, DataFileMeta, DataType, IntType};
+ use crate::table::data_file_reader::DataFileReader;
+ use crate::table::pk_vector_position_read::{
+ PKEY_VECTOR_POSITION_COLUMN, PKEY_VECTOR_SCORE_COLUMN,
+ };
+ use crate::table::schema_manager::SchemaManager;
+ use crate::table::source::{DataSplit, DataSplitBuilder, DeletionFile};
+ use arrow_array::{Array, Float32Array, Int32Array, Int64Array,
RecordBatch};
+ use bytes::Bytes;
+ use futures::TryStreamExt;
+ use paimon_mosaic_core::spec::COMPRESSION_NONE;
+ use paimon_mosaic_core::writer::{MosaicWriter, OutputFile, WriterOptions};
+ use roaring::RoaringBitmap;
+ use std::io;
+ use std::sync::Arc;
+
+ struct MemOutputFile {
+ data: Vec<u8>,
+ }
+ impl MemOutputFile {
+ fn new() -> Self {
+ Self { data: Vec::new() }
+ }
+ }
+ impl OutputFile for MemOutputFile {
+ fn write(&mut self, data: &[u8]) -> io::Result<()> {
+ self.data.extend_from_slice(data);
+ Ok(())
+ }
+ fn flush(&mut self) -> io::Result<()> {
+ Ok(())
+ }
+ fn pos(&self) -> u64 {
+ self.data.len() as u64
+ }
+ }
+
+ fn id_field() -> DataField {
+ DataField::new(0, "id".to_string(), DataType::Int(IntType::new()))
+ }
+ fn id_fields() -> Vec<DataField> {
+ vec![id_field()]
+ }
+ fn id_batch(ids: Vec<i32>) -> RecordBatch {
+ let schema = build_target_arrow_schema(&id_fields()).unwrap();
+ RecordBatch::try_new(schema,
vec![Arc::new(Int32Array::from(ids))]).unwrap()
+ }
+
+ fn data_file(file_name: &str, file_size: i64, row_count: i64) ->
DataFileMeta {
+ DataFileMeta {
+ file_name: file_name.to_string(),
+ file_size,
+ row_count,
+ min_key: Vec::new(),
+ max_key: Vec::new(),
+ key_stats: BinaryTableStats::empty(),
+ value_stats: BinaryTableStats::empty(),
+ min_sequence_number: 0,
+ max_sequence_number: 0,
+ schema_id: 1,
+ level: 0,
+ extra_files: Vec::new(),
+ creation_time: None,
+ delete_row_count: None,
+ embedded_index: None,
+ file_source: None,
+ value_stats_cols: None,
+ external_path: None,
+ first_row_id: Some(0),
+ write_cols: None,
+ }
+ }
+
+ fn write_mosaic_single_group(batch: &RecordBatch) -> Bytes {
+ let out = MemOutputFile::new();
+ let mut writer = MosaicWriter::new(
+ out,
+ batch.schema().as_ref(),
+ WriterOptions {
+ compression: COMPRESSION_NONE,
+ num_buckets: 2,
+ row_group_max_size: u64::MAX,
+ ..Default::default()
+ },
+ )
+ .unwrap();
+ writer.write_batch(batch).unwrap();
+ writer.close().unwrap();
+ Bytes::from(writer.output().data.to_vec())
+ }
+
+ fn write_mosaic_multi_group(batches: &[RecordBatch]) -> Bytes {
+ let out = MemOutputFile::new();
+ let mut writer = MosaicWriter::new(
+ out,
+ batches[0].schema().as_ref(),
+ WriterOptions {
+ compression: COMPRESSION_NONE,
+ num_buckets: 2,
+ row_group_max_size: 1,
+ ..Default::default()
+ },
+ )
+ .unwrap();
+ for batch in batches {
+ writer.write_batch(batch).unwrap();
+ }
+ writer.close().unwrap();
+ Bytes::from(writer.output().data.to_vec())
+ }
+
+ async fn write_deletion_file(
+ file_io: &crate::io::FileIO,
+ path: &str,
+ deleted_rows: &[u32],
+ ) -> DeletionFile {
+ 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::from(blob))
+ .await
+ .unwrap();
+ DeletionFile::new(
+ path.to_string(),
+ 0,
+ bitmap_length as i64,
+ Some(deleted_rows.len() as i64),
+ )
+ }
+
+ /// Build a predicate-free `DataFileReader` over an in-memory mosaic file
plus
+ /// the matching split. `deleted_rows`, when non-empty, writes a DV into
the split.
+ async fn build_reader_and_split(
+ table_path: &str,
+ data: &Bytes,
+ row_count: i64,
+ deleted_rows: &[u32],
+ ) -> (DataFileReader, DataSplit) {
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let bucket_path = format!("{table_path}/bucket-0");
+ let file_name = "part-0.mosaic";
+ file_io
+ .new_output(&format!("{bucket_path}/{file_name}"))
+ .unwrap()
+ .write(data.clone())
+ .await
+ .unwrap();
+
+ let mut split_builder = DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(crate::spec::BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path(bucket_path)
+ .with_total_buckets(1)
+ .with_data_files(vec![data_file(file_name, data.len() as i64,
row_count)]);
+ if !deleted_rows.is_empty() {
+ let df =
+ write_deletion_file(&file_io,
&format!("{table_path}/index/dv-0"), deleted_rows)
+ .await;
+ split_builder =
split_builder.with_data_deletion_files(vec![Some(df)]);
+ }
+ let split = split_builder.build().unwrap();
+
+ let schema_manager = SchemaManager::new(file_io.clone(),
table_path.to_string());
+ let reader = DataFileReader::new(
+ file_io,
+ schema_manager,
+ 1,
+ id_fields(),
+ id_fields(),
+ Vec::new(),
+ );
+ (reader, split)
+ }
+
+ fn column_by_name<'a>(batch: &'a RecordBatch, name: &str) -> Option<&'a
Arc<dyn Array>> {
+ batch
+ .schema()
+ .index_of(name)
+ .ok()
+ .map(|idx| batch.column(idx))
+ }
+ fn collect_i32(batches: &[RecordBatch], name: &str) -> Vec<i32> {
+ batches
+ .iter()
+ .flat_map(|b| {
+ column_by_name(b, name)
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ })
+ .collect()
+ }
+ fn collect_i64(batches: &[RecordBatch], name: &str) -> Vec<i64> {
+ batches
+ .iter()
+ .flat_map(|b| {
+ column_by_name(b, name)
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ })
+ .collect()
+ }
+ fn collect_f32(batches: &[RecordBatch], name: &str) -> Vec<f32> {
+ batches
+ .iter()
+ .flat_map(|b| {
+ column_by_name(b, name)
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Float32Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ })
+ .collect()
+ }
+
+ #[tokio::test]
+ async fn reads_ranges_with_position_column_no_scores() {
+ // rows id=[10,11,12], ranges [0..=0, 2..=2] -> ids [10,12], positions
[0,2];
+ // no _ROW_ID leak, no _PKEY_VECTOR_SCORE.
+ let data = write_mosaic_single_group(&id_batch(vec![10, 11, 12]));
+ let (reader, split) = build_reader_and_split("memory:/pkvisr_basic",
&data, 3, &[]).await;
+ let indexed = PkVectorIndexedSplit {
+ split,
+ row_ranges: vec![RowRange::new(0, 0), RowRange::new(2, 2)],
+ scores: None,
+ };
+
+ let batches = PkVectorIndexedSplitRead::new(reader)
+ .read(&indexed)
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ assert_eq!(collect_i32(&batches, "id"), vec![10, 12]);
+ assert_eq!(
+ collect_i64(&batches, PKEY_VECTOR_POSITION_COLUMN),
+ vec![0, 2]
+ );
+ for batch in &batches {
+ assert!(
+ column_by_name(batch, "_ROW_ID").is_none(),
+ "_ROW_ID must not leak"
+ );
+ assert!(column_by_name(batch, PKEY_VECTOR_SCORE_COLUMN).is_none());
+ }
+ }
+
+ #[tokio::test]
+ async fn reads_ranges_with_aligned_scores() {
+ let data = write_mosaic_single_group(&id_batch(vec![10, 11, 12, 13]));
+ let (reader, split) = build_reader_and_split("memory:/pkvisr_scores",
&data, 4, &[]).await;
+ let indexed = PkVectorIndexedSplit {
+ split,
+ row_ranges: vec![RowRange::new(0, 0), RowRange::new(2, 3)],
+ scores: Some(vec![0.9, 0.5, 0.1]),
+ };
+
+ let batches = PkVectorIndexedSplitRead::new(reader)
+ .read(&indexed)
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ assert_eq!(collect_i32(&batches, "id"), vec![10, 12, 13]);
+ assert_eq!(
+ collect_i64(&batches, PKEY_VECTOR_POSITION_COLUMN),
+ vec![0, 2, 3]
+ );
+ assert_eq!(
+ collect_f32(&batches, PKEY_VECTOR_SCORE_COLUMN),
+ vec![0.9, 0.5, 0.1]
+ );
+ }
+
+ #[tokio::test]
+ async fn deletion_vector_drops_selected_position_and_its_score() {
+ // select positions [0,1,2,3] via ranges, scores aligned; DV deletes
position 1.
+ // -> returned positions [0,2,3], scores [.4,.2,.1]; input scores
still had key 1.
+ let data = write_mosaic_single_group(&id_batch(vec![10, 11, 12, 13]));
+ let (reader, split) = build_reader_and_split("memory:/pkvisr_dv",
&data, 4, &[1]).await;
+ let indexed = PkVectorIndexedSplit {
+ split,
+ row_ranges: vec![RowRange::new(0, 3)],
+ scores: Some(vec![0.4, 0.3, 0.2, 0.1]),
+ };
+
+ let batches = PkVectorIndexedSplitRead::new(reader)
+ .read(&indexed)
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ assert_eq!(collect_i32(&batches, "id"), vec![10, 12, 13]);
+ assert_eq!(
+ collect_i64(&batches, PKEY_VECTOR_POSITION_COLUMN),
+ vec![0, 2, 3]
+ );
+ assert_eq!(
+ collect_f32(&batches, PKEY_VECTOR_SCORE_COLUMN),
+ vec![0.4, 0.2, 0.1]
+ );
+ }
+
+ #[tokio::test]
+ async fn alignment_holds_across_multiple_batches() {
+ // Three row groups [10,11] [12,13] [14,15] -> reader yields >1 batch.
+ // ranges [1..=2, 4..=4] span batch boundaries.
+ let data = write_mosaic_multi_group(&[
+ id_batch(vec![10, 11]),
+ id_batch(vec![12, 13]),
+ id_batch(vec![14, 15]),
+ ]);
+ let (reader, split) =
+ build_reader_and_split("memory:/pkvisr_multibatch", &data, 6,
&[]).await;
+ let indexed = PkVectorIndexedSplit {
+ split,
+ row_ranges: vec![RowRange::new(1, 2), RowRange::new(4, 4)],
+ scores: Some(vec![0.9, 0.5, 0.1]),
+ };
+
+ let batches = PkVectorIndexedSplitRead::new(reader)
+ .read(&indexed)
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ assert!(
+ batches.len() > 1,
+ "expected multiple batches, got {}",
+ batches.len()
+ );
+ assert_eq!(collect_i32(&batches, "id"), vec![11, 12, 14]);
+ assert_eq!(
+ collect_i64(&batches, PKEY_VECTOR_POSITION_COLUMN),
+ vec![1, 2, 4]
+ );
+ assert_eq!(
+ collect_f32(&batches, PKEY_VECTOR_SCORE_COLUMN),
+ vec![0.9, 0.5, 0.1]
+ );
+ }
+
+ #[tokio::test]
+ async fn multi_file_split_errors_at_call_site() {
+ // Eager validation: multi-file split errors from read(...) itself,
not the stream.
+ let data = write_mosaic_single_group(&id_batch(vec![10, 11]));
+ let (reader, _split) =
+ build_reader_and_split("memory:/pkvisr_multifile", &data, 2,
&[]).await;
+ let split = DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(crate::spec::BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path("memory:/pkvisr_multifile/bucket-0".to_string())
+ .with_total_buckets(1)
+ .with_data_files(vec![
+ data_file("part-0.mosaic", data.len() as i64, 2),
+ data_file("part-1.mosaic", data.len() as i64, 2),
+ ])
+ .build()
+ .unwrap();
+ let indexed = PkVectorIndexedSplit {
+ split,
+ row_ranges: vec![RowRange::new(0, 0)],
+ scores: None,
+ };
+
+ let err = PkVectorIndexedSplitRead::new(reader)
+ .read(&indexed)
+ .err()
+ .expect("multi-file split must error eagerly");
+ assert!(
+ format!("{err:?}").contains("exactly one data file"),
+ "got: {err:?}"
+ );
+ }
+}
diff --git a/crates/paimon/src/table/pk_vector_orchestrator.rs
b/crates/paimon/src/table/pk_vector_orchestrator.rs
new file mode 100644
index 0000000..b8664cc
--- /dev/null
+++ b/crates/paimon/src/table/pk_vector_orchestrator.rs
@@ -0,0 +1,1181 @@
+// 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 read orchestration: Rust equivalent of Java
+//! `PrimaryKeyVectorRead` + `PrimaryKeyVectorResult.splits()`.
+//!
+//! Per-bucket search via `bucket_search`, cross-bucket global Top-K merge,
+//! grouping survivors by data file into `PkVectorIndexedSplit`s, and lazy
+//! materialization via `PkVectorIndexedSplitRead`. Inputs are per-bucket
+//! `PkVectorSearchSplit`s supplied by the caller.
+
+// The read path that drives this module lands in a later change, so under
+// clippy -D warnings the module reads as dead_code until then. Suppress at the
+// module boundary.
+#![allow(dead_code)]
+
+use std::cmp::Ordering;
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use futures::StreamExt;
+
+use crate::deletion_vector::DeletionVector;
+use crate::spec::BinaryRow;
+use crate::table::data_file_reader::DataFileReader;
+use crate::table::pk_vector_indexed_split_read::{PkVectorIndexedSplit,
PkVectorIndexedSplitRead};
+use crate::table::source::{DataSplit, DataSplitBuilder, RowRange};
+use crate::table::ArrowRecordBatchStream;
+use crate::vindex::pkvector::ann::PkVectorAnnSearcher;
+use crate::vindex::pkvector::bucket::{bucket_search, BucketActiveFile,
BucketAnnSegment};
+use crate::vindex::pkvector::metric::{java_float_compare, VectorSearchMetric};
+use crate::vindex::pkvector::reader::PkVectorReader;
+use crate::vindex::pkvector::result::PkVectorSearchResult;
+
+fn data_invalid(message: impl Into<String>) -> crate::Error {
+ crate::Error::DataInvalid {
+ message: message.into(),
+ source: None,
+ }
+}
+
+/// Per-bucket search input. Rust equivalent of Java `BucketVectorSearchSplit`.
+/// A `PrimaryKeyVectorScan` mirror constructs these from a snapshot/manifest
plan.
+pub(crate) struct PkVectorSearchSplit {
+ /// The bucket's combined data split (>= 1 data file); source of the
+ /// partition/bucket/bucket_path/snapshot, the per-file `DataFileMeta`,
and the
+ /// deletion files. Its `data_files()` is the authority for re-associating
a hit's
+ /// file name back to a `DataFileMeta`.
+ pub data_split: DataSplit,
+ /// ANN payload segments for this bucket.
+ pub ann_segments: Vec<BucketAnnSegment>,
+ /// Files eligible for exact fallback.
+ pub active_files: Vec<BucketActiveFile>,
+}
+
+/// A `bucket_search` hit tagged with its source bucket. Rust equivalent of
Java
+/// `PrimaryKeyVectorRead.Candidate`. `partition`/`bucket` are the cross-bucket
+/// merge dimensions a lone `PkVectorSearchResult` lacks; `split_index` is the
+/// re-association handle back to `splits[split_index].data_split`.
+struct Candidate {
+ split_index: usize,
+ partition: BinaryRow,
+ bucket: i32,
+ data_file_name: String,
+ row_position: i64,
+ distance: f32,
+}
+
+/// 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
+/// partition's serialized bytes; Rust `Vec<u8>::cmp` is unsigned lexicographic
+/// then shorter-is-less, exactly the spec's contract (`[0x7f] < [0x80] <
[0xff]`).
+fn candidate_cmp(a: &Candidate, b: &Candidate) -> Ordering {
+ java_float_compare(a.distance, b.distance)
+ .then_with(|| {
+ a.partition
+ .to_serialized_bytes()
+ .cmp(&b.partition.to_serialized_bytes())
+ })
+ .then_with(|| a.bucket.cmp(&b.bucket))
+ .then_with(|| a.data_file_name.cmp(&b.data_file_name))
+ .then_with(|| a.row_position.cmp(&b.row_position))
+}
+
+/// Collect all candidates, order BEST_FIRST, keep the best `limit`.
+fn global_top_k(mut candidates: Vec<Candidate>, limit: usize) ->
Vec<Candidate> {
+ candidates.sort_by(candidate_cmp);
+ candidates.truncate(limit);
+ candidates
+}
+
+/// 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
+/// emitted in ascending group-key order (deterministic file/position output
+/// order). Mirrors Java `PrimaryKeyVectorResult.splits()`.
+fn build_indexed_splits(
+ survivors: Vec<Candidate>,
+ splits: &[PkVectorSearchSplit],
+ metric: VectorSearchMetric,
+) -> crate::Result<Vec<PkVectorIndexedSplit>> {
+ // Group key: (partition bytes, bucket, file name). BTreeMap keeps
ascending
+ // group order deterministically. Value: (split_index, Vec<(position,
distance)>).
+ use std::collections::BTreeMap;
+ type GroupKey = (Vec<u8>, i32, String);
+ let mut groups: BTreeMap<GroupKey, (usize, Vec<(i64, f32)>)> =
BTreeMap::new();
+ for c in survivors {
+ let key = (
+ c.partition.to_serialized_bytes(),
+ c.bucket,
+ c.data_file_name.clone(),
+ );
+ let entry = groups
+ .entry(key)
+ .or_insert_with(|| (c.split_index, Vec::new()));
+ // A (partition, bucket, file_name) group must map to a single source
+ // split. Two candidates sharing the group key but tagged with
different
+ // split_index means malformed input (e.g. duplicate buckets);
+ // silently merging them would materialize against the wrong split, so
+ // fail loud instead.
+ if entry.0 != c.split_index {
+ return Err(data_invalid(format!(
+ "vector search hits for {} map to different splits ({} and
{})",
+ c.data_file_name, entry.0, c.split_index
+ )));
+ }
+ entry.1.push((c.row_position, c.distance));
+ }
+
+ let mut out = Vec::with_capacity(groups.len());
+ for ((_partition, _bucket, file_name), (split_index, mut hits)) in groups {
+ // Sort positions ascending; reject duplicate (file, position).
+ hits.sort_by_key(|(pos, _)| *pos);
+ for pair in hits.windows(2) {
+ if pair[0].0 == pair[1].0 {
+ return Err(data_invalid(format!(
+ "duplicate (file, position) in vector search result: {} @
{}",
+ file_name, pair[0].0
+ )));
+ }
+ }
+
+ // Re-associate the file to its DataFileMeta + aligned deletion file.
+ let source = &splits[split_index].data_split;
+ let file_idx = source
+ .data_files()
+ .iter()
+ .position(|f| f.file_name == file_name)
+ .ok_or_else(|| {
+ data_invalid(format!(
+ "vector search hit references data file {file_name} not
present in its bucket split"
+ ))
+ })?;
+ let file_meta = source.data_files()[file_idx].clone();
+ let deletion_file = source
+ .data_deletion_files()
+ .and_then(|dfs| dfs.get(file_idx).cloned().flatten());
+
+ // Coalesce ascending positions into inclusive ranges; scores aligned
to
+ // ascending-position order.
+ let mut row_ranges: Vec<RowRange> = Vec::new();
+ let mut scores: Vec<f32> = Vec::with_capacity(hits.len());
+ let mut start = hits[0].0;
+ let mut end = hits[0].0;
+ scores.push(metric.distance_to_score(hits[0].1));
+ for &(pos, distance) in &hits[1..] {
+ if pos == end + 1 {
+ end = pos;
+ } else {
+ row_ranges.push(RowRange::new(start, end));
+ start = pos;
+ end = pos;
+ }
+ scores.push(metric.distance_to_score(distance));
+ }
+ row_ranges.push(RowRange::new(start, end));
+
+ let mut builder = DataSplitBuilder::new()
+ .with_snapshot(source.snapshot_id())
+ .with_partition(source.partition().clone())
+ .with_bucket(source.bucket())
+ .with_bucket_path(source.bucket_path().to_string())
+ .with_total_buckets(source.total_buckets())
+ .with_data_files(vec![file_meta]);
+ if let Some(df) = deletion_file {
+ builder = builder.with_data_deletion_files(vec![Some(df)]);
+ }
+ let split = builder.build()?;
+
+ out.push(PkVectorIndexedSplit {
+ split,
+ row_ranges,
+ scores: Some(scores),
+ });
+ }
+ Ok(out)
+}
+
+/// Build one bucket's DV map: keys are the union of active-file names and all
+/// ANN-source file names, so an ANN-source file not in `active_files` still
gets
+/// its DV. Uses one split-level factory. (Search-time DV; materialization
loads
+/// its own DV again — an accepted redundancy.)
+async fn build_bucket_dv_map(
+ reader: &DataFileReader,
+ split: &PkVectorSearchSplit,
+) -> crate::Result<HashMap<String, Arc<DeletionVector>>> {
+ let factory = reader.build_split_dv_factory(&split.data_split).await?;
+ let mut names: Vec<&str> = split
+ .active_files
+ .iter()
+ .map(|f| f.file_name.as_str())
+ .collect();
+ for segment in &split.ann_segments {
+ for source in segment.source_meta.source_files() {
+ names.push(source.file_name());
+ }
+ }
+ let mut dvs = HashMap::new();
+ for name in names {
+ if dvs.contains_key(name) {
+ continue;
+ }
+ if let Some(dv) =
DataFileReader::deletion_vector_for_file(factory.as_ref(), name) {
+ dvs.insert(name.to_string(), dv);
+ }
+ }
+ Ok(dvs)
+}
+
+/// Read orchestrator for the PK-table vector search path. Mirrors Java
+/// `PrimaryKeyVectorRead` + `PrimaryKeyVectorResult.splits()`.
+pub(crate) struct PkVectorOrchestrator {
+ reader: DataFileReader,
+}
+
+impl PkVectorOrchestrator {
+ pub(crate) fn new(reader: DataFileReader) -> Self {
+ Self { reader }
+ }
+
+ /// Run the eager per-bucket search, then lazily materialize the surviving
+ /// rows. `async` because the eager search phase is genuine async IO
+ /// needing the borrowed `exact_reader_factory` / `ann_searcher`; the
returned
+ /// stream owns only the built splits + a reader clone (so it is
`'static`).
+ #[allow(clippy::too_many_arguments)]
+ pub(crate) async fn read(
+ &self,
+ splits: &[PkVectorSearchSplit],
+ query: &[f32],
+ metric: VectorSearchMetric,
+ limit: usize,
+ ann_searcher: Option<&dyn PkVectorAnnSearcher>,
+ exact_reader_factory: &mut dyn FnMut(
+ &BucketActiveFile,
+ ) -> crate::Result<Box<dyn PkVectorReader>>,
+ search_options: &HashMap<String, String>,
+ ) -> crate::Result<ArrowRecordBatchStream> {
+ // Eager input-shape validation (Java checkArgument parity).
+ if limit == 0 {
+ return Err(data_invalid("vector search limit must be positive"));
+ }
+ if query.is_empty() {
+ return Err(data_invalid("vector search query must not be empty"));
+ }
+
+ // Eager per-bucket search -> tagged candidates.
+ let mut candidates: Vec<Candidate> = Vec::new();
+ for (split_index, split) in splits.iter().enumerate() {
+ let dvs = build_bucket_dv_map(&self.reader, split).await?;
+ let results = bucket_search(
+ ann_searcher,
+ &split.ann_segments,
+ &split.active_files,
+ &dvs,
+ exact_reader_factory,
+ query,
+ metric,
+ limit,
+ search_options,
+ )?;
+ for PkVectorSearchResult {
+ data_file_name,
+ row_position,
+ distance,
+ } in results
+ {
+ candidates.push(Candidate {
+ split_index,
+ partition: split.data_split.partition().clone(),
+ bucket: split.data_split.bucket(),
+ data_file_name,
+ row_position,
+ distance,
+ });
+ }
+ }
+
+ // Eager global merge + grouping + split construction.
+ let survivors = global_top_k(candidates, limit);
+ let indexed_splits = build_indexed_splits(survivors, splits, metric)?;
+
+ // Lazy materialization: own the splits + a reader clone.
+ let reader = self.reader.clone();
+ let stream = async_stream::try_stream! {
+ for indexed in indexed_splits {
+ let inner =
PkVectorIndexedSplitRead::new(reader.clone()).read(&indexed)?;
+ futures::pin_mut!(inner);
+ while let Some(batch) = inner.next().await {
+ yield batch?;
+ }
+ }
+ };
+ Ok(Box::pin(stream))
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::spec::stats::BinaryTableStats;
+ use crate::spec::DataFileMeta;
+
+ fn data_file(name: &str, row_count: i64) -> DataFileMeta {
+ DataFileMeta {
+ file_name: name.to_string(),
+ file_size: 1,
+ row_count,
+ min_key: Vec::new(),
+ max_key: Vec::new(),
+ key_stats: BinaryTableStats::empty(),
+ value_stats: BinaryTableStats::empty(),
+ min_sequence_number: 0,
+ max_sequence_number: 0,
+ schema_id: 1,
+ level: 0,
+ extra_files: Vec::new(),
+ creation_time: None,
+ delete_row_count: None,
+ embedded_index: None,
+ file_source: None,
+ value_stats_cols: None,
+ external_path: None,
+ first_row_id: Some(0),
+ write_cols: None,
+ }
+ }
+
+ fn bucket_split(bucket: i32, files: Vec<DataFileMeta>) -> DataSplit {
+ DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(BinaryRow::new(0))
+ .with_bucket(bucket)
+ .with_bucket_path(format!("memory:/pkvo/bucket-{bucket}"))
+ .with_total_buckets(1)
+ .with_data_files(files)
+ .build()
+ .unwrap()
+ }
+
+ fn search_split(bucket: i32, files: Vec<DataFileMeta>) ->
PkVectorSearchSplit {
+ PkVectorSearchSplit {
+ data_split: bucket_split(bucket, files),
+ ann_segments: Vec::new(),
+ active_files: Vec::new(),
+ }
+ }
+
+ // Candidate carrying an empty (arity-0) partition, matching
bucket_split's partition.
+ fn cand(split_index: usize, bucket: i32, file: &str, pos: i64, distance:
f32) -> Candidate {
+ Candidate {
+ split_index,
+ partition: BinaryRow::new(0),
+ bucket,
+ data_file_name: file.to_string(),
+ row_position: pos,
+ distance,
+ }
+ }
+
+ fn candidate(
+ split_index: usize,
+ partition_bytes: Vec<u8>,
+ bucket: i32,
+ file: &str,
+ pos: i64,
+ distance: f32,
+ ) -> Candidate {
+ Candidate {
+ split_index,
+ partition: BinaryRow::from_bytes(1, partition_bytes),
+ bucket,
+ data_file_name: file.to_string(),
+ row_position: pos,
+ distance,
+ }
+ }
+
+ fn ids(c: &[Candidate]) -> Vec<(i32, String, i64)> {
+ c.iter()
+ .map(|c| (c.bucket, c.data_file_name.clone(), c.row_position))
+ .collect()
+ }
+
+ #[test]
+ fn merges_global_top_k_with_deterministic_ties() {
+ // Java
PrimaryKeyVectorReadTest.testMergesGlobalTopKWithDeterministicTies:
+ // (b1,file-c,pos0,d=2), (b1,file-b,pos1,d=1), (b0,file-a,pos2,d=1),
limit=2.
+ // Same partition -> drop file-c (d=2); two d=1 tie on bucket: 0 < 1.
+ // Result: [(0,"file-a",2), (1,"file-b",1)].
+ let part = vec![0x00];
+ let survivors = global_top_k(
+ vec![
+ candidate(0, part.clone(), 1, "file-c", 0, 2.0),
+ candidate(0, part.clone(), 1, "file-b", 1, 1.0),
+ candidate(1, part.clone(), 0, "file-a", 2, 1.0),
+ ],
+ 2,
+ );
+ assert_eq!(
+ ids(&survivors),
+ vec![(0, "file-a".to_string(), 2), (1, "file-b".to_string(), 1)]
+ );
+ }
+
+ #[test]
+ fn orders_partition_bytes_as_unsigned() {
+ // Guards against signed-byte comparison: 0x7f < 0x80 < 0xff
(unsigned).
+ // Equal distance so level 2 (partition bytes) decides.
+ let survivors = global_top_k(
+ vec![
+ candidate(2, vec![0xff], 0, "f", 0, 1.0),
+ candidate(0, vec![0x7f], 0, "f", 0, 1.0),
+ candidate(1, vec![0x80], 0, "f", 0, 1.0),
+ ],
+ 3,
+ );
+ assert_eq!(
+ survivors
+ .iter()
+ .map(|c| c.partition.to_serialized_bytes().pop().unwrap())
+ .collect::<Vec<u8>>(),
+ vec![0x7f, 0x80, 0xff]
+ );
+ }
+
+ #[test]
+ fn truncates_to_limit() {
+ let part = vec![0x00];
+ let survivors = global_top_k(
+ vec![
+ candidate(0, part.clone(), 0, "f", 0, 3.0),
+ candidate(0, part.clone(), 0, "f", 1, 1.0),
+ candidate(0, part.clone(), 0, "f", 2, 2.0),
+ ],
+ 1,
+ );
+ assert_eq!(ids(&survivors), vec![(0, "f".to_string(), 1)]); //
smallest distance
+ }
+
+ #[test]
+ fn empty_candidates_yield_empty() {
+ let survivors = global_top_k(Vec::new(), 5);
+ assert!(survivors.is_empty());
+ }
+
+ #[test]
+ fn builds_two_splits_with_ascending_position_ordered_scores() {
+ // One bucket, two files. file-a hits at global order [pos=10, pos=2];
+ // build must reorder to positions [2,10] with scores [score(2),
score(10)].
+ let splits = vec![search_split(
+ 0,
+ vec![data_file("file-a", 20), data_file("file-b", 20)],
+ )];
+ let survivors = vec![
+ cand(0, 0, "file-a", 10, 3.0),
+ cand(0, 0, "file-b", 5, 2.0),
+ cand(0, 0, "file-a", 2, 1.0),
+ ];
+ let built = build_indexed_splits(survivors, &splits,
VectorSearchMetric::L2).unwrap();
+ assert_eq!(built.len(), 2);
+
+ // Group order is ascending (partition, bucket, name): file-a before
file-b.
+ let a = &built[0];
+ assert_eq!(a.split.data_files()[0].file_name, "file-a");
+ assert_eq!(
+ a.row_ranges,
+ vec![RowRange::new(2, 2), RowRange::new(10, 10)]
+ );
+ // scores in ascending-position order: score(d=1.0) for pos2,
score(d=3.0) for pos10.
+ assert_eq!(
+ a.scores.as_deref(),
+ Some(
+ [
+ VectorSearchMetric::L2.distance_to_score(1.0),
+ VectorSearchMetric::L2.distance_to_score(3.0),
+ ]
+ .as_slice()
+ )
+ );
+
+ let b = &built[1];
+ assert_eq!(b.split.data_files()[0].file_name, "file-b");
+ assert_eq!(b.row_ranges, vec![RowRange::new(5, 5)]);
+ }
+
+ #[test]
+ fn coalesces_consecutive_positions_into_one_range() {
+ let splits = vec![search_split(0, vec![data_file("f", 10)])];
+ let survivors = vec![
+ cand(0, 0, "f", 0, 1.0),
+ cand(0, 0, "f", 1, 1.0),
+ cand(0, 0, "f", 2, 1.0),
+ cand(0, 0, "f", 5, 1.0),
+ ];
+ let built = build_indexed_splits(survivors, &splits,
VectorSearchMetric::L2).unwrap();
+ assert_eq!(
+ built[0].row_ranges,
+ vec![RowRange::new(0, 2), RowRange::new(5, 5)]
+ );
+ }
+
+ #[test]
+ fn rejects_file_absent_from_bucket_split() {
+ let splits = vec![search_split(0, vec![data_file("known", 10)])];
+ let survivors = vec![cand(0, 0, "unknown", 0, 1.0)];
+ // PkVectorIndexedSplit has no Debug; map the Ok value away so
expect_err's
+ // `T: Debug` bound is satisfied without touching the shared type.
+ let err = build_indexed_splits(survivors, &splits,
VectorSearchMetric::L2)
+ .map(|_| ())
+ .expect_err("unknown file must error");
+ assert!(
+ format!("{err:?}").contains("unknown") ||
format!("{err:?}").contains("not"),
+ "got: {err:?}"
+ );
+ }
+
+ #[test]
+ fn rejects_duplicate_file_position() {
+ let splits = vec![search_split(0, vec![data_file("f", 10)])];
+ let survivors = vec![cand(0, 0, "f", 3, 1.0), cand(0, 0, "f", 3, 2.0)];
+ let err = build_indexed_splits(survivors, &splits,
VectorSearchMetric::L2)
+ .map(|_| ())
+ .expect_err("duplicate (file,pos) must error");
+ assert!(format!("{err:?}").contains("duplicate"), "got: {err:?}");
+ }
+
+ #[test]
+ fn rejects_same_group_key_from_different_splits() {
+ // Two buckets share (partition, bucket, file_name) but sit at
different
+ // split_index. Silently merging them would materialize against the
wrong
+ // split; fail loud instead (defensive guard against malformed
+ // input). Both search_splits have the empty partition + bucket 0 +
file "f".
+ let splits = vec![
+ search_split(0, vec![data_file("f", 10)]),
+ search_split(0, vec![data_file("f", 10)]),
+ ];
+ let survivors = vec![cand(0, 0, "f", 1, 1.0), cand(1, 0, "f", 2, 1.0)];
+ let err = build_indexed_splits(survivors, &splits,
VectorSearchMetric::L2)
+ .map(|_| ())
+ .expect_err("same group key from different splits must error");
+ assert!(
+ format!("{err:?}").contains("different splits")
+ || format!("{err:?}").contains("distinct splits"),
+ "got: {err:?}"
+ );
+ }
+}
+
+#[cfg(test)]
+mod e2e_tests {
+ use super::*;
+ use crate::arrow::build_target_arrow_schema;
+ use crate::io::{FileIO, FileIOBuilder};
+ use crate::spec::stats::BinaryTableStats;
+ use crate::spec::{
+ DataField, DataFileMeta, DataType, IntType, PkVectorSourceFile,
PkVectorSourceMeta,
+ };
+ use crate::table::pk_vector_position_read::{
+ PKEY_VECTOR_POSITION_COLUMN, PKEY_VECTOR_SCORE_COLUMN,
+ };
+ use crate::table::schema_manager::SchemaManager;
+ use crate::table::source::DeletionFile;
+ use crate::vindex::pkvector::reader::test_support::ArrayReader;
+ use arrow_array::{Array, Float32Array, Int32Array, Int64Array,
RecordBatch};
+ use bytes::Bytes;
+ use futures::TryStreamExt;
+ use paimon_mosaic_core::spec::COMPRESSION_NONE;
+ use paimon_mosaic_core::writer::{MosaicWriter, OutputFile, WriterOptions};
+ use roaring::RoaringBitmap;
+ use std::collections::HashSet;
+ use std::io;
+
+ struct MemOutputFile {
+ data: Vec<u8>,
+ }
+ impl MemOutputFile {
+ fn new() -> Self {
+ Self { data: Vec::new() }
+ }
+ }
+ impl OutputFile for MemOutputFile {
+ fn write(&mut self, data: &[u8]) -> io::Result<()> {
+ self.data.extend_from_slice(data);
+ Ok(())
+ }
+ fn flush(&mut self) -> io::Result<()> {
+ Ok(())
+ }
+ fn pos(&self) -> u64 {
+ self.data.len() as u64
+ }
+ }
+
+ fn id_field() -> DataField {
+ DataField::new(0, "id".to_string(), DataType::Int(IntType::new()))
+ }
+ fn id_fields() -> Vec<DataField> {
+ vec![id_field()]
+ }
+ fn id_batch(ids: Vec<i32>) -> RecordBatch {
+ let schema = build_target_arrow_schema(&id_fields()).unwrap();
+ RecordBatch::try_new(schema,
vec![Arc::new(Int32Array::from(ids))]).unwrap()
+ }
+
+ fn data_file(file_name: &str, file_size: i64, row_count: i64) ->
DataFileMeta {
+ DataFileMeta {
+ file_name: file_name.to_string(),
+ file_size,
+ row_count,
+ min_key: Vec::new(),
+ max_key: Vec::new(),
+ key_stats: BinaryTableStats::empty(),
+ value_stats: BinaryTableStats::empty(),
+ min_sequence_number: 0,
+ max_sequence_number: 0,
+ schema_id: 1,
+ level: 0,
+ extra_files: Vec::new(),
+ creation_time: None,
+ delete_row_count: None,
+ embedded_index: None,
+ file_source: None,
+ value_stats_cols: None,
+ external_path: None,
+ first_row_id: Some(0),
+ write_cols: None,
+ }
+ }
+
+ fn write_mosaic_single_group(batch: &RecordBatch) -> Bytes {
+ let out = MemOutputFile::new();
+ let mut writer = MosaicWriter::new(
+ out,
+ batch.schema().as_ref(),
+ WriterOptions {
+ compression: COMPRESSION_NONE,
+ num_buckets: 2,
+ row_group_max_size: u64::MAX,
+ ..Default::default()
+ },
+ )
+ .unwrap();
+ writer.write_batch(batch).unwrap();
+ writer.close().unwrap();
+ Bytes::from(writer.output().data.to_vec())
+ }
+
+ async fn write_deletion_file(
+ file_io: &FileIO,
+ path: &str,
+ deleted_rows: &[u32],
+ ) -> DeletionFile {
+ 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::from(blob))
+ .await
+ .unwrap();
+ DeletionFile::new(
+ path.to_string(),
+ 0,
+ bitmap_length as i64,
+ Some(deleted_rows.len() as i64),
+ )
+ }
+
+ fn make_reader(file_io: FileIO, table_path: &str) -> DataFileReader {
+ let schema_manager = SchemaManager::new(file_io.clone(),
table_path.to_string());
+ DataFileReader::new(
+ file_io,
+ schema_manager,
+ 1,
+ id_fields(),
+ id_fields(),
+ Vec::new(),
+ )
+ }
+
+ /// Write one mosaic data file into a bucket path and return its
`DataFileMeta`.
+ async fn write_file(
+ file_io: &FileIO,
+ bucket_path: &str,
+ file_name: &str,
+ ids: Vec<i32>,
+ ) -> DataFileMeta {
+ let row_count = ids.len() as i64;
+ let data = write_mosaic_single_group(&id_batch(ids));
+ file_io
+ .new_output(&format!("{bucket_path}/{file_name}"))
+ .unwrap()
+ .write(data.clone())
+ .await
+ .unwrap();
+ data_file(file_name, data.len() as i64, row_count)
+ }
+
+ fn ann_segment(sources: &[(&str, i64)]) -> BucketAnnSegment {
+ BucketAnnSegment {
+ source_meta: PkVectorSourceMeta::new(
+ sources
+ .iter()
+ .map(|(n, r)| PkVectorSourceFile::new((*n).to_string(),
*r).unwrap())
+ .collect(),
+ )
+ .unwrap(),
+ }
+ }
+
+ fn active(name: &str, rows: i64) -> BucketActiveFile {
+ BucketActiveFile {
+ file_name: name.to_string(),
+ row_count: rows,
+ }
+ }
+
+ fn column_by_name<'a>(batch: &'a RecordBatch, name: &str) -> Option<&'a
Arc<dyn Array>> {
+ batch
+ .schema()
+ .index_of(name)
+ .ok()
+ .map(|idx| batch.column(idx))
+ }
+ fn collect_i32(batches: &[RecordBatch], name: &str) -> Vec<i32> {
+ batches
+ .iter()
+ .flat_map(|b| {
+ column_by_name(b, name)
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ })
+ .collect()
+ }
+ fn collect_i64(batches: &[RecordBatch], name: &str) -> Vec<i64> {
+ batches
+ .iter()
+ .flat_map(|b| {
+ column_by_name(b, name)
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ })
+ .collect()
+ }
+ fn collect_f32(batches: &[RecordBatch], name: &str) -> Vec<f32> {
+ batches
+ .iter()
+ .flat_map(|b| {
+ column_by_name(b, name)
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Float32Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ })
+ .collect()
+ }
+
+ fn l2_score(distance: f32) -> f32 {
+ VectorSearchMetric::L2.distance_to_score(distance)
+ }
+
+ // Fake ANN searcher returning preset hits.
+ struct FakeAnn {
+ hits: Vec<PkVectorSearchResult>,
+ }
+ impl PkVectorAnnSearcher for FakeAnn {
+ 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.hits.clone())
+ }
+ }
+
+ #[tokio::test]
+ async fn eager_rejects_zero_limit() {
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let reader = make_reader(file_io, "memory:/pkvo_zero");
+ let splits: Vec<PkVectorSearchSplit> = Vec::new();
+ let mut factory = |_: &BucketActiveFile| -> crate::Result<Box<dyn
PkVectorReader>> {
+ unreachable!("no bucket search on eager-rejected input")
+ };
+ let opts = HashMap::new();
+ let err = PkVectorOrchestrator::new(reader)
+ .read(
+ &splits,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 0,
+ None,
+ &mut factory,
+ &opts,
+ )
+ .await
+ .map(|_| ())
+ .expect_err("limit == 0 must be rejected eagerly");
+ assert!(format!("{err:?}").contains("positive"), "got: {err:?}");
+ }
+
+ #[tokio::test]
+ async fn eager_rejects_empty_query() {
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let reader = make_reader(file_io, "memory:/pkvo_empty_query");
+ let splits: Vec<PkVectorSearchSplit> = Vec::new();
+ let mut factory = |_: &BucketActiveFile| -> crate::Result<Box<dyn
PkVectorReader>> {
+ unreachable!("no bucket search on eager-rejected input")
+ };
+ let opts = HashMap::new();
+ let err = PkVectorOrchestrator::new(reader)
+ .read(
+ &splits,
+ &[],
+ VectorSearchMetric::L2,
+ 5,
+ None,
+ &mut factory,
+ &opts,
+ )
+ .await
+ .map(|_| ())
+ .expect_err("empty query must be rejected eagerly");
+ assert!(format!("{err:?}").contains("empty"), "got: {err:?}");
+ }
+
+ #[tokio::test]
+ async fn
single_bucket_ann_plus_exact_merge_materializes_with_position_and_score() {
+ // One bucket, two files: "ann.mosaic" is ANN-covered (FakeAnn hit at
pos 1,
+ // distance 0.25); "exact.mosaic" is exact fallback (ArrayReader).
Covered
+ // files are NOT re-scanned by the exact fallback.
+ let table_path = "memory:/pkvo_single";
+ let bucket_path = format!("{table_path}/bucket-0");
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let ann_meta = write_file(&file_io, &bucket_path, "ann.mosaic",
vec![100, 101, 102]).await;
+ let exact_meta = write_file(&file_io, &bucket_path, "exact.mosaic",
vec![200, 201]).await;
+
+ 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![ann_meta, exact_meta])
+ .build()
+ .unwrap();
+ let split = PkVectorSearchSplit {
+ data_split,
+ ann_segments: vec![ann_segment(&[("ann.mosaic", 3)])],
+ active_files: vec![active("ann.mosaic", 3), active("exact.mosaic",
2)],
+ };
+
+ let ann = FakeAnn {
+ hits: vec![PkVectorSearchResult {
+ data_file_name: "ann.mosaic".to_string(),
+ row_position: 1,
+ distance: 0.25,
+ }],
+ };
+ // Exact fallback scans only "exact.mosaic": pos0 {1,0} d=1.0, pos1
{2,0} d=4.0.
+ let mut factory = |f: &BucketActiveFile| -> crate::Result<Box<dyn
PkVectorReader>> {
+ let vectors = match f.file_name.as_str() {
+ "exact.mosaic" => vec![Some(vec![1.0, 0.0]), Some(vec![2.0,
0.0])],
+ other => panic!("unexpected exact scan of covered file
{other}"),
+ };
+ Ok(Box::new(ArrayReader::new(2, vectors)))
+ };
+ let opts = HashMap::new();
+ let batches = PkVectorOrchestrator::new(make_reader(file_io,
table_path))
+ .read(
+ &[split],
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 3,
+ Some(&ann),
+ &mut factory,
+ &opts,
+ )
+ .await
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ // Output is ascending group (file name) then ascending position:
+ // ann.mosaic pos1 -> id 101; exact.mosaic pos0,1 -> ids 200,201.
+ assert_eq!(collect_i32(&batches, "id"), vec![101, 200, 201]);
+ assert_eq!(
+ collect_i64(&batches, PKEY_VECTOR_POSITION_COLUMN),
+ vec![1, 0, 1]
+ );
+ assert_eq!(
+ collect_f32(&batches, PKEY_VECTOR_SCORE_COLUMN),
+ vec![l2_score(0.25), l2_score(1.0), l2_score(4.0)]
+ );
+ for batch in &batches {
+ assert!(
+ column_by_name(batch, "_ROW_ID").is_none(),
+ "_ROW_ID must not leak"
+ );
+ }
+ }
+
+ #[tokio::test]
+ async fn multi_bucket_merge_keeps_global_top_k() {
+ // Two buckets, exact-only. limit=3 < 6 total hits; surviving rows are
the
+ // global-best 3 by distance across both buckets.
+ let table_path = "memory:/pkvo_multi";
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+
+ let b0_path = format!("{table_path}/bucket-0");
+ let b0_meta = write_file(&file_io, &b0_path, "b0.mosaic", vec![10, 11,
12]).await;
+ let split0 = PkVectorSearchSplit {
+ data_split: DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path(b0_path)
+ .with_total_buckets(2)
+ .with_data_files(vec![b0_meta])
+ .build()
+ .unwrap(),
+ ann_segments: Vec::new(),
+ active_files: vec![active("b0.mosaic", 3)],
+ };
+
+ let b1_path = format!("{table_path}/bucket-1");
+ let b1_meta = write_file(&file_io, &b1_path, "b1.mosaic", vec![20, 21,
22]).await;
+ let split1 = PkVectorSearchSplit {
+ data_split: DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(BinaryRow::new(0))
+ .with_bucket(1)
+ .with_bucket_path(b1_path)
+ .with_total_buckets(2)
+ .with_data_files(vec![b1_meta])
+ .build()
+ .unwrap(),
+ ann_segments: Vec::new(),
+ active_files: vec![active("b1.mosaic", 3)],
+ };
+
+ // b0: x = 1,4,6 -> d = 1,16,36. b1: x = 2,3,5 -> d = 4,9,25.
+ // Global best 3: d1 (b0 pos0 id10), d4 (b1 pos0 id20), d9 (b1 pos1
id21).
+ let mut factory = |f: &BucketActiveFile| -> crate::Result<Box<dyn
PkVectorReader>> {
+ let vectors = match f.file_name.as_str() {
+ "b0.mosaic" => vec![
+ Some(vec![1.0, 0.0]),
+ Some(vec![4.0, 0.0]),
+ Some(vec![6.0, 0.0]),
+ ],
+ "b1.mosaic" => vec![
+ Some(vec![2.0, 0.0]),
+ Some(vec![3.0, 0.0]),
+ Some(vec![5.0, 0.0]),
+ ],
+ other => panic!("unexpected file {other}"),
+ };
+ Ok(Box::new(ArrayReader::new(2, vectors)))
+ };
+ let opts = HashMap::new();
+ let batches = PkVectorOrchestrator::new(make_reader(file_io,
table_path))
+ .read(
+ &[split0, split1],
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 3,
+ None,
+ &mut factory,
+ &opts,
+ )
+ .await
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ // Ascending group order: bucket0 "b0.mosaic" pos0 -> 10; bucket1
"b1.mosaic"
+ // pos0,1 -> 20,21.
+ assert_eq!(collect_i32(&batches, "id"), vec![10, 20, 21]);
+ assert_eq!(
+ collect_i64(&batches, PKEY_VECTOR_POSITION_COLUMN),
+ vec![0, 0, 1]
+ );
+ assert_eq!(
+ collect_f32(&batches, PKEY_VECTOR_SCORE_COLUMN),
+ vec![l2_score(1.0), l2_score(4.0), l2_score(9.0)]
+ );
+ }
+
+ #[tokio::test]
+ async fn dv_deleted_exact_position_is_absent_from_output() {
+ // Exact fallback over "d.mosaic" with a DV deleting position 1. The
deleted
+ // position is absent; remaining position/score alignment holds.
+ let table_path = "memory:/pkvo_dv";
+ let bucket_path = format!("{table_path}/bucket-0");
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let meta = write_file(&file_io, &bucket_path, "d.mosaic", vec![30, 31,
32, 33]).await;
+ let df = write_deletion_file(&file_io,
&format!("{table_path}/index/dv-0"), &[1]).await;
+ 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![meta])
+ .with_data_deletion_files(vec![Some(df)])
+ .build()
+ .unwrap();
+ let split = PkVectorSearchSplit {
+ data_split,
+ ann_segments: Vec::new(),
+ active_files: vec![active("d.mosaic", 4)],
+ };
+
+ // pos0 {1,0} d=1, pos1 {2,0} d=4 (DELETED), pos2 {3,0} d=9, pos3
{0,0} d=0.
+ let mut factory = |f: &BucketActiveFile| -> crate::Result<Box<dyn
PkVectorReader>> {
+ let vectors = match f.file_name.as_str() {
+ "d.mosaic" => vec![
+ Some(vec![1.0, 0.0]),
+ Some(vec![2.0, 0.0]),
+ Some(vec![3.0, 0.0]),
+ Some(vec![0.0, 0.0]),
+ ],
+ other => panic!("unexpected file {other}"),
+ };
+ Ok(Box::new(ArrayReader::new(2, vectors)))
+ };
+ let opts = HashMap::new();
+ let batches = PkVectorOrchestrator::new(make_reader(file_io,
table_path))
+ .read(
+ &[split],
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 4,
+ None,
+ &mut factory,
+ &opts,
+ )
+ .await
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ // Position 1 (id 31) is absent. Remaining ascending positions 0,2,3.
+ assert_eq!(collect_i32(&batches, "id"), vec![30, 32, 33]);
+ assert_eq!(
+ collect_i64(&batches, PKEY_VECTOR_POSITION_COLUMN),
+ vec![0, 2, 3]
+ );
+ assert_eq!(
+ collect_f32(&batches, PKEY_VECTOR_SCORE_COLUMN),
+ vec![l2_score(1.0), l2_score(9.0), l2_score(0.0)]
+ );
+ }
+
+ #[tokio::test]
+ async fn output_is_file_position_order_not_best_first() {
+ // Best-first order (by distance) differs from file/position order.
Output must
+ // be ascending file/position order.
+ let table_path = "memory:/pkvo_order";
+ let bucket_path = format!("{table_path}/bucket-0");
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let meta = write_file(&file_io, &bucket_path, "o.mosaic", vec![40, 41,
42]).await;
+ let split = PkVectorSearchSplit {
+ data_split: DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path(bucket_path)
+ .with_total_buckets(1)
+ .with_data_files(vec![meta])
+ .build()
+ .unwrap(),
+ ann_segments: Vec::new(),
+ active_files: vec![active("o.mosaic", 3)],
+ };
+
+ // pos0 {3,0} d=9, pos1 {1,0} d=1, pos2 {2,0} d=4. Best-first =
[1,2,0].
+ let mut factory = |f: &BucketActiveFile| -> crate::Result<Box<dyn
PkVectorReader>> {
+ let vectors = match f.file_name.as_str() {
+ "o.mosaic" => vec![
+ Some(vec![3.0, 0.0]),
+ Some(vec![1.0, 0.0]),
+ Some(vec![2.0, 0.0]),
+ ],
+ other => panic!("unexpected file {other}"),
+ };
+ Ok(Box::new(ArrayReader::new(2, vectors)))
+ };
+ let opts = HashMap::new();
+ let batches = PkVectorOrchestrator::new(make_reader(file_io,
table_path))
+ .read(
+ &[split],
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 3,
+ None,
+ &mut factory,
+ &opts,
+ )
+ .await
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ // Ascending physical position order, not best-first distance order.
+ assert_eq!(collect_i32(&batches, "id"), vec![40, 41, 42]);
+ assert_eq!(
+ collect_i64(&batches, PKEY_VECTOR_POSITION_COLUMN),
+ vec![0, 1, 2]
+ );
+ // Scores aligned to ascending position: d=9,1,4.
+ assert_eq!(
+ collect_f32(&batches, PKEY_VECTOR_SCORE_COLUMN),
+ vec![l2_score(9.0), l2_score(1.0), l2_score(4.0)]
+ );
+ }
+}
diff --git a/crates/paimon/src/table/pk_vector_position_read.rs
b/crates/paimon/src/table/pk_vector_position_read.rs
new file mode 100644
index 0000000..6a5ec45
--- /dev/null
+++ b/crates/paimon/src/table/pk_vector_position_read.rs
@@ -0,0 +1,1079 @@
+// 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 position read (reader-kernel subset of
apache/paimon#8576).
+//!
+//! Materializes the selected physical rows of one data file and appends
+//! `_PKEY_VECTOR_POSITION` (+ optional `_PKEY_VECTOR_SCORE`) metadata columns.
+//! Rust equivalent of Java `PrimaryKeyVectorPositionReader`. This is the
lowest
+//! layer of the PK-vector read kernel; the sibling
`pk_vector_indexed_split_read`
+//! and `pk_vector_orchestrator` modules build the indexed-split contract and
+//! cross-bucket merge on top of it.
+
+// The read path that would call these items lives in a later change, so under
+// clippy -D warnings they read as dead_code. Suppress at the module boundary.
+#![allow(dead_code)]
+
+use std::collections::BTreeMap;
+use std::sync::Arc;
+
+use arrow_array::{Array, Float32Array, Int64Array, RecordBatch};
+use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Fields,
Schema as ArrowSchema};
+use futures::StreamExt;
+
+use crate::deletion_vector::DeletionVector;
+use crate::spec::{
+ BigIntType, DataField, DataFileMeta, DataType, ROW_ID_FIELD_ID,
ROW_ID_FIELD_NAME,
+};
+use crate::table::data_file_reader::DataFileReader;
+use crate::table::source::DataSplit;
+use crate::table::{ArrowRecordBatchStream, RowRange};
+
+pub(crate) const PKEY_VECTOR_POSITION_COLUMN: &str = "_PKEY_VECTOR_POSITION";
+pub(crate) const PKEY_VECTOR_SCORE_COLUMN: &str = "_PKEY_VECTOR_SCORE";
+
+fn data_invalid(message: impl Into<String>) -> crate::Error {
+ crate::Error::DataInvalid {
+ message: message.into(),
+ source: None,
+ }
+}
+
+/// Coalesce sorted, de-duplicated 0-based physical positions into contiguous
+/// global `RowRange`s (each bound offset by `first_row_id`). The read path
+/// converts these back to local ranges via
`to_local_row_ranges(first_row_id)`,
+/// so this round-trips to exactly the requested physical positions.
+fn positions_to_global_ranges(
+ sorted_positions: &[i64],
+ first_row_id: i64,
+) -> crate::Result<Vec<RowRange>> {
+ let mut ranges = Vec::new();
+ let mut iter = sorted_positions.iter().copied();
+ let Some(first) = iter.next() else {
+ return Ok(ranges);
+ };
+ let mut push_range = |start: i64, end: i64| -> crate::Result<()> {
+ let from = start
+ .checked_add(first_row_id)
+ .ok_or_else(|| data_invalid("vector position offset overflows
i64"))?;
+ let to = end
+ .checked_add(first_row_id)
+ .ok_or_else(|| data_invalid("vector position offset overflows
i64"))?;
+ ranges.push(RowRange::new(from, to));
+ Ok(())
+ };
+ let mut start = first;
+ let mut end = first;
+ for pos in iter {
+ if end.checked_add(1) == Some(pos) {
+ end = pos;
+ } else {
+ push_range(start, end)?;
+ start = pos;
+ end = pos;
+ }
+ }
+ push_range(start, end)?;
+ Ok(ranges)
+}
+
+/// Reads selected physical rows of one data file, appending position (+
optional
+/// score) metadata columns. Rust equivalent of Java
+/// `PrimaryKeyVectorPositionReader` (apache/paimon#8576, reader-kernel
subset).
+pub(crate) struct PkVectorPositionRead<'a> {
+ reader: &'a DataFileReader,
+}
+
+impl<'a> PkVectorPositionRead<'a> {
+ pub(crate) fn new(reader: &'a DataFileReader) -> Self {
+ Self { reader }
+ }
+
+ pub(crate) fn read(
+ &self,
+ split: &DataSplit,
+ file_meta: DataFileMeta,
+ data_fields: Option<Vec<DataField>>,
+ dv: Option<Arc<DeletionVector>>,
+ positions: impl IntoIterator<Item = i64>,
+ scores: Option<BTreeMap<i64, f32>>,
+ ) -> crate::Result<ArrowRecordBatchStream> {
+ // (1) normalize + validate positions
+ let mut sorted: Vec<i64> = positions.into_iter().collect();
+ sorted.sort_unstable();
+ sorted.dedup();
+ if sorted.is_empty() {
+ return Err(data_invalid("Selected row positions must not be
empty"));
+ }
+ if let Some(&first) = sorted.first() {
+ if first < 0 {
+ return Err(data_invalid(format!(
+ "Vector position must not be negative: {first}"
+ )));
+ }
+ }
+ if let Some(&last) = sorted.last() {
+ if last >= file_meta.row_count {
+ return Err(data_invalid(format!(
+ "Vector position {last} is outside data file row count {}",
+ file_meta.row_count
+ )));
+ }
+ }
+
+ // (2) scores contract on the input
+ if let Some(scores) = scores.as_ref() {
+ let key_ok =
+ scores.len() == sorted.len() &&
scores.keys().copied().eq(sorted.iter().copied());
+ if !key_ok {
+ return Err(data_invalid(
+ "Scores keys must exactly match the selected row
positions",
+ ));
+ }
+ }
+
+ // (3) reserved-column-name check against the requested output fields
+ for field in self.reader.read_type() {
+ if field.name() == PKEY_VECTOR_POSITION_COLUMN
+ || field.name() == PKEY_VECTOR_SCORE_COLUMN
+ {
+ return Err(data_invalid(format!(
+ "Reserved metadata column name conflicts with a table
column: {}",
+ field.name()
+ )));
+ }
+ if field.name() == ROW_ID_FIELD_NAME {
+ return Err(data_invalid(
+ "PK vector position read does not support a requested
_ROW_ID column; \
+ it is used internally for physical-position recovery",
+ ));
+ }
+ }
+
+ // (4) predicate guard: a row-filtering predicate would desync
positional
+ // row-id recovery, so reject a reader that carries one.
+ if self.reader.has_row_filtering_predicate() {
+ return Err(data_invalid(
+ "PK vector position read requires a predicate-free reader",
+ ));
+ }
+
+ // (5) first_row_id guard
+ let first_row_id = file_meta.first_row_id.ok_or_else(|| {
+ data_invalid("PK vector position read requires a data file with
first_row_id")
+ })?;
+
+ // Build a reader whose read_type includes _ROW_ID so the lower-level
read
+ // emits real global row-ids we can convert to physical positions. A
+ // caller-requested _ROW_ID was already rejected in (3), so it is
always
+ // absent here and appended fresh.
+ let mut inner_read_type = self.reader.read_type().to_vec();
+ inner_read_type.push(DataField::new(
+ ROW_ID_FIELD_ID,
+ ROW_ID_FIELD_NAME.to_string(),
+ DataType::BigInt(BigIntType::new()),
+ ));
+ let inner_reader = self.reader.clone().with_read_type(inner_read_type);
+
+ let ranges = positions_to_global_ranges(&sorted, first_row_id)?;
+ let inner = inner_reader.read_single_file_stream(
+ split,
+ file_meta.clone(),
+ data_fields,
+ dv,
+ Some(ranges),
+ )?;
+
+ let want_score_col = scores.is_some();
+
+ let stream = async_stream::try_stream! {
+ futures::pin_mut!(inner);
+ while let Some(batch) = inner.next().await {
+ let batch = batch?;
+ let out = append_metadata_columns(
+ batch,
+ first_row_id,
+ want_score_col,
+ scores.as_ref(),
+ )?;
+ yield out;
+ }
+ };
+ Ok(Box::pin(stream))
+ }
+}
+
+/// Split `batch` (which contains a `_ROW_ID` column) into an output batch with
+/// `_ROW_ID` removed and `_PKEY_VECTOR_POSITION` (+ optional
`_PKEY_VECTOR_SCORE`)
+/// appended. Positions are `_ROW_ID - first_row_id`; scores are looked up by
the
+/// returned position. Returns the new batch.
+fn append_metadata_columns(
+ batch: RecordBatch,
+ first_row_id: i64,
+ want_score_col: bool,
+ scores: Option<&BTreeMap<i64, f32>>,
+) -> crate::Result<RecordBatch> {
+ let schema = batch.schema();
+ let row_id_idx = schema
+ .index_of(ROW_ID_FIELD_NAME)
+ .map_err(|_| data_invalid("internal: _ROW_ID column missing from
position read"))?;
+ let row_ids = batch
+ .column(row_id_idx)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .ok_or_else(|| data_invalid("internal: _ROW_ID column is not Int64"))?;
+
+ let mut positions = Vec::with_capacity(row_ids.len());
+ let mut score_vals = if want_score_col {
+ Some(Vec::with_capacity(row_ids.len()))
+ } else {
+ None
+ };
+ for i in 0..row_ids.len() {
+ let position = row_ids.value(i) - first_row_id;
+ positions.push(position);
+ if let Some(sv) = score_vals.as_mut() {
+ let score = scores
+ .and_then(|m| m.get(&position).copied())
+ .ok_or_else(|| {
+ data_invalid(format!(
+ "internal: no score for returned position {position}"
+ ))
+ })?;
+ sv.push(score);
+ }
+ }
+
+ // Rebuild columns/fields excluding _ROW_ID, then append metadata columns.
+ let mut fields: Vec<ArrowField> = Vec::new();
+ let mut columns: Vec<Arc<dyn Array>> = Vec::new();
+ for (i, f) in schema.fields().iter().enumerate() {
+ if i == row_id_idx {
+ continue;
+ }
+ fields.push(f.as_ref().clone());
+ columns.push(batch.column(i).clone());
+ }
+ fields.push(ArrowField::new(
+ PKEY_VECTOR_POSITION_COLUMN,
+ ArrowDataType::Int64,
+ false,
+ ));
+ columns.push(Arc::new(Int64Array::from(positions)));
+ if let Some(sv) = score_vals {
+ fields.push(ArrowField::new(
+ PKEY_VECTOR_SCORE_COLUMN,
+ ArrowDataType::Float32,
+ false,
+ ));
+ columns.push(Arc::new(Float32Array::from(sv)));
+ }
+
+ let out_schema = Arc::new(ArrowSchema::new(Fields::from(fields)));
+ RecordBatch::try_new(out_schema, columns).map_err(|e|
crate::Error::UnexpectedError {
+ message: format!("Failed to build position-read batch: {e}"),
+ source: Some(Box::new(e)),
+ })
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::arrow::build_target_arrow_schema;
+ use crate::deletion_vector::DeletionVectorFactory;
+ use crate::io::FileIOBuilder;
+ use crate::spec::stats::BinaryTableStats;
+ use crate::spec::{
+ DataFileMeta, DataType, Datum, IntType, PredicateBuilder,
ROW_ID_FIELD_NAME,
+ };
+ use crate::table::data_file_reader::DataFileReader;
+ use crate::table::schema_manager::SchemaManager;
+ use crate::table::source::{DataSplit, DataSplitBuilder, DeletionFile};
+ use arrow_array::{Array, Float32Array, Int32Array, Int64Array,
RecordBatch};
+ use bytes::Bytes;
+ use futures::TryStreamExt;
+ use paimon_mosaic_core::spec::COMPRESSION_NONE;
+ use paimon_mosaic_core::writer::{MosaicWriter, OutputFile, WriterOptions};
+ use roaring::RoaringBitmap;
+ use std::collections::BTreeMap;
+ use std::io;
+ use std::sync::Arc;
+
+ struct MemOutputFile {
+ data: Vec<u8>,
+ }
+
+ impl MemOutputFile {
+ fn new() -> Self {
+ Self { data: Vec::new() }
+ }
+ }
+
+ impl OutputFile for MemOutputFile {
+ fn write(&mut self, data: &[u8]) -> io::Result<()> {
+ self.data.extend_from_slice(data);
+ Ok(())
+ }
+
+ fn flush(&mut self) -> io::Result<()> {
+ Ok(())
+ }
+
+ fn pos(&self) -> u64 {
+ self.data.len() as u64
+ }
+ }
+
+ fn id_field() -> DataField {
+ DataField::new(0, "id".to_string(), DataType::Int(IntType::new()))
+ }
+
+ /// A single `id: Int32` field is the reader's read-type in every test.
+ fn id_fields() -> Vec<DataField> {
+ vec![id_field()]
+ }
+
+ fn id_batch(ids: Vec<i32>) -> RecordBatch {
+ let schema = build_target_arrow_schema(&id_fields()).unwrap();
+ RecordBatch::try_new(schema,
vec![Arc::new(Int32Array::from(ids))]).unwrap()
+ }
+
+ /// `DataFileMeta` with a specified `first_row_id` (the read path requires
it).
+ fn data_file(
+ file_name: &str,
+ file_size: i64,
+ row_count: i64,
+ schema_id: i64,
+ first_row_id: Option<i64>,
+ ) -> DataFileMeta {
+ DataFileMeta {
+ file_name: file_name.to_string(),
+ file_size,
+ row_count,
+ min_key: Vec::new(),
+ max_key: Vec::new(),
+ key_stats: BinaryTableStats::empty(),
+ value_stats: BinaryTableStats::empty(),
+ min_sequence_number: 0,
+ max_sequence_number: 0,
+ schema_id,
+ level: 0,
+ extra_files: Vec::new(),
+ creation_time: None,
+ delete_row_count: None,
+ embedded_index: None,
+ file_source: None,
+ value_stats_cols: None,
+ external_path: None,
+ first_row_id,
+ write_cols: None,
+ }
+ }
+
+ fn write_mosaic_single_group(batch: &RecordBatch) -> Bytes {
+ let out = MemOutputFile::new();
+ let mut writer = MosaicWriter::new(
+ out,
+ batch.schema().as_ref(),
+ WriterOptions {
+ compression: COMPRESSION_NONE,
+ num_buckets: 2,
+ row_group_max_size: u64::MAX,
+ ..Default::default()
+ },
+ )
+ .unwrap();
+ writer.write_batch(batch).unwrap();
+ writer.close().unwrap();
+ Bytes::from(writer.output().data.to_vec())
+ }
+
+ /// Writes one row group per batch, so the reader yields one Arrow batch
per
+ /// input batch — used for the cross-batch alignment test.
+ fn write_mosaic_multi_group(batches: &[RecordBatch]) -> Bytes {
+ let out = MemOutputFile::new();
+ let mut writer = MosaicWriter::new(
+ out,
+ batches[0].schema().as_ref(),
+ WriterOptions {
+ compression: COMPRESSION_NONE,
+ num_buckets: 2,
+ row_group_max_size: 1,
+ ..Default::default()
+ },
+ )
+ .unwrap();
+ for batch in batches {
+ writer.write_batch(batch).unwrap();
+ }
+ writer.close().unwrap();
+ Bytes::from(writer.output().data.to_vec())
+ }
+
+ async fn write_deletion_file(
+ file_io: &crate::io::FileIO,
+ path: &str,
+ deleted_rows: &[u32],
+ ) -> DeletionFile {
+ 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::from(blob))
+ .await
+ .unwrap();
+ DeletionFile::new(
+ path.to_string(),
+ 0,
+ bitmap_length as i64,
+ Some(deleted_rows.len() as i64),
+ )
+ }
+
+ /// Build a `DataFileReader` over an in-memory mosaic file plus the
matching
+ /// `DataSplit`. `read_type`/`predicates` override the reader's projection
and
+ /// filter; `deleted_rows`, when non-empty, writes a DV into the split.
+ async fn build_reader_and_split(
+ table_path: &str,
+ data: &Bytes,
+ row_count: i64,
+ read_type: Vec<DataField>,
+ predicates: Vec<crate::spec::Predicate>,
+ deleted_rows: &[u32],
+ ) -> (DataFileReader, DataSplit, Option<Arc<DeletionVector>>) {
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let bucket_path = format!("{table_path}/bucket-0");
+ let file_name = "part-0.mosaic";
+ file_io
+ .new_output(&format!("{bucket_path}/{file_name}"))
+ .unwrap()
+ .write(data.clone())
+ .await
+ .unwrap();
+
+ let schema_id = 1;
+ let mut split_builder = DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(crate::spec::BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path(bucket_path)
+ .with_total_buckets(1)
+ .with_data_files(vec![data_file(
+ file_name,
+ data.len() as i64,
+ row_count,
+ schema_id,
+ Some(0),
+ )]);
+ let mut dv = None;
+ if !deleted_rows.is_empty() {
+ let df =
+ write_deletion_file(&file_io,
&format!("{table_path}/index/dv-0"), deleted_rows)
+ .await;
+ dv = Some(Arc::new(
+ DeletionVectorFactory::read(&file_io, &df).await.unwrap(),
+ ));
+ split_builder =
split_builder.with_data_deletion_files(vec![Some(df)]);
+ }
+ let split = split_builder.build().unwrap();
+
+ let schema_manager = SchemaManager::new(file_io.clone(),
table_path.to_string());
+ let reader = DataFileReader::new(
+ file_io,
+ schema_manager,
+ schema_id,
+ id_fields(),
+ read_type,
+ predicates,
+ );
+ (reader, split, dv)
+ }
+
+ fn column_by_name<'a>(batch: &'a RecordBatch, name: &str) -> Option<&'a
Arc<dyn Array>> {
+ batch
+ .schema()
+ .index_of(name)
+ .ok()
+ .map(|idx| batch.column(idx))
+ }
+
+ fn collect_i32(batches: &[RecordBatch], name: &str) -> Vec<i32> {
+ batches
+ .iter()
+ .flat_map(|b| {
+ column_by_name(b, name)
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ })
+ .collect()
+ }
+
+ fn collect_i64(batches: &[RecordBatch], name: &str) -> Vec<i64> {
+ batches
+ .iter()
+ .flat_map(|b| {
+ column_by_name(b, name)
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ })
+ .collect()
+ }
+
+ fn collect_f32(batches: &[RecordBatch], name: &str) -> Vec<f32> {
+ batches
+ .iter()
+ .flat_map(|b| {
+ column_by_name(b, name)
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Float32Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ })
+ .collect()
+ }
+
+ #[tokio::test]
+ async fn test_reads_selected_positions_with_position_column() {
+ // rows id=[10,11,12,13,14], first_row_id=0, select positions [0,2,4]
+ // -> output ids [10,12,14], _PKEY_VECTOR_POSITION [0,2,4], ascending;
+ // no _ROW_ID leak and no _PKEY_VECTOR_SCORE column.
+ let data = write_mosaic_single_group(&id_batch(vec![10, 11, 12, 13,
14]));
+ let (reader, split, _dv) = build_reader_and_split(
+ "memory:/pkvpr_basic",
+ &data,
+ 5,
+ id_fields(),
+ Vec::new(),
+ &[],
+ )
+ .await;
+
+ let batches = PkVectorPositionRead::new(&reader)
+ .read(
+ &split,
+ split.data_files()[0].clone(),
+ None,
+ None,
+ vec![0, 2, 4],
+ None,
+ )
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ assert_eq!(collect_i32(&batches, "id"), vec![10, 12, 14]);
+ assert_eq!(
+ collect_i64(&batches, PKEY_VECTOR_POSITION_COLUMN),
+ vec![0, 2, 4]
+ );
+ for batch in &batches {
+ assert!(
+ column_by_name(batch, ROW_ID_FIELD_NAME).is_none(),
+ "_ROW_ID must not leak into output"
+ );
+ assert!(
+ column_by_name(batch, PKEY_VECTOR_SCORE_COLUMN).is_none(),
+ "_PKEY_VECTOR_SCORE must be absent when no scores are supplied"
+ );
+ }
+ }
+
+ #[tokio::test]
+ async fn test_score_alignment_non_contiguous() {
+ // select [0,2,5] with scores {0:0.9, 2:0.5, 5:0.1}
+ // -> _PKEY_VECTOR_SCORE aligned by returned position: [0.9,0.5,0.1].
+ let data = write_mosaic_single_group(&id_batch(vec![10, 11, 12, 13,
14, 15]));
+ let (reader, split, _dv) = build_reader_and_split(
+ "memory:/pkvpr_scores",
+ &data,
+ 6,
+ id_fields(),
+ Vec::new(),
+ &[],
+ )
+ .await;
+
+ let scores = BTreeMap::from([(0, 0.9f32), (2, 0.5), (5, 0.1)]);
+ let batches = PkVectorPositionRead::new(&reader)
+ .read(
+ &split,
+ split.data_files()[0].clone(),
+ None,
+ None,
+ vec![0, 2, 5],
+ Some(scores),
+ )
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ assert_eq!(collect_i32(&batches, "id"), vec![10, 12, 15]);
+ assert_eq!(
+ collect_i64(&batches, PKEY_VECTOR_POSITION_COLUMN),
+ vec![0, 2, 5]
+ );
+ assert_eq!(
+ collect_f32(&batches, PKEY_VECTOR_SCORE_COLUMN),
+ vec![0.9, 0.5, 0.1]
+ );
+ }
+
+ #[tokio::test]
+ async fn test_score_omitted_no_score_column() {
+ let data = write_mosaic_single_group(&id_batch(vec![10, 11, 12]));
+ let (reader, split, _dv) = build_reader_and_split(
+ "memory:/pkvpr_noscore",
+ &data,
+ 3,
+ id_fields(),
+ Vec::new(),
+ &[],
+ )
+ .await;
+
+ let batches = PkVectorPositionRead::new(&reader)
+ .read(
+ &split,
+ split.data_files()[0].clone(),
+ None,
+ None,
+ vec![0, 1],
+ None,
+ )
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ for batch in &batches {
+ assert!(column_by_name(batch, PKEY_VECTOR_SCORE_COLUMN).is_none());
+ }
+ }
+
+ #[tokio::test]
+ async fn test_deletion_vector_skips_positions_and_keeps_alignment() {
+ // select [0,1,2,3] scores {0:.4,1:.3,2:.2,3:.1}, DV deletes position
1.
+ // -> returned positions [0,2,3], scores [.4,.2,.1]; position 1 absent.
+ // scores STILL contains key 1 (input contract), though row 1 is
deleted.
+ let data = write_mosaic_single_group(&id_batch(vec![10, 11, 12, 13]));
+ let (reader, split, dv) =
+ build_reader_and_split("memory:/pkvpr_dv", &data, 4, id_fields(),
Vec::new(), &[1])
+ .await;
+
+ let scores = BTreeMap::from([(0, 0.4f32), (1, 0.3), (2, 0.2), (3,
0.1)]);
+ let batches = PkVectorPositionRead::new(&reader)
+ .read(
+ &split,
+ split.data_files()[0].clone(),
+ None,
+ dv,
+ vec![0, 1, 2, 3],
+ Some(scores),
+ )
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ assert_eq!(collect_i32(&batches, "id"), vec![10, 12, 13]);
+ assert_eq!(
+ collect_i64(&batches, PKEY_VECTOR_POSITION_COLUMN),
+ vec![0, 2, 3]
+ );
+ assert_eq!(
+ collect_f32(&batches, PKEY_VECTOR_SCORE_COLUMN),
+ vec![0.4, 0.2, 0.1]
+ );
+ }
+
+ #[tokio::test]
+ async fn test_multi_batch_alignment() {
+ // Three row groups [10,11] [12,13] [14,15] -> reader yields >1 batch.
+ // Select [1,2,4] spanning batch boundaries; assert position/score
+ // alignment holds across batches (row_id_offset advances correctly).
+ let data = write_mosaic_multi_group(&[
+ id_batch(vec![10, 11]),
+ id_batch(vec![12, 13]),
+ id_batch(vec![14, 15]),
+ ]);
+ let (reader, split, _dv) = build_reader_and_split(
+ "memory:/pkvpr_multibatch",
+ &data,
+ 6,
+ id_fields(),
+ Vec::new(),
+ &[],
+ )
+ .await;
+
+ let scores = BTreeMap::from([(1, 0.9f32), (2, 0.5), (4, 0.1)]);
+ let batches = PkVectorPositionRead::new(&reader)
+ .read(
+ &split,
+ split.data_files()[0].clone(),
+ None,
+ None,
+ vec![1, 2, 4],
+ Some(scores),
+ )
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ assert!(
+ batches.len() > 1,
+ "expected multiple batches, got {}",
+ batches.len()
+ );
+ assert_eq!(collect_i32(&batches, "id"), vec![11, 12, 14]);
+ assert_eq!(
+ collect_i64(&batches, PKEY_VECTOR_POSITION_COLUMN),
+ vec![1, 2, 4]
+ );
+ assert_eq!(
+ collect_f32(&batches, PKEY_VECTOR_SCORE_COLUMN),
+ vec![0.9, 0.5, 0.1]
+ );
+ }
+
+ #[tokio::test]
+ async fn test_empty_positions_is_error() {
+ let data = write_mosaic_single_group(&id_batch(vec![10, 11, 12]));
+ let (reader, split, _dv) = build_reader_and_split(
+ "memory:/pkvpr_empty",
+ &data,
+ 3,
+ id_fields(),
+ Vec::new(),
+ &[],
+ )
+ .await;
+
+ let err = PkVectorPositionRead::new(&reader)
+ .read(
+ &split,
+ split.data_files()[0].clone(),
+ None,
+ None,
+ Vec::<i64>::new(),
+ None,
+ )
+ .err()
+ .expect("empty positions must be an error");
+ assert!(
+ format!("{err:?}").contains("must not be empty"),
+ "got: {err:?}"
+ );
+ }
+
+ #[tokio::test]
+ async fn test_position_out_of_range_is_error() {
+ let data = write_mosaic_single_group(&id_batch(vec![10, 11, 12]));
+ let (reader, split, _dv) =
+ build_reader_and_split("memory:/pkvpr_oor", &data, 3, id_fields(),
Vec::new(), &[])
+ .await;
+
+ let err = PkVectorPositionRead::new(&reader)
+ .read(
+ &split,
+ split.data_files()[0].clone(),
+ None,
+ None,
+ vec![5],
+ None,
+ )
+ .err()
+ .expect("out-of-range position must be an error");
+ let msg = format!("{err:?}");
+ assert!(msg.contains('5') && msg.contains('3'), "got: {msg}");
+ }
+
+ #[tokio::test]
+ async fn test_negative_position_is_error() {
+ let data = write_mosaic_single_group(&id_batch(vec![10, 11, 12]));
+ let (reader, split, _dv) =
+ build_reader_and_split("memory:/pkvpr_neg", &data, 3, id_fields(),
Vec::new(), &[])
+ .await;
+
+ let err = PkVectorPositionRead::new(&reader)
+ .read(
+ &split,
+ split.data_files()[0].clone(),
+ None,
+ None,
+ vec![-1],
+ None,
+ )
+ .err()
+ .expect("negative position must be an error");
+ assert!(
+ format!("{err:?}").contains("must not be negative"),
+ "got: {err:?}"
+ );
+ }
+
+ #[tokio::test]
+ async fn test_scores_key_mismatch_is_error() {
+ let data = write_mosaic_single_group(&id_batch(vec![10, 11, 12]));
+ let (reader, split, _dv) = build_reader_and_split(
+ "memory:/pkvpr_scoremismatch",
+ &data,
+ 3,
+ id_fields(),
+ Vec::new(),
+ &[],
+ )
+ .await;
+
+ // select [0,1] but scores only has key 0 -> mismatch.
+ let scores = BTreeMap::from([(0, 0.5f32)]);
+ let err = PkVectorPositionRead::new(&reader)
+ .read(
+ &split,
+ split.data_files()[0].clone(),
+ None,
+ None,
+ vec![0, 1],
+ Some(scores),
+ )
+ .err()
+ .expect("score key mismatch must be an error");
+ assert!(format!("{err:?}").contains("Scores keys"), "got: {err:?}");
+ }
+
+ #[tokio::test]
+ async fn test_missing_first_row_id_is_error() {
+ // first_row_id = None -> Err mentioning first_row_id. Positions must
be
+ // valid so validation reaches the first_row_id guard.
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let schema_manager = SchemaManager::new(file_io.clone(),
"memory:/pkvpr_frid".to_string());
+ let reader = DataFileReader::new(
+ file_io,
+ schema_manager,
+ 1,
+ id_fields(),
+ id_fields(),
+ Vec::new(),
+ );
+ let split = DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(crate::spec::BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path("memory:/pkvpr_frid/bucket-0".to_string())
+ .with_total_buckets(1)
+ .with_data_files(vec![data_file("part-0.mosaic", 1, 5, 1, None)])
+ .build()
+ .unwrap();
+
+ let err = PkVectorPositionRead::new(&reader)
+ .read(
+ &split,
+ split.data_files()[0].clone(),
+ None,
+ None,
+ vec![0],
+ None,
+ )
+ .err()
+ .expect("missing first_row_id must be an error");
+ assert!(format!("{err:?}").contains("first_row_id"), "got: {err:?}");
+ }
+
+ #[tokio::test]
+ async fn test_predicate_reader_is_rejected() {
+ // A row-filtering predicate on the reader must be rejected.
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let schema_manager = SchemaManager::new(file_io.clone(),
"memory:/pkvpr_pred".to_string());
+ let predicate = PredicateBuilder::new(&id_fields())
+ .equal("id", Datum::Int(10))
+ .unwrap();
+ let reader = DataFileReader::new(
+ file_io,
+ schema_manager,
+ 1,
+ id_fields(),
+ id_fields(),
+ vec![predicate],
+ );
+ let split = DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(crate::spec::BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path("memory:/pkvpr_pred/bucket-0".to_string())
+ .with_total_buckets(1)
+ .with_data_files(vec![data_file("part-0.mosaic", 1, 5, 1,
Some(0))])
+ .build()
+ .unwrap();
+
+ let err = PkVectorPositionRead::new(&reader)
+ .read(
+ &split,
+ split.data_files()[0].clone(),
+ None,
+ None,
+ vec![0],
+ None,
+ )
+ .err()
+ .expect("predicate reader must be rejected");
+ assert!(
+ format!("{err:?}").contains("predicate-free"),
+ "got: {err:?}"
+ );
+ }
+
+ #[tokio::test]
+ async fn test_reserved_column_name_conflict_is_error() {
+ // A reader whose read_type contains "_PKEY_VECTOR_POSITION" must be
rejected.
+ let reserved = DataField::new(
+ 0,
+ PKEY_VECTOR_POSITION_COLUMN.to_string(),
+ DataType::Int(IntType::new()),
+ );
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let schema_manager =
+ SchemaManager::new(file_io.clone(),
"memory:/pkvpr_reserved".to_string());
+ let reader = DataFileReader::new(
+ file_io,
+ schema_manager,
+ 1,
+ id_fields(),
+ vec![reserved],
+ Vec::new(),
+ );
+ let split = DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(crate::spec::BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path("memory:/pkvpr_reserved/bucket-0".to_string())
+ .with_total_buckets(1)
+ .with_data_files(vec![data_file("part-0.mosaic", 1, 5, 1,
Some(0))])
+ .build()
+ .unwrap();
+
+ let err = PkVectorPositionRead::new(&reader)
+ .read(
+ &split,
+ split.data_files()[0].clone(),
+ None,
+ None,
+ vec![0],
+ None,
+ )
+ .err()
+ .expect("reserved column name conflict must be an error");
+ assert!(
+ format!("{err:?}").contains("Reserved metadata column name"),
+ "got: {err:?}"
+ );
+ }
+
+ #[tokio::test]
+ async fn test_requested_row_id_column_is_rejected() {
+ // A reader whose read_type requests _ROW_ID must be rejected loudly
rather
+ // than have the column silently stripped from the output.
+ let requested_row_id = DataField::new(
+ ROW_ID_FIELD_ID,
+ ROW_ID_FIELD_NAME.to_string(),
+ DataType::BigInt(BigIntType::new()),
+ );
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let schema_manager =
+ SchemaManager::new(file_io.clone(),
"memory:/pkvpr_reqrowid".to_string());
+ let reader = DataFileReader::new(
+ file_io,
+ schema_manager,
+ 1,
+ id_fields(),
+ vec![id_field(), requested_row_id],
+ Vec::new(),
+ );
+ let split = DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(crate::spec::BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path("memory:/pkvpr_reqrowid/bucket-0".to_string())
+ .with_total_buckets(1)
+ .with_data_files(vec![data_file("part-0.mosaic", 1, 5, 1,
Some(0))])
+ .build()
+ .unwrap();
+
+ let err = PkVectorPositionRead::new(&reader)
+ .read(
+ &split,
+ split.data_files()[0].clone(),
+ None,
+ None,
+ vec![0],
+ None,
+ )
+ .err()
+ .expect("requested _ROW_ID must be an error");
+ assert!(format!("{err:?}").contains("_ROW_ID"), "got: {err:?}");
+ }
+
+ #[test]
+ fn test_positions_to_global_ranges_coalesces_and_offsets() {
+ // positions [0,1,2,4,5] with first_row_id 100 -> global ranges
+ // [100..=102] and [104..=105] (two coalesced ranges, offset by 100).
+ let ranges = positions_to_global_ranges(&[0, 1, 2, 4, 5],
100).unwrap();
+ assert_eq!(ranges.len(), 2);
+ assert_eq!((ranges[0].from(), ranges[0].to()), (100, 102));
+ assert_eq!((ranges[1].from(), ranges[1].to()), (104, 105));
+ }
+
+ #[test]
+ fn test_positions_to_global_ranges_single() {
+ let ranges = positions_to_global_ranges(&[3], 0).unwrap();
+ assert_eq!(ranges.len(), 1);
+ assert_eq!((ranges[0].from(), ranges[0].to()), (3, 3));
+ }
+
+ #[test]
+ fn test_positions_to_global_ranges_overflow_is_error() {
+ // position 1 offset by i64::MAX overflows -> Err.
+ let err = positions_to_global_ranges(&[1], i64::MAX)
+ .expect_err("offset overflow must be an error");
+ assert!(format!("{err:?}").contains("overflows"), "got: {err:?}");
+ }
+}
diff --git a/crates/paimon/src/vindex/pkvector/ann.rs
b/crates/paimon/src/vindex/pkvector/ann.rs
index f87e7f9..b5f26fd 100644
--- a/crates/paimon/src/vindex/pkvector/ann.rs
+++ b/crates/paimon/src/vindex/pkvector/ann.rs
@@ -20,7 +20,7 @@ use std::sync::Arc;
use super::bucket::BucketAnnSegment;
use super::data_invalid;
-use super::metric::VectorSearchMetric;
+use super::metric::{java_float_compare, VectorSearchMetric};
use super::result::PkVectorSearchResult;
use crate::deletion_vector::DeletionVector;
use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta};
@@ -119,8 +119,7 @@ pub(crate) fn map_ann_results(
});
}
results.sort_by(|a, b| {
- a.distance
- .total_cmp(&b.distance)
+ java_float_compare(a.distance, b.distance)
.then_with(|| a.data_file_name.cmp(&b.data_file_name))
.then_with(|| a.row_position.cmp(&b.row_position))
});
diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs
b/crates/paimon/src/vindex/pkvector/bucket.rs
index 23a2b1c..18fd8f8 100644
--- a/crates/paimon/src/vindex/pkvector/bucket.rs
+++ b/crates/paimon/src/vindex/pkvector/bucket.rs
@@ -22,7 +22,7 @@ use std::sync::Arc;
use super::ann::PkVectorAnnSearcher;
use super::data_invalid;
use super::exact::exact_search;
-use super::metric::VectorSearchMetric;
+use super::metric::{java_float_compare, VectorSearchMetric};
use super::reader::PkVectorReader;
use super::result::PkVectorSearchResult;
use crate::deletion_vector::DeletionVector;
@@ -43,10 +43,10 @@ pub(crate) struct BucketActiveFile {
}
/// 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.
+/// then row_position ASC. `java_float_compare` sorts NaN distances last (never
+/// best), matching Java `Float.compare`, and is panic-free.
fn best_first(a: &PkVectorSearchResult, b: &PkVectorSearchResult) -> Ordering {
- a.distance
- .total_cmp(&b.distance)
+ java_float_compare(a.distance, b.distance)
.then_with(|| a.data_file_name.cmp(&b.data_file_name))
.then_with(|| a.row_position.cmp(&b.row_position))
}
@@ -314,6 +314,50 @@ mod tests {
);
}
+ #[test]
+ fn nan_ann_hit_never_evicts_finite_candidate_from_top1() {
+ // The core failure mode from review: an ANN hit with a negative-NaN
+ // distance must not win the single bucket Top-1 slot over a finite
hit.
+ // Under f32::total_cmp the -NaN would rank best and evict the finite
+ // candidate here in the bucket heap, before any cross-bucket merge.
+ let negative_nan = f32::from_bits(0xffc00000);
+ assert!(negative_nan.is_nan());
+ let segment = BucketAnnSegment {
+ source_meta: meta(&[("data-1", 2)]),
+ };
+ let ann = FakeAnnSearcher {
+ result: vec![
+ PkVectorSearchResult {
+ data_file_name: "data-1".into(),
+ row_position: 0,
+ distance: negative_nan,
+ },
+ PkVectorSearchResult {
+ data_file_name: "data-1".into(),
+ row_position: 1,
+ distance: -1.0,
+ },
+ ],
+ };
+ let mut factory =
+ |_: &BucketActiveFile| -> crate::Result<Box<dyn PkVectorReader>> {
unreachable!() };
+ let results = bucket_search(
+ Some(&ann),
+ &[segment],
+ &[active("data-1", 2)],
+ &HashMap::new(),
+ &mut factory,
+ &[0.0, 0.0],
+ VectorSearchMetric::L2,
+ 1,
+ &HashMap::new(),
+ )
+ .unwrap();
+ assert_eq!(results.len(), 1);
+ assert_eq!(results[0].row_position, 1);
+ assert_eq!(results[0].distance, -1.0);
+ }
+
#[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
diff --git a/crates/paimon/src/vindex/pkvector/exact.rs
b/crates/paimon/src/vindex/pkvector/exact.rs
index c9c0480..e5d043c 100644
--- a/crates/paimon/src/vindex/pkvector/exact.rs
+++ b/crates/paimon/src/vindex/pkvector/exact.rs
@@ -19,14 +19,15 @@ use std::cmp::Ordering;
use std::collections::BinaryHeap;
use super::data_invalid;
-use super::metric::VectorSearchMetric;
+use super::metric::{java_float_compare, 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).
+/// top therefore evicts the least-wanted candidate. Uses `java_float_compare`
+/// for a deterministic total order over f32 that ranks NaN distances as worst
+/// (largest), so a NaN is evicted before any finite candidate (no panic).
struct WorstFirst(PkVectorSearchResult);
impl PartialEq for WorstFirst {
@@ -42,9 +43,7 @@ impl PartialOrd for WorstFirst {
}
impl Ord for WorstFirst {
fn cmp(&self, other: &Self) -> Ordering {
- self.0
- .distance
- .total_cmp(&other.0.distance)
+ java_float_compare(self.0.distance, other.0.distance)
.then_with(|| self.0.row_position.cmp(&other.0.row_position))
}
}
@@ -52,9 +51,7 @@ impl Ord for WorstFirst {
/// 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)
+ java_float_compare(candidate.distance, weakest.distance)
.then_with(|| candidate.row_position.cmp(&weakest.row_position))
== Ordering::Less
}
@@ -117,9 +114,7 @@ pub(crate) fn exact_search(
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))
+ java_float_compare(a.distance, b.distance).then_with(||
a.row_position.cmp(&b.row_position))
});
Ok(results)
}
@@ -156,6 +151,28 @@ mod tests {
}
}
+ #[test]
+ fn nan_distance_candidate_never_beats_finite_top1() {
+ // Row 0's stored vector contains NaN → its inner-product distance is a
+ // negative NaN (sign flipped by the metric's negation), which
+ // `f32::total_cmp` would rank as best (smallest) and select as Top-1.
+ // `java_float_compare` ranks NaN worst, so the finite row 1 wins the
+ // single Top-1 slot instead.
+ let mut reader = ArrayReader::new(2, vec![Some(vec![f32::NAN, 0.0]),
Some(vec![1.0, 0.0])]);
+ let results = exact_search(
+ "data-file",
+ &mut reader,
+ &[1.0, 0.0],
+ VectorSearchMetric::InnerProduct,
+ 1,
+ &no_exclusion(),
+ )
+ .unwrap();
+ assert_eq!(results.len(), 1);
+ assert_eq!(results[0].row_position, 1);
+ assert_eq!(results[0].distance, -1.0);
+ }
+
#[test]
fn test_rejects_dimension_mismatch() {
let mut reader = ArrayReader::new(2, vec![Some(vec![1.0, 0.0])]);
diff --git a/crates/paimon/src/vindex/pkvector/metric.rs
b/crates/paimon/src/vindex/pkvector/metric.rs
index 1dda56c..799d9b8 100644
--- a/crates/paimon/src/vindex/pkvector/metric.rs
+++ b/crates/paimon/src/vindex/pkvector/metric.rs
@@ -16,6 +16,23 @@
// under the License.
use super::data_invalid;
+use std::cmp::Ordering;
+
+/// Order two distances the way Java `Float.compare` does: every NaN sorts
after
+/// all numeric values (so a NaN distance is always ranked worst, never best),
+/// while non-NaN values keep `total_cmp`'s IEEE total order (which also
matches
+/// Java in ranking `-0.0` before `+0.0`). `f32::total_cmp` alone is unsuitable
+/// here because it orders a negative NaN before every finite value, which
would
+/// let a NaN distance (e.g. from a non-finite stored vector under inner
product)
+/// win Top-1.
+pub(crate) fn java_float_compare(a: f32, b: f32) -> Ordering {
+ match (a.is_nan(), b.is_nan()) {
+ (true, true) => Ordering::Equal,
+ (true, false) => Ordering::Greater,
+ (false, true) => Ordering::Less,
+ (false, false) => a.total_cmp(&b),
+ }
+}
/// Normalize a metric name: lowercase and `-` → `_`. NO trim (deliberately
/// stricter than the build-side `vindex::normalize_metric`, to match Java
@@ -82,6 +99,21 @@ impl VectorSearchMetric {
Self::InnerProduct => -score,
}
}
+
+ /// Convert a lower-is-better canonical distance (as produced by
+ /// `bucket_search`) to a higher-is-better score. Mirrors Java
+ /// `PrimaryKeyVectorResult.score(distance)`. No clamping — natural f32
+ /// behavior (L2 with `distance=inf` -> `0.0`), consistent with the sibling
+ /// `score_to_distance`.
+ /// NOTE: cosine here is `1 - distance` applied directly to the canonical
+ /// distance; it does NOT reuse `score_to_distance`'s clamp.
+ pub(crate) fn distance_to_score(&self, distance: f32) -> f32 {
+ match self {
+ Self::L2 => 1.0 / (1.0 + distance),
+ Self::Cosine => 1.0 - distance,
+ Self::InnerProduct => -distance,
+ }
+ }
}
fn squared_l2(query: &[f32], stored: &[f32]) -> f32 {
@@ -126,6 +158,51 @@ fn cosine_distance(similarity: f32) -> f32 {
mod tests {
use super::*;
+ /// A negative NaN (sign bit set): `f32::total_cmp` would order this
*before*
+ /// every finite value; `java_float_compare` must order it *after*.
+ const NEGATIVE_NAN: f32 = f32::from_bits(0xffc00000);
+
+ #[test]
+ fn java_float_compare_sorts_all_nan_after_finite() {
+ // Positive NaN after a finite value.
+ assert_eq!(java_float_compare(f32::NAN, 1.0), Ordering::Greater);
+ assert_eq!(java_float_compare(1.0, f32::NAN), Ordering::Less);
+ // Negative NaN also after a finite value (the total_cmp trap).
+ assert!(NEGATIVE_NAN.is_nan());
+ assert_eq!(java_float_compare(NEGATIVE_NAN, 1.0), Ordering::Greater);
+ assert_eq!(java_float_compare(1.0, NEGATIVE_NAN), Ordering::Less);
+ // NaN after +inf too.
+ assert_eq!(
+ java_float_compare(f32::NAN, f32::INFINITY),
+ Ordering::Greater
+ );
+ // Two NaNs compare equal.
+ assert_eq!(java_float_compare(f32::NAN, NEGATIVE_NAN),
Ordering::Equal);
+ }
+
+ #[test]
+ fn java_float_compare_keeps_java_numeric_order() {
+ // Ascending finite order preserved.
+ assert_eq!(java_float_compare(1.0, 2.0), Ordering::Less);
+ assert_eq!(java_float_compare(2.0, 1.0), Ordering::Greater);
+ assert_eq!(java_float_compare(1.0, 1.0), Ordering::Equal);
+ // -0.0 sorts before +0.0, matching Java Float.compare.
+ assert_eq!(java_float_compare(-0.0, 0.0), Ordering::Less);
+ assert_eq!(java_float_compare(0.0, -0.0), Ordering::Greater);
+ }
+
+ #[test]
+ fn java_float_compare_ranks_negative_nan_worst_when_sorting() {
+ // A -NaN distance must never sort ahead of a finite one: sorting
ascending
+ // (smallest = best) puts finite first, NaN last.
+ let mut xs = [NEGATIVE_NAN, 3.0, 1.0, 2.0];
+ xs.sort_by(|a, b| java_float_compare(*a, *b));
+ assert_eq!(xs[0], 1.0);
+ assert_eq!(xs[1], 2.0);
+ assert_eq!(xs[2], 3.0);
+ assert!(xs[3].is_nan());
+ }
+
#[test]
fn test_normalize_lowercases_and_replaces_hyphens_without_trimming() {
assert_eq!(normalize_metric("Inner-Product"), "inner_product");
@@ -208,4 +285,33 @@ mod tests {
);
assert_eq!(VectorSearchMetric::Cosine.score_to_distance(1.0), 0.0);
}
+
+ #[test]
+ fn test_distance_to_score_per_metric_formula() {
+ // L2: 1/(1+d); Cosine: 1-d; InnerProduct: -d.
+ assert_eq!(VectorSearchMetric::L2.distance_to_score(0.0), 1.0);
+ assert_eq!(VectorSearchMetric::L2.distance_to_score(1.0), 0.5);
+ assert_eq!(VectorSearchMetric::L2.distance_to_score(3.0), 0.25);
+ assert_eq!(VectorSearchMetric::Cosine.distance_to_score(0.0), 1.0);
+ assert_eq!(VectorSearchMetric::Cosine.distance_to_score(1.0), 0.0);
+ assert_eq!(VectorSearchMetric::InnerProduct.distance_to_score(0.0),
0.0);
+ assert_eq!(
+ VectorSearchMetric::InnerProduct.distance_to_score(3.0),
+ -3.0
+ );
+ }
+
+ #[test]
+ fn test_distance_to_score_inverts_score_to_distance_off_boundary() {
+ // Avoid boundary values (no L2 score=0/distance=inf, no cosine/ip
infinities):
+ // distance_to_score(score_to_distance(s)) == s for representative s.
+ for (metric, s) in [
+ (VectorSearchMetric::L2, 0.5f32),
+ (VectorSearchMetric::Cosine, 0.25f32),
+ (VectorSearchMetric::InnerProduct, 2.0f32),
+ ] {
+ let d = metric.score_to_distance(s);
+ assert_eq!(metric.distance_to_score(d), s, "metric {metric:?}");
+ }
+ }
}
