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 22a13b6 fix(table): align pk-vector source meta with Java (#531)
22a13b6 is described below
commit 22a13b676280ad18ec4af7a492867dbbf58c9ba4
Author: Junrui Lee <[email protected]>
AuthorDate: Fri Jul 17 08:35:20 2026 +0800
fix(table): align pk-vector source meta with Java (#531)
---
crates/paimon/src/spec/pk_vector_source.rs | 78 ++++++++++----
crates/paimon/src/table/pk_vector_orchestrator.rs | 1 +
crates/paimon/src/table/pk_vector_scan.rs | 124 ++++++++++++++++++++--
crates/paimon/src/table/vector_search_builder.rs | 12 +--
crates/paimon/src/vindex/pkvector/ann.rs | 19 ++--
crates/paimon/src/vindex/pkvector/bucket.rs | 1 +
crates/paimon/tests/pk_vector_baseline_test.rs | 13 ++-
7 files changed, 203 insertions(+), 45 deletions(-)
diff --git a/crates/paimon/src/spec/pk_vector_source.rs
b/crates/paimon/src/spec/pk_vector_source.rs
index 2c859e0..802cbf7 100644
--- a/crates/paimon/src/spec/pk_vector_source.rs
+++ b/crates/paimon/src/spec/pk_vector_source.rs
@@ -76,15 +76,28 @@ impl PkVectorSourceFile {
/// Mirrors Java `org.apache.paimon.index.pkvector.PkVectorSourceMeta`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PkVectorSourceMeta {
+ data_level: i32,
source_files: Vec<PkVectorSourceFile>,
}
impl PkVectorSourceMeta {
- pub fn new(source_files: Vec<PkVectorSourceFile>) -> crate::Result<Self> {
+ pub fn new(data_level: i32, source_files: Vec<PkVectorSourceFile>) ->
crate::Result<Self> {
+ if data_level <= 0 {
+ return Err(data_invalid(format!(
+ "source meta data level must be positive: {data_level}"
+ )));
+ }
if source_files.is_empty() {
return Err(data_invalid("a vector index must reference source
files"));
}
- Ok(Self { source_files })
+ Ok(Self {
+ data_level,
+ source_files,
+ })
+ }
+
+ pub fn data_level(&self) -> i32 {
+ self.data_level
}
pub fn source_files(&self) -> &[PkVectorSourceFile] {
@@ -136,6 +149,12 @@ impl PkVectorSourceMeta {
"unsupported vector source version: {version}"
)));
}
+ let data_level = cursor.read_i32_be()?;
+ if data_level <= 0 {
+ return Err(data_invalid(format!(
+ "source meta data level must be positive: {data_level}"
+ )));
+ }
let count = cursor.read_i32_be()?;
if count <= 0 {
return Err(data_invalid("a vector index must reference source
files"));
@@ -152,7 +171,7 @@ impl PkVectorSourceMeta {
"unexpected trailing bytes in vector source metadata",
));
}
- Self::new(source_files)
+ Self::new(data_level, source_files)
}
}
@@ -270,9 +289,10 @@ mod tests {
out
}
- fn frame(files: &[(&str, i64)]) -> Vec<u8> {
+ fn frame(data_level: i32, files: &[(&str, i64)]) -> Vec<u8> {
let mut out = Vec::new();
out.extend_from_slice(&1i32.to_be_bytes()); // version
+ out.extend_from_slice(&data_level.to_be_bytes());
out.extend_from_slice(&(files.len() as i32).to_be_bytes());
for (name, rows) in files {
out.extend_from_slice(&java_write_utf(name));
@@ -319,8 +339,9 @@ mod tests {
#[test]
fn deserialize_single_source_file() {
- let bytes = frame(&[("data-abc.parquet", 100)]);
+ let bytes = frame(1, &[("data-abc.parquet", 100)]);
let meta = PkVectorSourceMeta::deserialize(&bytes).unwrap();
+ assert_eq!(meta.data_level(), 1);
assert_eq!(meta.source_files().len(), 1);
assert_eq!(meta.source_files()[0].file_name(), "data-abc.parquet");
assert_eq!(meta.source_files()[0].row_count(), 100);
@@ -328,54 +349,67 @@ mod tests {
#[test]
fn deserialize_multi_source_files() {
- let bytes = frame(&[("f0", 3), ("f1", 5)]);
+ let bytes = frame(2, &[("f0", 3), ("f1", 5)]);
let meta = PkVectorSourceMeta::deserialize(&bytes).unwrap();
+ assert_eq!(meta.data_level(), 2);
assert_eq!(meta.source_files().len(), 2);
assert_eq!(meta.source_files()[1].row_count(), 5);
}
#[test]
fn deserialize_rejects_bad_version() {
- let mut bytes = frame(&[("f0", 1)]);
+ let mut bytes = frame(1, &[("f0", 1)]);
bytes[0..4].copy_from_slice(&2i32.to_be_bytes());
assert!(PkVectorSourceMeta::deserialize(&bytes).is_err());
}
+ #[test]
+ fn deserialize_rejects_zero_or_negative_data_level() {
+ assert!(PkVectorSourceMeta::deserialize(&frame(0, &[("f0",
1)])).is_err());
+ assert!(PkVectorSourceMeta::deserialize(&frame(-1, &[("f0",
1)])).is_err());
+ }
+
#[test]
fn deserialize_rejects_zero_count() {
let mut out = Vec::new();
out.extend_from_slice(&1i32.to_be_bytes());
+ out.extend_from_slice(&1i32.to_be_bytes());
out.extend_from_slice(&0i32.to_be_bytes());
assert!(PkVectorSourceMeta::deserialize(&out).is_err());
}
#[test]
fn deserialize_rejects_trailing_bytes() {
- let mut bytes = frame(&[("f0", 1)]);
+ let mut bytes = frame(1, &[("f0", 1)]);
bytes.push(0xFF);
assert!(PkVectorSourceMeta::deserialize(&bytes).is_err());
}
#[test]
fn deserialize_rejects_truncated_input() {
- let bytes = frame(&[("f0", 1)]);
+ let bytes = frame(1, &[("f0", 1)]);
assert!(PkVectorSourceMeta::deserialize(&bytes[..bytes.len() -
2]).is_err());
}
#[test]
fn deserialize_rejects_negative_row_count() {
- let bytes = frame(&[("f0", -1)]);
+ let bytes = frame(1, &[("f0", -1)]);
assert!(PkVectorSourceMeta::deserialize(&bytes).is_err());
}
#[test]
fn new_rejects_empty() {
- assert!(PkVectorSourceMeta::new(Vec::new()).is_err());
+ assert!(PkVectorSourceMeta::new(1, Vec::new()).is_err());
+ assert!(PkVectorSourceMeta::new(
+ 0,
+ vec![PkVectorSourceFile::new("f0".to_string(), 1).unwrap()]
+ )
+ .is_err());
}
#[test]
fn resolve_single_file() {
- let meta = PkVectorSourceMeta::deserialize(&frame(&[("f0",
3)])).unwrap();
+ let meta = PkVectorSourceMeta::deserialize(&frame(1, &[("f0",
3)])).unwrap();
assert_eq!(meta.resolve(0).unwrap(), ("f0".to_string(), 0));
assert_eq!(meta.resolve(2).unwrap(), ("f0".to_string(), 2));
}
@@ -383,7 +417,7 @@ mod tests {
#[test]
fn resolve_multi_file_prefix_sum_boundaries() {
// f0 owns ordinals 0..=2, f1 owns 3..=7.
- let meta = PkVectorSourceMeta::deserialize(&frame(&[("f0", 3), ("f1",
5)])).unwrap();
+ let meta = PkVectorSourceMeta::deserialize(&frame(1, &[("f0", 3),
("f1", 5)])).unwrap();
assert_eq!(meta.resolve(2).unwrap(), ("f0".to_string(), 2)); // last
of f0
assert_eq!(meta.resolve(3).unwrap(), ("f1".to_string(), 0)); // first
of f1
assert_eq!(meta.resolve(7).unwrap(), ("f1".to_string(), 4)); // last
of f1
@@ -391,13 +425,13 @@ mod tests {
#[test]
fn resolve_rejects_negative_ordinal() {
- let meta = PkVectorSourceMeta::deserialize(&frame(&[("f0",
3)])).unwrap();
+ let meta = PkVectorSourceMeta::deserialize(&frame(1, &[("f0",
3)])).unwrap();
assert!(meta.resolve(-1).is_err());
}
#[test]
fn resolve_rejects_ordinal_at_or_past_total() {
- let meta = PkVectorSourceMeta::deserialize(&frame(&[("f0",
3)])).unwrap();
+ let meta = PkVectorSourceMeta::deserialize(&frame(1, &[("f0",
3)])).unwrap();
assert!(meta.resolve(3).is_err()); // total == 3, valid range 0..=2
}
@@ -422,9 +456,10 @@ mod tests {
index_field_id: 0,
extra_field_ids: None,
index_meta: None,
- source_meta: Some(frame(&[("f0", 3)])),
+ source_meta: Some(frame(1, &[("f0", 3)])),
};
let parsed =
PkVectorSourceMeta::from_global_index_meta(&meta).unwrap();
+ assert_eq!(parsed.data_level(), 1);
assert_eq!(parsed.source_files()[0].file_name(), "f0");
}
@@ -432,10 +467,13 @@ mod tests {
fn resolve_rejects_row_count_overflow() {
// Two individually-valid row counts whose prefix sum overflows i64.
// Resolving past the first file forces the checked_add on the second.
- let meta = PkVectorSourceMeta::new(vec![
- PkVectorSourceFile::new("f0".to_string(), i64::MAX).unwrap(),
- PkVectorSourceFile::new("f1".to_string(), 1).unwrap(),
- ])
+ let meta = PkVectorSourceMeta::new(
+ 1,
+ vec![
+ PkVectorSourceFile::new("f0".to_string(), i64::MAX).unwrap(),
+ PkVectorSourceFile::new("f1".to_string(), 1).unwrap(),
+ ],
+ )
.unwrap();
assert!(meta.resolve(i64::MAX).is_err());
}
diff --git a/crates/paimon/src/table/pk_vector_orchestrator.rs
b/crates/paimon/src/table/pk_vector_orchestrator.rs
index 66a7543..59ee64c 100644
--- a/crates/paimon/src/table/pk_vector_orchestrator.rs
+++ b/crates/paimon/src/table/pk_vector_orchestrator.rs
@@ -764,6 +764,7 @@ mod e2e_tests {
fn ann_segment(sources: &[(&str, i64)]) -> BucketAnnSegment {
BucketAnnSegment::for_test(
PkVectorSourceMeta::new(
+ 1,
sources
.iter()
.map(|(n, r)| PkVectorSourceFile::new((*n).to_string(),
*r).unwrap())
diff --git a/crates/paimon/src/table/pk_vector_scan.rs
b/crates/paimon/src/table/pk_vector_scan.rs
index d5547e5..a8219fd 100644
--- a/crates/paimon/src/table/pk_vector_scan.rs
+++ b/crates/paimon/src/table/pk_vector_scan.rs
@@ -23,7 +23,8 @@
use std::collections::{BTreeMap, HashSet};
use crate::spec::{
- BinaryRow, DataFileMeta, FileKind, GlobalIndexMeta, IndexManifest,
PkVectorSourceMeta,
+ BinaryRow, DataFileMeta, FileKind, GlobalIndexMeta, IndexManifest,
PkVectorSourceFile,
+ PkVectorSourceMeta,
};
use crate::table::pk_vector_orchestrator::PkVectorSearchSplit;
use crate::table::source::{DataSplit, DataSplitBuilder, DeletionFile};
@@ -46,6 +47,56 @@ fn should_read_pk_index_source(file: &DataFileMeta) -> bool {
matches!(file.file_source, Some(src) if src == FILE_SOURCE_COMPACT) &&
file.level > 0
}
+fn source_files_unique(files: &[PkVectorSourceFile]) -> bool {
+ let mut seen = HashSet::new();
+ files.iter().all(|file| seen.insert(file.file_name()))
+}
+
+fn current_ann_segments(
+ active_data_files: &[DataFileMeta],
+ ann_segments: Vec<BucketAnnSegment>,
+) -> crate::Result<Vec<BucketAnnSegment>> {
+ let mut sources_by_level: BTreeMap<i32, Vec<PkVectorSourceFile>> =
BTreeMap::new();
+ for file in active_data_files {
+ if should_read_pk_index_source(file) {
+ sources_by_level
+ .entry(file.level)
+ .or_default()
+ .push(PkVectorSourceFile::new(
+ file.file_name.clone(),
+ file.row_count,
+ )?);
+ }
+ }
+ for sources in sources_by_level.values_mut() {
+ sources.sort_by(|a, b| a.file_name().cmp(b.file_name()));
+ }
+
+ let mut segments_by_level: BTreeMap<i32, Vec<BucketAnnSegment>> =
BTreeMap::new();
+ for segment in ann_segments {
+ let source_meta = &segment.source_meta;
+ let Some(desired) = sources_by_level.get(&source_meta.data_level())
else {
+ continue;
+ };
+ if source_files_unique(source_meta.source_files())
+ && desired.as_slice() == source_meta.source_files()
+ {
+ segments_by_level
+ .entry(source_meta.data_level())
+ .or_default()
+ .push(segment);
+ }
+ }
+
+ let mut current = Vec::new();
+ for mut level_segments in segments_by_level.into_values() {
+ if level_segments.len() == 1 {
+ current.push(level_segments.remove(0));
+ }
+ }
+ Ok(current)
+}
+
/// Combines one bucket's data splits into a single split, keeping data files
and
/// deletion files in strict parallel order and rejecting duplicate file names.
struct BucketAccumulator {
@@ -284,8 +335,11 @@ fn plan_from_inputs(
// Phase C: assemble one split per bucket that has data.
let mut out = Vec::new();
for (key, acc) in accum_by_bucket {
- let ann_segments = segments_by_bucket.remove(&key).unwrap_or_default();
let data_split = acc.build()?;
+ let ann_segments = current_ann_segments(
+ data_split.data_files(),
+ segments_by_bucket.remove(&key).unwrap_or_default(),
+ )?;
let active_files: Vec<BucketActiveFile> = data_split
.data_files()
.iter()
@@ -369,9 +423,10 @@ mod tests {
/// Build a `_SOURCE_META` blob the way `PkVectorSourceMeta::deserialize`
/// expects it. There is no public serializer, so we mirror the frame used
by
/// `pk_vector_source.rs`'s own round-trip tests.
- fn source_meta_bytes(files: &[(&str, i64)]) -> Vec<u8> {
+ fn source_meta_bytes(data_level: i32, files: &[(&str, i64)]) -> Vec<u8> {
let mut out = Vec::new();
out.extend_from_slice(&1i32.to_be_bytes()); // version
+ out.extend_from_slice(&data_level.to_be_bytes());
out.extend_from_slice(&(files.len() as i32).to_be_bytes());
for (name, rows) in files {
out.extend_from_slice(&java_write_utf(name));
@@ -380,14 +435,32 @@ mod tests {
out
}
- fn gim(field_id: i32, source_files: &[(&str, i64)]) -> GlobalIndexMeta {
+ fn gim(field_id: i32, data_level: i32, source_files: &[(&str, i64)]) ->
GlobalIndexMeta {
GlobalIndexMeta {
row_range_start: 0,
row_range_end: 0,
index_field_id: field_id,
extra_field_ids: None,
index_meta: Some(vec![1, 2, 3]),
- source_meta: Some(source_meta_bytes(source_files)),
+ source_meta: Some(source_meta_bytes(data_level, source_files)),
+ }
+ }
+
+ fn ann_segment(data_level: i32, path: &str, source_files: &[(&str, i64)])
-> BucketAnnSegment {
+ BucketAnnSegment {
+ source_meta: PkVectorSourceMeta::new(
+ data_level,
+ source_files
+ .iter()
+ .map(|(name, rows)| {
+ PkVectorSourceFile::new((*name).to_string(),
*rows).unwrap()
+ })
+ .collect(),
+ )
+ .unwrap(),
+ path: path.to_string(),
+ file_size: 1,
+ index_meta: Vec::new(),
}
}
@@ -397,7 +470,7 @@ mod tests {
let entries = vec![(
BinaryRow::new(0),
0,
- gim(2, &[("d0", 3)]),
+ gim(2, 5, &[("d0", 3)]),
"idx/seg0".to_string(),
10u64,
"seg0".to_string(),
@@ -411,7 +484,7 @@ mod tests {
let entries = vec![(
BinaryRow::new(0),
0,
- gim(2, &[("d0", 3)]),
+ gim(2, 5, &[("d0", 3)]),
"idx/seg0".to_string(),
10u64,
"seg0".to_string(),
@@ -436,6 +509,43 @@ mod tests {
assert_eq!(splits[0].active_files[0].file_name, "d0");
}
+ #[test]
+ fn current_segments_require_exact_level_source_set() {
+ let active = vec![
+ dfm("b", 2, 5, Some(1)),
+ dfm("a", 1, 5, Some(1)),
+ dfm("c", 3, 6, Some(1)),
+ ];
+ let current = current_ann_segments(
+ &active,
+ vec![
+ // Matches level 5 after active files are sorted by file name.
+ ann_segment(5, "current-l5", &[("a", 1), ("b", 2)]),
+ // Wrong level for the same source files -> stale.
+ ann_segment(4, "wrong-level", &[("a", 1), ("b", 2)]),
+ // Incomplete level 6 coverage -> stale.
+ ann_segment(6, "partial-l6", &[("c", 2)]),
+ ],
+ )
+ .unwrap();
+ assert_eq!(current.len(), 1);
+ assert_eq!(current[0].path, "current-l5");
+ }
+
+ #[test]
+ fn current_segments_drop_all_when_level_has_multiple_matches() {
+ let active = vec![dfm("a", 1, 5, Some(1))];
+ let current = current_ann_segments(
+ &active,
+ vec![
+ ann_segment(5, "first", &[("a", 1)]),
+ ann_segment(5, "second", &[("a", 1)]),
+ ],
+ )
+ .unwrap();
+ assert!(current.is_empty());
+ }
+
#[test]
fn rejects_data_split_with_wrong_snapshot() {
let data = DataSplitBuilder::new()
diff --git a/crates/paimon/src/table/vector_search_builder.rs
b/crates/paimon/src/table/vector_search_builder.rs
index 97f8b2e..122af04 100644
--- a/crates/paimon/src/table/vector_search_builder.rs
+++ b/crates/paimon/src/table/vector_search_builder.rs
@@ -2653,13 +2653,11 @@ mod tests {
/// A `PkVectorSearchSplit` carrying a single ANN segment addressed by
`path`.
fn pk_split_with_segment(path: &str) -> PkVectorSearchSplit {
let mut split = pk_search_split(0, vec![pk_data_file("file-a", 3,
Some(0))]);
- let source_meta =
-
crate::spec::PkVectorSourceMeta::new(vec![crate::spec::PkVectorSourceFile::new(
- "file-a".to_string(),
- 3,
- )
- .unwrap()])
- .unwrap();
+ let source_meta = crate::spec::PkVectorSourceMeta::new(
+ 1,
+ vec![crate::spec::PkVectorSourceFile::new("file-a".to_string(),
3).unwrap()],
+ )
+ .unwrap();
let mut segment = BucketAnnSegment::for_test(source_meta);
segment.path = path.to_string();
split.ann_segments = vec![segment];
diff --git a/crates/paimon/src/vindex/pkvector/ann.rs
b/crates/paimon/src/vindex/pkvector/ann.rs
index 0d1559b..aad463f 100644
--- a/crates/paimon/src/vindex/pkvector/ann.rs
+++ b/crates/paimon/src/vindex/pkvector/ann.rs
@@ -221,7 +221,7 @@ mod tests {
.iter()
.map(|(name, rows)| PkVectorSourceFile::new((*name).to_string(),
*rows).unwrap())
.collect();
- PkVectorSourceMeta::new(files).unwrap()
+ PkVectorSourceMeta::new(1, files).unwrap()
}
fn dv(deleted: &[u32]) -> Arc<DeletionVector> {
@@ -379,10 +379,13 @@ mod tests {
);
let segment = BucketAnnSegment::for_test({
use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta};
- PkVectorSourceMeta::new(vec![
- PkVectorSourceFile::new("f0".into(), 3).unwrap(),
- PkVectorSourceFile::new("f1".into(), 5).unwrap(),
- ])
+ PkVectorSourceMeta::new(
+ 1,
+ vec![
+ PkVectorSourceFile::new("f0".into(), 3).unwrap(),
+ PkVectorSourceFile::new("f1".into(), 5).unwrap(),
+ ],
+ )
.unwrap()
});
let mut dvs = HashMap::new();
@@ -416,7 +419,8 @@ mod tests {
);
let segment = BucketAnnSegment::for_test({
use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta};
- PkVectorSourceMeta::new(vec![PkVectorSourceFile::new("f0".into(),
1).unwrap()]).unwrap()
+ PkVectorSourceMeta::new(1,
vec![PkVectorSourceFile::new("f0".into(), 1).unwrap()])
+ .unwrap()
});
let err = searcher
.search(
@@ -440,7 +444,8 @@ mod tests {
);
let segment = BucketAnnSegment::for_test({
use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta};
- PkVectorSourceMeta::new(vec![PkVectorSourceFile::new("f0".into(),
1).unwrap()]).unwrap()
+ PkVectorSourceMeta::new(1,
vec![PkVectorSourceFile::new("f0".into(), 1).unwrap()])
+ .unwrap()
});
let results = searcher
.search(
diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs
b/crates/paimon/src/vindex/pkvector/bucket.rs
index 337e922..fcd7d1f 100644
--- a/crates/paimon/src/vindex/pkvector/bucket.rs
+++ b/crates/paimon/src/vindex/pkvector/bucket.rs
@@ -289,6 +289,7 @@ mod tests {
fn meta(files: &[(&str, i64)]) -> PkVectorSourceMeta {
PkVectorSourceMeta::new(
+ 1,
files
.iter()
.map(|(n, r)| PkVectorSourceFile::new((*n).into(),
*r).unwrap())
diff --git a/crates/paimon/tests/pk_vector_baseline_test.rs
b/crates/paimon/tests/pk_vector_baseline_test.rs
index 1cb0399..e1e3177 100644
--- a/crates/paimon/tests/pk_vector_baseline_test.rs
+++ b/crates/paimon/tests/pk_vector_baseline_test.rs
@@ -196,11 +196,13 @@ fn java_write_utf(s: &str) -> Vec<u8> {
/// Assemble the `_SOURCE_META` frame the way Java `PkVectorSourceMeta` writes
it
/// and `PkVectorSourceMeta::deserialize` expects: `i32-BE version=1`, `i32-BE
-/// count`, then per source file a `writeUTF` name and an `i64-BE` row count.
No
-/// trailing bytes. Source files are listed in global ordinal order.
-fn source_meta_bytes(files: &[(&str, i64)]) -> Vec<u8> {
+/// data_level`, `i32-BE count`, then per source file a `writeUTF` name and an
+/// `i64-BE` row count. No trailing bytes. Source files are listed in global
+/// ordinal order.
+fn source_meta_bytes(data_level: i32, files: &[(&str, i64)]) -> Vec<u8> {
let mut out = Vec::new();
out.extend_from_slice(&1i32.to_be_bytes()); // version
+ out.extend_from_slice(&data_level.to_be_bytes());
out.extend_from_slice(&(files.len() as i32).to_be_bytes());
for (name, rows) in files {
out.extend_from_slice(&java_write_utf(name));
@@ -396,7 +398,10 @@ async fn build_table(
row_range_end: row_count - 1,
index_field_id: vector_field_id,
extra_field_ids: None,
- source_meta: Some(source_meta_bytes(&[(&data_file_name,
row_count)])),
+ source_meta: Some(source_meta_bytes(
+ indexed_meta.level,
+ &[(&data_file_name, row_count)],
+ )),
index_meta: None,
}),
};
