From 52ed1d7e69df460a3af24e1f838cbbfc8e6d096c Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Tue, 14 Jul 2026 03:28:51 +0800 Subject: [PATCH 01/10] feat(spec): add primary-key vector index config options --- crates/paimon/src/spec/core_options.rs | 151 +++++++++++++++++++++++++ 1 file changed, 151 insertions(+) diff --git a/crates/paimon/src/spec/core_options.rs b/crates/paimon/src/spec/core_options.rs index 60d33ff5b..277f04cf5 100644 --- a/crates/paimon/src/spec/core_options.rs +++ b/crates/paimon/src/spec/core_options.rs @@ -117,6 +117,7 @@ pub(crate) const BLOB_FIELD_OPTION: &str = "blob-field"; pub(crate) const BLOB_DESCRIPTOR_FIELD_OPTION: &str = "blob-descriptor-field"; pub(crate) const BLOB_VIEW_FIELD_OPTION: &str = "blob-view-field"; pub const BLOB_VIEW_RESOLVE_ENABLED_OPTION: &str = "blob-view.resolve.enabled"; +const PK_VECTOR_INDEX_COLUMNS_OPTION: &str = "pk-vector.index.columns"; /// Merge engine for primary-key tables. /// @@ -1067,6 +1068,72 @@ impl<'a> CoreOptions<'a> { }) .unwrap_or_default() } + + /// True when the PK-vector index column option key is present (regardless of value). + pub fn primary_key_vector_index_enabled(&self) -> bool { + self.options.contains_key(PK_VECTOR_INDEX_COLUMNS_OPTION) + } + + /// The configured PK-vector index columns, split on ',' and trimmed. Errors when + /// the key is present but resolves to no non-blank column. + pub fn primary_key_vector_index_columns(&self) -> crate::Result> { + let raw = self + .options + .get(PK_VECTOR_INDEX_COLUMNS_OPTION) + .ok_or_else(|| crate::Error::ConfigInvalid { + message: "pk-vector.index.columns is not set".to_string(), + })?; + let columns: Vec = raw + .split(',') + .map(|c| c.trim().to_string()) + .filter(|c| !c.is_empty()) + .collect(); + if columns.is_empty() { + return Err(crate::Error::ConfigInvalid { + message: "pk-vector.index.columns is set but names no column".to_string(), + }); + } + Ok(columns) + } + + /// The single PK-vector index column. The first release supports exactly one. + pub fn primary_key_vector_index_column(&self) -> crate::Result { + let mut columns = self.primary_key_vector_index_columns()?; + if columns.len() != 1 { + return Err(crate::Error::ConfigInvalid { + message: format!( + "pk-vector.index.columns must name exactly one column, got {}", + columns.len() + ), + }); + } + Ok(columns.remove(0)) + } + + /// The index type for a PK-vector column. Required — planning and the index reader + /// both need it, so an absent value is a hard error rather than a guessed default. + pub fn primary_key_vector_index_type(&self, col: &str) -> crate::Result { + self.options + .get(&format!("fields.{col}.pk-vector.index.type")) + .map(|v| v.trim().to_string()) + .filter(|v| !v.is_empty()) + .ok_or_else(|| crate::Error::ConfigInvalid { + message: format!("fields.{col}.pk-vector.index.type is required but not set"), + }) + } + + /// The distance metric name for a PK-vector column, defaulting to inner_product. + /// Validated against the supported metrics; an unknown value is a hard error. + pub fn primary_key_vector_distance_metric(&self, col: &str) -> crate::Result { + let raw = self + .options + .get(&format!("fields.{col}.pk-vector.distance.metric")) + .map(|v| v.trim().to_string()) + .unwrap_or_else(|| "inner_product".to_string()); + // Validate now (fail-loud) without exposing the crate-private metric enum. + crate::vindex::pkvector::metric::VectorSearchMetric::parse(&raw)?; + Ok(raw) + } } /// Parse a memory size string to bytes using binary (1024-based) semantics. @@ -1856,4 +1923,88 @@ mod tests { let opts = CoreOptions::new(&options); assert!(!opts.ignore_update_before()); } + + #[test] + fn test_pk_vector_index_disabled_by_default() { + let opts = HashMap::new(); + assert!(!CoreOptions::new(&opts).primary_key_vector_index_enabled()); + } + + #[test] + fn test_pk_vector_single_column_and_type_and_metric() { + let opts = HashMap::from([ + ( + "pk-vector.index.columns".to_string(), + " embedding ".to_string(), + ), + ( + "fields.embedding.pk-vector.index.type".to_string(), + "ivf-flat".to_string(), + ), + ( + "fields.embedding.pk-vector.distance.metric".to_string(), + "Inner-Product".to_string(), + ), + ]); + let co = CoreOptions::new(&opts); + assert!(co.primary_key_vector_index_enabled()); + assert_eq!(co.primary_key_vector_index_column().unwrap(), "embedding"); + assert_eq!( + co.primary_key_vector_index_type("embedding").unwrap(), + "ivf-flat" + ); + assert_eq!( + co.primary_key_vector_distance_metric("embedding").unwrap(), + "Inner-Product" + ); + } + + #[test] + fn test_pk_vector_metric_defaults_to_inner_product() { + let opts = HashMap::from([("pk-vector.index.columns".to_string(), "e".to_string())]); + assert_eq!( + CoreOptions::new(&opts) + .primary_key_vector_distance_metric("e") + .unwrap(), + "inner_product" + ); + } + + #[test] + fn test_pk_vector_unknown_metric_errors() { + let opts = HashMap::from([ + ("pk-vector.index.columns".to_string(), "e".to_string()), + ( + "fields.e.pk-vector.distance.metric".to_string(), + "manhattan".to_string(), + ), + ]); + assert!(CoreOptions::new(&opts) + .primary_key_vector_distance_metric("e") + .is_err()); + } + + #[test] + fn test_pk_vector_empty_columns_errors() { + let opts = HashMap::from([("pk-vector.index.columns".to_string(), " , ".to_string())]); + let co = CoreOptions::new(&opts); + assert!(co.primary_key_vector_index_enabled()); // key present + assert!(co.primary_key_vector_index_columns().is_err()); + } + + #[test] + fn test_pk_vector_multiple_columns_unsupported() { + let opts = HashMap::from([("pk-vector.index.columns".to_string(), "a,b".to_string())]); + assert!(CoreOptions::new(&opts) + .primary_key_vector_index_column() + .is_err()); + } + + #[test] + fn test_pk_vector_index_type_absent_errors() { + let opts = HashMap::from([("pk-vector.index.columns".to_string(), "e".to_string())]); + assert!(CoreOptions::new(&opts) + .primary_key_vector_index_type("e") + .is_err()); + } } From 8f52528617ac518696f89010a8ff50e1ab3ee7e7 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Tue, 14 Jul 2026 18:58:40 +0800 Subject: [PATCH 02/10] feat(vindex): thread ANN segment and fast-mode skip through bucket search Pass the ANN segment to the scorer seam so a real scorer can select per-segment index bytes, and add a fast-mode flag to bucket_search that skips the exact data-file fallback (ANN-only search) while leaving the default behavior unchanged. --- .../src/table/pk_vector_orchestrator.rs | 7 +- crates/paimon/src/vindex/pkvector/ann.rs | 51 +++---- crates/paimon/src/vindex/pkvector/bucket.rs | 130 ++++++++++++------ 3 files changed, 113 insertions(+), 75 deletions(-) diff --git a/crates/paimon/src/table/pk_vector_orchestrator.rs b/crates/paimon/src/table/pk_vector_orchestrator.rs index b8664cc9f..e6cef37f5 100644 --- a/crates/paimon/src/table/pk_vector_orchestrator.rs +++ b/crates/paimon/src/table/pk_vector_orchestrator.rs @@ -292,6 +292,7 @@ impl PkVectorOrchestrator { metric, limit, search_options, + false, )?; for PkVectorSearchResult { data_file_name, @@ -743,15 +744,15 @@ mod e2e_tests { } fn ann_segment(sources: &[(&str, i64)]) -> BucketAnnSegment { - BucketAnnSegment { - source_meta: PkVectorSourceMeta::new( + BucketAnnSegment::for_test( + PkVectorSourceMeta::new( sources .iter() .map(|(n, r)| PkVectorSourceFile::new((*n).to_string(), *r).unwrap()) .collect(), ) .unwrap(), - } + ) } fn active(name: &str, rows: i64) -> BucketActiveFile { diff --git a/crates/paimon/src/vindex/pkvector/ann.rs b/crates/paimon/src/vindex/pkvector/ann.rs index b5f26fda4..5c650afb8 100644 --- a/crates/paimon/src/vindex/pkvector/ann.rs +++ b/crates/paimon/src/vindex/pkvector/ann.rs @@ -152,7 +152,8 @@ pub(crate) trait PkVectorAnnSearcher { /// with a segment's index bytes; tests inject a synthetic scorer. The adapter's /// own logic (live-row masking, ordinal mapping, deletion checks, ordering) is /// exercised independently of the scorer. -type Scorer = Box crate::Result>>>; +pub(crate) type Scorer = + Box crate::Result>>>; /// Structural vindex-backed `PkVectorAnnSearcher`. Composes the pure helpers /// (`build_live_row_ids`, `map_ann_results`) around the scorer seam. @@ -188,7 +189,7 @@ impl PkVectorAnnSearcher for VindexAnnSearcher { { search = search.with_include_row_ids(live); } - let scored = match (self.scorer)(&search)? { + let scored = match (self.scorer)(segment, &search)? { Some(map) => map, None => return Ok(Vec::new()), }; @@ -360,7 +361,7 @@ mod tests { let scorer_has_filter = Rc::clone(&seen_has_filter); let searcher = VindexAnnSearcher::new( "embedding".to_string(), - Box::new(move |search: &VectorSearch| { + Box::new(move |_segment: &BucketAnnSegment, search: &VectorSearch| { *scorer_limit.borrow_mut() = search.limit; *scorer_has_filter.borrow_mut() = search.include_row_ids.is_some(); let mut scores = HashMap::new(); @@ -369,16 +370,14 @@ mod tests { Ok(Some(scores)) }), ); - let segment = BucketAnnSegment { - source_meta: { - use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta}; - PkVectorSourceMeta::new(vec![ - PkVectorSourceFile::new("f0".into(), 3).unwrap(), - PkVectorSourceFile::new("f1".into(), 5).unwrap(), - ]) - .unwrap() - }, - }; + let segment = BucketAnnSegment::for_test({ + use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta}; + PkVectorSourceMeta::new(vec![ + PkVectorSourceFile::new("f0".into(), 3).unwrap(), + PkVectorSourceFile::new("f1".into(), 5).unwrap(), + ]) + .unwrap() + }); let mut dvs = HashMap::new(); dvs.insert("f0".to_string(), dv(&[1])); let results = searcher @@ -406,15 +405,12 @@ mod tests { fn test_vindex_adapter_rejects_non_positive_limit() { let searcher = VindexAnnSearcher::new( "embedding".to_string(), - Box::new(|_: &VectorSearch| Ok(None)), + Box::new(|_: &BucketAnnSegment, _: &VectorSearch| Ok(None)), ); - let segment = BucketAnnSegment { - source_meta: { - use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta}; - PkVectorSourceMeta::new(vec![PkVectorSourceFile::new("f0".into(), 1).unwrap()]) - .unwrap() - }, - }; + let segment = BucketAnnSegment::for_test({ + use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta}; + PkVectorSourceMeta::new(vec![PkVectorSourceFile::new("f0".into(), 1).unwrap()]).unwrap() + }); let err = searcher .search( &segment, @@ -433,15 +429,12 @@ mod tests { fn test_vindex_adapter_empty_scorer_result_is_empty() { let searcher = VindexAnnSearcher::new( "embedding".to_string(), - Box::new(|_: &VectorSearch| Ok(None)), + Box::new(|_: &BucketAnnSegment, _: &VectorSearch| Ok(None)), ); - let segment = BucketAnnSegment { - source_meta: { - use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta}; - PkVectorSourceMeta::new(vec![PkVectorSourceFile::new("f0".into(), 1).unwrap()]) - .unwrap() - }, - }; + let segment = BucketAnnSegment::for_test({ + use crate::spec::{PkVectorSourceFile, PkVectorSourceMeta}; + PkVectorSourceMeta::new(vec![PkVectorSourceFile::new("f0".into(), 1).unwrap()]).unwrap() + }); let results = searcher .search( &segment, diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs b/crates/paimon/src/vindex/pkvector/bucket.rs index 18fd8f8d4..01da95f58 100644 --- a/crates/paimon/src/vindex/pkvector/bucket.rs +++ b/crates/paimon/src/vindex/pkvector/bucket.rs @@ -28,11 +28,32 @@ use super::result::PkVectorSearchResult; use crate::deletion_vector::DeletionVector; use crate::spec::PkVectorSourceMeta; -/// One ANN segment to be searched by the bucket kernel: the source metadata -/// resolving segment ordinals back to physical `(data file, position)`. Only -/// `source_meta` is needed for ordinal mapping and live-row masking. +/// One ANN segment to be searched by the bucket kernel. `source_meta` resolves +/// segment ordinals back to physical `(data file, position)` and drives live-row +/// masking; the remaining fields address the segment's index file for the ANN +/// scorer that reads it. pub(crate) struct BucketAnnSegment { pub source_meta: PkVectorSourceMeta, + pub file_name: String, + /// Resolved index-file path (globally unique; the scorer's preload key). + pub path: String, + pub file_size: u64, + pub index_meta: Vec, +} + +#[cfg(test)] +impl BucketAnnSegment { + /// Build a segment with dummy index-file fields for tests that exercise only + /// `source_meta`-driven logic. + pub(crate) fn for_test(source_meta: PkVectorSourceMeta) -> Self { + Self { + source_meta, + file_name: "seg".to_string(), + path: "seg".to_string(), + file_size: 0, + index_meta: Vec::new(), + } + } } /// A data file participating in the bucket search, with its row count. Used by @@ -107,6 +128,7 @@ pub(crate) fn bucket_search( metric: VectorSearchMetric, limit: usize, search_options: &HashMap, + skip_exact_fallback: bool, ) -> crate::Result> { if limit == 0 { return Err(data_invalid("vector search limit must be positive")); @@ -170,29 +192,31 @@ pub(crate) fn bucket_search( } } - for file in active_files { - if covered.contains(&file.file_name) { - continue; - } - let dv = deletion_vectors.get(&file.file_name).cloned(); - let is_excluded = move |position: i64| -> bool { - match &dv { - Some(dv) => u64::try_from(position) - .map(|p| dv.is_deleted(p)) - .unwrap_or(false), - None => false, + if !skip_exact_fallback { + for file in active_files { + if covered.contains(&file.file_name) { + continue; + } + let dv = deletion_vectors.get(&file.file_name).cloned(); + let is_excluded = move |position: i64| -> bool { + match &dv { + Some(dv) => u64::try_from(position) + .map(|p| dv.is_deleted(p)) + .unwrap_or(false), + None => false, + } + }; + let mut reader = exact_reader_factory(file)?; + for result in exact_search( + &file.file_name, + reader.as_mut(), + query, + metric, + limit, + &is_excluded, + )? { + add_candidate(&mut heap, result, limit); } - }; - let mut reader = exact_reader_factory(file)?; - for result in exact_search( - &file.file_name, - reader.as_mut(), - query, - metric, - limit, - &is_excluded, - )? { - add_candidate(&mut heap, result, limit); } } @@ -260,6 +284,7 @@ mod tests { VectorSearchMetric::L2, 0, &HashMap::new(), + false, ) .unwrap_err(); assert!(err.to_string().contains("positive")); @@ -271,9 +296,7 @@ mod tests { // BEST_FIRST tie-break (data_file_name ASC, then row_position ASC). Feed // more than `limit` ANN hits and assert the kept set is the smallest // (file, position) pairs in that order. Locks the bounded-heap merge. - let segment = BucketAnnSegment { - source_meta: meta(&[("data-1", 3)]), - }; + let segment = BucketAnnSegment::for_test(meta(&[("data-1", 3)])); let hit = |file: &str, pos: i64| PkVectorSearchResult { data_file_name: file.into(), row_position: pos, @@ -301,6 +324,7 @@ mod tests { VectorSearchMetric::L2, 3, &HashMap::new(), + false, ) .unwrap(); // Top-3 BEST_FIRST: (data-1,0), (data-1,1), (data-1,2) — the larger @@ -322,9 +346,7 @@ mod tests { // 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 segment = BucketAnnSegment::for_test(meta(&[("data-1", 2)])); let ann = FakeAnnSearcher { result: vec![ PkVectorSearchResult { @@ -351,6 +373,7 @@ mod tests { VectorSearchMetric::L2, 1, &HashMap::new(), + false, ) .unwrap(); assert_eq!(results.len(), 1); @@ -362,9 +385,7 @@ mod tests { fn test_merges_ann_and_exact_without_rescanning_covered_files() { // data-1 is ANN-covered; data-2 is exact fallback. Factory must never be // called for data-1. - let segment = BucketAnnSegment { - source_meta: meta(&[("data-1", 2)]), - }; + let segment = BucketAnnSegment::for_test(meta(&[("data-1", 2)])); let ann = FakeAnnSearcher { result: vec![PkVectorSearchResult { data_file_name: "data-1".into(), @@ -391,6 +412,7 @@ mod tests { VectorSearchMetric::L2, 2, &HashMap::new(), + false, ) .unwrap(); assert_eq!( @@ -439,6 +461,7 @@ mod tests { VectorSearchMetric::L2, 2, &HashMap::new(), + false, ) .unwrap(); // Candidates: data-2 pos0 {1,0} dist 1.0; data-1 pos1 {2,0} dist 4.0. @@ -474,6 +497,7 @@ mod tests { VectorSearchMetric::L2, 1, &HashMap::new(), + false, ) .unwrap_err(); assert!(err.to_string().contains("duplicate") || err.to_string().contains("Duplicate")); @@ -484,9 +508,7 @@ mod tests { let ann = FakeAnnSearcher { result: vec![] }; // Segment references data-1 with 2 rows, but the active file has 3 rows. // An active source with a mismatched row count is still a hard error. - let segment = BucketAnnSegment { - source_meta: meta(&[("data-1", 2)]), - }; + let segment = BucketAnnSegment::for_test(meta(&[("data-1", 2)])); let mut factory = |_: &BucketActiveFile| -> crate::Result> { unreachable!() }; let err = bucket_search( @@ -499,6 +521,7 @@ mod tests { VectorSearchMetric::L2, 1, &HashMap::new(), + false, ) .unwrap_err(); assert!( @@ -513,9 +536,7 @@ mod tests { // instead of failing the whole query; data-2 is neither covered (so it // is not treated as ANN-covered) nor an active file (so it is not exact // scanned). The ANN searcher still runs for the segment. - let segment = BucketAnnSegment { - source_meta: meta(&[("data-1", 2), ("data-2", 2)]), - }; + let segment = BucketAnnSegment::for_test(meta(&[("data-1", 2), ("data-2", 2)])); let ann = FakeAnnSearcher { result: vec![PkVectorSearchResult { data_file_name: "data-1".into(), @@ -538,6 +559,7 @@ mod tests { VectorSearchMetric::L2, 2, &HashMap::new(), + false, ) .unwrap(); assert_eq!( @@ -554,9 +576,7 @@ mod tests { #[test] fn test_rejects_segments_without_ann_searcher() { - let segment = BucketAnnSegment { - source_meta: meta(&[("data-1", 2)]), - }; + let segment = BucketAnnSegment::for_test(meta(&[("data-1", 2)])); let mut factory = |_: &BucketActiveFile| -> crate::Result> { unreachable!() }; let err = bucket_search( @@ -569,6 +589,7 @@ mod tests { VectorSearchMetric::L2, 1, &HashMap::new(), + false, ) .unwrap_err(); assert!( @@ -577,6 +598,28 @@ mod tests { ); } + #[test] + fn test_skip_exact_fallback_does_not_call_factory() { + // No ANN segments, two active files. With skip_exact_fallback = true the + // factory must never be called and the result is empty. + let mut factory = + |_: &BucketActiveFile| -> crate::Result> { unreachable!() }; + let results = bucket_search( + None, + &[], + &[active("data-1", 2), active("data-2", 2)], + &HashMap::new(), + &mut factory, + &[0.0, 0.0], + VectorSearchMetric::L2, + 2, + &HashMap::new(), + true, // skip_exact_fallback + ) + .unwrap(); + assert!(results.is_empty()); + } + #[test] fn test_negative_active_row_count_rejected() { let mut factory = @@ -591,6 +634,7 @@ mod tests { VectorSearchMetric::L2, 1, &HashMap::new(), + false, ) .unwrap_err(); assert!(err.to_string().contains("row count") || err.to_string().contains("-1")); From 5259b179773bee9cd776a985d9ace9ca1c7c03f2 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Tue, 14 Jul 2026 18:59:14 +0800 Subject: [PATCH 03/10] feat(table): add primary-key vector read building blocks Expose best-first search candidates from the orchestrator, add PkVectorScan planning (snapshot resolution, index-manifest scan, per-bucket split assembly), and add the positional exact data-file reader that reuses the shared single-file read stream. --- crates/paimon/src/table/mod.rs | 2 + .../src/table/pk_vector_data_file_reader.rs | 451 ++++++++++++++++ .../src/table/pk_vector_orchestrator.rs | 210 ++++++-- crates/paimon/src/table/pk_vector_scan.rs | 502 ++++++++++++++++++ 4 files changed, 1134 insertions(+), 31 deletions(-) create mode 100644 crates/paimon/src/table/pk_vector_data_file_reader.rs create mode 100644 crates/paimon/src/table/pk_vector_scan.rs diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index cf581aa57..5c34918d2 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -56,9 +56,11 @@ mod lumina_index_build_builder; pub(crate) mod merge_tree_split_generator; mod partition_filter; mod partition_stat; +mod pk_vector_data_file_reader; mod pk_vector_indexed_split_read; mod pk_vector_orchestrator; mod pk_vector_position_read; +mod pk_vector_scan; mod postpone_file_writer; mod prepared_files; mod read_builder; diff --git a/crates/paimon/src/table/pk_vector_data_file_reader.rs b/crates/paimon/src/table/pk_vector_data_file_reader.rs new file mode 100644 index 000000000..b0ba71c8c --- /dev/null +++ b/crates/paimon/src/table/pk_vector_data_file_reader.rs @@ -0,0 +1,451 @@ +// 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. + +//! Exact sequential vector reader over one data file's vector column. Mirrors +//! Java `org.apache.paimon.index.pkvector.PkVectorDataFileReader`. +//! +//! The factory projects the single vector column, reads the whole file in +//! physical order, and preloads it into memory as `Vec>>` +//! (a NULL row is `None`). Deletion vectors are deliberately NOT applied here: +//! physical position must stay in lockstep with the segment ordinal so the +//! bucket search can address rows by position. The returned reader then serves +//! vectors from memory one physical row at a time. + +use arrow_array::{Array, FixedSizeListArray, Float32Array, ListArray}; +use futures::TryStreamExt; + +use crate::spec::{DataField, DataType}; +use crate::table::data_file_reader::DataFileReader; +use crate::table::source::DataSplit; +use crate::vindex::pkvector::bucket::BucketActiveFile; +use crate::vindex::pkvector::reader::PkVectorReader; + +fn data_invalid(message: impl Into) -> crate::Error { + crate::Error::DataInvalid { + message: message.into(), + source: None, + } +} + +/// Builds an exact [`PkVectorReader`] over one data file's vector column. +/// +/// `reader` is configured (via [`DataFileReader::with_read_type`]) to project +/// only the vector column, so each read returns a single-column batch. Mirrors +/// Java `PkVectorDataFileReader` (as a factory owning the projected reader). +#[allow(dead_code)] +pub(crate) struct DataFilePkVectorReaderFactory { + reader: DataFileReader, + data_split: DataSplit, + vector_field: DataField, + dimension: usize, +} + +impl DataFilePkVectorReaderFactory { + /// Configure `reader` to project the vector column only and capture the + /// vector dimension from the schema field. The field must be a fixed-length + /// `Vector` type; anything else is rejected as invalid. + #[allow(dead_code)] + pub(crate) fn new( + reader: DataFileReader, + data_split: DataSplit, + vector_field: DataField, + ) -> crate::Result { + let dimension = match vector_field.data_type() { + DataType::Vector(vector_type) => vector_type.length() as usize, + other => { + return Err(data_invalid(format!( + "PK-vector reader requires a fixed-length Vector field, got {other:?}" + ))); + } + }; + let reader = reader.with_read_type(vec![vector_field.clone()]); + Ok(Self { + reader, + data_split, + vector_field, + dimension, + }) + } + + /// Preload the whole vector column of `file` into memory and return a + /// sequential reader over it. `file` must name a data file present in this + /// factory's split. The drained row count is checked against the file's + /// `DataFileMeta.row_count`. + #[allow(dead_code)] + pub(crate) async fn create( + &self, + file: &BucketActiveFile, + ) -> crate::Result> { + let file_meta = self + .data_split + .data_files() + .iter() + .find(|meta| meta.file_name == file.file_name) + .cloned() + .ok_or_else(|| { + data_invalid(format!( + "data file '{}' not found in split for PK-vector read", + file.file_name + )) + })?; + let row_count = file_meta.row_count; + + let data_fields = self.reader.derive_data_fields(&file_meta).await?; + let mut stream = self.reader.read_single_file_stream( + &self.data_split, + file_meta, + data_fields, + None, + None, + )?; + + let mut vectors: Vec>> = Vec::new(); + while let Some(batch) = stream.try_next().await? { + append_batch_vectors( + &batch, + self.vector_field.name(), + self.dimension, + &mut vectors, + )?; + } + + let drained = vectors.len() as i64; + if drained > row_count { + return Err(data_invalid( + "data file produced more rows than DataFileMeta.row_count", + )); + } + if drained < row_count { + return Err(data_invalid( + "data file ended before DataFileMeta.row_count", + )); + } + + Ok(Box::new(DataFilePkVectorReader { + dimension: self.dimension, + row_count, + vectors, + position: 0, + })) + } +} + +/// Extract one batch's vector column into `out`, one entry per row (NULL row = +/// `None`). The column must be a `FixedSizeList`/`List` of `Float32`; every +/// non-null row's child slice must have exactly `dimension` elements. Mirrors +/// the layout handling in `vector_search_builder`. +fn append_batch_vectors( + batch: &arrow_array::RecordBatch, + field_name: &str, + dimension: usize, + out: &mut Vec>>, +) -> crate::Result<()> { + let index = batch + .schema() + .index_of(field_name) + .map_err(|e| data_invalid(format!("vector column '{field_name}' not found: {e}")))?; + let column = batch.column(index); + + enum VectorLayout<'a> { + List(&'a ListArray), + Fixed(&'a FixedSizeListArray), + } + let layout = if let Some(a) = column.as_any().downcast_ref::() { + VectorLayout::List(a) + } else if let Some(a) = column.as_any().downcast_ref::() { + VectorLayout::Fixed(a) + } else { + return Err(data_invalid( + "PK-vector read requires Arrow List or FixedSizeList", + )); + }; + + let values = match layout { + VectorLayout::List(a) => a.values(), + VectorLayout::Fixed(a) => a.values(), + } + .as_any() + .downcast_ref::() + .ok_or_else(|| data_invalid("PK-vector read requires Float32 vector elements"))?; + + for row in 0..batch.num_rows() { + let is_null = match layout { + VectorLayout::List(a) => a.is_null(row), + VectorLayout::Fixed(a) => a.is_null(row), + }; + if is_null { + out.push(None); + continue; + } + let (start, end) = match layout { + VectorLayout::List(a) => { + let offsets = a.value_offsets(); + (offsets[row] as usize, offsets[row + 1] as usize) + } + VectorLayout::Fixed(a) => { + let len = a.value_length() as usize; + (row * len, (row + 1) * len) + } + }; + if end - start != dimension { + return Err(data_invalid(format!( + "vector row has {} elements, expected dimension {dimension}", + end - start + ))); + } + let mut vector = Vec::with_capacity(dimension); + for i in start..end { + vector.push(values.value(i)); + } + out.push(Some(vector)); + } + Ok(()) +} + +/// In-memory sequential reader over one file's preloaded vector column. Each +/// [`read_next_vector`](PkVectorReader::read_next_vector) advances exactly one +/// physical row; a NULL row returns `false` but still advances the position. +struct DataFilePkVectorReader { + dimension: usize, + row_count: i64, + /// Preloaded whole-file column in physical order; `None` = NULL row. + vectors: Vec>>, + position: usize, +} + +impl PkVectorReader for DataFilePkVectorReader { + fn dimension(&self) -> usize { + self.dimension + } + + fn row_count(&self) -> i64 { + self.row_count + } + + fn read_next_vector(&mut self, reuse: &mut [f32]) -> crate::Result { + if reuse.len() != self.dimension { + return Err(data_invalid(format!( + "reuse buffer length {} does not match vector dimension {}", + reuse.len(), + self.dimension + ))); + } + if self.position as i64 >= self.row_count { + return Err(data_invalid("read past row count")); + } + let entry = &self.vectors[self.position]; + self.position += 1; + match entry { + Some(vector) => { + reuse.copy_from_slice(vector); + Ok(true) + } + None => Ok(false), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn null_row_consumes_ordinal() { + let mut r = DataFilePkVectorReader { + dimension: 2, + row_count: 3, + vectors: vec![Some(vec![1.0, 2.0]), None, Some(vec![3.0, 4.0])], + position: 0, + }; + let mut buf = [0.0f32; 2]; + assert!(r.read_next_vector(&mut buf).unwrap()); + assert_eq!(buf, [1.0, 2.0]); + assert!(!r.read_next_vector(&mut buf).unwrap()); // null: false, ordinal advanced + assert!(r.read_next_vector(&mut buf).unwrap()); + assert_eq!(buf, [3.0, 4.0]); + assert_eq!(r.row_count(), 3); + assert_eq!(r.dimension(), 2); + } + + #[test] + fn read_past_row_count_errors() { + let mut r = DataFilePkVectorReader { + dimension: 1, + row_count: 1, + vectors: vec![Some(vec![1.0])], + position: 0, + }; + let mut buf = [0.0f32; 1]; + assert!(r.read_next_vector(&mut buf).unwrap()); + assert!(r.read_next_vector(&mut buf).is_err()); // past row_count + } + + #[test] + fn reuse_len_mismatch_errors() { + let mut r = DataFilePkVectorReader { + dimension: 2, + row_count: 1, + vectors: vec![Some(vec![1.0, 2.0])], + position: 0, + }; + let mut buf = [0.0f32; 1]; + assert!(r.read_next_vector(&mut buf).is_err()); + } +} + +#[cfg(test)] +mod integration_tests { + use super::*; + use crate::arrow::build_target_arrow_schema; + use crate::arrow::format::{FormatFileWriter, ParquetFormatWriter}; + use crate::io::FileIOBuilder; + use crate::spec::stats::BinaryTableStats; + use crate::spec::{DataFileMeta, FloatType, VectorType}; + use crate::table::schema_manager::SchemaManager; + use crate::table::source::DataSplitBuilder; + use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; + use arrow_array::RecordBatch; + use arrow_schema::{DataType as ArrowDataType, Field as ArrowField}; + use std::sync::Arc; + + fn vector_field() -> DataField { + let vector_type = VectorType::try_new(true, 2, DataType::Float(FloatType::new())).unwrap(); + DataField::new(0, "embedding".to_string(), DataType::Vector(vector_type)) + } + + fn data_file(file_name: &str, file_size: i64, row_count: i64, schema_id: 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: None, + write_cols: None, + } + } + + /// Write a FixedSizeList vector column + /// (`[1,2]`, NULL, `[3,4]`) as a parquet data file, build the factory over + /// its split, preload via `create`, and assert the whole-file sequential + /// read plus the "file not in split" error. + #[tokio::test] + async fn create_preloads_and_reads_whole_file() { + let field = vector_field(); + let read_fields = vec![field.clone()]; + let arrow_schema = build_target_arrow_schema(&read_fields).unwrap(); + + let mut builder = FixedSizeListBuilder::new(Float32Builder::new(), 2).with_field(Arc::new( + ArrowField::new("element", ArrowDataType::Float32, true), + )); + builder.values().append_value(1.0); + builder.values().append_value(2.0); + builder.append(true); + builder.values().append_value(0.0); + builder.values().append_value(0.0); + builder.append(false); // NULL vector row + builder.values().append_value(3.0); + builder.values().append_value(4.0); + builder.append(true); + let vec_array = builder.finish(); + let batch = RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(vec_array)]).unwrap(); + + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let table_path = "memory:/pk_vector_data_file_reader"; + let bucket_path = format!("{table_path}/bucket-0"); + let file_name = "part-0.parquet"; + let file_path = format!("{bucket_path}/{file_name}"); + let output = file_io.new_output(&file_path).unwrap(); + let mut writer: Box = Box::new( + ParquetFormatWriter::new(&output, arrow_schema.clone(), "zstd", 1) + .await + .unwrap(), + ); + writer.write(&batch).await.unwrap(); + let file_size = writer.close().await.unwrap(); + + let table_schema_id = 1; + let data_split = 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, + file_size as i64, + 3, + table_schema_id, + )]) + .build() + .unwrap(); + + let schema_manager = SchemaManager::new(file_io.clone(), table_path.to_string()); + let reader = DataFileReader::new( + file_io, + schema_manager, + table_schema_id, + read_fields.clone(), + read_fields.clone(), + Vec::new(), + ); + + let factory = + DataFilePkVectorReaderFactory::new(reader, data_split, field.clone()).unwrap(); + + let present = BucketActiveFile { + file_name: file_name.to_string(), + row_count: 3, + }; + let mut pk_reader = factory.create(&present).await.unwrap(); + assert_eq!(pk_reader.dimension(), 2); + assert_eq!(pk_reader.row_count(), 3); + + let mut buf = [0.0f32; 2]; + assert!(pk_reader.read_next_vector(&mut buf).unwrap()); + assert_eq!(buf, [1.0, 2.0]); + assert!(!pk_reader.read_next_vector(&mut buf).unwrap()); // NULL row + assert!(pk_reader.read_next_vector(&mut buf).unwrap()); + assert_eq!(buf, [3.0, 4.0]); + assert!(pk_reader.read_next_vector(&mut buf).is_err()); // past row count + + // A file name absent from the split is rejected as invalid. + let missing = BucketActiveFile { + file_name: "absent.parquet".to_string(), + row_count: 3, + }; + let err = factory + .create(&missing) + .await + .err() + .expect("absent file must be rejected"); + assert!(matches!(err, crate::Error::DataInvalid { .. })); + } +} diff --git a/crates/paimon/src/table/pk_vector_orchestrator.rs b/crates/paimon/src/table/pk_vector_orchestrator.rs index e6cef37f5..ac05367d2 100644 --- a/crates/paimon/src/table/pk_vector_orchestrator.rs +++ b/crates/paimon/src/table/pk_vector_orchestrator.rs @@ -67,17 +67,17 @@ pub(crate) struct PkVectorSearchSplit { pub active_files: Vec, } -/// 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, +/// A `bucket_search` hit tagged with its source bucket. `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`. +pub(crate) struct PkVectorCandidate { + pub split_index: usize, + pub partition: BinaryRow, + pub bucket: i32, + pub data_file_name: String, + pub row_position: i64, + pub distance: f32, } /// 5-level BEST_FIRST (smallest = best) key. Level 1 orders distance with @@ -85,7 +85,7 @@ struct Candidate { /// under inner product) sorts last rather than winning Top-1. Level 2 uses the /// partition's serialized bytes; Rust `Vec::cmp` is unsigned lexicographic /// then shorter-is-less, exactly the spec's contract (`[0x7f] < [0x80] < [0xff]`). -fn candidate_cmp(a: &Candidate, b: &Candidate) -> Ordering { +fn candidate_cmp(a: &PkVectorCandidate, b: &PkVectorCandidate) -> Ordering { java_float_compare(a.distance, b.distance) .then_with(|| { a.partition @@ -98,7 +98,7 @@ fn candidate_cmp(a: &Candidate, b: &Candidate) -> Ordering { } /// Collect all candidates, order BEST_FIRST, keep the best `limit`. -fn global_top_k(mut candidates: Vec, limit: usize) -> Vec { +fn global_top_k(mut candidates: Vec, limit: usize) -> Vec { candidates.sort_by(candidate_cmp); candidates.truncate(limit); candidates @@ -110,7 +110,7 @@ fn global_top_k(mut candidates: Vec, limit: usize) -> Vec /// emitted in ascending group-key order (deterministic file/position output /// order). Mirrors Java `PrimaryKeyVectorResult.splits()`. fn build_indexed_splits( - survivors: Vec, + survivors: Vec, splits: &[PkVectorSearchSplit], metric: VectorSearchMetric, ) -> crate::Result> { @@ -253,12 +253,14 @@ impl PkVectorOrchestrator { 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`). + /// Run the eager per-bucket search + cross-bucket global Top-K and return the + /// best-first survivors (through the full 5-level tie-break, raw distance + /// preserved). The exact-reader factory is split-scoped: it receives the + /// current split index and split so a caller can build a reader keyed to the + /// specific split/file. `skip_exact_fallback` forwards to `bucket_search`. #[allow(clippy::too_many_arguments)] - pub(crate) async fn read( + #[allow(clippy::type_complexity)] + pub(crate) async fn search_candidates( &self, splits: &[PkVectorSearchSplit], query: &[f32], @@ -266,10 +268,13 @@ impl PkVectorOrchestrator { limit: usize, ann_searcher: Option<&dyn PkVectorAnnSearcher>, exact_reader_factory: &mut dyn FnMut( + usize, + &PkVectorSearchSplit, &BucketActiveFile, ) -> crate::Result>, search_options: &HashMap, - ) -> crate::Result { + skip_exact_fallback: bool, + ) -> crate::Result> { // Eager input-shape validation (Java checkArgument parity). if limit == 0 { return Err(data_invalid("vector search limit must be positive")); @@ -279,20 +284,23 @@ impl PkVectorOrchestrator { } // Eager per-bucket search -> tagged candidates. - let mut candidates: Vec = Vec::new(); + let mut candidates: Vec = Vec::new(); for (split_index, split) in splits.iter().enumerate() { let dvs = build_bucket_dv_map(&self.reader, split).await?; + // Wrap the split-scoped factory into bucket_search's per-file signature. + let mut bucket_factory = + |file: &BucketActiveFile| exact_reader_factory(split_index, split, file); let results = bucket_search( ann_searcher, &split.ann_segments, &split.active_files, &dvs, - exact_reader_factory, + &mut bucket_factory, query, metric, limit, search_options, - false, + skip_exact_fallback, )?; for PkVectorSearchResult { data_file_name, @@ -300,7 +308,7 @@ impl PkVectorOrchestrator { distance, } in results { - candidates.push(Candidate { + candidates.push(PkVectorCandidate { split_index, partition: split.data_split.partition().clone(), bucket: split.data_split.bucket(), @@ -311,8 +319,41 @@ impl PkVectorOrchestrator { } } - // Eager global merge + grouping + split construction. - let survivors = global_top_k(candidates, limit); + Ok(global_top_k(candidates, limit)) + } + + /// See spec §3. `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>, + search_options: &HashMap, + ) -> crate::Result { + // Wrap the per-file factory into the split-scoped shape search_candidates + // expects; the split index/split are unused on this back-compat path. + let mut wrapped = + |_: usize, _: &PkVectorSearchSplit, f: &BucketActiveFile| exact_reader_factory(f); + let survivors = self + .search_candidates( + splits, + query, + metric, + limit, + ann_searcher, + &mut wrapped, + search_options, + false, + ) + .await?; let indexed_splits = build_indexed_splits(survivors, splits, metric)?; // Lazy materialization: own the splits + a reader clone. @@ -382,8 +423,14 @@ mod tests { } // 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 { + fn cand( + split_index: usize, + bucket: i32, + file: &str, + pos: i64, + distance: f32, + ) -> PkVectorCandidate { + PkVectorCandidate { split_index, partition: BinaryRow::new(0), bucket, @@ -400,8 +447,8 @@ mod tests { file: &str, pos: i64, distance: f32, - ) -> Candidate { - Candidate { + ) -> PkVectorCandidate { + PkVectorCandidate { split_index, partition: BinaryRow::from_bytes(1, partition_bytes), bucket, @@ -411,7 +458,7 @@ mod tests { } } - fn ids(c: &[Candidate]) -> Vec<(i32, String, i64)> { + fn ids(c: &[PkVectorCandidate]) -> Vec<(i32, String, i64)> { c.iter() .map(|c| (c.bucket, c.data_file_name.clone(), c.row_position)) .collect() @@ -1179,4 +1226,105 @@ mod e2e_tests { vec![l2_score(9.0), l2_score(1.0), l2_score(4.0)] ); } + + #[tokio::test] + async fn search_candidates_returns_best_first_survivors() { + // One bucket, exact-only, three rows; limit 2. Best-first by distance. + let table_path = "memory:/pkvo_candidates"; + let bucket_path = format!("{table_path}/bucket-0"); + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let meta = write_file(&file_io, &bucket_path, "c.mosaic", vec![1, 2, 3]).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("c.mosaic", 3)], + }; + // pos0 {3,0} d=9, pos1 {1,0} d=1, pos2 {2,0} d=4. + let mut factory = |_: usize, + _: &PkVectorSearchSplit, + f: &BucketActiveFile| + -> crate::Result> { + assert_eq!(f.file_name, "c.mosaic"); + Ok(Box::new(ArrayReader::new( + 2, + vec![ + Some(vec![3.0, 0.0]), + Some(vec![1.0, 0.0]), + Some(vec![2.0, 0.0]), + ], + ))) + }; + let opts = HashMap::new(); + let cands = PkVectorOrchestrator::new(make_reader(file_io, table_path)) + .search_candidates( + &[split], + &[0.0, 0.0], + VectorSearchMetric::L2, + 2, + None, + &mut factory, + &opts, + false, + ) + .await + .unwrap(); + // Best-first: pos1 (d=1), pos2 (d=4). + assert_eq!( + cands + .iter() + .map(|c| (c.row_position, c.distance)) + .collect::>(), + vec![(1, 1.0), (2, 4.0)] + ); + } + + #[tokio::test] + async fn search_candidates_fast_mode_skips_exact_factory() { + let table_path = "memory:/pkvo_fast"; + let bucket_path = format!("{table_path}/bucket-0"); + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let meta = write_file(&file_io, &bucket_path, "f.mosaic", vec![1, 2]).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("f.mosaic", 2)], + }; + let mut factory = |_: usize, + _: &PkVectorSearchSplit, + _: &BucketActiveFile| + -> crate::Result> { + unreachable!("fast mode must not read exact") + }; + let opts = HashMap::new(); + let cands = PkVectorOrchestrator::new(make_reader(file_io, table_path)) + .search_candidates( + &[split], + &[0.0, 0.0], + VectorSearchMetric::L2, + 2, + None, + &mut factory, + &opts, + true, + ) + .await + .unwrap(); + assert!(cands.is_empty()); + } } diff --git a/crates/paimon/src/table/pk_vector_scan.rs b/crates/paimon/src/table/pk_vector_scan.rs new file mode 100644 index 000000000..e635269e1 --- /dev/null +++ b/crates/paimon/src/table/pk_vector_scan.rs @@ -0,0 +1,502 @@ +// 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 search planning: resolve a snapshot, plan its data splits, +//! scan the index manifest for this column's ANN segments, and accumulate one +//! search split per bucket. Mirror of Java `PrimaryKeyVectorScan` and +//! `PrimaryKeyIndexSourcePolicy`. + +use std::collections::{BTreeMap, HashSet}; + +use crate::spec::{ + BinaryRow, DataFileMeta, FileKind, GlobalIndexMeta, IndexManifest, PkVectorSourceMeta, +}; +use crate::table::pk_vector_orchestrator::PkVectorSearchSplit; +use crate::table::source::{DataSplit, DataSplitBuilder, DeletionFile}; +use crate::table::Table; +use crate::vindex::pkvector::bucket::{BucketActiveFile, BucketAnnSegment}; + +const INDEX_DIR: &str = "index"; +const FILE_SOURCE_COMPACT: i32 = 1; + +fn data_invalid(message: impl Into) -> crate::Error { + crate::Error::DataInvalid { + message: message.into(), + source: None, + } +} + +/// Mirror of `PrimaryKeyIndexSourcePolicy.shouldRead`: only compacted, non-level-0 +/// files back the PK-vector index; an absent file source reads as false. +fn should_read_pk_index_source(file: &DataFileMeta) -> bool { + matches!(file.file_source, Some(src) if src == FILE_SOURCE_COMPACT) && file.level > 0 +} + +/// 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 { + snapshot_id: i64, + partition: BinaryRow, + bucket: i32, + bucket_path: Option, + total_buckets: Option, + data_files: Vec, + deletion_files: Vec>, + seen: HashSet, + any_deletion: bool, +} + +impl BucketAccumulator { + fn new(snapshot_id: i64, partition: BinaryRow, bucket: i32) -> Self { + Self { + snapshot_id, + partition, + bucket, + bucket_path: None, + total_buckets: None, + data_files: Vec::new(), + deletion_files: Vec::new(), + seen: HashSet::new(), + any_deletion: false, + } + } + + fn add(&mut self, split: &DataSplit) -> crate::Result<()> { + if split.snapshot_id() != self.snapshot_id { + return Err(data_invalid( + "data split snapshot id does not match plan snapshot", + )); + } + if split.partition().to_serialized_bytes() != self.partition.to_serialized_bytes() { + return Err(data_invalid( + "data split partition does not match bucket group", + )); + } + if split.bucket() != self.bucket { + return Err(data_invalid( + "data split bucket does not match bucket group", + )); + } + match &self.bucket_path { + Some(p) if p != split.bucket_path() => { + return Err(data_invalid("inconsistent bucket path within bucket group")) + } + None => self.bucket_path = Some(split.bucket_path().to_string()), + _ => {} + } + match self.total_buckets { + Some(tb) if tb != split.total_buckets() => { + return Err(data_invalid( + "inconsistent total buckets within bucket group", + )) + } + None => self.total_buckets = Some(split.total_buckets()), + _ => {} + } + let dvs = split.data_deletion_files(); + for (i, file) in split.data_files().iter().enumerate() { + if !self.seen.insert(file.file_name.clone()) { + return Err(data_invalid(format!( + "duplicate data file in bucket group: {}", + file.file_name + ))); + } + self.data_files.push(file.clone()); + let df = dvs.and_then(|d| d.get(i).cloned().flatten()); + if df.is_some() { + self.any_deletion = true; + } + self.deletion_files.push(df); + } + Ok(()) + } + + fn build(self) -> crate::Result { + let mut builder = DataSplitBuilder::new() + .with_snapshot(self.snapshot_id) + .with_partition(self.partition) + .with_bucket(self.bucket) + .with_bucket_path( + self.bucket_path + .ok_or_else(|| data_invalid("bucket group has no bucket path"))?, + ) + .with_total_buckets(self.total_buckets.unwrap_or(1)) + .with_data_files(self.data_files) + .with_raw_convertible(false); + if self.any_deletion { + builder = builder.with_data_deletion_files(self.deletion_files); + } + builder.build() + } +} + +/// The snapshot id plus the per-bucket search splits produced by planning. +#[allow(dead_code)] +pub(crate) struct PkVectorScanPlan { + pub snapshot_id: i64, + pub splits: Vec, +} + +#[allow(dead_code)] +pub(crate) struct PkVectorScan<'a> { + table: &'a Table, + vector_field_id: i32, + index_type: String, +} + +#[allow(dead_code)] +impl<'a> PkVectorScan<'a> { + pub(crate) fn new(table: &'a Table, vector_field_id: i32, index_type: String) -> Self { + Self { + table, + vector_field_id, + index_type, + } + } + + pub(crate) async fn plan(&self) -> crate::Result { + let snapshot_manager = self.table.snapshot_manager(); + let snapshot = match snapshot_manager.get_latest_snapshot().await? { + Some(s) => s, + None => { + return Ok(PkVectorScanPlan { + snapshot_id: -1, + splits: Vec::new(), + }) + } + }; + let snapshot_id = snapshot.id(); + + // Data splits (scan all files). + let builder = self.table.new_read_builder(); + let data_splits = builder + .new_scan() + .with_scan_all_files() + .plan() + .await? + .splits() + .to_vec(); + + // Index-manifest scan into filtered ANN payload tuples. + let table_path = self.table.location().trim_end_matches('/'); + let mut entries = Vec::new(); + if let Some(name) = snapshot.index_manifest() { + let path = snapshot_manager.manifest_path(name); + for entry in IndexManifest::read(self.table.file_io(), &path).await? { + if entry.kind != FileKind::Add { + continue; + } + if entry.index_file.index_type != self.index_type { + continue; + } + let Some(gim) = entry.index_file.global_index_meta.clone() else { + continue; + }; + if gim.index_field_id != self.vector_field_id { + continue; + } + if gim.source_meta.is_none() { + continue; + } + let partition = BinaryRow::from_serialized_bytes(&entry.partition)?; + let resolved_path = + format!("{table_path}/{INDEX_DIR}/{}", entry.index_file.file_name); + let file_size = u64::try_from(entry.index_file.file_size) + .map_err(|_| data_invalid("index file size must not be negative"))?; + entries.push(( + partition, + entry.bucket, + gim, + resolved_path, + file_size, + entry.index_file.file_name.clone(), + )); + } + } + + let splits = plan_from_inputs(snapshot_id, data_splits, entries)?; + Ok(PkVectorScanPlan { + snapshot_id, + splits, + }) + } +} + +/// Pure planning core, drivable without a live snapshot: group ANN payloads and +/// data splits by `(partition, bucket)`, then assemble one search split per +/// bucket that has data. Index-only buckets are dropped, not errored. +#[allow(dead_code)] +#[allow(clippy::type_complexity)] +fn plan_from_inputs( + snapshot_id: i64, + data_splits: Vec, + index_entries: Vec<(BinaryRow, i32, GlobalIndexMeta, String, u64, String)>, +) -> crate::Result> { + type Key = (Vec, i32); + + // Phase A: group ANN payloads by (partition, bucket). + let mut segments_by_bucket: BTreeMap> = BTreeMap::new(); + for (partition, bucket, gim, path, file_size, file_name) in index_entries { + let source_meta = PkVectorSourceMeta::from_global_index_meta(&gim) + .map_err(|_| data_invalid(format!("index file {file_name} is not active")))?; + let key = (partition.to_serialized_bytes(), bucket); + segments_by_bucket + .entry(key) + .or_default() + .push(BucketAnnSegment { + source_meta, + file_name, + path, + file_size, + index_meta: gim.index_meta.clone().unwrap_or_default(), + }); + } + + // Phase B: group data splits by (partition, bucket). + let mut accum_by_bucket: BTreeMap = BTreeMap::new(); + for split in &data_splits { + let key = (split.partition().to_serialized_bytes(), split.bucket()); + let acc = accum_by_bucket.entry(key).or_insert_with(|| { + BucketAccumulator::new(snapshot_id, split.partition().clone(), split.bucket()) + }); + acc.add(split)?; + } + + // 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 active_files: Vec = data_split + .data_files() + .iter() + .filter(|f| should_read_pk_index_source(f)) + .map(|f| BucketActiveFile { + file_name: f.file_name.clone(), + row_count: f.row_count, + }) + .collect(); + out.push(PkVectorSearchSplit { + data_split, + ann_segments, + active_files, + }); + } + // Index-only buckets left in segments_by_bucket are intentionally dropped. + Ok(out) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::spec::stats::BinaryTableStats; + use crate::spec::{BinaryRow, DataFileMeta, GlobalIndexMeta}; + use crate::table::source::{DataSplitBuilder, DeletionFile}; + + fn dfm(name: &str, rows: i64, level: i32, file_source: Option) -> DataFileMeta { + DataFileMeta { + file_name: name.into(), + file_size: 1, + row_count: rows, + 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, + extra_files: Vec::new(), + creation_time: None, + delete_row_count: None, + embedded_index: None, + file_source, + value_stats_cols: None, + external_path: None, + first_row_id: Some(0), + write_cols: None, + } + } + + #[test] + fn should_read_matches_java_policy() { + assert!(should_read_pk_index_source(&dfm("a", 1, 1, Some(1)))); // COMPACT + level>0 + assert!(!should_read_pk_index_source(&dfm("a", 1, 0, Some(1)))); // COMPACT + level==0 + assert!(!should_read_pk_index_source(&dfm("a", 1, 3, Some(0)))); // APPEND + assert!(!should_read_pk_index_source(&dfm("a", 1, 3, None))); // absent -> false + } + + /// Build one Java `DataOutput#writeUTF` value (u16-BE length + modified + /// UTF-8) for the ASCII test file names used here. + fn java_write_utf(s: &str) -> Vec { + let mut body = Vec::new(); + for c in s.encode_utf16() { + if (0x0001..=0x007F).contains(&c) { + body.push(c as u8); + } else if c > 0x07FF { + body.push(0xE0 | (c >> 12) as u8); + body.push(0x80 | ((c >> 6) & 0x3F) as u8); + body.push(0x80 | (c & 0x3F) as u8); + } else { + body.push(0xC0 | (c >> 6) as u8); + body.push(0x80 | (c & 0x3F) as u8); + } + } + let mut out = (body.len() as u16).to_be_bytes().to_vec(); + out.extend_from_slice(&body); + out + } + + /// 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 { + let mut out = Vec::new(); + out.extend_from_slice(&1i32.to_be_bytes()); // version + out.extend_from_slice(&(files.len() as i32).to_be_bytes()); + for (name, rows) in files { + out.extend_from_slice(&java_write_utf(name)); + out.extend_from_slice(&rows.to_be_bytes()); + } + out + } + + fn gim(field_id: 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)), + } + } + + #[test] + fn drops_index_only_bucket_without_error() { + // Payload for (part=[], bucket 0) but NO data split -> no split, no error. + let entries = vec![( + BinaryRow::new(0), + 0, + gim(2, &[("d0", 3)]), + "idx/seg0".to_string(), + 10u64, + "seg0".to_string(), + )]; + let splits = plan_from_inputs(1, Vec::new(), entries).unwrap(); + assert!(splits.is_empty()); + } + + #[test] + fn builds_one_split_per_bucket_with_data() { + let entries = vec![( + BinaryRow::new(0), + 0, + gim(2, &[("d0", 3)]), + "idx/seg0".to_string(), + 10u64, + "seg0".to_string(), + )]; + let data = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path("memory:/t/bucket-0".to_string()) + .with_total_buckets(1) + .with_data_files(vec![dfm("d0", 3, 5, Some(1))]) + .build() + .unwrap(); + let splits = plan_from_inputs(1, vec![data], entries).unwrap(); + assert_eq!(splits.len(), 1); + assert_eq!(splits[0].ann_segments.len(), 1); + let seg = &splits[0].ann_segments[0]; + assert_eq!(seg.file_name, "seg0"); + assert_eq!(seg.path, "idx/seg0"); + assert_eq!(seg.file_size, 10); + assert_eq!(seg.source_meta.resolve(0).unwrap(), ("d0".to_string(), 0)); + assert_eq!(splits[0].active_files.len(), 1); // d0 is COMPACT + level>0 + assert_eq!(splits[0].active_files[0].file_name, "d0"); + } + + #[test] + fn rejects_data_split_with_wrong_snapshot() { + let data = DataSplitBuilder::new() + .with_snapshot(2) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path("memory:/t/bucket-0".to_string()) + .with_total_buckets(1) + .with_data_files(vec![dfm("d0", 3, 5, Some(1))]) + .build() + .unwrap(); + assert!(plan_from_inputs(1, vec![data], Vec::new()).is_err()); + } + + #[test] + fn accumulator_rejects_duplicate_file_name() { + // Two splits in the SAME (partition, bucket) carrying a data file with the + // same name must fail loud via the accumulator's duplicate-file guard. + let split_a = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path("memory:/t/bucket-0".to_string()) + .with_total_buckets(1) + .with_data_files(vec![dfm("dup", 3, 5, Some(1))]) + .build() + .unwrap(); + let split_b = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path("memory:/t/bucket-0".to_string()) + .with_total_buckets(1) + .with_data_files(vec![dfm("dup", 3, 5, Some(1))]) + .build() + .unwrap(); + assert!(plan_from_inputs(1, vec![split_a, split_b], Vec::new()).is_err()); + } + + #[test] + fn accumulator_keeps_deletion_files_in_parallel_order() { + // One split, two data files; only the second carries a deletion file. The + // built split must preserve the [None, Some] alignment parallel to + // data_files. + let dv = DeletionFile::new("dv".to_string(), 0, 1, Some(1)); + let data = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path("memory:/t/bucket-0".to_string()) + .with_total_buckets(1) + .with_data_files(vec![dfm("d0", 3, 5, Some(1)), dfm("d1", 3, 5, Some(1))]) + .with_data_deletion_files(vec![None, Some(dv)]) + .build() + .unwrap(); + let splits = plan_from_inputs(1, vec![data], Vec::new()).unwrap(); + assert_eq!(splits.len(), 1); + let dvs = splits[0] + .data_split + .data_deletion_files() + .expect("deletion files preserved"); + assert_eq!(dvs.len(), 2); + assert!(dvs[0].is_none()); + assert!(dvs[1].is_some()); + // Both files are COMPACT + level>0, so both appear as active files. + assert_eq!(splits[0].active_files.len(), 2); + } +} From c4c7625009ae270a891af877007b5f9de7cb98b9 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Tue, 14 Jul 2026 18:59:33 +0800 Subject: [PATCH 04/10] feat(table): route primary-key vector search to the bucket-local reader Add the PK-vector branch to VectorSearchBuilder: resolve the configured column, plan the per-bucket splits, build the real vindex ANN scorer and exact-fallback readers, verify the segment metric matches the configured metric (fail loud on mismatch), and convert best-first candidates into a SearchResult. Drop the dead-code allowances and unused helpers now that the search path is wired. --- crates/paimon/src/table/data_file_reader.rs | 11 +- .../src/table/pk_vector_data_file_reader.rs | 18 +- .../src/table/pk_vector_indexed_split_read.rs | 14 +- .../src/table/pk_vector_orchestrator.rs | 31 +- .../src/table/pk_vector_position_read.rs | 24 +- crates/paimon/src/table/pk_vector_scan.rs | 24 +- .../paimon/src/table/vector_search_builder.rs | 617 +++++++++++++++++- crates/paimon/src/vindex/pkvector/bucket.rs | 2 - crates/paimon/src/vindex/pkvector/metric.rs | 80 +-- crates/paimon/src/vindex/pkvector/mod.rs | 5 - 10 files changed, 715 insertions(+), 111 deletions(-) diff --git a/crates/paimon/src/table/data_file_reader.rs b/crates/paimon/src/table/data_file_reader.rs index 3395f8dd8..791ca0b0f 100644 --- a/crates/paimon/src/table/data_file_reader.rs +++ b/crates/paimon/src/table/data_file_reader.rs @@ -71,20 +71,16 @@ 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) -> 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. + /// Exposed for the sibling `pk_vector_position_read` module, which drives the + /// PK-vector materialization read path (no production caller yet). #[allow(dead_code)] pub(super) fn read_type(&self) -> &[DataField] { &self.read_type @@ -92,7 +88,8 @@ impl DataFileReader { /// 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. + /// matching `reject_row_id_with_predicate`'s notion. Consumed by + /// `pk_vector_position_read` (materialization read path; no production caller yet). #[allow(dead_code)] pub(super) fn has_row_filtering_predicate(&self) -> bool { self.predicates diff --git a/crates/paimon/src/table/pk_vector_data_file_reader.rs b/crates/paimon/src/table/pk_vector_data_file_reader.rs index b0ba71c8c..f2c3a6ddb 100644 --- a/crates/paimon/src/table/pk_vector_data_file_reader.rs +++ b/crates/paimon/src/table/pk_vector_data_file_reader.rs @@ -46,7 +46,6 @@ fn data_invalid(message: impl Into) -> crate::Error { /// `reader` is configured (via [`DataFileReader::with_read_type`]) to project /// only the vector column, so each read returns a single-column batch. Mirrors /// Java `PkVectorDataFileReader` (as a factory owning the projected reader). -#[allow(dead_code)] pub(crate) struct DataFilePkVectorReaderFactory { reader: DataFileReader, data_split: DataSplit, @@ -58,7 +57,6 @@ impl DataFilePkVectorReaderFactory { /// Configure `reader` to project the vector column only and capture the /// vector dimension from the schema field. The field must be a fixed-length /// `Vector` type; anything else is rejected as invalid. - #[allow(dead_code)] pub(crate) fn new( reader: DataFileReader, data_split: DataSplit, @@ -85,7 +83,6 @@ impl DataFilePkVectorReaderFactory { /// sequential reader over it. `file` must name a data file present in this /// factory's split. The drained row count is checked against the file's /// `DataFileMeta.row_count`. - #[allow(dead_code)] pub(crate) async fn create( &self, file: &BucketActiveFile, @@ -384,12 +381,19 @@ mod integration_tests { let file_path = format!("{bucket_path}/{file_name}"); let output = file_io.new_output(&file_path).unwrap(); let mut writer: Box = Box::new( - ParquetFormatWriter::new(&output, arrow_schema.clone(), "zstd", 1) - .await - .unwrap(), + ParquetFormatWriter::new( + &output, + arrow_schema.clone(), + "zstd", + 1, + None, + &std::collections::HashMap::new(), + ) + .await + .unwrap(), ); writer.write(&batch).await.unwrap(); - let file_size = writer.close().await.unwrap(); + let file_size = writer.close().await.unwrap().file_size; let table_schema_id = 1; let data_split = DataSplitBuilder::new() diff --git a/crates/paimon/src/table/pk_vector_indexed_split_read.rs b/crates/paimon/src/table/pk_vector_indexed_split_read.rs index 00bd1ae52..54475ddb6 100644 --- a/crates/paimon/src/table/pk_vector_indexed_split_read.rs +++ b/crates/paimon/src/table/pk_vector_indexed_split_read.rs @@ -15,21 +15,19 @@ // specific language governing permissions and limitations // under the License. -//! Primary-key vector indexed-split read-path contract (read-path subset of -//! apache/paimon#8576). +//! Primary-key vector indexed-split read-path contract. //! //! `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 +//! It is a pure consumer: no bucket/ANN search, no cross-bucket merge, 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. +// This materialization read path is exercised only by its own tests: the wired +// vector search returns matched row-ids and scores directly, so no production +// caller drives row materialization yet. Suppress dead_code at the module +// boundary until a materializing caller exists. #![allow(dead_code)] use std::collections::BTreeMap; diff --git a/crates/paimon/src/table/pk_vector_orchestrator.rs b/crates/paimon/src/table/pk_vector_orchestrator.rs index ac05367d2..21adb41a7 100644 --- a/crates/paimon/src/table/pk_vector_orchestrator.rs +++ b/crates/paimon/src/table/pk_vector_orchestrator.rs @@ -15,18 +15,12 @@ // specific language governing permissions and limitations // under the License. -//! Primary-key vector read orchestration: Rust equivalent of Java -//! `PrimaryKeyVectorRead` + `PrimaryKeyVectorResult.splits()`. +//! 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)] +//! materialization via `PkVectorIndexedSplitRead`. use std::cmp::Ordering; use std::collections::HashMap; @@ -53,8 +47,9 @@ fn data_invalid(message: impl Into) -> crate::Error { } } -/// Per-bucket search input. Rust equivalent of Java `BucketVectorSearchSplit`. -/// A `PrimaryKeyVectorScan` mirror constructs these from a snapshot/manifest plan. +/// One bucket's search input. Rust equivalent of Java +/// `BucketVectorSearchSplit`. Constructed from a snapshot/manifest plan by +/// `PkVectorScan`. 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 @@ -109,6 +104,9 @@ fn global_top_k(mut candidates: Vec, limit: usize) -> Vec, splits: &[PkVectorSearchSplit], @@ -214,7 +212,8 @@ fn build_indexed_splits( /// 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.) +/// its own DV again — an accepted redundancy between the search and +/// materialization phases.) async fn build_bucket_dv_map( reader: &DataFileReader, split: &PkVectorSearchSplit, @@ -322,9 +321,15 @@ impl PkVectorOrchestrator { Ok(global_top_k(candidates, limit)) } - /// See spec §3. `async` because the eager search phase is genuine async IO + /// Run the eager per-bucket search phase, then return a stream that lazily + /// materializes 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`). + /// + /// The wired vector search returns matched row-ids and scores via + /// `search_candidates`; no production caller drives row materialization yet. + #[allow(dead_code)] #[allow(clippy::too_many_arguments)] pub(crate) async fn read( &self, diff --git a/crates/paimon/src/table/pk_vector_position_read.rs b/crates/paimon/src/table/pk_vector_position_read.rs index 6a5ec454a..e8fedb896 100644 --- a/crates/paimon/src/table/pk_vector_position_read.rs +++ b/crates/paimon/src/table/pk_vector_position_read.rs @@ -15,17 +15,19 @@ // specific language governing permissions and limitations // under the License. -//! Primary-key vector position read (reader-kernel subset of apache/paimon#8576). +//! Primary-key vector position read (Rust equivalent of Java +//! `PrimaryKeyVectorPositionReader`). //! //! 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. +//! 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. + +// This materialization read path is exercised only by its own tests: the wired +// vector search returns matched row-ids and scores directly, so no production +// caller drives row materialization yet. Suppress dead_code at the module +// boundary until a materializing caller exists. #![allow(dead_code)] use std::collections::BTreeMap; @@ -93,7 +95,7 @@ fn positions_to_global_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). +/// `PrimaryKeyVectorPositionReader`. pub(crate) struct PkVectorPositionRead<'a> { reader: &'a DataFileReader, } @@ -164,8 +166,8 @@ impl<'a> PkVectorPositionRead<'a> { } } - // (4) predicate guard: a row-filtering predicate would desync positional - // row-id recovery, so reject a reader that carries one. + // (4) predicate guard: a residual row-filtering predicate would desync + // positional row-id assignment, so reject it here. if self.reader.has_row_filtering_predicate() { return Err(data_invalid( "PK vector position read requires a predicate-free reader", diff --git a/crates/paimon/src/table/pk_vector_scan.rs b/crates/paimon/src/table/pk_vector_scan.rs index e635269e1..54900ad5b 100644 --- a/crates/paimon/src/table/pk_vector_scan.rs +++ b/crates/paimon/src/table/pk_vector_scan.rs @@ -144,21 +144,17 @@ impl BucketAccumulator { } } -/// The snapshot id plus the per-bucket search splits produced by planning. -#[allow(dead_code)] +/// The per-bucket search splits produced by planning. pub(crate) struct PkVectorScanPlan { - pub snapshot_id: i64, pub splits: Vec, } -#[allow(dead_code)] pub(crate) struct PkVectorScan<'a> { table: &'a Table, vector_field_id: i32, index_type: String, } -#[allow(dead_code)] impl<'a> PkVectorScan<'a> { pub(crate) fn new(table: &'a Table, vector_field_id: i32, index_type: String) -> Self { Self { @@ -172,12 +168,7 @@ impl<'a> PkVectorScan<'a> { let snapshot_manager = self.table.snapshot_manager(); let snapshot = match snapshot_manager.get_latest_snapshot().await? { Some(s) => s, - None => { - return Ok(PkVectorScanPlan { - snapshot_id: -1, - splits: Vec::new(), - }) - } + None => return Ok(PkVectorScanPlan { splits: Vec::new() }), }; let snapshot_id = snapshot.id(); @@ -229,17 +220,13 @@ impl<'a> PkVectorScan<'a> { } let splits = plan_from_inputs(snapshot_id, data_splits, entries)?; - Ok(PkVectorScanPlan { - snapshot_id, - splits, - }) + Ok(PkVectorScanPlan { splits }) } } /// Pure planning core, drivable without a live snapshot: group ANN payloads and /// data splits by `(partition, bucket)`, then assemble one search split per /// bucket that has data. Index-only buckets are dropped, not errored. -#[allow(dead_code)] #[allow(clippy::type_complexity)] fn plan_from_inputs( snapshot_id: i64, @@ -259,9 +246,11 @@ fn plan_from_inputs( .or_default() .push(BucketAnnSegment { source_meta, - file_name, path, file_size, + // Not consumed on the search path: the vindex reader loads its + // metadata from the index file bytes and ignores this field, so an + // absent value defaulting to an empty vec is acceptable. index_meta: gim.index_meta.clone().unwrap_or_default(), }); } @@ -424,7 +413,6 @@ mod tests { assert_eq!(splits.len(), 1); assert_eq!(splits[0].ann_segments.len(), 1); let seg = &splits[0].ann_segments[0]; - assert_eq!(seg.file_name, "seg0"); assert_eq!(seg.path, "idx/seg0"); assert_eq!(seg.file_size, 10); assert_eq!(seg.source_meta.resolve(0).unwrap(), ("d0".to_string(), 0)); diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index c5bdc9643..4c15d1ae0 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -22,13 +22,23 @@ use crate::spec::{ CoreOptions, DataField, FileKind, GlobalIndexSearchMode, IndexFileMeta, IndexManifest, IndexManifestEntry, ROW_ID_FIELD_NAME, }; +use crate::table::data_file_reader::DataFileReader; use crate::table::global_index_scanner::{ deleted_row_ranges_for_data_evolution_dvs, search_limit_with_deleted_rows, unindexed_ranges_for_global_index_entries, RowRangeIndex, }; +use crate::table::pk_vector_data_file_reader::DataFilePkVectorReaderFactory; +use crate::table::pk_vector_orchestrator::{ + PkVectorCandidate, PkVectorOrchestrator, PkVectorSearchSplit, +}; +use crate::table::pk_vector_scan::PkVectorScan; use crate::table::{find_field_id_by_name, merge_row_ranges, RowRange, Table}; use crate::vector_search::{GlobalIndexIOMeta, SearchResult, VectorSearch}; use crate::vindex::is_vindex_index_type; +use crate::vindex::pkvector::ann::VindexAnnSearcher; +use crate::vindex::pkvector::bucket::{BucketActiveFile, BucketAnnSegment}; +use crate::vindex::pkvector::metric::VectorSearchMetric; +use crate::vindex::pkvector::reader::PkVectorReader; use crate::vindex::reader::VindexVectorGlobalIndexReader; use arrow_array::{Array, FixedSizeListArray, Float32Array, Int64Array, ListArray, RecordBatch}; use futures::TryStreamExt; @@ -119,7 +129,8 @@ impl<'a> VectorSearchBuilder<'a> { pub async fn execute_scored(&self) -> crate::Result { // Fail closed: returns data-derived row ranges outside `TableScan`/`TableRead`. - CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?; + let core = CoreOptions::new(self.table.schema().options()); + core.ensure_read_authorized()?; let vector_column = self.vector_column .as_deref() @@ -136,6 +147,30 @@ impl<'a> VectorSearchBuilder<'a> { message: "Limit must be set via with_limit()".to_string(), })?; + // Primary-key vector search branch: mirrors Java `PrimaryKeyVectorRead`. + // Only taken when the table enables the PK-vector index AND this query + // targets a configured PK-vector column; otherwise fall through to the + // data-evolution (DE) global-index path below. + // + // Membership is resolved first via the non-erroring columns accessor so a + // malformed PK-vector config (e.g. more than one column, or a blank list) + // cannot abort an unrelated DE query. The exactly-one-column rule is + // enforced only once this query is known to target a PK-vector column, + // keeping fail-loud behavior for a genuinely-broken config on the path + // where erroring is correct. + if core.primary_key_vector_index_enabled() { + let targets_pk_column = core + .primary_key_vector_index_columns() + .ok() + .is_some_and(|cols| cols.iter().any(|c| c == vector_column)); + if targets_pk_column { + let pk_col = core.primary_key_vector_index_column()?; + return self + .execute_primary_key_vector_search(&core, &pk_col, query_vector, limit) + .await; + } + } + let mut batch_builder = BatchVectorSearchBuilder::new(self.table); let mut results = batch_builder .with_vector_column(vector_column) @@ -148,6 +183,143 @@ impl<'a> VectorSearchBuilder<'a> { debug_assert_eq!(results.len(), 1); Ok(results.remove(0)) } + + /// Run the primary-key bucket-local vector search: plan the per-bucket splits, + /// build the real vindex ANN scorer and (outside FAST mode) the exact-fallback + /// readers, run the orchestrator, and convert the best-first candidates into a + /// `SearchResult`. Mirrors Java `PrimaryKeyVectorRead`. + async fn execute_primary_key_vector_search( + &self, + core: &CoreOptions<'_>, + pk_col: &str, + query_vector: &[f32], + limit: usize, + ) -> crate::Result { + // Residual-filter guard: PK vector search accepts partition filters only. + // This builder exposes no data-predicate setter, so there is nothing to + // reject here; the guard mirrors Java `checkArgument(filter == null)` and, + // if a filter setter is ever added, it must error rather than be ignored. + + // `primary_key_vector_distance_metric` returns a validated name; re-parse + // into the enum for the numeric semantics. + let metric = VectorSearchMetric::parse(&core.primary_key_vector_distance_metric(pk_col)?)?; + let index_type = core.primary_key_vector_index_type(pk_col)?; + let field_id = + find_field_id_by_name(self.table.schema().fields(), pk_col).ok_or_else(|| { + crate::Error::DataInvalid { + message: format!("PK-vector column '{pk_col}' not found in schema"), + source: None, + } + })?; + let vector_field = self + .table + .schema() + .fields() + .iter() + .find(|f| f.name() == pk_col) + .cloned() + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("PK-vector column '{pk_col}' not found in schema"), + source: None, + })?; + + let search_mode = core.global_index_search_mode()?; + let skip_exact_fallback = search_mode == GlobalIndexSearchMode::Fast; + + let plan = PkVectorScan::new(self.table, field_id, index_type) + .plan() + .await?; + if plan.splits.is_empty() { + return Ok(SearchResult::empty()); + } + + // Production data-file reader, mirroring `table_read.rs::new_data_file_reader` + // but projecting only the vector column with no predicates. + let reader = DataFileReader::new( + self.table.file_io().clone(), + self.table.schema_manager().clone(), + self.table.schema().id(), + self.table.schema().fields().to_vec(), + vec![vector_field.clone()], + Vec::new(), + ); + + // Real ANN scorer: preload each segment's bytes (keyed by resolved, + // globally unique path) and drive the vindex reader from memory. + let segment_bytes = preload_segment_bytes(self.table.file_io(), &plan.splits).await?; + // Fail loud on a config/segment metric mismatch before scoring, mirroring + // Java `PkVectorAnnSegmentSearcher.search`. + verify_pk_vector_segment_metrics(&plan.splits, &segment_bytes, metric)?; + let options = { + let mut o = self.table.schema().options().clone(); + o.extend(self.options.clone()); + o + }; + let search_options = options.clone(); + let field_name = pk_col.to_string(); + let scorer: crate::vindex::pkvector::ann::Scorer = + Box::new(move |segment: &BucketAnnSegment, search: &VectorSearch| { + let data = segment_bytes + .get(&segment.path) + .ok_or_else(|| crate::Error::DataInvalid { + message: "missing preloaded ANN bytes for segment".to_string(), + source: None, + })? + .clone(); + let io_meta = GlobalIndexIOMeta::new( + segment.path.clone(), + segment.file_size, + segment.index_meta.clone(), + ); + let mut reader = VindexVectorGlobalIndexReader::new(io_meta, options.clone()); + reader.visit_vector_search(search, |_| Ok(Cursor::new(data))) + }); + let ann_searcher = VindexAnnSearcher::new(field_name, scorer); + + // Exact-fallback readers, keyed by (split_index, file_name). In FAST mode + // the kernel never invokes the factory, so skip the in-memory column read + // entirely. + let mut exact_readers: HashMap<(usize, String), Box> = HashMap::new(); + if !skip_exact_fallback { + for (split_index, split) in plan.splits.iter().enumerate() { + let factory = DataFilePkVectorReaderFactory::new( + reader.clone(), + split.data_split.clone(), + vector_field.clone(), + )?; + for active in &split.active_files { + let r = factory.create(active).await?; + exact_readers.insert((split_index, active.file_name.clone()), r); + } + } + } + let mut factory = |split_index: usize, + _split: &PkVectorSearchSplit, + file: &BucketActiveFile| + -> crate::Result> { + exact_readers + .remove(&(split_index, file.file_name.clone())) + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("no preloaded exact reader for {}", file.file_name), + source: None, + }) + }; + + let candidates = PkVectorOrchestrator::new(reader) + .search_candidates( + &plan.splits, + query_vector, + metric, + limit, + Some(&ann_searcher), + &mut factory, + &search_options, + skip_exact_fallback, + ) + .await?; + + candidates_to_search_result(&candidates, &plan.splits, metric) + } } impl<'a> BatchVectorSearchBuilder<'a> { @@ -519,6 +691,135 @@ fn is_vector_global_index_file(index_file: &IndexFileMeta) -> bool { VectorIndexBackend::from_index_type(&index_file.index_type).is_some() } +/// Preload every ANN segment's bytes into a map keyed by the resolved (globally +/// unique) segment path. The scorer closure reads from this map so the vindex +/// reader is driven from memory without per-search IO. +async fn preload_segment_bytes( + file_io: &FileIO, + splits: &[PkVectorSearchSplit], +) -> crate::Result>> { + let mut out = HashMap::new(); + for split in splits { + for segment in &split.ann_segments { + if out.contains_key(&segment.path) { + continue; + } + let input = file_io.new_input(&segment.path)?; + let bytes = input.read().await.map_err(|e| crate::Error::DataInvalid { + message: format!("failed to read ANN index file '{}': {e}", segment.path), + source: None, + })?; + out.insert(segment.path.clone(), bytes.to_vec()); + } + } + Ok(out) +} + +/// Fail loud when an ANN segment was trained with a metric other than the +/// configured one, mirroring the search-time `checkArgument` in Java +/// `PkVectorAnnSegmentSearcher.search`. Opens each distinct segment's preloaded +/// bytes once and compares its trained metric against `configured`. +fn verify_pk_vector_segment_metrics( + splits: &[PkVectorSearchSplit], + segment_bytes: &HashMap>, + configured: VectorSearchMetric, +) -> crate::Result<()> { + let mut checked: HashSet<&str> = HashSet::new(); + for split in splits { + for segment in &split.ann_segments { + if !checked.insert(segment.path.as_str()) { + continue; + } + let bytes = + segment_bytes + .get(&segment.path) + .ok_or_else(|| crate::Error::DataInvalid { + message: format!( + "missing preloaded ANN bytes for segment '{}'", + segment.path + ), + source: None, + })?; + let reader = VIndexReader::open(Cursor::new(bytes.clone())).map_err(|e| { + crate::Error::DataInvalid { + message: format!( + "failed to open ANN index file '{}' for metric check: {e}", + segment.path + ), + source: Some(Box::new(e)), + } + })?; + let segment_metric = reader.metadata().metric; + if VectorSearchMetric::from_vindex(segment_metric) != configured { + return Err(crate::Error::DataInvalid { + message: format!( + "ANN segment metric {} does not match configured metric {}", + segment_metric.as_str(), + configured.as_str() + ), + source: None, + }); + } + } + } + Ok(()) +} + +/// candidate order (no re-sort). Each candidate's global row id is +/// `first_row_id + row_position` of the data file it references; the score is +/// derived from the raw distance via the metric. A candidate referencing a file +/// absent from its split, or a file with no `first_row_id`, fails loud. +fn candidates_to_search_result( + candidates: &[PkVectorCandidate], + splits: &[PkVectorSearchSplit], + metric: VectorSearchMetric, +) -> crate::Result { + let mut row_ids = Vec::with_capacity(candidates.len()); + let mut scores = Vec::with_capacity(candidates.len()); + for c in candidates { + let split = splits + .get(c.split_index) + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("candidate split_index {} out of range", c.split_index), + source: None, + })?; + let file_meta = split + .data_split + .data_files() + .iter() + .find(|f| f.file_name == c.data_file_name) + .ok_or_else(|| crate::Error::DataInvalid { + message: format!( + "candidate references data file {} not present in its split", + c.data_file_name + ), + source: None, + })?; + let first_row_id = file_meta + .first_row_id + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("data file {} has no first_row_id", c.data_file_name), + source: None, + })?; + let global = + first_row_id + .checked_add(c.row_position) + .ok_or_else(|| crate::Error::DataInvalid { + message: "global row id overflows i64".to_string(), + source: None, + })?; + row_ids.push( + u64::try_from(global).map_err(|_| crate::Error::DataInvalid { + message: format!("negative global row id {global}"), + source: None, + })?, + ); + scores.push(metric.distance_to_score(c.distance)); + } + // Order preserved: best-first, as produced by the orchestrator. + Ok(SearchResult::new(row_ids, scores)) +} + fn indexed_search_limit(limit: usize, refine_factor: usize) -> crate::Result { if refine_factor == 0 { return Ok(limit); @@ -1324,10 +1625,12 @@ mod tests { use crate::catalog::Identifier; use crate::io::FileIOBuilder; use crate::lumina::{LEGACY_LUMINA_VECTOR_ANN_IDENTIFIER, LUMINA_IDENTIFIER}; + use crate::spec::stats::BinaryTableStats; use crate::spec::{ - ArrayType, DataType, FloatType, GlobalIndexMeta, IndexFileMeta, IndexManifestEntry, - IntType, Schema, TableSchema, + ArrayType, BinaryRow, DataFileMeta, DataType, FloatType, GlobalIndexMeta, IndexFileMeta, + IndexManifestEntry, IntType, Schema, TableSchema, }; + use crate::table::source::DataSplitBuilder; use crate::vindex::IVF_FLAT_IDENTIFIER; use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; use arrow_array::ArrayRef; @@ -1876,6 +2179,314 @@ mod tests { ); } + fn pk_data_file(name: &str, row_count: i64, first_row_id: Option) -> 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, + write_cols: None, + } + } + + fn pk_search_split(bucket: i32, files: Vec) -> PkVectorSearchSplit { + PkVectorSearchSplit { + data_split: DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(bucket) + .with_bucket_path(format!("memory:/t/bucket-{bucket}")) + .with_total_buckets(1) + .with_data_files(files) + .build() + .unwrap(), + ann_segments: Vec::new(), + active_files: Vec::new(), + } + } + + fn pk_candidate( + split_index: usize, + bucket: i32, + file: &str, + pos: i64, + distance: f32, + ) -> PkVectorCandidate { + PkVectorCandidate { + split_index, + partition: BinaryRow::new(0), + bucket, + data_file_name: file.to_string(), + row_position: pos, + distance, + } + } + + #[test] + fn candidates_to_search_result_global_row_id_and_best_first_order() { + // Two files in one split with different first_row_id. The helper is a pure + // order-preserving map: the orchestrator already established best-first + // order upstream, so the candidate INPUT order here is deliberately NOT in + // score order and NOT in (file, position) order. This proves the helper + // preserves the given sequence rather than sorting. + let splits = vec![pk_search_split( + 0, + vec![ + pk_data_file("file-a", 100, Some(1000)), + pk_data_file("file-b", 100, Some(5000)), + ], + )]; + // Input sequence (NOT sorted by score, NOT sorted by file/position): + // c0: file-b pos5 d=2.0 -> WORST distance, appears FIRST + // c1: file-b pos1 d=1.0 -> tie with c2 + // c2: file-a pos2 d=1.0 -> tie with c1 + // A score-based best-first re-sort would produce [c1, c2, c0] (worst last); + // a (file, position) re-sort would produce [c2 (file-a), c1, c0]. Both + // differ from the input order, so the exact assertion below discriminates. + let candidates = vec![ + pk_candidate(0, 0, "file-b", 5, 2.0), + pk_candidate(0, 0, "file-b", 1, 1.0), + pk_candidate(0, 0, "file-a", 2, 1.0), + ]; + let result = candidates_to_search_result(&candidates, &splits, VectorSearchMetric::L2) + .expect("conversion succeeds"); + // global_row_id = first_row_id + position; INPUT order preserved (not sorted). + assert_eq!(result.row_ids, vec![5005, 5001, 1002]); + assert_eq!( + result.scores, + vec![ + VectorSearchMetric::L2.distance_to_score(2.0), + VectorSearchMetric::L2.distance_to_score(1.0), + VectorSearchMetric::L2.distance_to_score(1.0), + ] + ); + } + + #[test] + fn candidates_to_search_result_absent_first_row_id_fails_loud() { + let splits = vec![pk_search_split(0, vec![pk_data_file("file-a", 100, None)])]; + let candidates = vec![pk_candidate(0, 0, "file-a", 0, 1.0)]; + let err = candidates_to_search_result(&candidates, &splits, VectorSearchMetric::L2) + .expect_err("absent first_row_id must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("first_row_id")), + "unexpected error: {err:?}" + ); + } + + /// Build a real vindex IVF-flat segment trained with `metric`, returning the + /// serialized bytes. `nlist = 1` keeps training trivial and deterministic; the + /// only thing the metric check cares about is the persisted metadata metric. + fn build_vindex_segment_bytes(metric: &str) -> Vec { + use paimon_vindex_core::index::{VectorIndexConfig, VectorIndexTrainer, VectorIndexWriter}; + use paimon_vindex_core::io::PosWriter; + + const DIM: usize = 2; + let vectors: Vec = vec![1.0, 0.0, 0.0, 1.0, 1.0, 1.0]; + let n = vectors.len() / DIM; + let ids: Vec = (0..n as i64).collect(); + let options = HashMap::from([ + ("index.type".to_string(), "ivf_flat".to_string()), + ("dimension".to_string(), DIM.to_string()), + ("nlist".to_string(), "1".to_string()), + ("metric".to_string(), metric.to_string()), + ]); + let config = VectorIndexConfig::from_options(&options).unwrap(); + let training = VectorIndexTrainer::train(config, &vectors, n).unwrap(); + let mut writer = VectorIndexWriter::new(training); + writer.add_vectors(&ids, &vectors, n).unwrap(); + let mut bytes = Vec::new(); + { + let mut output = PosWriter::new(&mut bytes); + writer.write(&mut output).unwrap(); + } + bytes + } + + /// 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 mut segment = BucketAnnSegment::for_test(source_meta); + segment.path = path.to_string(); + split.ann_segments = vec![segment]; + split + } + + #[test] + fn verify_pk_vector_segment_metrics_accepts_matching_metric() { + // Real IVF segment trained with L2; configured metric L2 => Ok. + let bytes = build_vindex_segment_bytes("l2"); + let splits = vec![pk_split_with_segment("seg-l2")]; + let segment_bytes = HashMap::from([("seg-l2".to_string(), bytes)]); + verify_pk_vector_segment_metrics(&splits, &segment_bytes, VectorSearchMetric::L2) + .expect("matching metric must pass"); + } + + #[test] + fn verify_pk_vector_segment_metrics_rejects_mismatched_metric() { + // Real IVF segment trained with L2; configured metric Cosine => fail loud. + let bytes = build_vindex_segment_bytes("l2"); + let splits = vec![pk_split_with_segment("seg-l2")]; + let segment_bytes = HashMap::from([("seg-l2".to_string(), bytes)]); + let err = + verify_pk_vector_segment_metrics(&splits, &segment_bytes, VectorSearchMetric::Cosine) + .expect_err("mismatched metric must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } + if message.contains("does not match configured metric") + && message.contains("l2") + && message.contains("cosine")), + "unexpected error: {err:?}" + ); + } + + #[test] + fn candidates_to_search_result_missing_file_fails_loud() { + let splits = vec![pk_search_split( + 0, + vec![pk_data_file("known", 100, Some(0))], + )]; + let candidates = vec![pk_candidate(0, 0, "unknown", 0, 1.0)]; + let err = candidates_to_search_result(&candidates, &splits, VectorSearchMetric::L2) + .expect_err("missing file must fail loud"); + assert!(matches!(err, crate::Error::DataInvalid { .. })); + } + + fn pk_vector_table(options: &[(&str, &str)]) -> Table { + let mut builder = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column( + "embedding", + DataType::Array(ArrayType::new(DataType::Float(FloatType::new()))), + ); + for (k, v) in options { + builder = builder.option(*k, *v); + } + let schema = builder.build().unwrap(); + Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "pk_vector_test"), + "memory:/pk_vector_test".to_string(), + TableSchema::new(0, &schema), + None, + ) + } + + #[tokio::test] + async fn pk_branch_disabled_falls_through_to_de_path() { + // No pk-vector.index.columns: behaves exactly as the DE path. With no + // snapshot the DE path returns an empty result; the PK branch must not + // intercept it. + let table = pk_vector_table(&[]); + let result = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0]) + .with_limit(5) + .execute_scored() + .await + .unwrap(); + assert!(result.is_empty()); + } + + #[tokio::test] + async fn pk_branch_enabled_empty_plan_returns_empty() { + // pk-vector.index.columns set, but no snapshot -> empty plan -> empty result. + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.embedding.pk-vector.distance.metric", "l2"), + ]); + let result = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0]) + .with_limit(5) + .execute_scored() + .await + .unwrap(); + assert!(result.is_empty()); + } + + #[tokio::test] + async fn pk_branch_other_column_falls_through_to_de_path() { + // pk-vector index configured for "embedding", but the query targets a + // different column -> the PK branch must not intercept; DE path (no + // snapshot) yields empty. Discriminator: the PK column carries a + // DELIBERATELY INVALID distance metric, which the PK branch parses eagerly + // (`VectorSearchMetric::parse`) and would fail on. So a regression that + // dropped the `pk_col == vector_column` guard and ran the PK branch for + // "other" would surface as Err here, not Ok(empty) -- the assertion + // therefore proves the DE path ran, not merely that the result is empty. + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ( + "fields.embedding.pk-vector.distance.metric", + "not-a-real-metric", + ), + ]); + let result = table + .new_vector_search_builder() + .with_vector_column("other") + .with_query_vector(vec![1.0]) + .with_limit(5) + .execute_scored() + .await + .unwrap(); + assert!(result.is_empty()); + } + + #[tokio::test] + async fn pk_branch_multi_column_config_does_not_break_unrelated_de_query() { + // A malformed multi-column PK-vector config ("a,b") must not abort an + // unrelated DE vector query. The query targets a column NOT among the + // configured PK-vector columns, so membership resolution short-circuits + // before the exactly-one-column rule fires -- the query falls through to + // the DE path (no snapshot -> empty) instead of surfacing the "must name + // exactly one column" error. + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "a,b"), + ("fields.a.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.a.pk-vector.distance.metric", "l2"), + ]); + let result = table + .new_vector_search_builder() + .with_vector_column("other") + .with_query_vector(vec![1.0]) + .with_limit(5) + .execute_scored() + .await; + match result { + Ok(search) => assert!(search.is_empty()), + Err(err) => panic!( + "unrelated DE query must not error on a malformed multi-column PK config: {err}" + ), + } + } + fn make_lumina_entry( file_name: &str, index_type: &str, diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs b/crates/paimon/src/vindex/pkvector/bucket.rs index 01da95f58..ccd57dbf6 100644 --- a/crates/paimon/src/vindex/pkvector/bucket.rs +++ b/crates/paimon/src/vindex/pkvector/bucket.rs @@ -34,7 +34,6 @@ use crate::spec::PkVectorSourceMeta; /// scorer that reads it. pub(crate) struct BucketAnnSegment { pub source_meta: PkVectorSourceMeta, - pub file_name: String, /// Resolved index-file path (globally unique; the scorer's preload key). pub path: String, pub file_size: u64, @@ -48,7 +47,6 @@ impl BucketAnnSegment { pub(crate) fn for_test(source_meta: PkVectorSourceMeta) -> Self { Self { source_meta, - file_name: "seg".to_string(), path: "seg".to_string(), file_size: 0, index_meta: Vec::new(), diff --git a/crates/paimon/src/vindex/pkvector/metric.rs b/crates/paimon/src/vindex/pkvector/metric.rs index 799d9b836..7d654e585 100644 --- a/crates/paimon/src/vindex/pkvector/metric.rs +++ b/crates/paimon/src/vindex/pkvector/metric.rs @@ -41,14 +41,6 @@ pub(crate) fn normalize_metric(metric: &str) -> String { metric.to_ascii_lowercase().replace('-', "_") } -/// True if the (normalized) metric is one of the three supported metrics. -pub(crate) fn is_supported_metric(metric: &str) -> bool { - matches!( - normalize_metric(metric).as_str(), - "l2" | "cosine" | "inner_product" - ) -} - /// Numeric semantics for a supported vector search metric. Mirrors Java /// `org.apache.paimon.globalindex.VectorSearchMetric`. #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -59,6 +51,17 @@ pub(crate) enum VectorSearchMetric { } impl VectorSearchMetric { + /// Map a vindex-core metric to this enum. Mirrors the build-side + /// `RawVectorMetric::from_vindex`; lets the read path compare the metric a + /// segment was trained with against the configured metric. + pub(crate) fn from_vindex(metric: paimon_vindex_core::distance::MetricType) -> Self { + match metric { + paimon_vindex_core::distance::MetricType::L2 => Self::L2, + paimon_vindex_core::distance::MetricType::Cosine => Self::Cosine, + paimon_vindex_core::distance::MetricType::InnerProduct => Self::InnerProduct, + } + } + /// Normalize, validate, and map to the enum. Errors on an unsupported metric. pub(crate) fn parse(metric: &str) -> crate::Result { match normalize_metric(metric).as_str() { @@ -71,12 +74,13 @@ impl VectorSearchMetric { } } - /// Higher-is-better score for exact vector search. - pub(crate) fn compute_score(&self, query: &[f32], stored: &[f32]) -> f32 { + /// Canonical lowercase name, matching the vindex-core `MetricType::as_str` + /// spelling so mismatch diagnostics read consistently on both sides. + pub(crate) fn as_str(&self) -> &'static str { match self { - Self::L2 => 1.0 / (1.0 + squared_l2(query, stored)), - Self::Cosine => cosine_similarity(query, stored), - Self::InnerProduct => inner_product(query, stored), + Self::L2 => "l2", + Self::Cosine => "cosine", + Self::InnerProduct => "inner_product", } } @@ -211,15 +215,6 @@ mod tests { assert_eq!(normalize_metric(" l2 "), " l2 "); } - #[test] - fn test_is_supported_only_for_three_metrics() { - assert!(is_supported_metric("L2")); - assert!(is_supported_metric("cosine")); - assert!(is_supported_metric("inner-product")); - assert!(!is_supported_metric("manhattan")); - assert!(!is_supported_metric(" l2 ")); - } - #[test] fn test_parse_rejects_unsupported_metric() { assert!(VectorSearchMetric::parse("l2").is_ok()); @@ -242,20 +237,10 @@ mod tests { ); } - #[test] - fn test_compute_score_higher_is_better() { - let q = [2.0f32, 0.0]; - let s = [1.0f32, 0.0]; - assert_eq!(VectorSearchMetric::L2.compute_score(&q, &s), 0.5); // 1/(1+1) - assert_eq!(VectorSearchMetric::Cosine.compute_score(&q, &s), 1.0); // parallel - assert_eq!(VectorSearchMetric::InnerProduct.compute_score(&q, &s), 2.0); - } - #[test] fn test_cosine_zero_norm_similarity_is_zero() { let zero = [0.0f32, 0.0]; let s = [1.0f32, 0.0]; - assert_eq!(VectorSearchMetric::Cosine.compute_score(&zero, &s), 0.0); assert_eq!(VectorSearchMetric::Cosine.compute_distance(&zero, &s), 1.0); } @@ -264,15 +249,19 @@ mod tests { // query [0,3] -> norm 9, stored [1,2] -> norm 5; sqrt(5) is irrational // so the f32-sqrt path (each sqrt taken in f32, then widened) and the // f64-sqrt path (widen first, sqrt in f64) produce different f32 bits: - // buggy score 0.8944271 vs correct 0.8944272. Encode the f64 contract in - // the expected value (dot / (sqrt(9.0f64) * sqrt(5.0f64)) as f32) rather - // than a magic literal, so this pins the Java-matching f64 arithmetic. + // buggy similarity 0.8944271 vs correct 0.8944272. Encode the f64 contract + // in the expected value (dot / (sqrt(9.0f64) * sqrt(5.0f64)) as f32) rather + // than a magic literal, so this pins the Java-matching f64 arithmetic. The + // cosine distance is `1 - similarity`, so the sqrt path is exercised here. let q = [0.0f32, 3.0]; let s = [1.0f32, 2.0]; let dot = 6.0f32; let denominator = ((9.0f64).sqrt() * (5.0f64).sqrt()) as f32; - let expected = dot / denominator; - assert_eq!(VectorSearchMetric::Cosine.compute_score(&q, &s), expected); + let expected_similarity = dot / denominator; + assert_eq!( + VectorSearchMetric::Cosine.compute_distance(&q, &s), + 1.0 - expected_similarity + ); } #[test] @@ -314,4 +303,21 @@ mod tests { assert_eq!(metric.distance_to_score(d), s, "metric {metric:?}"); } } + + #[test] + fn test_from_vindex_maps_every_variant() { + use paimon_vindex_core::distance::MetricType; + assert_eq!( + VectorSearchMetric::from_vindex(MetricType::L2), + VectorSearchMetric::L2 + ); + assert_eq!( + VectorSearchMetric::from_vindex(MetricType::Cosine), + VectorSearchMetric::Cosine + ); + assert_eq!( + VectorSearchMetric::from_vindex(MetricType::InnerProduct), + VectorSearchMetric::InnerProduct + ); + } } diff --git a/crates/paimon/src/vindex/pkvector/mod.rs b/crates/paimon/src/vindex/pkvector/mod.rs index 3a685ca38..a1a6446bf 100644 --- a/crates/paimon/src/vindex/pkvector/mod.rs +++ b/crates/paimon/src/vindex/pkvector/mod.rs @@ -20,11 +20,6 @@ //! Read-only bucket-local approximate-nearest-neighbour search over the //! primary-key vector index. -// The kernel is crate-private and has no production caller yet, so its items -// are unreachable outside their own tests. Suppress the resulting dead_code -// lint at the module boundary until the read path wires it in. -#![allow(dead_code)] - pub(crate) mod ann; pub(crate) mod bucket; pub(crate) mod exact; From 40c4af21fbdce2cdfbe2a7d36ae74dc819bf0af1 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Tue, 14 Jul 2026 18:59:50 +0800 Subject: [PATCH 05/10] feat(table): materialize primary-key vector search rows best-first Add VectorSearchBuilder::execute_read (+with_projection): plan and search as execute_scored does, then materialize the surviving rows in best-first order, dropping the internal position column and hiding _ROW_ID. Output is the projected columns plus _PKEY_VECTOR_SCORE. Materialization is done per split so rows can be mapped back to their search rank; the file/position-order convenience wrapper is removed. Covered by a self-contained end-to-end test that builds a real table (vindex IVF segment, compacted data file, source metadata) and asserts row ids, scores, best-first order, and hidden columns. --- crates/paimon/src/table/data_file_reader.rs | 6 +- .../src/table/pk_vector_indexed_split_read.rs | 6 - .../src/table/pk_vector_orchestrator.rs | 223 +++--- .../src/table/pk_vector_position_read.rs | 6 - .../paimon/src/table/vector_search_builder.rs | 603 +++++++++++++++- .../paimon/tests/pk_vector_baseline_test.rs | 676 ++++++++++++++++++ 6 files changed, 1372 insertions(+), 148 deletions(-) create mode 100644 crates/paimon/tests/pk_vector_baseline_test.rs diff --git a/crates/paimon/src/table/data_file_reader.rs b/crates/paimon/src/table/data_file_reader.rs index 791ca0b0f..771f69bc1 100644 --- a/crates/paimon/src/table/data_file_reader.rs +++ b/crates/paimon/src/table/data_file_reader.rs @@ -80,8 +80,7 @@ impl DataFileReader { /// The effective read-type (requested output fields) of this reader. /// Exposed for the sibling `pk_vector_position_read` module, which drives the - /// PK-vector materialization read path (no production caller yet). - #[allow(dead_code)] + /// PK-vector materialization read path. pub(super) fn read_type(&self) -> &[DataField] { &self.read_type } @@ -89,8 +88,7 @@ impl DataFileReader { /// 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. Consumed by - /// `pk_vector_position_read` (materialization read path; no production caller yet). - #[allow(dead_code)] + /// `pk_vector_position_read` (materialization read path). pub(super) fn has_row_filtering_predicate(&self) -> bool { self.predicates .iter() diff --git a/crates/paimon/src/table/pk_vector_indexed_split_read.rs b/crates/paimon/src/table/pk_vector_indexed_split_read.rs index 54475ddb6..668957a6a 100644 --- a/crates/paimon/src/table/pk_vector_indexed_split_read.rs +++ b/crates/paimon/src/table/pk_vector_indexed_split_read.rs @@ -24,12 +24,6 @@ //! It is a pure consumer: no bucket/ANN search, no cross-bucket merge, no //! serialization. -// This materialization read path is exercised only by its own tests: the wired -// vector search returns matched row-ids and scores directly, so no production -// caller drives row materialization yet. Suppress dead_code at the module -// boundary until a materializing caller exists. -#![allow(dead_code)] - use std::collections::BTreeMap; use futures::StreamExt; diff --git a/crates/paimon/src/table/pk_vector_orchestrator.rs b/crates/paimon/src/table/pk_vector_orchestrator.rs index 21adb41a7..140dd02ac 100644 --- a/crates/paimon/src/table/pk_vector_orchestrator.rs +++ b/crates/paimon/src/table/pk_vector_orchestrator.rs @@ -26,14 +26,11 @@ 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::pk_vector_indexed_split_read::PkVectorIndexedSplit; 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}; @@ -104,10 +101,7 @@ fn global_top_k(mut candidates: Vec, limit: usize) -> Vec, splits: &[PkVectorSearchSplit], metric: VectorSearchMetric, @@ -320,60 +314,6 @@ impl PkVectorOrchestrator { Ok(global_top_k(candidates, limit)) } - - /// Run the eager per-bucket search phase, then return a stream that lazily - /// materializes 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`). - /// - /// The wired vector search returns matched row-ids and scores via - /// `search_candidates`; no production caller drives row materialization yet. - #[allow(dead_code)] - #[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>, - search_options: &HashMap, - ) -> crate::Result { - // Wrap the per-file factory into the split-scoped shape search_candidates - // expects; the split index/split are unused on this back-compat path. - let mut wrapped = - |_: usize, _: &PkVectorSearchSplit, f: &BucketActiveFile| exact_reader_factory(f); - let survivors = self - .search_candidates( - splits, - query, - metric, - limit, - ann_searcher, - &mut wrapped, - search_options, - false, - ) - .await?; - 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)] @@ -643,6 +583,7 @@ mod e2e_tests { use crate::spec::{ DataField, DataFileMeta, DataType, IntType, PkVectorSourceFile, PkVectorSourceMeta, }; + use crate::table::pk_vector_indexed_split_read::PkVectorIndexedSplitRead; use crate::table::pk_vector_position_read::{ PKEY_VECTOR_POSITION_COLUMN, PKEY_VECTOR_SCORE_COLUMN, }; @@ -887,17 +828,54 @@ mod e2e_tests { } } + /// Run the eager per-bucket search + global Top-K, group survivors into indexed + /// splits, then materialize each split in file/position order. This is the + /// materialization path the production best-first read reorders on top of; the + /// tests below drive it directly through its `pub(crate)` components. + #[allow(clippy::too_many_arguments)] + async fn materialize_via_splits( + reader: DataFileReader, + splits: &[PkVectorSearchSplit], + query: &[f32], + metric: VectorSearchMetric, + limit: usize, + ann: Option<&dyn PkVectorAnnSearcher>, + factory: &mut dyn FnMut(&BucketActiveFile) -> crate::Result>, + opts: &HashMap, + ) -> crate::Result> { + let orch = PkVectorOrchestrator::new(reader.clone()); + // Wrap the per-file factory into the split-scoped shape search_candidates + // expects; the split index/split are unused here. + let mut wrapped = |_: usize, _: &PkVectorSearchSplit, f: &BucketActiveFile| factory(f); + let survivors = orch + .search_candidates(splits, query, metric, limit, ann, &mut wrapped, opts, false) + .await?; + let indexed_splits = build_indexed_splits(survivors, splits, metric)?; + let mut out = Vec::new(); + for indexed in indexed_splits { + let batches: Vec = PkVectorIndexedSplitRead::new(reader.clone()) + .read(&indexed)? + .try_collect() + .await?; + out.extend(batches); + } + Ok(out) + } + #[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 = Vec::new(); - let mut factory = |_: &BucketActiveFile| -> crate::Result> { + let mut factory = |_: usize, + _: &PkVectorSearchSplit, + _: &BucketActiveFile| + -> crate::Result> { unreachable!("no bucket search on eager-rejected input") }; let opts = HashMap::new(); let err = PkVectorOrchestrator::new(reader) - .read( + .search_candidates( &splits, &[0.0, 0.0], VectorSearchMetric::L2, @@ -905,6 +883,7 @@ mod e2e_tests { None, &mut factory, &opts, + false, ) .await .map(|_| ()) @@ -917,12 +896,15 @@ mod e2e_tests { let file_io = FileIOBuilder::new("memory").build().unwrap(); let reader = make_reader(file_io, "memory:/pkvo_empty_query"); let splits: Vec = Vec::new(); - let mut factory = |_: &BucketActiveFile| -> crate::Result> { + let mut factory = |_: usize, + _: &PkVectorSearchSplit, + _: &BucketActiveFile| + -> crate::Result> { unreachable!("no bucket search on eager-rejected input") }; let opts = HashMap::new(); let err = PkVectorOrchestrator::new(reader) - .read( + .search_candidates( &splits, &[], VectorSearchMetric::L2, @@ -930,6 +912,7 @@ mod e2e_tests { None, &mut factory, &opts, + false, ) .await .map(|_| ()) @@ -979,21 +962,18 @@ mod e2e_tests { 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::>() - .await - .unwrap(); + let batches = materialize_via_splits( + make_reader(file_io, table_path), + &[split], + &[0.0, 0.0], + VectorSearchMetric::L2, + 3, + Some(&ann), + &mut factory, + &opts, + ) + .await + .unwrap(); // Output is ascending group (file name) then ascending position: // ann.mosaic pos1 -> id 101; exact.mosaic pos0,1 -> ids 200,201. @@ -1072,21 +1052,18 @@ mod e2e_tests { 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::>() - .await - .unwrap(); + let batches = materialize_via_splits( + make_reader(file_io, table_path), + &[split0, split1], + &[0.0, 0.0], + VectorSearchMetric::L2, + 3, + None, + &mut factory, + &opts, + ) + .await + .unwrap(); // Ascending group order: bucket0 "b0.mosaic" pos0 -> 10; bucket1 "b1.mosaic" // pos0,1 -> 20,21. @@ -1140,21 +1117,18 @@ mod e2e_tests { 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::>() - .await - .unwrap(); + let batches = materialize_via_splits( + make_reader(file_io, table_path), + &[split], + &[0.0, 0.0], + VectorSearchMetric::L2, + 4, + None, + &mut factory, + &opts, + ) + .await + .unwrap(); // Position 1 (id 31) is absent. Remaining ascending positions 0,2,3. assert_eq!(collect_i32(&batches, "id"), vec![30, 32, 33]); @@ -1203,21 +1177,18 @@ mod e2e_tests { 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::>() - .await - .unwrap(); + let batches = materialize_via_splits( + make_reader(file_io, table_path), + &[split], + &[0.0, 0.0], + VectorSearchMetric::L2, + 3, + None, + &mut factory, + &opts, + ) + .await + .unwrap(); // Ascending physical position order, not best-first distance order. assert_eq!(collect_i32(&batches, "id"), vec![40, 41, 42]); diff --git a/crates/paimon/src/table/pk_vector_position_read.rs b/crates/paimon/src/table/pk_vector_position_read.rs index e8fedb896..9662c8b44 100644 --- a/crates/paimon/src/table/pk_vector_position_read.rs +++ b/crates/paimon/src/table/pk_vector_position_read.rs @@ -24,12 +24,6 @@ //! `pk_vector_indexed_split_read` and `pk_vector_orchestrator` modules build the //! indexed-split contract and cross-bucket merge on top of it. -// This materialization read path is exercised only by its own tests: the wired -// vector search returns matched row-ids and scores directly, so no production -// caller drives row materialization yet. Suppress dead_code at the module -// boundary until a materializing caller exists. -#![allow(dead_code)] - use std::collections::BTreeMap; use std::sync::Arc; diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index 4c15d1ae0..e4fa06b5f 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -28,11 +28,18 @@ use crate::table::global_index_scanner::{ unindexed_ranges_for_global_index_entries, RowRangeIndex, }; use crate::table::pk_vector_data_file_reader::DataFilePkVectorReaderFactory; +use crate::table::pk_vector_indexed_split_read::PkVectorIndexedSplitRead; use crate::table::pk_vector_orchestrator::{ - PkVectorCandidate, PkVectorOrchestrator, PkVectorSearchSplit, + build_indexed_splits, PkVectorCandidate, PkVectorOrchestrator, PkVectorSearchSplit, +}; +use crate::table::pk_vector_position_read::{ + PKEY_VECTOR_POSITION_COLUMN, PKEY_VECTOR_SCORE_COLUMN, +}; +use crate::table::pk_vector_scan::{PkVectorScan, PkVectorScanPlan}; +use crate::table::read_builder::resolve_projected_fields; +use crate::table::{ + find_field_id_by_name, merge_row_ranges, ArrowRecordBatchStream, RowRange, Table, }; -use crate::table::pk_vector_scan::PkVectorScan; -use crate::table::{find_field_id_by_name, merge_row_ranges, RowRange, Table}; use crate::vector_search::{GlobalIndexIOMeta, SearchResult, VectorSearch}; use crate::vindex::is_vindex_index_type; use crate::vindex::pkvector::ann::VindexAnnSearcher; @@ -41,7 +48,8 @@ use crate::vindex::pkvector::metric::VectorSearchMetric; use crate::vindex::pkvector::reader::PkVectorReader; use crate::vindex::reader::VindexVectorGlobalIndexReader; use arrow_array::{Array, FixedSizeListArray, Float32Array, Int64Array, ListArray, RecordBatch}; -use futures::TryStreamExt; +use arrow_select::interleave::interleave_record_batch; +use futures::{stream, TryStreamExt}; use paimon_vindex_core::distance::MetricType; use paimon_vindex_core::index::VectorIndexReader as VIndexReader; use roaring::RoaringTreemap; @@ -82,6 +90,7 @@ pub struct VectorSearchBuilder<'a> { query_vector: Option>, limit: Option, options: HashMap, + projection: Option>, } pub struct BatchVectorSearchBuilder<'a> { @@ -100,6 +109,7 @@ impl<'a> VectorSearchBuilder<'a> { query_vector: None, limit: None, options: HashMap::new(), + projection: None, } } @@ -123,6 +133,15 @@ impl<'a> VectorSearchBuilder<'a> { self } + /// Restrict the columns materialized by [`execute_read`](Self::execute_read) + /// to `cols` (plus the always-appended `_PKEY_VECTOR_SCORE`). Without this + /// call `execute_read` materializes every user table column. Only affects + /// `execute_read`; the search-only paths ignore it. + pub fn with_projection(&mut self, cols: &[&str]) -> &mut Self { + self.projection = Some(cols.iter().map(|c| c.to_string()).collect()); + self + } + pub async fn execute(&self) -> crate::Result> { self.execute_scored().await?.to_row_ranges() } @@ -184,6 +203,55 @@ impl<'a> VectorSearchBuilder<'a> { Ok(results.remove(0)) } + /// Run the vector search and materialize the matching rows as Arrow batches, + /// ordered best-first. Only supported for primary-key vector indexes; a + /// data-evolution table or a query targeting a non-PK-vector column fails + /// loud. Output columns are the projected user table columns (all user + /// columns by default, or those set via + /// [`with_projection`](Self::with_projection)) plus `_PKEY_VECTOR_SCORE`; + /// `_ROW_ID` and `_PKEY_VECTOR_POSITION` are always hidden. + pub async fn execute_read(&self) -> crate::Result { + // Fail closed: returns data outside `TableScan`/`TableRead`. + let core = CoreOptions::new(self.table.schema().options()); + core.ensure_read_authorized()?; + let vector_column = + self.vector_column + .as_deref() + .ok_or_else(|| crate::Error::ConfigInvalid { + message: "Vector column must be set via with_vector_column()".to_string(), + })?; + let query_vector = + self.query_vector + .as_ref() + .ok_or_else(|| crate::Error::ConfigInvalid { + message: "Query vector must be set via with_query_vector()".to_string(), + })?; + let limit = self.limit.ok_or_else(|| crate::Error::ConfigInvalid { + message: "Limit must be set via with_limit()".to_string(), + })?; + + // Only the primary-key vector path can materialize rows. The data-evolution + // (global-index) path returns data-derived row-ids, not table rows, so a + // read against it (or against a non-PK-vector column) fails loud. + if core.primary_key_vector_index_enabled() { + let targets_pk_column = core + .primary_key_vector_index_columns() + .ok() + .is_some_and(|cols| cols.iter().any(|c| c == vector_column)); + if targets_pk_column { + let pk_col = core.primary_key_vector_index_column()?; + return self + .execute_primary_key_vector_read(&core, &pk_col, query_vector, limit) + .await; + } + } + + Err(crate::Error::DataInvalid { + message: "vector search read is only supported for primary-key vector indexes".into(), + source: None, + }) + } + /// Run the primary-key bucket-local vector search: plan the per-bucket splits, /// build the real vindex ANN scorer and (outside FAST mode) the exact-fallback /// readers, run the orchestrator, and convert the best-first candidates into a @@ -195,6 +263,26 @@ impl<'a> VectorSearchBuilder<'a> { query_vector: &[f32], limit: usize, ) -> crate::Result { + let (candidates, plan, metric) = self + .plan_and_search_pk_candidates(core, pk_col, query_vector, limit) + .await?; + candidates_to_search_result(&candidates, &plan.splits, metric) + } + + /// Shared PK-vector search core for both the search-only and search-and-read + /// paths: plan the per-bucket splits, verify the configured metric against each + /// ANN segment, build the real vindex ANN scorer and (outside FAST mode) the + /// exact-fallback readers, and run the orchestrator. Returns the best-first + /// candidates together with the plan and resolved metric so the caller can + /// either serialize them to a `SearchResult` or materialize their rows. An + /// empty plan yields empty candidates. + async fn plan_and_search_pk_candidates( + &self, + core: &CoreOptions<'_>, + pk_col: &str, + query_vector: &[f32], + limit: usize, + ) -> crate::Result<(Vec, PkVectorScanPlan, VectorSearchMetric)> { // Residual-filter guard: PK vector search accepts partition filters only. // This builder exposes no data-predicate setter, so there is nothing to // reject here; the guard mirrors Java `checkArgument(filter == null)` and, @@ -230,7 +318,7 @@ impl<'a> VectorSearchBuilder<'a> { .plan() .await?; if plan.splits.is_empty() { - return Ok(SearchResult::empty()); + return Ok((Vec::new(), plan, metric)); } // Production data-file reader, mirroring `table_read.rs::new_data_file_reader` @@ -318,7 +406,132 @@ impl<'a> VectorSearchBuilder<'a> { ) .await?; - candidates_to_search_result(&candidates, &plan.splits, metric) + Ok((candidates, plan, metric)) + } + + /// Materialize the best-first PK-vector search hits into Arrow rows. Mirrors + /// Java `PrimaryKeyVectorRead` feeding its result splits into an ordinary table + /// read: the search decides which rows, a subsequent read decides which + /// columns. + /// + /// Output columns are the projected user table columns (all user columns when + /// [`with_projection`](Self::with_projection) was not called) plus + /// `_PKEY_VECTOR_SCORE`; `_ROW_ID` and `_PKEY_VECTOR_POSITION` are always + /// hidden. Rows are emitted best-first (the candidate order), which differs + /// from the file/position order the orchestrator materializes in. + async fn execute_primary_key_vector_read( + &self, + core: &CoreOptions<'_>, + pk_col: &str, + query_vector: &[f32], + limit: usize, + ) -> crate::Result { + let (candidates, plan, metric) = self + .plan_and_search_pk_candidates(core, pk_col, query_vector, limit) + .await?; + + // Resolve the materialization read-type up front so an invalid projection + // (unknown column, or a reserved metadata / row-id name) fails loud + // unconditionally, even when the plan is empty and no rows will be read. + // Default (no `with_projection`) is every user table column. + let read_type = self.resolve_materialize_read_type()?; + + if candidates.is_empty() { + return Ok(Box::pin(stream::empty())); + } + + // A separate, predicate-free materialization reader projecting the user + // columns (the search reader projects only the vector column). Mirrors + // `table_read.rs::new_data_file_reader` with an empty predicate list. + let materialize_reader = DataFileReader::new( + self.table.file_io().clone(), + self.table.schema_manager().clone(), + self.table.schema().id(), + self.table.schema().fields().to_vec(), + read_type, + Vec::new(), + ); + + // Rank each candidate by its best-first position, then reduce the physical + // materialization order back to best-first. The orchestrator emits rows in + // ascending (partition, bucket, file, position); the rank map keyed by + // (partition bytes, bucket, file, position) recovers the candidate order. + let mut rank_of: HashMap<(Vec, i32, String, i64), usize> = HashMap::new(); + for (rank, c) in candidates.iter().enumerate() { + rank_of.insert( + ( + c.partition.to_serialized_bytes(), + c.bucket, + c.data_file_name.clone(), + c.row_position, + ), + rank, + ); + } + + let indexed_splits = build_indexed_splits(candidates, &plan.splits, metric)?; + + // Materialize every indexed split, retaining each batch and, per row, the + // (rank, batch_index, row_index) tuple so we can reorder to best-first. + // Top-K is small, so full in-memory collection is acceptable. + let mut batches: Vec = Vec::new(); + let mut ranked: Vec = Vec::new(); + for indexed in indexed_splits { + let partition_bytes = indexed.split.partition().to_serialized_bytes(); + let bucket = indexed.split.bucket(); + let file_name = indexed.split.data_files()[0].file_name.clone(); + let mut stream = + PkVectorIndexedSplitRead::new(materialize_reader.clone()).read(&indexed)?; + while let Some(batch) = stream.try_next().await? { + let batch_index = batches.len(); + collect_ranked_rows( + &batch, + batch_index, + &partition_bytes, + bucket, + &file_name, + &rank_of, + &mut ranked, + )?; + batches.push(batch); + } + } + + // Reorder to best-first and drop the position column. + let output = reorder_and_strip_position(&batches, ranked)?; + Ok(Box::pin(stream::iter(output.into_iter().map(Ok)))) + } + + /// Resolve the projected fields for the materialization read-type. Default + /// (no projection set) is all user table fields; otherwise the requested + /// names resolved via `resolve_projected_fields`. Rejects reserved metadata + /// names and `_ROW_ID` so a user cannot request a hidden column. + fn resolve_materialize_read_type(&self) -> crate::Result> { + let fields = match &self.projection { + None => self.table.schema().fields().to_vec(), + Some(names) => { + for name in names { + if name == PKEY_VECTOR_POSITION_COLUMN + || name == PKEY_VECTOR_SCORE_COLUMN + || name == ROW_ID_FIELD_NAME + { + return Err(crate::Error::DataInvalid { + message: format!( + "vector search read projection must not request reserved column '{name}'" + ), + source: None, + }); + } + } + resolve_projected_fields( + self.table.identifier().full_name(), + self.table.schema().fields(), + names, + true, + )? + } + }; + Ok(fields) } } @@ -820,6 +1033,111 @@ fn candidates_to_search_result( Ok(SearchResult::new(row_ids, scores)) } +/// One materialized row tagged with its best-first `rank` and its `(batch_index, +/// row_index)` location in the retained materialization batches. +struct RankedRow { + rank: usize, + batch_index: usize, + row_index: usize, +} + +/// For each row in a materialized batch, look up its best-first rank via the +/// `(partition bytes, bucket, file, position)` key and record its location. The +/// `_PKEY_VECTOR_POSITION` column supplies the physical position; every row must +/// map to a candidate rank (the batch came from that candidate's file), so a miss +/// fails loud rather than silently dropping a row. +#[allow(clippy::too_many_arguments)] +fn collect_ranked_rows( + batch: &RecordBatch, + batch_index: usize, + partition_bytes: &[u8], + bucket: i32, + file_name: &str, + rank_of: &HashMap<(Vec, i32, String, i64), usize>, + out: &mut Vec, +) -> crate::Result<()> { + let position_idx = batch + .schema() + .index_of(PKEY_VECTOR_POSITION_COLUMN) + .map_err(|_| crate::Error::DataInvalid { + message: format!("materialized batch missing {PKEY_VECTOR_POSITION_COLUMN} column"), + source: None, + })?; + let positions = batch + .column(position_idx) + .as_any() + .downcast_ref::() + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("{PKEY_VECTOR_POSITION_COLUMN} column is not Int64"), + source: None, + })?; + for row_index in 0..batch.num_rows() { + let position = positions.value(row_index); + let key = ( + partition_bytes.to_vec(), + bucket, + file_name.to_string(), + position, + ); + let rank = *rank_of.get(&key).ok_or_else(|| crate::Error::DataInvalid { + message: format!( + "materialized row (file {file_name}, position {position}) has no matching search candidate" + ), + source: None, + })?; + out.push(RankedRow { + rank, + batch_index, + row_index, + }); + } + Ok(()) +} + +/// Reorder the materialized rows into best-first order and drop the internal +/// `_PKEY_VECTOR_POSITION` column, yielding a single output batch (empty input +/// yields no batches). The projected user columns and `_PKEY_VECTOR_SCORE` are +/// retained. +fn reorder_and_strip_position( + batches: &[RecordBatch], + mut ranked: Vec, +) -> crate::Result> { + if ranked.is_empty() { + return Ok(Vec::new()); + } + ranked.sort_by_key(|r| r.rank); + let indices: Vec<(usize, usize)> = ranked + .iter() + .map(|r| (r.batch_index, r.row_index)) + .collect(); + let refs: Vec<&RecordBatch> = batches.iter().collect(); + let reordered = + interleave_record_batch(&refs, &indices).map_err(|e| crate::Error::DataInvalid { + message: format!("failed to reorder vector search read rows: {e}"), + source: None, + })?; + + // Drop the internal position column; keep every other column (projected user + // columns + _PKEY_VECTOR_SCORE) in order. + let position_idx = reordered + .schema() + .index_of(PKEY_VECTOR_POSITION_COLUMN) + .map_err(|_| crate::Error::DataInvalid { + message: format!("reordered batch missing {PKEY_VECTOR_POSITION_COLUMN} column"), + source: None, + })?; + let keep: Vec = (0..reordered.num_columns()) + .filter(|i| *i != position_idx) + .collect(); + let projected = reordered + .project(&keep) + .map_err(|e| crate::Error::DataInvalid { + message: format!("failed to drop position column: {e}"), + source: None, + })?; + Ok(vec![projected]) +} + fn indexed_search_limit(limit: usize, refine_factor: usize) -> crate::Result { if refine_factor == 0 { return Ok(limit); @@ -1634,9 +1952,14 @@ mod tests { use crate::vindex::IVF_FLAT_IDENTIFIER; use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; use arrow_array::ArrayRef; + use arrow_array::Int32Array; use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; use std::sync::Arc; + fn l2_score(distance: f32) -> f32 { + VectorSearchMetric::L2.distance_to_score(distance) + } + fn make_field(id: i32, name: &str) -> DataField { DataField::new(id, name.to_string(), DataType::Int(IntType::default())) } @@ -2515,4 +2838,272 @@ mod tests { version: 1, } } + + // ---- Task B: search-and-read (`execute_read`) tests ---- + + /// Build a small materialization batch: user column `id: Int32`, the internal + /// `_PKEY_VECTOR_POSITION: Int64`, and `_PKEY_VECTOR_SCORE: Float32` (mirroring + /// what `PkVectorIndexedSplitRead` emits for a single file). + fn materialized_batch(rows: &[(i32, i64, f32)]) -> RecordBatch { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new(PKEY_VECTOR_POSITION_COLUMN, ArrowDataType::Int64, false), + ArrowField::new(PKEY_VECTOR_SCORE_COLUMN, ArrowDataType::Float32, false), + ])); + let ids = Int32Array::from(rows.iter().map(|(id, _, _)| *id).collect::>()); + let positions = Int64Array::from(rows.iter().map(|(_, pos, _)| *pos).collect::>()); + let scores = Float32Array::from(rows.iter().map(|(_, _, s)| *s).collect::>()); + RecordBatch::try_new( + schema, + vec![Arc::new(ids), Arc::new(positions), Arc::new(scores)], + ) + .unwrap() + } + + fn i32_col(batch: &RecordBatch, name: &str) -> Vec { + let idx = batch.schema().index_of(name).unwrap(); + batch + .column(idx) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + } + fn f32_col(batch: &RecordBatch, name: &str) -> Vec { + let idx = batch.schema().index_of(name).unwrap(); + batch + .column(idx) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + } + + #[test] + fn reorder_and_strip_position_recovers_best_first_and_drops_position() { + // Single file, one bucket. The materialization reader emits rows in + // ascending physical position [pos0, pos1, pos2] -> ids [40,41,42]. The + // search candidates ranked them best-first as pos1(rank0), pos2(rank1), + // pos0(rank2), which is NEITHER position order nor score order-by-batch. + // The reorder must yield ids [41,42,40] and drop _PKEY_VECTOR_POSITION. + let batch = materialized_batch(&[ + (40, 0, l2_score(9.0)), + (41, 1, l2_score(1.0)), + (42, 2, l2_score(4.0)), + ]); + let batches = vec![batch]; + let part = BinaryRow::new(0).to_serialized_bytes(); + let mut rank_of: HashMap<(Vec, i32, String, i64), usize> = HashMap::new(); + rank_of.insert((part.clone(), 0, "o.mosaic".to_string(), 1), 0); + rank_of.insert((part.clone(), 0, "o.mosaic".to_string(), 2), 1); + rank_of.insert((part.clone(), 0, "o.mosaic".to_string(), 0), 2); + + let mut ranked = Vec::new(); + collect_ranked_rows(&batches[0], 0, &part, 0, "o.mosaic", &rank_of, &mut ranked).unwrap(); + let out = reorder_and_strip_position(&batches, ranked).unwrap(); + assert_eq!(out.len(), 1); + let out = &out[0]; + + // Best-first row order, not ascending position order. + assert_eq!(i32_col(out, "id"), vec![41, 42, 40]); + // Score column preserved and aligned to the reordered rows. + assert_eq!( + f32_col(out, PKEY_VECTOR_SCORE_COLUMN), + vec![l2_score(1.0), l2_score(4.0), l2_score(9.0)] + ); + // Position column dropped; _ROW_ID never present. + assert!(out.schema().index_of(PKEY_VECTOR_POSITION_COLUMN).is_err()); + assert!(out.schema().index_of("_ROW_ID").is_err()); + } + + #[test] + fn reorder_and_strip_position_merges_rows_across_files() { + // Two files (two materialization batches). Best-first interleaves them: + // file-b pos0 (rank0), file-a pos1 (rank1), file-a pos0 (rank2). The + // reorder must pull rows from both batches into one best-first output. + let batch_a = materialized_batch(&[(10, 0, l2_score(9.0)), (11, 1, l2_score(1.0))]); + let batch_b = materialized_batch(&[(20, 0, l2_score(0.5))]); + let batches = vec![batch_a, batch_b]; + let part = BinaryRow::new(0).to_serialized_bytes(); + let mut rank_of: HashMap<(Vec, i32, String, i64), usize> = HashMap::new(); + rank_of.insert((part.clone(), 0, "b".to_string(), 0), 0); + rank_of.insert((part.clone(), 0, "a".to_string(), 1), 1); + rank_of.insert((part.clone(), 0, "a".to_string(), 0), 2); + + let mut ranked = Vec::new(); + collect_ranked_rows(&batches[0], 0, &part, 0, "a", &rank_of, &mut ranked).unwrap(); + collect_ranked_rows(&batches[1], 1, &part, 0, "b", &rank_of, &mut ranked).unwrap(); + let out = reorder_and_strip_position(&batches, ranked).unwrap(); + assert_eq!(i32_col(&out[0], "id"), vec![20, 11, 10]); + assert_eq!( + f32_col(&out[0], PKEY_VECTOR_SCORE_COLUMN), + vec![l2_score(0.5), l2_score(1.0), l2_score(9.0)] + ); + } + + #[test] + fn reorder_and_strip_position_empty_yields_no_batches() { + let out = reorder_and_strip_position(&[], Vec::new()).unwrap(); + assert!(out.is_empty()); + } + + #[test] + fn collect_ranked_rows_missing_candidate_fails_loud() { + // A materialized position with no candidate rank must fail loud rather than + // silently drop the row. + let batch = materialized_batch(&[(40, 7, l2_score(1.0))]); + let part = BinaryRow::new(0).to_serialized_bytes(); + let rank_of: HashMap<(Vec, i32, String, i64), usize> = HashMap::new(); + let mut ranked = Vec::new(); + let err = collect_ranked_rows(&batch, 0, &part, 0, "f", &rank_of, &mut ranked) + .expect_err("missing candidate must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } if message.contains("no matching search candidate")), + "unexpected error: {err:?}" + ); + } + + #[tokio::test] + async fn execute_read_de_table_fails_loud() { + // No pk-vector index configured: execute_read must fail loud (the DE path + // has no row materialization). + let table = pk_vector_table(&[]); + let err = table + .new_vector_search_builder() + .with_vector_column("embedding") + .with_query_vector(vec![1.0]) + .with_limit(5) + .execute_read() + .await + .map(|_| ()) + .expect_err("DE read must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } + if message.contains("only supported for primary-key")), + "unexpected error: {err:?}" + ); + } + + #[tokio::test] + async fn execute_read_non_pk_column_fails_loud() { + // pk-vector index configured for "embedding", but the query targets a + // different column -> read is unsupported. + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.embedding.pk-vector.distance.metric", "l2"), + ]); + let err = table + .new_vector_search_builder() + .with_vector_column("other") + .with_query_vector(vec![1.0]) + .with_limit(5) + .execute_read() + .await + .map(|_| ()) + .expect_err("non-PK column read must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } + if message.contains("only supported for primary-key")), + "unexpected error: {err:?}" + ); + } + + #[tokio::test] + async fn execute_read_empty_plan_reserved_projection_fails_loud() { + // Empty plan (no snapshot) must still fail loud on a reserved-name + // projection: projection validity does not depend on whether the search + // matched any rows. A regression that resolved the projection only after + // the `candidates.is_empty()` early return would yield an empty stream here + // instead of an error. + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.embedding.pk-vector.distance.metric", "l2"), + ]); + for reserved in [ + ROW_ID_FIELD_NAME, + PKEY_VECTOR_POSITION_COLUMN, + PKEY_VECTOR_SCORE_COLUMN, + ] { + let mut builder = table.new_vector_search_builder(); + builder + .with_vector_column("embedding") + .with_query_vector(vec![1.0]) + .with_limit(5) + .with_projection(&["id", reserved]); + let err = builder + .execute_read() + .await + .map(|_| ()) + .expect_err("empty plan + reserved projection must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } + if message.contains("reserved column")), + "unexpected error for {reserved}: {err:?}" + ); + } + } + + #[tokio::test] + async fn execute_read_projection_reserved_name_fails_loud() { + // Projecting a reserved metadata / row-id column must fail loud. The guard + // lives in `resolve_materialize_read_type`, which `execute_read` invokes + // before the empty-plan early return; assert on the resolver directly here. + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.embedding.pk-vector.distance.metric", "l2"), + ]); + for reserved in [ + ROW_ID_FIELD_NAME, + PKEY_VECTOR_POSITION_COLUMN, + PKEY_VECTOR_SCORE_COLUMN, + ] { + let mut builder = table.new_vector_search_builder(); + builder + .with_vector_column("embedding") + .with_query_vector(vec![1.0]) + .with_limit(5) + .with_projection(&["id", reserved]); + let err = builder + .resolve_materialize_read_type() + .expect_err("reserved projection must fail loud"); + assert!( + matches!(err, crate::Error::DataInvalid { ref message, .. } + if message.contains("reserved column")), + "unexpected error for {reserved}: {err:?}" + ); + } + } + + #[test] + fn resolve_materialize_read_type_default_is_all_user_columns() { + // No with_projection -> every user table column (id + embedding). + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.embedding.pk-vector.distance.metric", "l2"), + ]); + let builder = table.new_vector_search_builder(); + let fields = builder.resolve_materialize_read_type().unwrap(); + let names: Vec<&str> = fields.iter().map(|f| f.name()).collect(); + assert_eq!(names, vec!["id", "embedding"]); + } + + #[test] + fn resolve_materialize_read_type_projection_selects_named_columns() { + let table = pk_vector_table(&[ + ("pk-vector.index.columns", "embedding"), + ("fields.embedding.pk-vector.index.type", IVF_FLAT_IDENTIFIER), + ("fields.embedding.pk-vector.distance.metric", "l2"), + ]); + let mut builder = table.new_vector_search_builder(); + builder.with_projection(&["id"]); + let fields = builder.resolve_materialize_read_type().unwrap(); + let names: Vec<&str> = fields.iter().map(|f| f.name()).collect(); + assert_eq!(names, vec!["id"]); + } } diff --git a/crates/paimon/tests/pk_vector_baseline_test.rs b/crates/paimon/tests/pk_vector_baseline_test.rs new file mode 100644 index 000000000..b287257a5 --- /dev/null +++ b/crates/paimon/tests/pk_vector_baseline_test.rs @@ -0,0 +1,676 @@ +// 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. + +//! End-to-end acceptance gate for primary-key vector search. +//! +//! Builds a complete, self-contained primary-key table in a temporary directory +//! entirely from Rust — data file, a real vindex IVF-flat ANN index segment, and +//! the snapshot/manifest/index-manifest metadata — then reads it back through the +//! public `new_vector_search_builder()` API and asserts both the search result +//! (`execute_scored()` -> `row_ids`/`scores`) and the materialized rows +//! (`execute_read()` -> Arrow batches, best-first order, `_PKEY_VECTOR_SCORE`). +//! +//! Why Rust-built rather than a committed cross-language fixture: the Java +//! primary-key vector ANN segment is an opaque native Lumina format that cannot +//! be reproduced byte-for-byte here, whereas the Rust read path is backed by the +//! vindex (IVF) segment format. This test therefore validates the Rust read path +//! against real vindex IVF segment bytes it produces itself, with no committed +//! binaries and nothing skipped. +//! +//! Two constraints the primary-key read path enforces are satisfied by hand +//! (mirroring Java `PrimaryKeyIndexSourcePolicy` and `PkVectorSourceMeta`): +//! 1. Only a compacted (`file_source == COMPACT`), non-level-0 data file backs +//! the index, so the written file's meta is cloned with `level = 1` and +//! `file_source = Some(1)`. +//! 2. `GlobalIndexMeta.source_meta` must be the Java `PkVectorSourceMeta` frame +//! (big-endian ints/longs, `writeUTF` file names), assembled below. +//! +//! Determinism: every fixture uses `nlist = 1`, so the single IVF inverted list +//! is scanned exhaustively and the search is exact; datasets are chosen so the +//! top-k distances have strict gaps, making the best-first order unique and +//! immune to tie-breaks or IVF approximation. + +use std::collections::HashMap; +use std::io::Cursor; + +use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; +use arrow_array::{Array, ArrayRef, FixedSizeListArray, Float32Array, Int32Array, RecordBatch}; +use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; +use bytes::Bytes; +use futures::TryStreamExt; +use paimon::catalog::Identifier; +use paimon::io::{FileIO, FileIOBuilder}; +use paimon::spec::{ + DataFileMeta, DataType, FloatType, GlobalIndexMeta, IndexFileMeta, IntType, Schema, + TableSchema, VectorType, +}; +use paimon::table::{CommitMessage, SchemaManager, Table, TableCommit}; +use paimon_vindex_core::index::{VectorIndexConfig, VectorIndexTrainer, VectorIndexWriter}; +use paimon_vindex_core::io::PosWriter; +use std::sync::Arc; + +/// Vector dimension for the test datasets. +const DIM: usize = 4; +/// The primary-key vector column name. +const VECTOR_COLUMN: &str = "embedding"; +/// vindex index type (IVF-flat); matches `IndexFileMeta.index_type`. +const INDEX_TYPE: &str = "ivf-flat"; + +/// vindex L2 distance is the *squared* L2 (see paimon-vindex-core `fvec_l2sqr`), +/// and the primary-key vector metric is `l2`, whose `distance_to_score` is +/// `1 / (1 + distance)`. Kept in one place so both tests agree with the kernel. +fn l2_score(distance: f32) -> f32 { + 1.0 / (1.0 + distance) +} + +/// Brute-force exact squared-L2 top-k over the fixture rows, returning +/// `(physical_position, squared_l2_distance)` best-first. Physical position == +/// row index == global row id (the fixture pins `first_row_id = 0`). This is the +/// ground truth the fixture is validated against, derived from the data rather +/// than hand-tabulated, so the assertions cannot drift from the vectors. +fn analytic_topk(query: &[f32], vectors: &[[f32; DIM]], k: usize) -> Vec<(u64, f32)> { + let mut scored: Vec<(u64, f32)> = vectors + .iter() + .enumerate() + .map(|(pos, v)| { + let dist: f32 = v + .iter() + .zip(query.iter()) + .map(|(a, b)| (a - b) * (a - b)) + .sum(); + (pos as u64, dist) + }) + .collect(); + scored.sort_by(|a, b| a.1.total_cmp(&b.1)); + scored.truncate(k); + scored +} + +/// Table options that route searches into the primary-key vector branch +/// (`VectorSearchBuilder::execute_primary_key_vector_search`). Default search +/// mode is FAST, so only the ANN segment is consulted (no exact fallback). +fn table_options() -> Vec<(String, String)> { + vec![ + ("bucket".to_string(), "1".to_string()), + ( + "pk-vector.index.columns".to_string(), + VECTOR_COLUMN.to_string(), + ), + ( + format!("fields.{VECTOR_COLUMN}.pk-vector.index.type"), + INDEX_TYPE.to_string(), + ), + ( + format!("fields.{VECTOR_COLUMN}.pk-vector.distance.metric"), + "l2".to_string(), + ), + ] +} + +/// Primary-key schema `(id INT PRIMARY KEY, embedding VECTOR)`. +fn pk_vector_schema() -> TableSchema { + let mut builder = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column( + VECTOR_COLUMN, + DataType::Vector( + VectorType::try_new(true, DIM as u32, DataType::Float(FloatType::new())).unwrap(), + ), + ) + .primary_key(["id"]); + for (k, v) in table_options() { + builder = builder.option(k, v); + } + TableSchema::new(0, &builder.build().unwrap()) +} + +/// Arrow batch matching the table schema: `id` (== physical position) plus a +/// `FixedSizeList` vector column named to match paimon's target Arrow +/// schema (`element`). +fn data_batch(vectors: &[[f32; DIM]]) -> RecordBatch { + let ids: Vec = (0..vectors.len() as i32).collect(); + + let element_field = Arc::new(ArrowField::new("element", ArrowDataType::Float32, true)); + let mut vector_builder = FixedSizeListBuilder::new(Float32Builder::new(), DIM as i32) + .with_field(element_field.clone()); + for vector in vectors { + for &value in vector { + vector_builder.values().append_value(value); + } + vector_builder.append(true); + } + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new( + VECTOR_COLUMN, + ArrowDataType::FixedSizeList(element_field, DIM as i32), + true, + ), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(ids)) as ArrayRef, + Arc::new(vector_builder.finish()) as ArrayRef, + ], + ) + .unwrap() +} + +/// Encode one Java `DataOutput#writeUTF` value (u16-BE byte length + modified +/// UTF-8). ASCII file names are the common case; multibyte handling mirrors the +/// round-trip helper in `PkVectorSourceMeta`'s own tests. +fn java_write_utf(s: &str) -> Vec { + let mut body = Vec::new(); + for c in s.encode_utf16() { + if (0x0001..=0x007F).contains(&c) { + body.push(c as u8); + } else if c > 0x07FF { + body.push(0xE0 | (c >> 12) as u8); + body.push(0x80 | ((c >> 6) & 0x3F) as u8); + body.push(0x80 | (c & 0x3F) as u8); + } else { + body.push(0xC0 | (c >> 6) as u8); + body.push(0x80 | (c & 0x3F) as u8); + } + } + let mut out = (body.len() as u16).to_be_bytes().to_vec(); + out.extend_from_slice(&body); + out +} + +/// 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 { + let mut out = Vec::new(); + out.extend_from_slice(&1i32.to_be_bytes()); // version + out.extend_from_slice(&(files.len() as i32).to_be_bytes()); + for (name, rows) in files { + out.extend_from_slice(&java_write_utf(name)); + out.extend_from_slice(&rows.to_be_bytes()); + } + out +} + +/// Build a real vindex IVF-flat index segment over `vectors` (label == physical +/// position) and write it into `{table}/index/{file_name}`. Container format and +/// API usage mirror `VindexIndexBuildBuilder::build_index_file`, so the segment +/// is readable by `VindexVectorGlobalIndexReader::visit_vector_search`. +/// +/// `nlist = 1` keeps the search exact (a single inverted list scanned in full). +/// The vindex `metric = l2` matches the table's `pk-vector.distance.metric = l2`, +/// so distances agree. +async fn write_ann_segment( + file_io: &FileIO, + table_location: &str, + file_name: &str, + vectors: &[[f32; DIM]], +) -> u64 { + let n = vectors.len(); + let flat: Vec = vectors.iter().flat_map(|v| v.iter().copied()).collect(); + let ids: Vec = (0..n as i64).collect(); + + let native_options = HashMap::from([ + ("index.type".to_string(), "ivf_flat".to_string()), + ("dimension".to_string(), DIM.to_string()), + ("nlist".to_string(), "1".to_string()), + ("metric".to_string(), "l2".to_string()), + ]); + let config = VectorIndexConfig::from_options(&native_options).unwrap(); + + let training = VectorIndexTrainer::train(config, &flat, n).unwrap(); + let mut writer = VectorIndexWriter::new(training); + writer.add_vectors(&ids, &flat, n).unwrap(); + let mut bytes = Vec::new(); + { + let mut output = PosWriter::new(&mut bytes); + writer.write(&mut output).unwrap(); + } + + let index_dir = format!("{}/index", table_location.trim_end_matches('/')); + file_io.mkdirs(&index_dir).await.unwrap(); + let index_path = format!("{index_dir}/{file_name}"); + let file_size = bytes.len() as u64; + file_io + .new_output(&index_path) + .unwrap() + .write(Bytes::from(bytes)) + .await + .unwrap(); + file_size +} + +/// Round-trip the segment bytes through the reader in isolation, asserting the +/// analytic expectation. This proves the produced vindex bytes are readable and +/// the distances match before the full table read path is exercised. +fn assert_segment_reads_back(bytes: &[u8], query: &[f32], expected: &[(u64, f32)]) { + use paimon_vindex_core::index::{VectorIndexReader, VectorSearchParams}; + let mut reader = VectorIndexReader::open(Cursor::new(bytes.to_vec())).unwrap(); + reader.optimize_for_search().unwrap(); + let (labels, distances) = reader + .search(query, VectorSearchParams::new(expected.len(), 1)) + .unwrap(); + // Pair and sort best-first (smallest distance) to compare with the analytic + // expectation regardless of the reader's internal ordering. + let mut pairs: Vec<(i64, f32)> = labels.into_iter().zip(distances).collect(); + pairs.sort_by(|a, b| a.1.total_cmp(&b.1)); + let got: Vec<(u64, f32)> = pairs + .into_iter() + .map(|(label, distance)| (label as u64, distance)) + .collect(); + for ((got_id, got_d), (want_id, want_d)) in got.iter().zip(expected.iter()) { + assert_eq!(got_id, want_id, "segment label diverges from expected"); + assert!( + (got_d - want_d).abs() < 1e-3, + "segment distance diverges: got {got_d}, want {want_d}" + ); + } +} + +/// Open a table from the local filesystem, loading its latest schema. +async fn open_table(file_io: &FileIO, location: &str) -> Table { + let schema = SchemaManager::new(file_io.clone(), location.to_string()) + .latest() + .await + .expect("failed to list schemas") + .expect("table has no schema"); + Table::new( + file_io.clone(), + Identifier::new("default", "pkvector_baseline"), + location.to_string(), + (*schema).clone(), + None, + ) +} + +/// Build a complete self-contained primary-key vector table over `vectors` in a +/// fresh temp dir: persist the schema, write a real data file, apply the two +/// PK-vector constraints to its meta, build+commit a real vindex ANN segment, and +/// verify the segment reads back against `analytic_topk`. Returns the temp dir +/// (kept alive by the caller) and the opened table. +async fn build_table( + query: &[f32], + vectors: &[[f32; DIM]], + k: usize, +) -> (tempfile::TempDir, Table) { + let tmp = tempfile::tempdir().expect("create temp dir"); + let location = format!("file://{}", tmp.path().display()); + let file_io = FileIOBuilder::new("file").build().unwrap(); + + // Table layout dirs, then persist the schema. + for dir in ["schema", "snapshot", "manifest", "index"] { + file_io.mkdirs(&format!("{location}/{dir}")).await.unwrap(); + } + let schema = pk_vector_schema(); + file_io + .new_output(&format!("{location}/schema/schema-{}", schema.id())) + .unwrap() + .write(Bytes::from(serde_json::to_vec(&schema).unwrap())) + .await + .unwrap(); + + let table = open_table(&file_io, &location).await; + + // Write a real data file via the public write path to obtain a genuine + // DataFileMeta (real file name, row count, stats, file size). Its stats type + // is crate-private, so we reuse this meta rather than construct one. + let write_builder = table.new_write_builder(); + let mut writer = write_builder.new_write().unwrap(); + writer + .write_arrow_batch(&data_batch(vectors)) + .await + .unwrap(); + let write_messages = writer.prepare_commit().await.unwrap(); + assert_eq!( + write_messages.len(), + 1, + "single bucket -> one write message" + ); + let written = &write_messages[0]; + assert_eq!(written.new_files.len(), 1, "single data file expected"); + let base_meta = written.new_files[0].clone(); + let bucket = written.bucket; + let partition = written.partition.clone(); + let data_file_name = base_meta.file_name.clone(); + let row_count = base_meta.row_count; + + // Constraint 1 (PrimaryKeyIndexSourcePolicy.shouldRead): only a compacted, + // non-level-0 file backs the PK-vector index. Clone the real meta and set + // level > 0 + file_source == COMPACT (1). Pin first_row_id = 0 so the global + // row id equals the physical position. + let indexed_meta = DataFileMeta { + level: 1, + file_source: Some(1), + first_row_id: Some(0), + ..base_meta + }; + + // Build and persist the real vindex ANN segment; verify it reads back before + // wiring it into the table. + let index_file_name = "vector-ivf-flat-pkvector-baseline.index".to_string(); + let index_file_size = write_ann_segment(&file_io, &location, &index_file_name, vectors).await; + { + let bytes = file_io + .new_input(&format!("{location}/index/{index_file_name}")) + .unwrap() + .read() + .await + .unwrap(); + assert_segment_reads_back(&bytes, query, &analytic_topk(query, vectors, k)); + } + + // Constraint 2: GlobalIndexMeta.source_meta must be the Java PkVectorSourceMeta + // frame naming the backing data file(s) in ordinal order. Here one source file + // owns all rows, so ordinal == physical position. + let vector_field_id = schema + .fields() + .iter() + .find(|f| f.name() == VECTOR_COLUMN) + .expect("vector field present") + .id(); + let index_file = IndexFileMeta { + index_type: INDEX_TYPE.to_string(), + file_name: index_file_name, + file_size: i32::try_from(index_file_size).unwrap(), + row_count: i32::try_from(row_count).unwrap(), + deletion_vectors_ranges: None, + global_index_meta: Some(GlobalIndexMeta { + row_range_start: 0, + row_range_end: row_count - 1, + index_field_id: vector_field_id, + extra_field_ids: None, + source_meta: Some(source_meta_bytes(&[(&data_file_name, row_count)])), + index_meta: None, + }), + }; + + // Commit the indexed data file together with the ANN index segment in one + // snapshot. TableCommit writes the data manifest, index manifest, and snapshot. + let mut message = CommitMessage::new(partition, bucket, vec![indexed_meta]); + message.new_index_files = vec![index_file]; + TableCommit::new(table.clone(), "pkvector-baseline".to_string()) + .commit(vec![message]) + .await + .unwrap(); + + (tmp, table) +} + +/// Run `execute_read()` and flatten the stream into per-row `(id, score)` tuples +/// in emission order (best-first), returning the collected batches too for +/// schema / row-content assertions. +async fn read_id_and_scores( + table: &Table, + query: Vec, + limit: usize, + projection: Option<&[&str]>, +) -> (Vec, Vec, Vec) { + let mut builder = table.new_vector_search_builder(); + builder + .with_vector_column(VECTOR_COLUMN) + .with_query_vector(query) + .with_limit(limit); + if let Some(cols) = projection { + builder.with_projection(cols); + } + let batches = builder + .execute_read() + .await + .expect("primary-key vector read failed") + .try_collect::>() + .await + .expect("collecting read batches failed"); + + let ids: Vec = batches + .iter() + .flat_map(|b| { + let idx = b.schema().index_of("id").unwrap(); + b.column(idx) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + let scores: Vec = batches + .iter() + .flat_map(|b| { + let idx = b.schema().index_of("_PKEY_VECTOR_SCORE").unwrap(); + b.column(idx) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + (ids, scores, batches) +} + +/// Extract the materialized vector column across all batches, one `Vec` per +/// row in emission order. +fn collect_vectors(batches: &[RecordBatch]) -> Vec> { + let mut out = Vec::new(); + for batch in batches { + let idx = batch.schema().index_of(VECTOR_COLUMN).unwrap(); + let fsl = batch + .column(idx) + .as_any() + .downcast_ref::() + .expect("vector column must materialize as FixedSizeList"); + for row in 0..fsl.len() { + let values = fsl.value(row); + let floats = values + .as_any() + .downcast_ref::() + .expect("vector element must be Float32"); + out.push(floats.values().to_vec()); + } + } + out +} + +/// Fixture #1: distances 1 < 41 < 67 < 181, top-3 = rows 0, 4, 5. Here the +/// best-first order (0, 4, 5) happens to be ascending physical position, so this +/// case cannot by itself catch a "position order" regression — that gap is closed +/// by the discriminating fixture below. This case pins the search result and the +/// score alignment on the read path. +fn fixture_smoke() -> ([f32; DIM], Vec<[f32; DIM]>) { + let query = [9.0, 0.0, 0.0, 0.0]; + let vectors = vec![ + [10.0, 0.0, 0.0, 0.0], // row 0 -> (9-10)^2 = 1 + [0.0, 10.0, 0.0, 0.0], // row 1 -> 81 + 100 = 181 + [0.0, 0.0, 10.0, 0.0], // row 2 -> 181 + [0.0, 0.0, 0.0, 10.0], // row 3 -> 181 + [5.0, 5.0, 0.0, 0.0], // row 4 -> 16 + 25 = 41 + [1.0, 1.0, 1.0, 1.0], // row 5 -> 64 + 1 + 1 + 1 = 67 + ]; + (query, vectors) +} + +/// Fixture #2 (discriminating): the nearest neighbour sits at physical position +/// 5, the second nearest at position 1, the third at position 3, so the +/// best-first order [5, 1, 3] is NOT the ascending physical-position order +/// [1, 3, 5]. If the read path ever degraded to emitting rows in physical +/// position order, the ordering assertion below would fail. +/// +/// query [10,0,0,0] +/// pos0 [0,4,0,0] -> 100 + 16 = 116 +/// pos1 [8,0,0,0] -> 4 (2nd nearest) +/// pos2 [0,0,5,0] -> 100 + 25 = 125 +/// pos3 [7,0,0,0] -> 9 (3rd nearest) +/// pos4 [0,0,0,6] -> 100 + 36 = 136 +/// pos5 [9,0,0,0] -> 1 (nearest) +/// Strict gaps 1 < 4 < 9 < 116 < 125 < 136 make the top-3 order unique. +fn fixture_discriminating() -> ([f32; DIM], Vec<[f32; DIM]>) { + let query = [10.0, 0.0, 0.0, 0.0]; + let vectors = vec![ + [0.0, 4.0, 0.0, 0.0], // pos 0 + [8.0, 0.0, 0.0, 0.0], // pos 1 + [0.0, 0.0, 5.0, 0.0], // pos 2 + [7.0, 0.0, 0.0, 0.0], // pos 3 + [0.0, 0.0, 0.0, 6.0], // pos 4 + [9.0, 0.0, 0.0, 0.0], // pos 5 + ]; + (query, vectors) +} + +#[tokio::test] +async fn pk_vector_end_to_end_returns_expected_row_ids_and_scores() { + let (query, vectors) = fixture_smoke(); + let (_tmp, table) = build_table(&query, &vectors, 3).await; + + let expected = analytic_topk(&query, &vectors, 3); + let expected_row_ids: Vec = expected.iter().map(|(id, _)| *id).collect(); + let expected_scores: Vec = expected.iter().map(|(_, d)| l2_score(*d)).collect(); + + // Search path: execute_scored() -> row ids + scores. + let result = table + .new_vector_search_builder() + .with_vector_column(VECTOR_COLUMN) + .with_query_vector(query.to_vec()) + .with_limit(3) + .execute_scored() + .await + .expect("primary-key vector search failed"); + + assert_eq!( + result.row_ids, expected_row_ids, + "row ids diverge from the analytic expectation" + ); + assert_eq!( + result.scores.len(), + expected_scores.len(), + "score count diverges from the analytic expectation" + ); + for (got, want) in result.scores.iter().zip(&expected_scores) { + assert!( + (got - want).abs() < 1e-4, + "score diverges from the analytic expectation: got {got}, want {want}" + ); + } + + // Search-and-read: execute_read() materializes the matching rows best-first + // with a `_PKEY_VECTOR_SCORE` column, hiding `_ROW_ID`/`_PKEY_VECTOR_POSITION`. + // Projection ['id'] excludes the vector column. + let (ids, scores, batches) = read_id_and_scores(&table, query.to_vec(), 3, Some(&["id"])).await; + + let expected_ids: Vec = expected_row_ids.iter().map(|&id| id as i32).collect(); + assert_eq!(ids, expected_ids, "materialized rows must be best-first"); + assert_eq!(scores.len(), 3); + for (got, want) in scores.iter().zip(&expected_scores) { + assert!( + (got - want).abs() < 1e-4, + "materialized score diverges: got {got}, want {want}" + ); + } + for batch in &batches { + assert!( + batch.schema().index_of("_ROW_ID").is_err(), + "_ROW_ID must not leak into read output" + ); + assert!( + batch.schema().index_of("_PKEY_VECTOR_POSITION").is_err(), + "_PKEY_VECTOR_POSITION must not leak into read output" + ); + assert!( + batch.schema().index_of(VECTOR_COLUMN).is_err(), + "projection ['id'] must exclude the vector column" + ); + } +} + +/// Closes the "best-first == physical position order" discriminative gap: reads +/// back a fixture whose nearest neighbours are at physical positions 5, 1, 3 (in +/// that order) and asserts the materialized output is emitted best-first +/// [5, 1, 3], not in ascending physical position [1, 3, 5]. Also asserts the full +/// row content (id + vector values) and the aligned `_PKEY_VECTOR_SCORE`, with no +/// `_ROW_ID`/`_PKEY_VECTOR_POSITION` leaking. +#[tokio::test] +async fn pk_vector_read_orders_rows_best_first_not_by_position() { + let (query, vectors) = fixture_discriminating(); + let (_tmp, table) = build_table(&query, &vectors, 3).await; + + let expected = analytic_topk(&query, &vectors, 3); + let expected_ids: Vec = expected.iter().map(|(id, _)| *id as i32).collect(); + // The whole point of this fixture: best-first order != ascending position. + assert_eq!( + expected_ids, + vec![5, 1, 3], + "fixture must produce best-first order distinct from physical position order" + ); + let mut position_order = expected_ids.clone(); + position_order.sort_unstable(); + assert_ne!( + expected_ids, position_order, + "fixture is only discriminating if best-first != ascending position" + ); + + // Default projection (all user columns): id + vector column materialize. + let (ids, scores, batches) = read_id_and_scores(&table, query.to_vec(), 3, None).await; + + // Row order == best-first, NOT physical position order. A regression to + // position order would emit [1, 3, 5] and fail here. + assert_eq!( + ids, expected_ids, + "materialized rows must be best-first [5, 1, 3], not position order [1, 3, 5]" + ); + + // Row content: the materialized vector for each emitted row equals the source + // vector at that physical position. + let got_vectors = collect_vectors(&batches); + assert_eq!(got_vectors.len(), 3, "three rows expected"); + for (row_idx, (id, _)) in expected.iter().enumerate() { + assert_eq!( + got_vectors[row_idx], + vectors[*id as usize].to_vec(), + "materialized vector for row id {id} diverges from source data" + ); + } + + // Score alignment: `_PKEY_VECTOR_SCORE` matches metric.distance_to_score for + // each emitted row, in best-first order. + assert_eq!(scores.len(), 3); + for (got, (_, distance)) in scores.iter().zip(&expected) { + assert!( + (got - l2_score(*distance)).abs() < 1e-4, + "materialized score diverges: got {got}, want {}", + l2_score(*distance) + ); + } + + // Hidden metadata columns must not leak into the output. + for batch in &batches { + assert!( + batch.schema().index_of("_ROW_ID").is_err(), + "_ROW_ID must not leak into read output" + ); + assert!( + batch.schema().index_of("_PKEY_VECTOR_POSITION").is_err(), + "_PKEY_VECTOR_POSITION must not leak into read output" + ); + // Default projection keeps the user vector column. + assert!( + batch.schema().index_of(VECTOR_COLUMN).is_ok(), + "default projection must materialize the vector column" + ); + } +} From 762d3cd973bf7bedd5aa0414c5ca1e2e87052caa Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Thu, 16 Jul 2026 14:36:48 +0800 Subject: [PATCH 06/10] fix(vindex): require Send/Sync on pk-vector search seams for async spawn The public PK-vector search path is async and its future is spawned on a Send runtime by callers such as the DataFusion integration. The search seams held across .await were non-Send trait objects, so the future was not Send and the dependent crates failed to build. Require Send on PkVectorReader, Send + Sync on PkVectorAnnSearcher and the ANN Scorer closure, and Send on the exact-reader factory parameters. Test scorer switches Rc to Arc to satisfy the bound. --- .../src/table/pk_vector_orchestrator.rs | 8 ++-- crates/paimon/src/vindex/pkvector/ann.rs | 37 +++++++++++-------- crates/paimon/src/vindex/pkvector/reader.rs | 6 ++- 3 files changed, 32 insertions(+), 19 deletions(-) diff --git a/crates/paimon/src/table/pk_vector_orchestrator.rs b/crates/paimon/src/table/pk_vector_orchestrator.rs index 140dd02ac..7f9759780 100644 --- a/crates/paimon/src/table/pk_vector_orchestrator.rs +++ b/crates/paimon/src/table/pk_vector_orchestrator.rs @@ -260,11 +260,12 @@ impl PkVectorOrchestrator { metric: VectorSearchMetric, limit: usize, ann_searcher: Option<&dyn PkVectorAnnSearcher>, - exact_reader_factory: &mut dyn FnMut( + exact_reader_factory: &mut (dyn FnMut( usize, &PkVectorSearchSplit, &BucketActiveFile, - ) -> crate::Result>, + ) -> crate::Result> + + Send), search_options: &HashMap, skip_exact_fallback: bool, ) -> crate::Result> { @@ -840,7 +841,8 @@ mod e2e_tests { metric: VectorSearchMetric, limit: usize, ann: Option<&dyn PkVectorAnnSearcher>, - factory: &mut dyn FnMut(&BucketActiveFile) -> crate::Result>, + factory: &mut (dyn FnMut(&BucketActiveFile) -> crate::Result> + + Send), opts: &HashMap, ) -> crate::Result> { let orch = PkVectorOrchestrator::new(reader.clone()); diff --git a/crates/paimon/src/vindex/pkvector/ann.rs b/crates/paimon/src/vindex/pkvector/ann.rs index 5c650afb8..0d1559bb4 100644 --- a/crates/paimon/src/vindex/pkvector/ann.rs +++ b/crates/paimon/src/vindex/pkvector/ann.rs @@ -128,7 +128,11 @@ pub(crate) fn map_ann_results( /// One ANN segment's search dependency for the bucket kernel. Bucket tests fake /// this (mirroring Java's mock of `PkVectorAnnSegmentSearcher`). -pub(crate) trait PkVectorAnnSearcher { +/// +/// `Send + Sync` so a `&dyn PkVectorAnnSearcher` can be held across the `.await` +/// points of the async search path (the returned future is spawned on a `Send` +/// runtime by callers such as the DataFusion integration). +pub(crate) trait PkVectorAnnSearcher: Send + Sync { #[allow(clippy::too_many_arguments)] fn search( &self, @@ -152,8 +156,11 @@ pub(crate) trait PkVectorAnnSearcher { /// with a segment's index bytes; tests inject a synthetic scorer. The adapter's /// own logic (live-row masking, ordinal mapping, deletion checks, ordering) is /// exercised independently of the scorer. -pub(crate) type Scorer = - Box crate::Result>>>; +pub(crate) type Scorer = Box< + dyn Fn(&BucketAnnSegment, &VectorSearch) -> crate::Result>> + + Send + + Sync, +>; /// Structural vindex-backed `PkVectorAnnSearcher`. Composes the pure helpers /// (`build_live_row_ids`, `map_ann_results`) around the scorer seam. @@ -351,19 +358,19 @@ mod tests { #[test] fn test_vindex_adapter_composes_live_rows_and_maps_results() { // Scorer records the VectorSearch it received and returns synthetic ordinals. - // The scorer must be `'static`, so share the recording cells via `Rc` moved - // into the closure rather than borrowing locals. - use std::cell::RefCell; - use std::rc::Rc; - let seen_limit = Rc::new(RefCell::new(0usize)); - let seen_has_filter = Rc::new(RefCell::new(false)); - let scorer_limit = Rc::clone(&seen_limit); - let scorer_has_filter = Rc::clone(&seen_has_filter); + // The scorer must be `'static` and `Send + Sync`, so share the recording + // cells via `Arc>` moved into the closure rather than borrowing + // locals. + use std::sync::{Arc, Mutex}; + let seen_limit = Arc::new(Mutex::new(0usize)); + let seen_has_filter = Arc::new(Mutex::new(false)); + let scorer_limit = Arc::clone(&seen_limit); + let scorer_has_filter = Arc::clone(&seen_has_filter); let searcher = VindexAnnSearcher::new( "embedding".to_string(), Box::new(move |_segment: &BucketAnnSegment, search: &VectorSearch| { - *scorer_limit.borrow_mut() = search.limit; - *scorer_has_filter.borrow_mut() = search.include_row_ids.is_some(); + *scorer_limit.lock().unwrap() = search.limit; + *scorer_has_filter.lock().unwrap() = search.include_row_ids.is_some(); let mut scores = HashMap::new(); scores.insert(3u64, 0.5f32); // -> (f1, 0) scores.insert(0u64, 0.25f32); // -> (f0, 0), l2 dist 3.0 @@ -394,9 +401,9 @@ mod tests { // Sorted BEST_FIRST by distance: (f1,0) dist 1.0 then (f0,0) dist 3.0. assert_eq!(results[0].data_file_name, "f1"); assert_eq!(results[1].data_file_name, "f0"); - assert_eq!(*seen_limit.borrow(), 2); + assert_eq!(*seen_limit.lock().unwrap(), 2); assert!( - *seen_has_filter.borrow(), + *seen_has_filter.lock().unwrap(), "DV present -> include_row_ids set" ); } diff --git a/crates/paimon/src/vindex/pkvector/reader.rs b/crates/paimon/src/vindex/pkvector/reader.rs index 3298013af..f0f7d0aad 100644 --- a/crates/paimon/src/vindex/pkvector/reader.rs +++ b/crates/paimon/src/vindex/pkvector/reader.rs @@ -17,7 +17,11 @@ /// Sequential exact-scan source of vectors for one data file. Mirrors Java /// `org.apache.paimon.index.pkvector.PkVectorReader`. -pub(crate) trait PkVectorReader { +/// +/// `Send` so a boxed reader can be held across the `.await` points of the async +/// search path (the returned future is spawned on a `Send` runtime by callers +/// such as the DataFusion integration). +pub(crate) trait PkVectorReader: Send { fn dimension(&self) -> usize; fn row_count(&self) -> i64; From 5e72d42d95eddb4c841b745775d51f201c035e7d Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Thu, 16 Jul 2026 15:00:27 +0800 Subject: [PATCH 07/10] test(table): gate pk-vector baseline fixture tests off windows The baseline tests build the fixture table location as a file:// URL from a temp dir path, which FileIO cannot derive on Windows (see #397). The sibling rest_catalog_test gates its identical file:// tempdir tests with cfg(not(windows)); do the same here so the suite passes on Windows while still running on Linux and macOS. --- crates/paimon/tests/pk_vector_baseline_test.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/crates/paimon/tests/pk_vector_baseline_test.rs b/crates/paimon/tests/pk_vector_baseline_test.rs index b287257a5..1cb0399f0 100644 --- a/crates/paimon/tests/pk_vector_baseline_test.rs +++ b/crates/paimon/tests/pk_vector_baseline_test.rs @@ -533,6 +533,10 @@ fn fixture_discriminating() -> ([f32; DIM], Vec<[f32; DIM]>) { (query, vectors) } +// Gated off Windows: the fixture table location is a `file://` URL built from a +// temp dir path, which `FileIO` cannot derive on Windows (see #397); the sibling +// `rest_catalog_test` gates its identical `file://` tempdir tests the same way. +#[cfg(not(windows))] #[tokio::test] async fn pk_vector_end_to_end_returns_expected_row_ids_and_scores() { let (query, vectors) = fixture_smoke(); @@ -604,6 +608,8 @@ async fn pk_vector_end_to_end_returns_expected_row_ids_and_scores() { /// [5, 1, 3], not in ascending physical position [1, 3, 5]. Also asserts the full /// row content (id + vector values) and the aligned `_PKEY_VECTOR_SCORE`, with no /// `_ROW_ID`/`_PKEY_VECTOR_POSITION` leaking. +// Gated off Windows for the same `file://` tempdir reason as the test above. +#[cfg(not(windows))] #[tokio::test] async fn pk_vector_read_orders_rows_best_first_not_by_position() { let (query, vectors) = fixture_discriminating(); From 09d17373fde61d1a2a79e0375a7771d5d3fef6ea Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Thu, 16 Jul 2026 15:55:00 +0800 Subject: [PATCH 08/10] perf(table): preload exact readers only for ANN-uncovered files MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The non-fast search path preheated an exact PkVectorReader for every active file, reading each file's full vector column into memory up front. But the bucket search only runs the exact fallback on files not covered by an ANN segment, so covered files' columns were read and never used — IO/memory amplification on large compacted tables. Extract the covered-file rule into a shared covered_source_files helper (used by both bucket_search and the reader preload) and preload only the uncovered active files, matching Java, which creates a PkVectorReader lazily only for uncovered files. --- .../paimon/src/table/vector_search_builder.rs | 11 ++- crates/paimon/src/vindex/pkvector/bucket.rs | 81 ++++++++++++++++--- 2 files changed, 78 insertions(+), 14 deletions(-) diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index e4fa06b5f..6aeea1e77 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -43,7 +43,7 @@ use crate::table::{ use crate::vector_search::{GlobalIndexIOMeta, SearchResult, VectorSearch}; use crate::vindex::is_vindex_index_type; use crate::vindex::pkvector::ann::VindexAnnSearcher; -use crate::vindex::pkvector::bucket::{BucketActiveFile, BucketAnnSegment}; +use crate::vindex::pkvector::bucket::{covered_source_files, BucketActiveFile, BucketAnnSegment}; use crate::vindex::pkvector::metric::VectorSearchMetric; use crate::vindex::pkvector::reader::PkVectorReader; use crate::vindex::reader::VindexVectorGlobalIndexReader; @@ -366,16 +366,23 @@ impl<'a> VectorSearchBuilder<'a> { // Exact-fallback readers, keyed by (split_index, file_name). In FAST mode // the kernel never invokes the factory, so skip the in-memory column read - // entirely. + // entirely. Otherwise preload only the *uncovered* active files: files an + // ANN segment already covers never reach the exact fallback, so reading + // their vector column here would be wasted IO/memory. Mirrors Java, which + // creates a `PkVectorReader` lazily only for uncovered files. let mut exact_readers: HashMap<(usize, String), Box> = HashMap::new(); if !skip_exact_fallback { for (split_index, split) in plan.splits.iter().enumerate() { + let covered = covered_source_files(&split.ann_segments, &split.active_files); let factory = DataFilePkVectorReaderFactory::new( reader.clone(), split.data_split.clone(), vector_field.clone(), )?; for active in &split.active_files { + if covered.contains(&active.file_name) { + continue; + } let r = factory.create(active).await?; exact_readers.insert((split_index, active.file_name.clone()), r); } diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs b/crates/paimon/src/vindex/pkvector/bucket.rs index ccd57dbf6..d749e8124 100644 --- a/crates/paimon/src/vindex/pkvector/bucket.rs +++ b/crates/paimon/src/vindex/pkvector/bucket.rs @@ -108,6 +108,35 @@ fn add_candidate(heap: &mut BinaryHeap, candidate: PkVectorSearchRes } } +/// Active data files whose rows are already covered by an ANN segment's source +/// metadata, matched by both file name AND row count. The bucket exact fallback +/// skips these files, so a caller that preloads exact readers should preload +/// only the *uncovered* active files (`active_files` minus this set) rather than +/// reading every active file's vector column up front. A source naming an +/// inactive file, or one whose row count disagrees with the active file, is not +/// covered here; `bucket_search` rejects the row-count mismatch separately. +pub(crate) fn covered_source_files( + ann_segments: &[BucketAnnSegment], + active_files: &[BucketActiveFile], +) -> HashSet { + let row_counts: HashMap<&str, i64> = active_files + .iter() + .map(|f| (f.file_name.as_str(), f.row_count)) + .collect(); + let mut covered = HashSet::new(); + for segment in ann_segments { + for source in segment.source_meta.source_files() { + if row_counts + .get(source.file_name()) + .is_some_and(|&rc| rc == source.row_count()) + { + covered.insert(source.file_name().to_string()); + } + } + } + covered +} + /// ANN + exact data-file fallback search for one snapshot bucket. Mirrors Java /// `org.apache.paimon.index.pkvector.PrimaryKeyVectorBucketSearch.search`. /// @@ -154,26 +183,25 @@ pub(crate) fn bucket_search( let mut heap: BinaryHeap = BinaryHeap::with_capacity(limit + 1); let active_source_files: HashSet = files_by_name.keys().map(|name| name.to_string()).collect(); - let mut covered: HashSet = HashSet::new(); + // Active files whose rows an ANN segment already covers; the exact fallback + // skips them. Same rule the caller's exact-reader preload uses, so both agree + // on which files still need an exact reader. + let covered = covered_source_files(ann_segments, active_files); for segment in ann_segments { + // An active ANN source with a mismatched row count is corruption (the + // ordinal-to-position mapping would be wrong). An inactive source (no + // matching active file) is skipped: it was compacted away and its ordinal + // range is masked out of the ANN live-row bitmap. Mirrors Java master + // `PrimaryKeyVectorBucketSearch` (`file == null` -> continue). for source in segment.source_meta.source_files() { - // An ANN source that is no longer an active file (e.g. compacted away) - // is skipped, not rejected: its ordinal range is masked out of the ANN - // live-row bitmap and the remaining active sources are still searched. - // Mirrors Java master `PrimaryKeyVectorBucketSearch` (`file == null` - // -> continue). Active sources still require a row-count match. - match files_by_name.get(source.file_name()) { - Some(active) if active.row_count == source.row_count() => { - covered.insert(source.file_name().to_string()); - } - Some(_) => { + if let Some(active) = files_by_name.get(source.file_name()) { + if active.row_count != source.row_count() { return Err(data_invalid(format!( "ANN source {} does not match the active data file", source.file_name() ))); } - None => continue, } } let searcher = ann_searcher.ok_or_else(|| data_invalid("ANN search is not configured"))?; @@ -637,4 +665,33 @@ mod tests { .unwrap_err(); assert!(err.to_string().contains("row count") || err.to_string().contains("-1")); } + + #[test] + fn covered_source_files_matches_by_name_and_row_count() { + // "data-1" is an active ANN source with matching row count -> covered. + // "data-2" is active but its row count disagrees with the ANN source -> not + // covered (bucket_search rejects that separately). "data-3" is an active + // file with no ANN source -> not covered (it needs an exact reader). + let segment = BucketAnnSegment::for_test(meta(&[("data-1", 3), ("data-2", 9)])); + let active = vec![ + active("data-1", 3), + active("data-2", 2), + active("data-3", 5), + ]; + let covered = covered_source_files(&[segment], &active); + assert!(covered.contains("data-1")); + assert!(!covered.contains("data-2")); + assert!(!covered.contains("data-3")); + assert_eq!(covered.len(), 1); + } + + #[test] + fn covered_source_files_ignores_inactive_source() { + // ANN source names a file that is not active (compacted away) -> not + // covered, and no active file needs it. + let segment = BucketAnnSegment::for_test(meta(&[("gone", 4)])); + let active = vec![active("data-1", 3)]; + let covered = covered_source_files(&[segment], &active); + assert!(covered.is_empty()); + } } From b541e0b06910ddc30c4a3ffb34c5c949e55041d6 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Thu, 16 Jul 2026 16:37:17 +0800 Subject: [PATCH 09/10] fix(table): plan pk-vector scan on one snapshot and guard hit positions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Snapshot consistency: the scan resolved the snapshot twice — the index manifest and snapshot id from get_latest_snapshot(), the data splits from a separate (time-travel-aware) scan. A time-travel query then read data from the travelled snapshot but the index from latest and failed the snapshot-id guard, and a concurrent commit between the two resolutions could do the same. Derive the snapshot from the data scan's own splits and read the index manifest from that one snapshot, so both sides are consistent and honor time travel, matching Java PrimaryKeyVectorScan. Also fail loud on a non-ADD index-manifest entry (a malformed manifest) instead of silently skipping it, and validate each hit's physical row position is in-range for its data file (non-negative, < row count, fits i32) before building splits or global row ids, mirroring the bounds Java PrimaryKeyVectorResult enforces. --- .../src/table/pk_vector_orchestrator.rs | 37 +++++++++++++++++++ crates/paimon/src/table/pk_vector_scan.rs | 34 ++++++++++++----- .../paimon/src/table/vector_search_builder.rs | 4 +- 3 files changed, 65 insertions(+), 10 deletions(-) diff --git a/crates/paimon/src/table/pk_vector_orchestrator.rs b/crates/paimon/src/table/pk_vector_orchestrator.rs index 7f9759780..66a754334 100644 --- a/crates/paimon/src/table/pk_vector_orchestrator.rs +++ b/crates/paimon/src/table/pk_vector_orchestrator.rs @@ -44,6 +44,25 @@ fn data_invalid(message: impl Into) -> crate::Error { } } +/// Validate a hit's physical row position against its data file, mirroring the +/// bounds Java `PrimaryKeyVectorResult.splits()` enforces per candidate: the +/// position must be non-negative, within the file's row count, and fit in an +/// `i32`. A position outside this range means a corrupt ANN index or malformed +/// source metadata resolved to a bogus ordinal; fail loud rather than emit a +/// wrong row. +pub(crate) fn validate_row_position( + file_name: &str, + row_position: i64, + row_count: i64, +) -> crate::Result<()> { + if row_position < 0 || row_position >= row_count || row_position > i32::MAX as i64 { + return Err(data_invalid(format!( + "vector search hit position {row_position} out of range for {file_name} (row count {row_count})" + ))); + } + Ok(()) +} + /// One bucket's search input. Rust equivalent of Java /// `BucketVectorSearchSplit`. Constructed from a snapshot/manifest plan by /// `PkVectorScan`. @@ -163,6 +182,11 @@ pub(crate) fn build_indexed_splits( .data_deletion_files() .and_then(|dfs| dfs.get(file_idx).cloned().flatten()); + // Every hit's physical position must be in range for its data file. + for &(pos, _) in &hits { + validate_row_position(&file_name, pos, file_meta.row_count)?; + } + // Coalesce ascending positions into inclusive ranges; scores aligned to // ascending-position order. let mut row_ranges: Vec = Vec::new(); @@ -1305,4 +1329,17 @@ mod e2e_tests { .unwrap(); assert!(cands.is_empty()); } + + #[test] + fn validate_row_position_bounds() { + // In range. + assert!(validate_row_position("f", 0, 3).is_ok()); + assert!(validate_row_position("f", 2, 3).is_ok()); + // Negative, at/over row count, and past i32::MAX all fail loud. + assert!(validate_row_position("f", -1, 3).is_err()); + assert!(validate_row_position("f", 3, 3).is_err()); + assert!(validate_row_position("f", i32::MAX as i64 + 1, i64::MAX).is_err()); + let err = validate_row_position("data-1", 9, 3).unwrap_err(); + assert!(err.to_string().contains("out of range") && err.to_string().contains("data-1")); + } } diff --git a/crates/paimon/src/table/pk_vector_scan.rs b/crates/paimon/src/table/pk_vector_scan.rs index 54900ad5b..d5547e576 100644 --- a/crates/paimon/src/table/pk_vector_scan.rs +++ b/crates/paimon/src/table/pk_vector_scan.rs @@ -166,15 +166,18 @@ impl<'a> PkVectorScan<'a> { pub(crate) async fn plan(&self) -> crate::Result { let snapshot_manager = self.table.snapshot_manager(); - let snapshot = match snapshot_manager.get_latest_snapshot().await? { - Some(s) => s, - None => return Ok(PkVectorScanPlan { splits: Vec::new() }), - }; - let snapshot_id = snapshot.id(); - // Data splits (scan all files). - let builder = self.table.new_read_builder(); - let data_splits = builder + // Data splits first, via the table's own scan resolution (which honors + // time travel / scan.snapshot-id). Deriving the snapshot from the scan's + // own output — rather than resolving `get_latest_snapshot()` separately — + // keeps the index manifest and the data splits on ONE snapshot, matching + // Java `PrimaryKeyVectorScan` (resolve one snapshot up front, read data and + // index from it). It also avoids a time-travel mismatch (data from the + // travelled snapshot, index from latest) and a TOCTOU where a concurrent + // commit lands between two independent resolutions. + let data_splits = self + .table + .new_read_builder() .new_scan() .with_scan_all_files() .plan() @@ -182,14 +185,27 @@ impl<'a> PkVectorScan<'a> { .splits() .to_vec(); + // No data files -> nothing to search. + let Some(first_split) = data_splits.first() else { + return Ok(PkVectorScanPlan { splits: Vec::new() }); + }; + let snapshot_id = first_split.snapshot_id(); + let snapshot = snapshot_manager.get_snapshot(snapshot_id).await?; + // Index-manifest scan into filtered ANN payload tuples. let table_path = self.table.location().trim_end_matches('/'); let mut entries = Vec::new(); if let Some(name) = snapshot.index_manifest() { let path = snapshot_manager.manifest_path(name); for entry in IndexManifest::read(self.table.file_io(), &path).await? { + // The on-disk index manifest is combined to live ADD entries only. + // A non-ADD entry means a malformed manifest; fail loud rather than + // silently drop it (mirrors Java `checkArgument(kind == ADD)`). if entry.kind != FileKind::Add { - continue; + return Err(data_invalid(format!( + "index manifest entry {} is not active (kind {:?})", + entry.index_file.file_name, entry.kind + ))); } if entry.index_file.index_type != self.index_type { continue; diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index 6aeea1e77..97f8b2e51 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -30,7 +30,8 @@ use crate::table::global_index_scanner::{ use crate::table::pk_vector_data_file_reader::DataFilePkVectorReaderFactory; use crate::table::pk_vector_indexed_split_read::PkVectorIndexedSplitRead; use crate::table::pk_vector_orchestrator::{ - build_indexed_splits, PkVectorCandidate, PkVectorOrchestrator, PkVectorSearchSplit, + build_indexed_splits, validate_row_position, PkVectorCandidate, PkVectorOrchestrator, + PkVectorSearchSplit, }; use crate::table::pk_vector_position_read::{ PKEY_VECTOR_POSITION_COLUMN, PKEY_VECTOR_SCORE_COLUMN, @@ -1021,6 +1022,7 @@ fn candidates_to_search_result( message: format!("data file {} has no first_row_id", c.data_file_name), source: None, })?; + validate_row_position(&c.data_file_name, c.row_position, file_meta.row_count)?; let global = first_row_id .checked_add(c.row_position) From 59900be4318bea17c411cc41315115d6d2368c71 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Thu, 16 Jul 2026 23:03:23 +0800 Subject: [PATCH 10/10] fix(vindex): validate ANN segment uniqueness in bucket search --- crates/paimon/src/vindex/pkvector/bucket.rs | 101 ++++++++++++++++++++ 1 file changed, 101 insertions(+) diff --git a/crates/paimon/src/vindex/pkvector/bucket.rs b/crates/paimon/src/vindex/pkvector/bucket.rs index d749e8124..337e92255 100644 --- a/crates/paimon/src/vindex/pkvector/bucket.rs +++ b/crates/paimon/src/vindex/pkvector/bucket.rs @@ -180,6 +180,33 @@ pub(crate) fn bucket_search( } } + // Validate ANN segments mirror Java PkVectorBucketIndexState constructor checks: + // (1) payload file uniqueness, (2) no source file covered by multiple segments. + let mut segments_by_path: HashMap<&str, usize> = HashMap::new(); + let mut source_to_segment: HashMap<&str, &str> = HashMap::new(); + for (idx, segment) in ann_segments.iter().enumerate() { + if segments_by_path + .insert(segment.path.as_str(), idx) + .is_some() + { + return Err(data_invalid(format!( + "ANN segment payload {} appears more than once", + segment.path + ))); + } + for source in segment.source_meta.source_files() { + if let Some(&prior_segment_path) = source_to_segment.get(source.file_name()) { + return Err(data_invalid(format!( + "source data file {} is covered by both ANN segments {} and {}", + source.file_name(), + prior_segment_path, + segment.path + ))); + } + source_to_segment.insert(source.file_name(), segment.path.as_str()); + } + } + let mut heap: BinaryHeap = BinaryHeap::with_capacity(limit + 1); let active_source_files: HashSet = files_by_name.keys().map(|name| name.to_string()).collect(); @@ -646,6 +673,80 @@ mod tests { assert!(results.is_empty()); } + #[test] + fn test_rejects_duplicate_ann_segment_path() { + let seg1 = BucketAnnSegment { + source_meta: meta(&[("data-1", 2)]), + path: "duplicate-path".to_string(), + file_size: 100, + index_meta: vec![1, 2, 3], + }; + let seg2 = BucketAnnSegment { + source_meta: meta(&[("data-2", 2)]), + path: "duplicate-path".to_string(), + file_size: 200, + index_meta: vec![4, 5, 6], + }; + let ann = FakeAnnSearcher { result: vec![] }; + let mut factory = + |_: &BucketActiveFile| -> crate::Result> { unreachable!() }; + let err = bucket_search( + Some(&ann), + &[seg1, seg2], + &[active("data-1", 2), active("data-2", 2)], + &HashMap::new(), + &mut factory, + &[0.0, 0.0], + VectorSearchMetric::L2, + 1, + &HashMap::new(), + false, + ) + .unwrap_err(); + assert!( + err.to_string().contains("duplicate-path") + && err.to_string().contains("appears more than once") + ); + } + + #[test] + fn test_rejects_source_file_covered_by_multiple_segments() { + let seg1 = BucketAnnSegment { + source_meta: meta(&[("data-1", 2)]), + path: "segment-1".to_string(), + file_size: 100, + index_meta: vec![1, 2, 3], + }; + let seg2 = BucketAnnSegment { + source_meta: meta(&[("data-1", 2), ("data-2", 2)]), + path: "segment-2".to_string(), + file_size: 200, + index_meta: vec![4, 5, 6], + }; + let ann = FakeAnnSearcher { result: vec![] }; + let mut factory = + |_: &BucketActiveFile| -> crate::Result> { unreachable!() }; + let err = bucket_search( + Some(&ann), + &[seg1, seg2], + &[active("data-1", 2), active("data-2", 2)], + &HashMap::new(), + &mut factory, + &[0.0, 0.0], + VectorSearchMetric::L2, + 1, + &HashMap::new(), + false, + ) + .unwrap_err(); + assert!( + err.to_string().contains("data-1") + && err.to_string().contains("covered by both") + && err.to_string().contains("segment-1") + && err.to_string().contains("segment-2") + ); + } + #[test] fn test_negative_active_row_count_rejected() { let mut factory =