This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 120ff9ec feat(table): plan a PK-vector search from a decoded bucket
split (#757)
120ff9ec is described below
commit 120ff9ec27a994bbdf4b0de1f176f10b6ed5d1c2
Author: Junrui Lee <[email protected]>
AuthorDate: Tue Sep 1 20:38:14 2026 +0800
feat(table): plan a PK-vector search from a decoded bucket split (#757)
---
crates/paimon/src/arrow/format/mosaic.rs | 58 ++
crates/paimon/src/table/data_file_reader.rs | 40 +-
crates/paimon/src/table/pk_vector_bucket_split.rs | 89 +++
.../paimon/src/table/pk_vector_data_file_reader.rs | 220 +++++--
crates/paimon/src/table/pk_vector_scan.rs | 675 ++++++++++++++++++++-
crates/paimon/src/table/vector_search_builder.rs | 526 ++++++++++++++--
crates/paimon/src/vindex/pkvector/ann.rs | 72 ++-
7 files changed, 1579 insertions(+), 101 deletions(-)
diff --git a/crates/paimon/src/arrow/format/mosaic.rs
b/crates/paimon/src/arrow/format/mosaic.rs
index 1fdbbace..b10c756c 100644
--- a/crates/paimon/src/arrow/format/mosaic.rs
+++ b/crates/paimon/src/arrow/format/mosaic.rs
@@ -877,6 +877,33 @@ mod tests {
.await
}
+ /// The byte ranges a read of `data` requests when limited to
`row_selection`.
+ async fn read_ranges_with_row_selection(
+ data: Bytes,
+ read_fields: &[DataField],
+ row_selection: Option<Vec<RowRange>>,
+ ) -> crate::Result<Vec<Range<u64>>> {
+ let file_size = data.len() as u64;
+ let calls = Arc::new(Mutex::new(Vec::new()));
+ let _: Vec<RecordBatch> = MosaicFormatReader
+ .read_batch_stream(
+ Box::new(TrackingFileRead {
+ data,
+ calls: Arc::clone(&calls),
+ }),
+ file_size,
+ read_fields,
+ None,
+ None,
+ row_selection,
+ )
+ .await?
+ .try_collect()
+ .await?;
+ let ranges = calls.lock().unwrap().clone();
+ Ok(ranges)
+ }
+
async fn read_ranges_with_predicates(
data: Bytes,
read_fields: &[DataField],
@@ -1298,6 +1325,37 @@ mod tests {
);
}
+ #[tokio::test]
+ async fn test_row_selection_skips_unselected_row_group_reads() {
+ // Three row groups of two rows. A selection inside the last one must
not
+ // fetch the column data of the first two: this is what makes a narrow
+ // engine-supplied row range cheaper than reading the file and
discarding
+ // rows afterwards, and it is granular to a row group, not to a row.
+ let fields = data_fields();
+ let projected = vec![fields[0].clone()];
+ let data = multi_row_group_mosaic(vec!["id".to_string()]);
+
+ let all = read_ranges_with_row_selection(data.clone(), &projected,
None)
+ .await
+ .unwrap();
+ let last_only =
+ read_ranges_with_row_selection(data, &projected,
Some(vec![RowRange::new(4, 5)]))
+ .await
+ .unwrap();
+
+ assert!(
+ last_only.len() < all.len(),
+ "a selection in one row group must request fewer ranges than a
full read: \
+ {last_only:?} vs {all:?}"
+ );
+ let selected_bytes: u64 = last_only.iter().map(|r| r.end -
r.start).sum();
+ let all_bytes: u64 = all.iter().map(|r| r.end - r.start).sum();
+ assert!(
+ selected_bytes < all_bytes,
+ "and fewer bytes: {selected_bytes} vs {all_bytes}"
+ );
+ }
+
#[tokio::test]
async fn test_read_predicate_missing_stats_still_filters_rows() {
let fields = data_fields();
diff --git a/crates/paimon/src/table/data_file_reader.rs
b/crates/paimon/src/table/data_file_reader.rs
index 25e20549..8958b382 100644
--- a/crates/paimon/src/table/data_file_reader.rs
+++ b/crates/paimon/src/table/data_file_reader.rs
@@ -599,6 +599,36 @@ impl DataFileReader {
data_fields: Option<Vec<DataField>>,
dv: Option<Arc<DeletionVector>>,
local_positions: Vec<i64>,
+ ) -> crate::Result<ArrowRecordBatchStream> {
+ self.read_single_file_stream_local_ranges(
+ split,
+ file_meta,
+ data_fields,
+ dv,
+ coalesce_positions_to_local_ranges(&local_positions),
+ )
+ }
+
+ /// As [`Self::read_single_file_stream_local`], but the selection is
already a
+ /// list of file-local inclusive ranges: sorted ascending,
non-overlapping, and
+ /// within `[0, file_meta.row_count)`. A caller that already holds ranges
— an
+ /// engine-supplied bucket split does — hands them over directly rather
than
+ /// expanding them into positions this would only coalesce back.
+ ///
+ /// The emitted rows are always exactly the selected ones, but what that
saves is
+ /// the format's business, and it differs: mosaic skips a row group before
+ /// touching its column data, parquet skips pages through the offset index,
+ /// `.row` prunes blocks. Avro is the exception — its reader loads the
whole file
+ /// and deserializes every record before applying the selection, so there a
+ /// narrow selection saves only what comes after decoding: Arrow column
+ /// materialization, and whatever the caller does per row.
+ pub(super) fn read_single_file_stream_local_ranges(
+ &self,
+ split: &DataSplit,
+ file_meta: DataFileMeta,
+ data_fields: Option<Vec<DataField>>,
+ dv: Option<Arc<DeletionVector>>,
+ local_ranges: Vec<RowRange>,
) -> crate::Result<ArrowRecordBatchStream> {
// Local-position selection is only sound against a predicate-free
reader: a
// row-filtering predicate drops arbitrary selected rows and desyncs
the
@@ -666,12 +696,10 @@ impl DataFileReader {
}
};
- // Interpret `local_positions` directly as file-local ranges (no
- // `to_local_row_ranges`, no `first_row_id`), then fold the DV in.
- // `merge_row_selection` intersects the selection with the file's
- // non-deleted ranges, so the reader emits exactly the selected,
non-deleted
- // rows in ascending physical order.
- let local_ranges =
coalesce_positions_to_local_ranges(&local_positions);
+ // Interpret the ranges directly as file-local (no
`to_local_row_ranges`, no
+ // `first_row_id`), then fold the DV in. `merge_row_selection`
intersects the
+ // selection with the file's non-deleted ranges, so the reader emits
exactly
+ // the selected, non-deleted rows in ascending physical order.
let row_selection =
merge_row_selection(file_meta.row_count, dv.as_deref(),
Some(&local_ranges));
diff --git a/crates/paimon/src/table/pk_vector_bucket_split.rs
b/crates/paimon/src/table/pk_vector_bucket_split.rs
index 0c50b675..c8282675 100644
--- a/crates/paimon/src/table/pk_vector_bucket_split.rs
+++ b/crates/paimon/src/table/pk_vector_bucket_split.rs
@@ -126,6 +126,39 @@ impl BucketVectorPayload {
.as_deref()
.expect("a decoded payload always carries source metadata")
}
+
+ /// Consume the payload into the pieces a planner needs, so its decoded
metadata
+ /// moves out of the payload rather than being cloned out of it.
+ ///
+ /// Two decoded fields are deliberately left behind. `row_count` is the
payload's
+ /// own row count, which the read path derives from the source metadata
instead.
+ /// `deletion_vectors_ranges` belongs to deletion-vector index files --
Java
+ /// builds a vector payload through the overload that leaves it null, and
a read
+ /// takes its deletion vectors from the bucket's data split -- so a value
here
+ /// describes something this payload is not, and is ignored rather than
applied.
+ #[cfg_attr(not(test), allow(dead_code))]
+ pub(crate) fn into_parts(self) -> BucketVectorPayloadParts {
+ BucketVectorPayloadParts {
+ index_type: self.index_type,
+ file_name: self.file_name,
+ file_size: self.file_size,
+ external_path: self.external_path,
+ global_index_meta: self.global_index_meta,
+ }
+ }
+}
+
+/// The owned pieces of a [`BucketVectorPayload`], produced by
+/// [`BucketVectorPayload::into_parts`].
+#[cfg_attr(not(test), allow(dead_code))]
+pub(crate) struct BucketVectorPayloadParts {
+ pub(crate) index_type: String,
+ pub(crate) file_name: String,
+ /// As decoded: Java writes a signed length, so a negative value is
possible on
+ /// the wire and is rejected where it is converted, not here.
+ pub(crate) file_size: i64,
+ pub(crate) external_path: Option<String>,
+ pub(crate) global_index_meta: GlobalIndexMeta,
}
impl BucketVectorSearchSplit {
@@ -144,6 +177,19 @@ impl BucketVectorSearchSplit {
&self.row_ranges_by_file
}
+ /// Consume the split into its three parts, so a planner can take
ownership of
+ /// the data split, the payloads and the row ranges without cloning them.
+ #[cfg_attr(not(test), allow(dead_code))]
+ pub(crate) fn into_parts(
+ self,
+ ) -> (
+ DataSplit,
+ Vec<BucketVectorPayload>,
+ IndexMap<String, Vec<RowRange>>,
+ ) {
+ (self.data_split, self.payload_files, self.row_ranges_by_file)
+ }
+
/// Parse a Java `BucketVectorSearchSplit#serialize` message.
///
/// Integers are big-endian and file names are Java modified UTF-8,
following
@@ -454,6 +500,49 @@ fn read_count(cur: &mut &[u8], element: &str) ->
crate::Result<usize> {
Ok(count)
}
+#[cfg(test)]
+impl BucketVectorSearchSplit {
+ /// Assemble a split directly, for tests that need shapes the decoder will
not
+ /// produce -- a nested split that wrongly carries row ranges, two splits
for one
+ /// bucket, a negative payload size. Production splits always come from
+ /// [`Self::deserialize`].
+ pub(crate) fn new_for_test(
+ data_split: DataSplit,
+ payload_files: Vec<BucketVectorPayload>,
+ row_ranges_by_file: IndexMap<String, Vec<RowRange>>,
+ ) -> Self {
+ Self {
+ data_split,
+ payload_files,
+ row_ranges_by_file,
+ }
+ }
+}
+
+#[cfg(test)]
+impl BucketVectorPayload {
+ #[allow(clippy::too_many_arguments)]
+ pub(crate) fn new_for_test(
+ index_type: &str,
+ file_name: &str,
+ file_size: i64,
+ row_count: i64,
+ deletion_vectors_ranges: Option<IndexMap<String, DeletionVectorMeta>>,
+ external_path: Option<String>,
+ global_index_meta: GlobalIndexMeta,
+ ) -> Self {
+ Self {
+ index_type: index_type.to_string(),
+ file_name: file_name.to_string(),
+ file_size,
+ row_count,
+ deletion_vectors_ranges,
+ external_path,
+ global_index_meta,
+ }
+ }
+}
+
#[cfg(test)]
mod tests {
use super::*;
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 c3087beb..dc42f172 100644
--- a/crates/paimon/src/table/pk_vector_data_file_reader.rs
+++ b/crates/paimon/src/table/pk_vector_data_file_reader.rs
@@ -35,7 +35,7 @@ use futures::TryStreamExt;
use crate::spec::{DataField, DataType};
use crate::table::data_file_reader::DataFileReader;
-use crate::table::source::DataSplit;
+use crate::table::source::{DataSplit, RowRange};
use crate::vindex::pkvector::bucket::BucketActiveFile;
use crate::vindex::pkvector::exact::{drain_best_first, push_bounded,
validate_query, WorstFirst};
use crate::vindex::pkvector::metric::VectorSearchMetric;
@@ -95,8 +95,17 @@ impl DataFilePkVectorReaderFactory {
/// opened. Each surviving physical position (not NULL, not `is_excluded`)
is
/// scored against every query into that query's bounded heap; a NULL row
is
/// skipped but still advances the position so the position stays in
lockstep
- /// with `is_excluded`. The drained row count is checked against the file's
- /// `DataFileMeta.row_count` (both truncation and overrun fail loud).
+ /// with `is_excluded`. The read is checked against its selection:
emitting more
+ /// or fewer rows than were selected fails loud. A selected read can only
vouch
+ /// for the ranges it asked for, so unlike a full read it cannot notice a
file
+ /// truncated somewhere else.
+ ///
+ /// `allowed_rows`, when given, limits the read to those file-local
inclusive
+ /// ranges — normalized, as the plan carries them. An engine-supplied
bucket
+ /// split can restrict a large file to a handful of rows, and scoring is
not the
+ /// expensive part: reading the rest of the file to have `is_excluded`
reject it
+ /// afterwards is. `is_excluded` still applies on top, since it also folds
in a
+ /// residual predicate and the deletion vector.
pub(crate) async fn search_file(
&self,
file: &BucketActiveFile,
@@ -104,6 +113,7 @@ impl DataFilePkVectorReaderFactory {
metric: VectorSearchMetric,
exact_limit: usize,
is_excluded: &(dyn Fn(i64) -> bool + Sync),
+ allowed_rows: Option<&[RowRange]>,
) -> crate::Result<Vec<Vec<PkVectorSearchResult>>> {
if exact_limit == 0 {
return Err(data_invalid("vector search limit must be positive"));
@@ -129,21 +139,48 @@ impl DataFilePkVectorReaderFactory {
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,
- )?;
+ // An empty selection permits nothing, so there is nothing to open.
+ if allowed_rows.is_some_and(|ranges| ranges.is_empty()) {
+ return Ok(vec![Vec::new(); queries.len()]);
+ }
+ let mut stream = match allowed_rows {
+ Some(ranges) => self.reader.read_single_file_stream_local_ranges(
+ &self.data_split,
+ file_meta,
+ data_fields,
+ None,
+ ranges.to_vec(),
+ )?,
+ None => self.reader.read_single_file_stream(
+ &self.data_split,
+ file_meta,
+ data_fields,
+ None,
+ None,
+ )?,
+ };
+ // Rows arrive in ascending physical order and the read emits exactly
what
+ // was selected (no pushdown predicate, no deletion vector here), so
walking
+ // the selection in step with the rows gives each row its file-local
+ // position: the whole file when nothing restricted it.
+ let mut selected: Box<dyn Iterator<Item = i64> + Send> = match
allowed_rows {
+ Some(ranges) => {
+ let ranges: Vec<RowRange> = ranges.into();
+ Box::new(
+ ranges
+ .into_iter()
+ .flat_map(|range| range.from()..=range.to()),
+ )
+ }
+ None => Box::new(0..row_count.max(0)),
+ };
let mut heaps: Vec<BinaryHeap<WorstFirst>> = (0..queries.len())
.map(|_| BinaryHeap::with_capacity(exact_limit + 1))
.collect();
// One reused buffer per batch; a NULL row leaves it untouched (and is
not
- // scored). `position` is the monotonic physical row counter across
batches.
+ // scored) but still consumes its physical position.
let mut batch_vectors: Vec<Option<Vec<f32>>> = Vec::new();
- let mut position: i64 = 0;
while let Some(batch) = stream.try_next().await? {
batch_vectors.clear();
append_batch_vectors(
@@ -153,13 +190,11 @@ impl DataFilePkVectorReaderFactory {
&mut batch_vectors,
)?;
for entry in &batch_vectors {
- let pos = position;
- position += 1;
- if pos >= row_count {
+ let Some(pos) = selected.next() else {
return Err(data_invalid(
- "data file produced more rows than
DataFileMeta.row_count",
+ "data file produced more rows than the selection
allows",
));
- }
+ };
let Some(vector) = entry else {
continue; // NULL row: not scored, position already
advanced.
};
@@ -177,14 +212,10 @@ impl DataFilePkVectorReaderFactory {
}
}
- if position > row_count {
+ // The overrun side is caught above, when the selection runs dry
mid-batch.
+ if selected.next().is_some() {
return Err(data_invalid(
- "data file produced more rows than DataFileMeta.row_count",
- ));
- }
- if position < row_count {
- return Err(data_invalid(
- "data file ended before DataFileMeta.row_count",
+ "data file ended before the selection was exhausted",
));
}
@@ -431,6 +462,72 @@ mod integration_tests {
/// The streaming per-file search must produce candidates byte-identical to
/// the reference `exact_search` over an in-memory `ArrayReader` of the
same
/// data, including a NULL row and a residual/DV exclusion.
+ #[tokio::test]
+ async fn test_search_file_only_reads_the_rows_the_plan_allows() {
+ // Five rows; the plan allows physical positions 3-4 only. The nearest
rows to
+ // the origin are 0 and 1, so a result of 3 and 4 means the read never
saw
+ // them — the exact fallback does not need a data predicate to be
handed a
+ // narrow split, and reading the rest of the file to reject it
afterwards is
+ // what this avoids.
+ //
+ // A full read cannot pass either: positions come from the selection,
so five
+ // emitted rows against a two-row selection fails loudly.
+ let rows = vec![
+ Some(vec![0.0, 0.0]),
+ Some(vec![0.5, 0.0]),
+ Some(vec![5.0, 0.0]),
+ Some(vec![6.0, 0.0]),
+ Some(vec![7.0, 0.0]),
+ ];
+ let (factory, file_name) =
+ build_factory(&rows, rows.len() as i64,
"memory:/pkvdfr_plan_ranges").await;
+ let active = BucketActiveFile {
+ file_name: file_name.clone(),
+ row_count: rows.len() as i64,
+ };
+ let query = [0.0f32, 0.0];
+
+ let results = factory
+ .search_file(
+ &active,
+ &[&query],
+ VectorSearchMetric::L2,
+ 5,
+ &|_| false,
+ Some(&[RowRange::new(3, 4)]),
+ )
+ .await
+ .unwrap();
+ let positions: Vec<i64> = results[0].iter().map(|hit|
hit.row_position).collect();
+ assert_eq!(positions, vec![3, 4]);
+ }
+
+ #[tokio::test]
+ async fn test_search_file_reads_nothing_for_an_empty_selection() {
+ let rows = vec![Some(vec![0.0, 0.0]), Some(vec![1.0, 0.0])];
+ let (factory, file_name) =
+ build_factory(&rows, rows.len() as i64,
"memory:/pkvdfr_empty_selection").await;
+ let active = BucketActiveFile {
+ file_name,
+ row_count: rows.len() as i64,
+ };
+ let query = [0.0f32, 0.0];
+
+ let results = factory
+ .search_file(
+ &active,
+ &[&query],
+ VectorSearchMetric::L2,
+ 5,
+ &|_| false,
+ Some(&[]),
+ )
+ .await
+ .unwrap();
+ assert_eq!(results.len(), 1, "one query in, one result list out");
+ assert!(results[0].is_empty());
+ }
+
#[tokio::test]
async fn search_file_matches_exact_search_reference() {
use crate::vindex::pkvector::exact::exact_search;
@@ -454,7 +551,14 @@ mod integration_tests {
let query = [0.0f32, 0.0];
let streamed = factory
- .search_file(&active, &[&query], VectorSearchMetric::L2, 2,
&is_excluded)
+ .search_file(
+ &active,
+ &[&query],
+ VectorSearchMetric::L2,
+ 2,
+ &is_excluded,
+ None,
+ )
.await
.unwrap();
@@ -498,7 +602,14 @@ mod integration_tests {
let query = [0.0f32, 0.0];
let streamed = factory
- .search_file(&active, &[&query], VectorSearchMetric::L2, 2, &|_|
false)
+ .search_file(
+ &active,
+ &[&query],
+ VectorSearchMetric::L2,
+ 2,
+ &|_| false,
+ None,
+ )
.await
.unwrap();
@@ -549,6 +660,7 @@ mod integration_tests {
VectorSearchMetric::L2,
2,
&is_excluded,
+ None,
)
.await
.unwrap();
@@ -556,11 +668,25 @@ mod integration_tests {
// Each query searched alone must equal its slot in the batch.
let only_q0 = factory
- .search_file(&active, &[&q0], VectorSearchMetric::L2, 2,
&is_excluded)
+ .search_file(
+ &active,
+ &[&q0],
+ VectorSearchMetric::L2,
+ 2,
+ &is_excluded,
+ None,
+ )
.await
.unwrap();
let only_q1 = factory
- .search_file(&active, &[&q1], VectorSearchMetric::L2, 2,
&is_excluded)
+ .search_file(
+ &active,
+ &[&q1],
+ VectorSearchMetric::L2,
+ 2,
+ &is_excluded,
+ None,
+ )
.await
.unwrap();
assert_eq!(batch[0], only_q0[0]);
@@ -588,7 +714,14 @@ mod integration_tests {
};
let query = [0.0f32, 0.0];
let err = factory
- .search_file(&active, &[&query], VectorSearchMetric::L2, 2, &|_|
false)
+ .search_file(
+ &active,
+ &[&query],
+ VectorSearchMetric::L2,
+ 2,
+ &|_| false,
+ None,
+ )
.await
.expect_err("row-count truncation must fail loud");
assert!(err.to_string().contains("ended before"), "got: {err}");
@@ -609,7 +742,14 @@ mod integration_tests {
// Wrong dimension.
let bad_dim = [1.0f32];
let err = factory
- .search_file(&present, &[&bad_dim], VectorSearchMetric::L2, 2,
&|_| false)
+ .search_file(
+ &present,
+ &[&bad_dim],
+ VectorSearchMetric::L2,
+ 2,
+ &|_| false,
+ None,
+ )
.await
.expect_err("dimension mismatch must fail loud");
assert!(err.to_string().contains("dimension"), "got: {err}");
@@ -617,9 +757,14 @@ mod integration_tests {
// Non-finite element.
let bad_finite = [f32::NAN, 0.0];
let err = factory
- .search_file(&present, &[&bad_finite], VectorSearchMetric::L2, 2,
&|_| {
- false
- })
+ .search_file(
+ &present,
+ &[&bad_finite],
+ VectorSearchMetric::L2,
+ 2,
+ &|_| false,
+ None,
+ )
.await
.expect_err("non-finite query must fail loud");
assert!(err.to_string().contains("finite"), "got: {err}");
@@ -631,7 +776,14 @@ mod integration_tests {
};
let query = [0.0f32, 0.0];
let err = factory
- .search_file(&missing, &[&query], VectorSearchMetric::L2, 2, &|_|
false)
+ .search_file(
+ &missing,
+ &[&query],
+ VectorSearchMetric::L2,
+ 2,
+ &|_| false,
+ None,
+ )
.await
.expect_err("absent file must be rejected");
assert!(matches!(err, crate::Error::DataInvalid { .. }));
diff --git a/crates/paimon/src/table/pk_vector_scan.rs
b/crates/paimon/src/table/pk_vector_scan.rs
index 8521aad6..21d70d1d 100644
--- a/crates/paimon/src/table/pk_vector_scan.rs
+++ b/crates/paimon/src/table/pk_vector_scan.rs
@@ -20,15 +20,22 @@
//! search split per bucket. Mirror of Java `PrimaryKeyVectorScan` and
//! `PrimaryKeyIndexSourcePolicy`.
-use std::collections::{BTreeMap, HashSet};
+use std::collections::{BTreeMap, HashMap, HashSet};
+
+use indexmap::IndexMap;
+
+use roaring::RoaringTreemap;
use crate::spec::{
should_read_pk_index_source, BinaryRow, DataFileMeta, FileKind,
IndexManifest, Predicate,
PrimaryKeyIndexSourceFile, PrimaryKeyIndexSourceMeta,
};
+use crate::table::bucket_filter::split_partition_and_data_predicates;
use crate::table::index_file_path::IndexFileLocation;
+use crate::table::partition_filter::PartitionFilter;
+use crate::table::pk_vector_bucket_split::BucketVectorSearchSplit;
use crate::table::pk_vector_orchestrator::PkVectorSearchSplit;
-use crate::table::source::{DataSplit, DataSplitBuilder, DeletionFile};
+use crate::table::source::{merge_row_ranges, DataSplit, DataSplitBuilder,
DeletionFile, RowRange};
use crate::table::Table;
use crate::vindex::pkvector::bucket::{BucketActiveFile, BucketAnnSegment};
@@ -42,6 +49,41 @@ struct UnresolvedAnnSegment {
index_meta: Vec<u8>,
}
+/// A bucket's identity across planning inputs: the partition's serialized
bytes
+/// (`BinaryRow` is not hashable) paired with the bucket number.
+type BucketKey = (Vec<u8>, i32);
+
+/// Expand inclusive row ranges into the positions they allow. Only the search
+/// kernel's membership tests need this; a read is limited by the ranges
themselves.
+pub(super) fn positions_in_ranges(ranges: &[RowRange]) ->
crate::Result<RoaringTreemap> {
+ let mut positions = RoaringTreemap::new();
+ for range in ranges {
+ let from = u64::try_from(range.from())
+ .map_err(|_| data_invalid("row range bound must not be
negative"))?;
+ let to = u64::try_from(range.to())
+ .map_err(|_| data_invalid("row range bound must not be
negative"))?;
+ positions.insert_range(from..=to);
+ }
+ Ok(positions)
+}
+
+/// The whole of a file, for a file the message left unrestricted.
+///
+/// A zero-row file gets no range rather than an empty one: `RowRange` is
inclusive,
+/// so it cannot express "nothing". An unknown count
(`DataFileMeta::ROW_COUNT_UNKNOWN`,
+/// or anything else negative) is rejected rather than read as "nothing",
which would
+/// silently drop the file from the search. Decoding only checks the count of
a file
+/// the message lists ranges for, so this is where an omitted one is checked.
+fn whole_file_range(row_count: i64) -> crate::Result<Vec<RowRange>> {
+ match row_count {
+ 0 => Ok(Vec::new()),
+ count if count > 0 => Ok(vec![RowRange::new(0, count - 1)]),
+ count => Err(data_invalid(format!(
+ "data file row count must be known and non-negative, got {count}"
+ ))),
+ }
+}
+
fn data_invalid(message: impl Into<String>) -> crate::Error {
crate::Error::DataInvalid {
message: message.into(),
@@ -206,6 +248,17 @@ pub(crate) struct PkVectorScanPlan {
// at all (never written), which also yields empty `splits`.
pub snapshot_id: i64,
pub splits: Vec<PkVectorSearchSplit>,
+ // Per-split allow-list of physical rows, indexed parallel to `splits`:
only the
+ // rows listed for a data file may produce candidates from it. Ranges
rather than
+ // materialized positions, because this is what a read is limited to —
expanding
+ // a whole-file range of a large file into positions costs memory no
reader needs.
+ // Each list is normalized: sorted, non-overlapping, inclusive, file-local.
+ // Populated when the plan was built from engine-supplied bucket splits,
which
+ // carry row ranges the engine's own planner already resolved. `None` for
a plan
+ // read from this table's index manifest, which places no positional
restriction
+ // of its own -- distinct from `Some` of an empty allow-list, which permits
+ // nothing.
+ pub physical_row_ranges_by_split: Option<Vec<HashMap<String,
Vec<RowRange>>>>,
}
pub(crate) struct PkVectorScan<'a> {
@@ -271,6 +324,7 @@ impl<'a> PkVectorScan<'a> {
return Ok(PkVectorScanPlan {
snapshot_id: 0,
splits: Vec::new(),
+ physical_row_ranges_by_split: None,
});
};
let snapshot = snapshot_manager.get_snapshot(snapshot_id).await?;
@@ -344,8 +398,223 @@ impl<'a> PkVectorScan<'a> {
Ok(PkVectorScanPlan {
snapshot_id,
splits,
+ physical_row_ranges_by_split: None,
})
}
+
+ /// Build a plan from bucket splits an engine planned elsewhere, instead
of from
+ /// this table's index manifest.
+ ///
+ /// The splits are the planning input and are taken as authoritative: their
+ /// payload files, their per-file row ranges, and the snapshot they pin
are used
+ /// as given, and no index manifest is read. Only the partition conjuncts
of this
+ /// scan's filter are re-applied, because a caller may narrow the query
further
+ /// than the planner that produced the splits.
+ ///
+ /// Mirrors what Java's `PrimaryKeyVectorRead` does with a
+ /// `BucketVectorSearchSplit`: search the payloads the split names, over
the rows
+ /// the split allows.
+ // Entry point for engine-supplied splits; no in-tree caller reads a plan
from
+ // them yet, and the tests drive `plan_from_bucket_splits` directly.
+ #[allow(dead_code)]
+ pub(crate) fn plan_for_bucket_vector_splits(
+ &self,
+ splits: Vec<BucketVectorSearchSplit>,
+ ) -> crate::Result<PkVectorScanPlan> {
+ // Partition conjuncts only. Data conjuncts stay a per-row residual
applied
+ // during the search: pruning a whole bucket on them would drop rows
that
+ // still match.
+ let partition_filter = self.filter.as_ref().and_then(|filter| {
+ let (partition_predicate, _data_predicates) =
split_partition_and_data_predicates(
+ filter.clone(),
+ self.table.schema().fields(),
+ self.table.schema().partition_keys(),
+ );
+ partition_predicate.map(|predicate| {
+ PartitionFilter::from_predicate(predicate,
&self.table.schema().partition_fields())
+ })
+ });
+ plan_from_bucket_splits(
+ &self.index_type,
+ self.vector_field_id,
+ partition_filter.as_ref(),
+ self.table.location().trim_end_matches('/'),
+ self.table
+ .schema()
+ .core_options()
+ .index_file_in_data_file_dir(),
+ splits,
+ )
+ }
+}
+
+/// The `Table`-independent core of
[`PkVectorScan::plan_for_bucket_vector_splits`],
+/// so planning from engine-supplied splits is testable the same way planning
from a
+/// manifest is.
+#[cfg_attr(not(test), allow(dead_code))]
+fn plan_from_bucket_splits(
+ index_type: &str,
+ vector_field_id: i32,
+ partition_filter: Option<&PartitionFilter>,
+ table_path: &str,
+ index_file_in_data_file_dir: bool,
+ splits: Vec<BucketVectorSearchSplit>,
+) -> crate::Result<PkVectorScanPlan> {
+ // A plan's snapshot id stays authoritative even when nothing is
searchable, and
+ // empty input pins no snapshot to report. Reject rather than invent one.
+ if splits.is_empty() {
+ return Err(data_invalid(
+ "bucket-split planning requires at least one bucket split",
+ ));
+ }
+
+ let mut snapshot_id: Option<i64> = None;
+ let mut seen_buckets: HashSet<BucketKey> = HashSet::new();
+ let mut data_splits: Vec<DataSplit> = Vec::with_capacity(splits.len());
+ let mut index_entries: Vec<(BinaryRow, i32, UnresolvedAnnSegment)> =
Vec::new();
+ let mut listed_ranges: HashMap<BucketKey, IndexMap<String, Vec<RowRange>>>
= HashMap::new();
+
+ for split in splits {
+ let (data_split, payload_files, row_ranges_by_file) =
split.into_parts();
+
+ // Row ranges belong to the bucket form, one list per data file. A
nested
+ // split carrying its own would be a second authority over which
physical
+ // rows are readable, free to disagree with the first. Java's planner
+ // builds the nested split without them.
+ if data_split.row_ranges().is_some() {
+ return Err(data_invalid(
+ "a bucket split's nested data split must not carry row ranges",
+ ));
+ }
+
+ // One snapshot across every split: candidates found under different
+ // snapshots cannot be merged into a single Top-K. Checked before
pruning,
+ // so a mismatch is reported even when the offending split would have
been
+ // pruned away and the inconsistency left no trace.
+ match snapshot_id {
+ None => snapshot_id = Some(data_split.snapshot_id()),
+ Some(pinned) if pinned != data_split.snapshot_id() => {
+ return Err(data_invalid(format!(
+ "bucket splits pin different snapshots: {} and {}",
+ pinned,
+ data_split.snapshot_id()
+ )));
+ }
+ Some(_) => {}
+ }
+
+ // Java emits exactly one split per (partition, bucket). Buffers
decoded
+ // independently cannot enforce that between them, and two splits for
one
+ // bucket would search its rows twice.
+ let key: BucketKey = (
+ data_split.partition().to_serialized_bytes(),
+ data_split.bucket(),
+ );
+ if !seen_buckets.insert(key.clone()) {
+ return Err(data_invalid(format!(
+ "bucket splits repeat bucket {} of one partition",
+ data_split.bucket()
+ )));
+ }
+
+ if let Some(filter) = partition_filter {
+ if !filter.matches_entry(&key.0)? {
+ continue;
+ }
+ }
+
+ for payload in payload_files {
+ let parts = payload.into_parts();
+ // The same three filters the manifest route applies: this column's
+ // index type, this column's field id, and a payload that carries
the
+ // source metadata a search needs to map ordinals back to rows.
+ if parts.index_type != index_type
+ || parts.global_index_meta.index_field_id != vector_field_id
+ || parts.global_index_meta.source_meta.is_none()
+ {
+ continue;
+ }
+ // Java writes the size as a signed long, so the wire allows a
+ // negative value the segment addressing cannot represent.
+ let file_size = u64::try_from(parts.file_size)
+ .map_err(|_| data_invalid("index file size must not be
negative"))?;
+ let source_meta =
+
PrimaryKeyIndexSourceMeta::from_global_index_meta(&parts.global_index_meta)
+ .map_err(|_| {
+ data_invalid(format!("index file {} is not active",
parts.file_name))
+ })?;
+ let index_meta = parts
+ .global_index_meta
+ .index_meta
+ .clone()
+ .unwrap_or_default();
+ index_entries.push((
+ data_split.partition().clone(),
+ data_split.bucket(),
+ UnresolvedAnnSegment {
+ source_meta,
+ file_name: parts.file_name,
+ external_path: parts.external_path,
+ file_size,
+ index_meta,
+ },
+ ));
+ }
+
+ listed_ranges.insert(key, row_ranges_by_file);
+ data_splits.push(data_split);
+ }
+
+ // Non-empty input always pins one: the first split sets it and a mismatch
+ // returns early.
+ let snapshot_id = snapshot_id.expect("non-empty bucket-split input pins a
snapshot");
+
+ let splits = plan_from_inputs(
+ snapshot_id,
+ data_splits,
+ index_entries,
+ table_path,
+ index_file_in_data_file_dir,
+ )?;
+
+ // Normalize the row ranges against the planned splits, which are grouped
by
+ // bucket and so may be ordered differently from the input.
+ //
+ // A file the message lists is restricted to the positions it lists. A
file it
+ // omits is unrestricted: Java records ranges only for the files its own
+ // pre-filter narrowed, and leaves the rest out. The search kernel reads a
+ // missing entry as "no rows allowed", the opposite meaning, so the
omission
+ // has to be turned into an explicit full-file range here rather than
passed
+ // through.
+ let physical_row_ranges_by_split = splits
+ .iter()
+ .map(|split| {
+ let listed = listed_ranges.get(&(
+ split.data_split.partition().to_serialized_bytes(),
+ split.data_split.bucket(),
+ ));
+ split
+ .data_split
+ .data_files()
+ .iter()
+ .map(|file| {
+ let allowed = match listed.and_then(|ranges|
ranges.get(&file.file_name)) {
+ // Decoding checks each range's bounds but not their
order or
+ // whether they overlap, and a read needs them
normalized.
+ Some(ranges) => merge_row_ranges(ranges.clone()),
+ None => whole_file_range(file.row_count)?,
+ };
+ Ok((file.file_name.clone(), allowed))
+ })
+ .collect::<crate::Result<HashMap<String, Vec<RowRange>>>>()
+ })
+ .collect::<crate::Result<Vec<_>>>()?;
+
+ Ok(PkVectorScanPlan {
+ snapshot_id,
+ splits,
+ physical_row_ranges_by_split: Some(physical_row_ranges_by_split),
+ })
}
/// Pure planning core, drivable without a live snapshot: group ANN payloads
and
@@ -425,6 +694,8 @@ mod tests {
use super::*;
use crate::spec::stats::BinaryTableStats;
use crate::spec::{BinaryRow, DataFileMeta, GlobalIndexMeta};
+ use crate::spec::{DataField, DeletionVectorMeta};
+ use crate::table::pk_vector_bucket_split::BucketVectorPayload;
use crate::table::source::{DataSplitBuilder, DeletionFile};
fn dfm(name: &str, rows: i64, level: i32, file_source: Option<i32>) ->
DataFileMeta {
@@ -992,4 +1263,404 @@ mod tests {
);
}
}
+
+ // ---- planning from engine-supplied bucket splits ----
+
+ const BUCKET_SPLIT_GOLDEN: &[u8] =
include_bytes!("goldens/bucket_vector_search_split_v1.bin");
+
+ fn int_partition(value: i32) -> BinaryRow {
+ let mut builder = crate::spec::BinaryRowBuilder::new(1);
+ builder.write_int(0, value);
+ BinaryRow::from_serialized_bytes(&builder.build_serialized()).unwrap()
+ }
+
+ /// A bucket split as an engine would hand one over. Its data files are
COMPACT
+ /// above level 0, so they are exact-fallback eligible, and the caller's
payload
+ /// metadata is expected to name exactly that level's source set.
+ fn engine_split(
+ snapshot: i64,
+ bucket: i32,
+ partition: BinaryRow,
+ files: Vec<DataFileMeta>,
+ payloads: Vec<BucketVectorPayload>,
+ ranges: &[(&str, &[(i64, i64)])],
+ ) -> BucketVectorSearchSplit {
+ let data_split = DataSplitBuilder::new()
+ .with_snapshot(snapshot)
+ .with_partition(partition)
+ .with_bucket(bucket)
+ .with_bucket_path(format!("bucket-{bucket}"))
+ .with_total_buckets(1)
+ .with_data_files(files)
+ .build()
+ .unwrap();
+ BucketVectorSearchSplit::new_for_test(
+ data_split,
+ payloads,
+ ranges
+ .iter()
+ .map(|(name, bounds)| {
+ (
+ (*name).to_string(),
+ bounds
+ .iter()
+ .map(|(from, to)| RowRange::new(*from, *to))
+ .collect(),
+ )
+ })
+ .collect(),
+ )
+ }
+
+ fn engine_payload(meta: GlobalIndexMeta) -> BucketVectorPayload {
+ BucketVectorPayload::new_for_test("ivf-pq", "seg0", 1, 4, None, None,
meta)
+ }
+
+ /// One bucket holding `d0`, with a payload whose source set matches it.
+ fn one_file_split(
+ snapshot: i64,
+ bucket: i32,
+ ranges: &[(&str, &[(i64, i64)])],
+ ) -> BucketVectorSearchSplit {
+ engine_split(
+ snapshot,
+ bucket,
+ BinaryRow::new(0),
+ vec![dfm("d0", 4, 5, Some(1))],
+ vec![engine_payload(gim(2, 5, &[("d0", 4)]))],
+ ranges,
+ )
+ }
+
+ /// The positions a file's normalized ranges allow, for assertions that
read
+ /// better as a row list than as ranges.
+ fn allowed(map: &HashMap<String, Vec<RowRange>>, file: &str) -> Vec<u64> {
+ map.get(file)
+ .map(|ranges| {
+ positions_in_ranges(ranges)
+ .expect("planned ranges are in range")
+ .iter()
+ .collect()
+ })
+ .unwrap_or_default()
+ }
+
+ #[test]
+ fn plans_the_java_golden_bucket_split() {
+ let split =
BucketVectorSearchSplit::deserialize(BUCKET_SPLIT_GOLDEN).unwrap();
+ let plan = plan_from_bucket_splits("ivf-pq", 7, None, "/tbl", false,
vec![split]).unwrap();
+
+ // The snapshot the split pins, not one re-resolved from the table.
+ assert_eq!(plan.snapshot_id, 11);
+ assert_eq!(plan.splits.len(), 1);
+ let planned = &plan.splits[0];
+
+ // The payload's own external path wins over both directory layouts.
+ assert_eq!(planned.ann_segments.len(), 1);
+ assert_eq!(planned.ann_segments[0].path,
"s3://vector-bucket/ann-0.idx");
+ assert_eq!(planned.ann_segments[0].file_size, 5_000_000_000);
+ assert_eq!(planned.ann_segments[0].source_meta.data_level(), 1);
+
+ // `data-1.orc` is COMPACT above level 0, so exact fallback may read
it.
+ assert_eq!(planned.active_files.len(), 1);
+ assert_eq!(planned.active_files[0].file_name, "data-1.orc");
+
+ // The message allows rows 0-1 and 4-5 of a six-row file.
+ let ranges = plan
+ .physical_row_ranges_by_split
+ .expect("a split-driven plan restricts positions");
+ assert_eq!(allowed(&ranges[0], "data-1.orc"), vec![0, 1, 4, 5]);
+ }
+
+ #[test]
+ fn rejects_an_unknown_row_count_on_an_unlisted_file() {
+ // A file the message lists no ranges for is read as "the whole file",
which
+ // needs a real row count. `ROW_COUNT_UNKNOWN` is -1, and reading that
as "no
+ // rows" would drop the file from the search without a word; the
decoder only
+ // checks the count of files it does carry ranges for.
+ let error = whole_file_range(DataFileMeta::ROW_COUNT_UNKNOWN)
+ .map(|_| ())
+ .expect_err("an unknown row count cannot stand in for the whole
file");
+ assert!(error.to_string().contains("must be known"), "{error}");
+ assert!(whole_file_range(0).unwrap().is_empty());
+ assert_eq!(whole_file_range(3).unwrap(), vec![RowRange::new(0, 2)]);
+ }
+
+ #[test]
+ fn rejects_empty_bucket_split_input() {
+ let error = plan_from_bucket_splits("ivf-pq", 2, None, "/tbl", false,
Vec::new())
+ .map(|_| ())
+ .expect_err("empty input pins no snapshot to report");
+ assert!(
+ error.to_string().contains("at least one bucket split"),
+ "{error}"
+ );
+ }
+
+ #[test]
+ fn rejects_bucket_splits_pinning_different_snapshots() {
+ let error = plan_from_bucket_splits(
+ "ivf-pq",
+ 2,
+ None,
+ "/tbl",
+ false,
+ vec![one_file_split(11, 0, &[]), one_file_split(12, 1, &[])],
+ )
+ .map(|_| ())
+ .expect_err("candidates from two snapshots cannot merge into one
Top-K");
+ assert!(
+ error.to_string().contains("pin different snapshots"),
+ "{error}"
+ );
+ }
+
+ #[test]
+ fn rejects_two_splits_for_one_bucket() {
+ let error = plan_from_bucket_splits(
+ "ivf-pq",
+ 2,
+ None,
+ "/tbl",
+ false,
+ vec![one_file_split(11, 0, &[]), one_file_split(11, 0, &[])],
+ )
+ .map(|_| ())
+ .expect_err("one bucket twice would search its rows twice");
+ assert!(error.to_string().contains("repeat bucket 0"), "{error}");
+ }
+
+ #[test]
+ fn rejects_nested_data_split_carrying_row_ranges() {
+ let data_split = DataSplitBuilder::new()
+ .with_snapshot(11)
+ .with_partition(BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path("bucket-0".to_string())
+ .with_total_buckets(1)
+ .with_data_files(vec![dfm("d0", 4, 5, Some(1))])
+ .with_row_ranges(vec![RowRange::new(0, 1)])
+ .build()
+ .unwrap();
+ let split = BucketVectorSearchSplit::new_for_test(
+ data_split,
+ vec![engine_payload(gim(2, 5, &[("d0", 4)]))],
+ IndexMap::new(),
+ );
+ let error = plan_from_bucket_splits("ivf-pq", 2, None, "/tbl", false,
vec![split])
+ .map(|_| ())
+ .expect_err("two row-range authorities may disagree");
+ assert!(
+ error.to_string().contains("must not carry row ranges"),
+ "{error}"
+ );
+ }
+
+ #[test]
+ fn unlisted_file_is_unrestricted_and_an_empty_list_excludes_one() {
+ // Java records ranges only for the files its own pre-filter narrowed,
so an
+ // omitted file means "all rows". An explicitly empty list means "no
rows",
+ // and the two must not collapse into each other.
+ let split = engine_split(
+ 11,
+ 0,
+ BinaryRow::new(0),
+ vec![dfm("d0", 4, 5, Some(1)), dfm("d1", 3, 5, Some(1))],
+ vec![engine_payload(gim(2, 5, &[("d0", 4), ("d1", 3)]))],
+ &[("d0", &[])],
+ );
+ let plan = plan_from_bucket_splits("ivf-pq", 2, None, "/tbl", false,
vec![split]).unwrap();
+ let ranges = plan
+ .physical_row_ranges_by_split
+ .expect("split-driven plan");
+
+ assert!(allowed(&ranges[0], "d0").is_empty());
+ assert_eq!(allowed(&ranges[0], "d1"), vec![0, 1, 2]);
+ }
+
+ #[test]
+ fn rejects_negative_payload_file_size() {
+ let split = engine_split(
+ 11,
+ 0,
+ BinaryRow::new(0),
+ vec![dfm("d0", 4, 5, Some(1))],
+ vec![BucketVectorPayload::new_for_test(
+ "ivf-pq",
+ "seg0",
+ -1,
+ 4,
+ None,
+ None,
+ gim(2, 5, &[("d0", 4)]),
+ )],
+ &[],
+ );
+ let error = plan_from_bucket_splits("ivf-pq", 2, None, "/tbl", false,
vec![split])
+ .map(|_| ())
+ .expect_err("a signed wire size can be negative, segment
addressing cannot");
+ assert!(
+ error.to_string().contains("must not be negative"),
+ "{error}"
+ );
+ }
+
+ #[test]
+ fn ignores_payload_deletion_vector_ranges() {
+ // The field belongs to deletion-vector index files. Java builds a
vector
+ // payload through the overload that leaves it null, and a read takes
its
+ // deletion vectors from the bucket's data split, so a value here
describes
+ // something this payload is not.
+ let mut dv = IndexMap::new();
+ dv.insert(
+ "d0".to_string(),
+ DeletionVectorMeta {
+ offset: 0,
+ length: 8,
+ cardinality: Some(1),
+ },
+ );
+ let split = engine_split(
+ 11,
+ 0,
+ BinaryRow::new(0),
+ vec![dfm("d0", 4, 5, Some(1))],
+ vec![BucketVectorPayload::new_for_test(
+ "ivf-pq",
+ "seg0",
+ 1,
+ 4,
+ Some(dv),
+ None,
+ gim(2, 5, &[("d0", 4)]),
+ )],
+ &[],
+ );
+ let plan = plan_from_bucket_splits("ivf-pq", 2, None, "/tbl", false,
vec![split]).unwrap();
+ assert_eq!(plan.splits.len(), 1);
+ assert_eq!(plan.splits[0].ann_segments.len(), 1);
+ let ranges = plan
+ .physical_row_ranges_by_split
+ .expect("split-driven plan");
+ // Unaffected: the whole file stays readable.
+ assert_eq!(allowed(&ranges[0], "d0"), vec![0, 1, 2, 3]);
+ }
+
+ #[test]
+ fn skips_payloads_for_another_column_or_index_type() {
+ let split = engine_split(
+ 11,
+ 0,
+ BinaryRow::new(0),
+ vec![dfm("d0", 4, 5, Some(1))],
+ vec![
+ // Another column's vector index.
+ engine_payload(gim(99, 5, &[("d0", 4)])),
+ // This column, but another index type.
+ BucketVectorPayload::new_for_test(
+ "flat",
+ "seg1",
+ 1,
+ 4,
+ None,
+ None,
+ gim(2, 5, &[("d0", 4)]),
+ ),
+ ],
+ &[],
+ );
+ let plan = plan_from_bucket_splits("ivf-pq", 2, None, "/tbl", false,
vec![split]).unwrap();
+ assert_eq!(plan.splits.len(), 1);
+ assert!(plan.splits[0].ann_segments.is_empty());
+ // Still exact-fallback eligible: no ANN segment covers the file.
+ assert_eq!(plan.splits[0].active_files.len(), 1);
+ }
+
+ fn partition_filter_on_dt(keep: i32) -> PartitionFilter {
+ let fields = vec![DataField::new(
+ 0,
+ "dt".to_string(),
+ crate::spec::DataType::Int(crate::spec::IntType::new()),
+ )];
+ let builder = crate::spec::PredicateBuilder::new(&fields);
+ let predicate = builder.equal("dt",
crate::spec::Datum::Int(keep)).unwrap();
+ PartitionFilter::from_predicate(predicate, &fields)
+ }
+
+ #[test]
+ fn snapshot_mismatch_is_rejected_before_partition_pruning() {
+ // Both splits are pruned by this filter. The mismatch must still be
reported:
+ // pruning first would hide an inconsistent input behind an empty plan.
+ let filter = partition_filter_on_dt(3);
+ let error = plan_from_bucket_splits(
+ "ivf-pq",
+ 2,
+ Some(&filter),
+ "/tbl",
+ false,
+ vec![
+ engine_split(
+ 11,
+ 0,
+ int_partition(1),
+ vec![dfm("d0", 4, 5, Some(1))],
+ vec![engine_payload(gim(2, 5, &[("d0", 4)]))],
+ &[],
+ ),
+ engine_split(
+ 12,
+ 1,
+ int_partition(2),
+ vec![dfm("d0", 4, 5, Some(1))],
+ vec![engine_payload(gim(2, 5, &[("d0", 4)]))],
+ &[],
+ ),
+ ],
+ )
+ .map(|_| ())
+ .expect_err("a snapshot mismatch outranks pruning");
+ assert!(
+ error.to_string().contains("pin different snapshots"),
+ "{error}"
+ );
+ }
+
+ #[test]
+ fn pruning_every_split_keeps_the_pinned_snapshot() {
+ let filter = partition_filter_on_dt(3);
+ let plan = plan_from_bucket_splits(
+ "ivf-pq",
+ 2,
+ Some(&filter),
+ "/tbl",
+ false,
+ vec![engine_split(
+ 11,
+ 0,
+ int_partition(1),
+ vec![dfm("d0", 4, 5, Some(1))],
+ vec![engine_payload(gim(2, 5, &[("d0", 4)]))],
+ &[],
+ )],
+ )
+ .unwrap();
+ assert!(plan.splits.is_empty());
+ // Still authoritative with nothing left to search.
+ assert_eq!(plan.snapshot_id, 11);
+ assert_eq!(
+ plan.physical_row_ranges_by_split.as_deref(),
+ Some([].as_slice())
+ );
+ }
+
+ #[test]
+ fn resolves_a_payload_without_an_external_path_into_the_bucket_directory()
{
+ let split = one_file_split(11, 0, &[]);
+ let plan = plan_from_bucket_splits("ivf-pq", 2, None, "/tbl", true,
vec![split]).unwrap();
+ assert_eq!(plan.splits[0].ann_segments[0].path, "bucket-0/seg0");
+
+ let split = one_file_split(11, 0, &[]);
+ let plan = plan_from_bucket_splits("ivf-pq", 2, None, "/tbl", false,
vec![split]).unwrap();
+ assert_eq!(plan.splits[0].ann_segments[0].path, "/tbl/index/seg0");
+ }
}
diff --git a/crates/paimon/src/table/vector_search_builder.rs
b/crates/paimon/src/table/vector_search_builder.rs
index 6bf11330..447d6995 100644
--- a/crates/paimon/src/table/vector_search_builder.rs
+++ b/crates/paimon/src/table/vector_search_builder.rs
@@ -44,7 +44,7 @@ use crate::table::pk_vector_orchestrator::{
use crate::table::pk_vector_position_read::{
PkVectorPositionRead, PKEY_VECTOR_POSITION_COLUMN, SEARCH_SCORE_COLUMN,
};
-use crate::table::pk_vector_scan::{PkVectorScan, PkVectorScanPlan};
+use crate::table::pk_vector_scan::{positions_in_ranges, PkVectorScan,
PkVectorScanPlan};
use crate::table::read_builder::resolve_projected_fields;
use crate::table::row_id_predicate::intersect_sorted_ranges;
use crate::table::source::DataSplit;
@@ -762,7 +762,26 @@ pub(crate) fn ensure_no_reserved_read_columns(fields:
&[DataField]) -> crate::Re
/// vector, so it is computed once and the SAME slice is shared across all
queries.
/// Rerank stays per-query (each query reranks its own indexed list).
#[allow(clippy::too_many_arguments)]
-async fn plan_and_search_pk_candidates_batch(
+/// Query-level parameters for a primary-key vector search: everything
resolvable
+/// from the table schema, the options and the queries alone, independent of
which
+/// splits planning yields. Resolved before planning so a malformed query or
option
+/// fails loud even when the plan turns out empty.
+struct PkVectorSearchParams {
+ metric: VectorSearchMetric,
+ /// Fan-out limit for bucket orchestration plus ANN and exact-file leaves
(Java
+ /// `GLOBAL_INDEX_THREAD_NUM`); `1` reproduces strictly sequential
execution.
+ concurrency: usize,
+ index_type: String,
+ field_id: i32,
+ vector_field: DataField,
+ skip_exact_fallback: bool,
+ refine_factor: usize,
+ indexed_limit: usize,
+}
+
+/// Resolve the query-level parameters and reject a query the search cannot
answer
+/// correctly, before any planning or read happens.
+fn resolve_pk_vector_search_params(
table: &Table,
query_options: &HashMap<String, String>,
filter: Option<&Predicate>,
@@ -770,11 +789,7 @@ async fn plan_and_search_pk_candidates_batch(
pk_col: &str,
queries: &[&[f32]],
limit: usize,
-) -> crate::Result<(
- Vec<Vec<PkVectorCandidate>>,
- PkVectorScanPlan,
- VectorSearchMetric,
-)> {
+) -> crate::Result<PkVectorSearchParams> {
// Residual pre-filter guard, mirroring Java `PrimaryKeyVectorScan`. A DATA
// predicate set via `with_filter` is applied post-recall by re-reading
each
// candidate file's physical rows (see below). That physical-position
filtering
@@ -876,13 +891,128 @@ async fn plan_and_search_pk_candidates_batch(
}
}
- let plan = PkVectorScan::new(table, field_id, index_type.clone(),
filter.cloned())
- .plan()
- .await?;
+ Ok(PkVectorSearchParams {
+ metric,
+ concurrency,
+ index_type,
+ field_id,
+ vector_field,
+ skip_exact_fallback,
+ refine_factor,
+ indexed_limit,
+ })
+}
+
+/// Search an already-resolved plan across every query and return each query's
raw
+/// indexed and exact candidate lists, before any rerank or merge.
+///
+/// Plan-dependent concurrency — the vindex segment count, batch-index
parallelism
+/// and the range-read bound — is derived here from the plan that is actually
being
+/// searched, so a narrowed plan can never be searched under limits computed
for a
+/// wider one.
+/// Combine the two per-split row allow-lists a search can be handed: the
physical
+/// rows an engine-supplied plan restricts each file to, and the positions a
residual
+/// data predicate leaves behind.
+///
+/// Both sides list what is permitted, and both read a file's absence as "no
rows
+/// allowed", so combining them intersects files as well as positions. Either
side
+/// alone passes through unchanged; neither side means no positional
restriction.
+///
+/// This is where the plan's ranges become positions: the search kernel tests
+/// membership, while a read is limited by the ranges themselves. When the
residual
+/// was evaluated over those same ranges the intersection cannot remove
anything, and
+/// is kept as the invariant that says so.
+fn intersect_row_allow_lists(
+ physical: Option<&[HashMap<String, Vec<RowRange>>]>,
+ residual: Option<Vec<HashMap<String, RoaringTreemap>>>,
+ split_count: usize,
+) -> crate::Result<Option<Vec<HashMap<String, RoaringTreemap>>>> {
+ if let Some(maps) = physical {
+ if maps.len() != split_count {
+ return Err(crate::Error::DataInvalid {
+ message: format!(
+ "plan carries {} physical row allow-lists for
{split_count} splits",
+ maps.len()
+ ),
+ source: None,
+ });
+ }
+ }
+ match (physical, residual) {
+ (None, residual) => Ok(residual),
+ (Some(physical), None) => Ok(Some(
+ physical
+ .iter()
+ .map(|per_file| {
+ per_file
+ .iter()
+ .map(|(file, ranges)| Ok((file.clone(),
positions_in_ranges(ranges)?)))
+ .collect::<crate::Result<HashMap<String,
RoaringTreemap>>>()
+ })
+ .collect::<crate::Result<Vec<_>>>()?,
+ )),
+ (Some(physical), Some(residual)) => {
+ if residual.len() != split_count {
+ return Err(crate::Error::DataInvalid {
+ message: format!(
+ "residual carries {} row allow-lists for {split_count}
splits",
+ residual.len()
+ ),
+ source: None,
+ });
+ }
+ Ok(Some(
+ physical
+ .iter()
+ .zip(residual)
+ .map(|(physical, residual)| {
+ physical
+ .iter()
+ .filter(|(file, _)|
residual.contains_key(file.as_str()))
+ .map(|(file, ranges)| {
+ let allowed = positions_in_ranges(ranges)?;
+ let kept = &residual[file.as_str()];
+ Ok((file.clone(), allowed & kept))
+ })
+ .collect::<crate::Result<HashMap<String,
RoaringTreemap>>>()
+ })
+ .collect::<crate::Result<Vec<_>>>()?,
+ ))
+ }
+ }
+}
+
+#[allow(clippy::too_many_arguments)]
+async fn search_pk_raw_candidates_batch_with_plan(
+ table: &Table,
+ query_options: &HashMap<String, String>,
+ filter: Option<&Predicate>,
+ core: &CoreOptions<'_>,
+ pk_col: &str,
+ queries: &[&[f32]],
+ limit: usize,
+ plan: &PkVectorScanPlan,
+ params: &PkVectorSearchParams,
+) -> crate::Result<Vec<OrchestratorSearchResult>> {
+ // An empty plan has nothing to search. Returned before the backend is
resolved
+ // so a table with no searchable data never errors on an unrecognized
index type.
if plan.splits.is_empty() {
- return Ok((vec![Vec::new(); queries.len()], plan, metric));
+ return Ok(queries
+ .iter()
+ .map(|_| OrchestratorSearchResult {
+ indexed: Vec::new(),
+ exact: Vec::new(),
+ })
+ .collect());
}
+ let metric = params.metric;
+ let concurrency = params.concurrency;
+ let index_type = params.index_type.clone();
+ let vector_field = params.vector_field.clone();
+ let skip_exact_fallback = params.skip_exact_fallback;
+ let indexed_limit = params.indexed_limit;
+
// Resolve the vector index backend from the single configured index type.
// Java enforces one index type per PK table and Rust filters segments to
it,
// so one backend serves every segment. Computed after the empty-plan
return so
@@ -1073,13 +1203,21 @@ async fn plan_and_search_pk_candidates_batch(
Vec::new(),
);
let mut per_split = Vec::with_capacity(plan.splits.len());
- for split in &plan.splits {
+ for (index, split) in plan.splits.iter().enumerate() {
+ // The plan's selection for this split, so the residual is
+ // evaluated over the rows an engine-supplied split allows
rather
+ // than over the whole file.
+ let allowed_rows = plan
+ .physical_row_ranges_by_split
+ .as_ref()
+ .and_then(|per_split| per_split.get(index));
per_split.push(
residual_positions_by_file(
&residual_reader,
&split.data_split,
&split.active_files,
&file_predicates,
+ allowed_rows,
)
.await?,
);
@@ -1089,6 +1227,15 @@ async fn plan_and_search_pk_candidates_batch(
}
None => None,
};
+ // Fold the plan's own positional restriction into the same allow-list. A
plan
+ // built from engine-supplied bucket splits carries the physical positions
each
+ // file is limited to; a plan read from the index manifest carries none.
Both
+ // sides list what is permitted, so combining them is an intersection.
+ let residual_by_split = intersect_row_allow_lists(
+ plan.physical_row_ranges_by_split.as_deref(),
+ residual_by_split,
+ plan.splits.len(),
+ )?;
// Build the exact-fallback search on demand: the kernel calls this only
for a
// file it actually searches (uncovered by ANN, residual-allowed, and only
when
@@ -1098,8 +1245,12 @@ async fn plan_and_search_pk_candidates_batch(
// per-query bounded heaps (all queries share one stream).
let reader_for_factory = reader.clone();
let vector_field_for_factory = vector_field.clone();
+ // The plan's own per-file selection, so an exact fallback reads only the
rows an
+ // engine-supplied split allows. `is_excluded` still rejects on top of it,
but it
+ // cannot un-read a row.
+ let physical_for_factory = plan.physical_row_ranges_by_split.clone();
let factory = as_split_exact_file_search(
- move |_split_index: usize,
+ move |split_index: usize,
split: &PkVectorSearchSplit,
file: &BucketActiveFile,
queries: &[&[f32]],
@@ -1115,11 +1266,24 @@ async fn plan_and_search_pk_candidates_batch(
row_count: file.row_count,
};
let owned_queries: Vec<Vec<f32>> = queries.iter().map(|q|
q.to_vec()).collect();
+ let allowed_rows =
physical_for_factory.as_ref().and_then(|per_split| {
+ per_split
+ .get(split_index)
+ .and_then(|per_file| per_file.get(&active.file_name))
+ .cloned()
+ });
Box::pin(async move {
let factory = DataFilePkVectorReaderFactory::new(reader,
data_split, vector_field)?;
let query_refs: Vec<&[f32]> = owned_queries.iter().map(|q|
q.as_slice()).collect();
factory
- .search_file(&active, &query_refs, metric, exact_limit,
is_excluded)
+ .search_file(
+ &active,
+ &query_refs,
+ metric,
+ exact_limit,
+ is_excluded,
+ allowed_rows.as_deref(),
+ )
.await
})
},
@@ -1150,6 +1314,41 @@ async fn plan_and_search_pk_candidates_batch(
)
.await?;
+ Ok(searches)
+}
+
+/// Search an already-resolved plan and return one merged, best-first
candidate list
+/// per query: the raw layer above, followed by the optional exact rerank of
the
+/// approximate candidates and the merge with the exact-fallback candidates.
+#[allow(clippy::too_many_arguments)]
+async fn search_pk_candidates_batch_with_plan(
+ table: &Table,
+ query_options: &HashMap<String, String>,
+ filter: Option<&Predicate>,
+ core: &CoreOptions<'_>,
+ pk_col: &str,
+ queries: &[&[f32]],
+ limit: usize,
+ plan: &PkVectorScanPlan,
+ params: &PkVectorSearchParams,
+) -> crate::Result<Vec<Vec<PkVectorCandidate>>> {
+ let searches = search_pk_raw_candidates_batch_with_plan(
+ table,
+ query_options,
+ filter,
+ core,
+ pk_col,
+ queries,
+ limit,
+ plan,
+ params,
+ )
+ .await?;
+
+ let metric = params.metric;
+ let refine_factor = params.refine_factor;
+ let vector_field = params.vector_field.clone();
+
// Per query: exact rerank of the approximate candidates when a refine
factor is
// set (exact-fallback candidates are already exact and are not reranked),
then
// merge the (possibly reranked) indexed list with the exact list into one
@@ -1186,7 +1385,56 @@ async fn plan_and_search_pk_candidates_batch(
per_query_candidates.push(merge_candidates(indexed, search.exact,
limit));
}
- Ok((per_query_candidates, plan, metric))
+ Ok(per_query_candidates)
+}
+
+/// Plan the whole table and search it: resolve the query parameters, read the
index
+/// manifest into a plan, then search that plan. The plan and metric are
returned
+/// alongside the candidates because callers re-associate hits through the
plan.
+async fn plan_and_search_pk_candidates_batch(
+ table: &Table,
+ query_options: &HashMap<String, String>,
+ filter: Option<&Predicate>,
+ core: &CoreOptions<'_>,
+ pk_col: &str,
+ queries: &[&[f32]],
+ limit: usize,
+) -> crate::Result<(
+ Vec<Vec<PkVectorCandidate>>,
+ PkVectorScanPlan,
+ VectorSearchMetric,
+)> {
+ let params = resolve_pk_vector_search_params(
+ table,
+ query_options,
+ filter,
+ core,
+ pk_col,
+ queries,
+ limit,
+ )?;
+ let plan = PkVectorScan::new(
+ table,
+ params.field_id,
+ params.index_type.clone(),
+ filter.cloned(),
+ )
+ .plan()
+ .await?;
+ let metric = params.metric;
+ let candidates = search_pk_candidates_batch_with_plan(
+ table,
+ query_options,
+ filter,
+ core,
+ pk_col,
+ queries,
+ limit,
+ &plan,
+ ¶ms,
+ )
+ .await?;
+ Ok((candidates, plan, metric))
}
impl<'a> BatchVectorSearchBuilder<'a> {
@@ -2242,11 +2490,16 @@ fn is_vector_global_index_file(index_file:
&IndexFileMeta) -> bool {
/// row-collecting half of Java `PrimaryKeyVectorRead`'s `executeFilter`: the
/// predicate is NOT pushed down (a pushed filter would drop rows before their
/// position could be recovered). Instead `reader` projects only the residual
-/// columns and carries no pushdown predicate; every physical row is scanned in
-/// file order, the residual is evaluated here at the Arrow level, and each
-/// surviving row's file-local 0-based position is its running ordinal in the
scan.
-/// This needs no `_ROW_ID` and no `first_row_id` — real primary-key tables
never
-/// write one.
+/// columns and carries no pushdown predicate, the residual is evaluated here
at the
+/// Arrow level, and each surviving row's file-local 0-based position is
recovered
+/// from the selection the read was limited to. This needs no `_ROW_ID` and no
+/// `first_row_id` — real primary-key tables never write one.
+///
+/// `allowed_rows` is the plan's per-file physical selection, when it has one.
The
+/// residual is evaluated over exactly those rows: an engine-supplied bucket
split
+/// can restrict a huge file to a handful of ranges, and reading the whole
file only
+/// to discard everything outside them afterwards would defeat the split. With
no
+/// selection every physical row is scanned, as before.
///
/// Every *active* data file in the split gets an entry, possibly empty. The
/// bucket search treats an absent entry and an empty entry identically (the
file
@@ -2263,6 +2516,7 @@ async fn residual_positions_by_file(
split: &DataSplit,
active_files: &[BucketActiveFile],
residual: &FilePredicates,
+ allowed_rows: Option<&HashMap<String, Vec<RowRange>>>,
) -> crate::Result<HashMap<String, RoaringTreemap>> {
let scan_fields = reader.read_type().to_vec();
let active_names: HashSet<&str> = active_files.iter().map(|f|
f.file_name.as_str()).collect();
@@ -2273,16 +2527,48 @@ async fn residual_positions_by_file(
if !active_names.contains(file_meta.file_name.as_str()) {
continue;
}
+ let selection = match allowed_rows {
+ // A plan that lists nothing for a file permits nothing from it,
whether
+ // the list is empty or the file is absent: both sides of the
eventual
+ // intersection read absence that way. Registering it empty says
so and
+ // costs no read.
+ Some(by_file) => match by_file.get(&file_meta.file_name) {
+ Some(ranges) if !ranges.is_empty() => Some(ranges.clone()),
+ _ => {
+ out.entry(file_meta.file_name.clone()).or_default();
+ continue;
+ }
+ },
+ None => None,
+ };
let data_fields = reader.derive_data_fields(file_meta).await?;
- let mut stream =
- reader.read_single_file_stream(split, file_meta.clone(),
data_fields, None, None)?;
+ let mut stream = match selection.clone() {
+ Some(ranges) => reader.read_single_file_stream_local_ranges(
+ split,
+ file_meta.clone(),
+ data_fields,
+ None,
+ ranges,
+ )?,
+ None => {
+ reader.read_single_file_stream(split, file_meta.clone(),
data_fields, None, None)?
+ }
+ };
// Register the file up front so a file whose rows all fail the
residual
// still appears in the map (empty set).
let positions = out.entry(file_meta.file_name.clone()).or_default();
- // The scan has no row selection and no DV, so rows arrive in physical
file
- // order with no gaps: each row's file-local 0-based position is its
running
- // ordinal `base + row_index`.
- let mut base: u64 = 0;
+ // Rows arrive in ascending physical order, and the read emitted
exactly what
+ // was selected (no pushdown predicate, no deletion vector), so
walking the
+ // selection in step with the rows recovers each row's file-local
position.
+ let mut selected: Box<dyn Iterator<Item = u64> + Send> = match
&selection {
+ Some(ranges) => Box::new(
+ ranges
+ .clone()
+ .into_iter()
+ .flat_map(|range| (range.from() as u64)..=(range.to() as
u64)),
+ ),
+ None => Box::new(0..file_meta.row_count.max(0) as u64),
+ };
while let Some(batch) = stream.try_next().await? {
let num_rows = batch.num_rows();
let mask = evaluate_predicates_mask(
@@ -2291,24 +2577,34 @@ async fn residual_positions_by_file(
&residual.file_fields,
&scan_fields,
)?;
- match mask {
- Some(mask) => {
- for row_index in 0..num_rows {
- // NULL follows the same NULL -> false convention the
Arrow
- // filter kernel applies, so a null mask slot drops
the row.
- if mask.is_valid(row_index) && mask.value(row_index) {
- positions.insert(base + row_index as u64);
- }
- }
- }
- // No predicate contributed a mask (identity) -> keep every
row.
- None => {
- for row_index in 0..num_rows {
- positions.insert(base + row_index as u64);
- }
+ for row_index in 0..num_rows {
+ let position = selected.next().ok_or_else(||
crate::Error::DataInvalid {
+ message: format!(
+ "residual scan of '{}' emitted more rows than the
selection allows",
+ file_meta.file_name
+ ),
+ source: None,
+ })?;
+ let keep = match &mask {
+ // NULL follows the same NULL -> false convention the
Arrow filter
+ // kernel applies, so a null mask slot drops the row.
+ Some(mask) => mask.is_valid(row_index) &&
mask.value(row_index),
+ // No predicate contributed a mask (identity) -> keep
every row.
+ None => true,
+ };
+ if keep {
+ positions.insert(position);
}
}
- base += num_rows as u64;
+ }
+ if selected.next().is_some() {
+ return Err(crate::Error::DataInvalid {
+ message: format!(
+ "residual scan of '{}' emitted fewer rows than the
selection allows",
+ file_meta.file_name
+ ),
+ source: None,
+ });
}
}
Ok(out)
@@ -7771,19 +8067,74 @@ mod residual_positions_tests {
&[("part-0.mosaic", vec![1, 2, 3, 4, 5], 0)],
)
.await;
- let map = residual_positions_by_file(&reader, &split, &active,
&residual_id_gt(2))
+ let map = residual_positions_by_file(&reader, &split, &active,
&residual_id_gt(2), None)
.await
.unwrap();
assert_eq!(sorted(&map["part-0.mosaic"]), vec![2, 3, 4]);
}
+ #[tokio::test]
+ async fn test_residual_only_evaluates_the_rows_the_plan_allows() {
+ // ids [1,2,3,4,5]; the plan allows positions 3-4 only. `id > 2`
matches 2,3,4
+ // over the whole file, so a result of 3,4 is the plan's restriction
taking
+ // effect *before* evaluation: position 2 is never seen.
+ //
+ // This also cannot pass under a full read. The scan walks the
selection in
+ // step with the emitted rows, so a read that emitted all five would
run the
+ // selection dry and fail loudly rather than return a filtered answer.
+ let (reader, split, active) = build_reader_and_split(
+ "memory:/rpf_plan_ranges",
+ &[("part-0.mosaic", vec![1, 2, 3, 4, 5], 0)],
+ )
+ .await;
+ let allowed = HashMap::from([("part-0.mosaic".to_string(),
vec![RowRange::new(3, 4)])]);
+ let map = residual_positions_by_file(
+ &reader,
+ &split,
+ &active,
+ &residual_id_gt(2),
+ Some(&allowed),
+ )
+ .await
+ .unwrap();
+ assert_eq!(sorted(&map["part-0.mosaic"]), vec![3, 4]);
+ }
+
+ #[tokio::test]
+ async fn test_residual_does_not_read_a_file_the_plan_excludes() {
+ // A file the plan lists no rows for is registered empty and never
opened. The
+ // empty entry is what tells the search the file contributes nothing;
an
+ // absent one would mean the same, but then the map would not cover
the split.
+ let (reader, split, active) = build_reader_and_split(
+ "memory:/rpf_plan_excludes",
+ &[("part-0.mosaic", vec![1, 2, 3], 0)],
+ )
+ .await;
+ for allowed in [
+ HashMap::from([("part-0.mosaic".to_string(), Vec::new())]),
+ HashMap::new(),
+ ] {
+ let map = residual_positions_by_file(
+ &reader,
+ &split,
+ &active,
+ &residual_id_gt(0),
+ Some(&allowed),
+ )
+ .await
+ .unwrap();
+ assert!(map.contains_key("part-0.mosaic"));
+ assert!(sorted(&map["part-0.mosaic"]).is_empty());
+ }
+ }
+
#[tokio::test]
async fn test_residual_matches_none_yields_empty_entry() {
// id > 100 matches nothing; the file still gets a (present, empty)
entry.
let (reader, split, active) =
build_reader_and_split("memory:/rpf_none", &[("part-0.mosaic",
vec![1, 2, 3], 0)])
.await;
- let map = residual_positions_by_file(&reader, &split, &active,
&residual_id_gt(100))
+ let map = residual_positions_by_file(&reader, &split, &active,
&residual_id_gt(100), None)
.await
.unwrap();
assert!(map.contains_key("part-0.mosaic"));
@@ -7794,7 +8145,7 @@ mod residual_positions_tests {
async fn test_residual_matches_all_yields_full_set() {
let (reader, split, active) =
build_reader_and_split("memory:/rpf_all", &[("part-0.mosaic",
vec![1, 2, 3], 0)]).await;
- let map = residual_positions_by_file(&reader, &split, &active,
&residual_id_gt(0))
+ let map = residual_positions_by_file(&reader, &split, &active,
&residual_id_gt(0), None)
.await
.unwrap();
assert_eq!(sorted(&map["part-0.mosaic"]), vec![0, 1, 2]);
@@ -7812,7 +8163,7 @@ mod residual_positions_tests {
],
)
.await;
- let map = residual_positions_by_file(&reader, &split, &active,
&residual_id_gt(3))
+ let map = residual_positions_by_file(&reader, &split, &active,
&residual_id_gt(3), None)
.await
.unwrap();
assert_eq!(sorted(&map["part-0.mosaic"]), vec![3, 4]);
@@ -7855,7 +8206,7 @@ mod residual_positions_tests {
.with_data_files(metas)
.build()
.unwrap();
- let map = residual_positions_by_file(&reader, &split, &active,
&residual_id_gt(2))
+ let map = residual_positions_by_file(&reader, &split, &active,
&residual_id_gt(2), None)
.await
.unwrap();
assert_eq!(sorted(&map["part-0.mosaic"]), vec![2, 3, 4]);
@@ -7871,7 +8222,7 @@ mod residual_positions_tests {
// recovered from each row's ordinal in the scan, so the residual still
// works: ids [1,2,3] with id > 0 -> all match -> local positions
[0,1,2].
let (reader, split, active) =
build_reader_and_split_no_first_row_id().await;
- let map = residual_positions_by_file(&reader, &split, &active,
&residual_id_gt(0))
+ let map = residual_positions_by_file(&reader, &split, &active,
&residual_id_gt(0), None)
.await
.expect("missing first_row_id must not fail the residual read");
assert_eq!(sorted(&map["part-0.mosaic"]), vec![0, 1, 2]);
@@ -7914,4 +8265,83 @@ mod residual_positions_tests {
}];
(reader, split, active)
}
+
+ // ---- combining the plan's positional restriction with the residual ----
+
+ fn allow_list(entries: &[(&str, &[u64])]) -> HashMap<String,
RoaringTreemap> {
+ entries
+ .iter()
+ .map(|(file, positions)| ((*file).to_string(),
positions.iter().copied().collect()))
+ .collect()
+ }
+
+ /// The plan side carries ranges, so its fixtures are built from the
positions
+ /// each file allows and coalesced the way the planner normalizes them.
+ fn range_allow_list(entries: &[(&str, &[u64])]) -> HashMap<String,
Vec<RowRange>> {
+ entries
+ .iter()
+ .map(|(file, positions)| {
+ let ranges = positions
+ .iter()
+ .map(|p| RowRange::new(*p as i64, *p as i64))
+ .collect();
+ ((*file).to_string(), merge_row_ranges(ranges))
+ })
+ .collect()
+ }
+
+ fn listed(map: &HashMap<String, RoaringTreemap>, file: &str) -> Vec<u64> {
+ map.get(file)
+ .map(|positions| positions.iter().collect())
+ .unwrap_or_default()
+ }
+
+ #[test]
+ fn no_restriction_on_either_side_stays_unrestricted() {
+ assert!(intersect_row_allow_lists(None, None, 1).unwrap().is_none());
+ }
+
+ #[test]
+ fn one_side_alone_passes_through() {
+ let physical = vec![range_allow_list(&[("d0", &[1, 2])])];
+ let only_physical = intersect_row_allow_lists(Some(&physical), None, 1)
+ .unwrap()
+ .expect("a plan restriction survives on its own");
+ assert_eq!(listed(&only_physical[0], "d0"), vec![1, 2]);
+
+ let residual = vec![allow_list(&[("d0", &[3])])];
+ let only_residual = intersect_row_allow_lists(None, Some(residual), 1)
+ .unwrap()
+ .expect("a residual survives on its own");
+ assert_eq!(listed(&only_residual[0], "d0"), vec![3]);
+ }
+
+ #[test]
+ fn both_sides_intersect_and_a_file_either_omits_is_dropped() {
+ // `d0`: both list positions, so only the shared ones survive. `d1`:
the
+ // residual kept nothing there, and its absence means "no rows", so
the file
+ // must not come back unrestricted from the plan side.
+ let physical = vec![range_allow_list(&[("d0", &[1, 2, 3]), ("d1", &[0,
1])])];
+ let residual = vec![allow_list(&[("d0", &[2, 3, 4])])];
+ let combined = intersect_row_allow_lists(Some(&physical),
Some(residual), 1)
+ .unwrap()
+ .expect("both sides restrict");
+ assert_eq!(listed(&combined[0], "d0"), vec![2, 3]);
+ assert!(!combined[0].contains_key("d1"));
+ }
+
+ #[test]
+ fn rejects_allow_lists_that_do_not_cover_every_split() {
+ let physical = vec![range_allow_list(&[("d0", &[1])])];
+ let error = intersect_row_allow_lists(Some(&physical), None, 2)
+ .map(|_| ())
+ .expect_err("an allow-list per split is what makes the index
meaningful");
+ assert!(error.to_string().contains("for 2 splits"), "{error}");
+
+ let residual = vec![allow_list(&[("d0", &[1])])];
+ let error = intersect_row_allow_lists(Some(&physical), Some(residual),
2)
+ .map(|_| ())
+ .expect_err("the residual must cover every split too");
+ assert!(error.to_string().contains("for 2 splits"), "{error}");
+ }
}
diff --git a/crates/paimon/src/vindex/pkvector/ann.rs
b/crates/paimon/src/vindex/pkvector/ann.rs
index 76cf6aa4..e2be9f87 100644
--- a/crates/paimon/src/vindex/pkvector/ann.rs
+++ b/crates/paimon/src/vindex/pkvector/ann.rs
@@ -89,18 +89,30 @@ pub(crate) fn build_live_row_ids(
// file_offset). A missing/empty entry allows no rows.
Some(ranges) => {
if let Some(allowed) = ranges.get(source_file.file_name())
{
- for position in allowed.iter() {
- if position >= row_count {
- return Err(data_invalid(format!(
- "residual position {position} is out of
range for source file {} ({} rows)",
- source_file.file_name(),
- row_count
- )));
+ // A producer that restricts only some files leaves
the rest
+ // unrestricted, and an adapter has to spell that out
as an
+ // explicit whole-file allow-list. Insert it as one
range
+ // rather than walking every position, which would
cost one
+ // insert per row of the file. `len` plus a maximum of
+ // `row_count - 1` can only describe the full set, and
it
+ // subsumes the per-position bound check below.
+ if allowed.len() == row_count && allowed.max() ==
Some(row_count - 1) {
+ live.insert_range(file_offset..end);
+ } else {
+ for position in allowed.iter() {
+ if position >= row_count {
+ return Err(data_invalid(format!(
+ "residual position {position} is out
of range for source file {} ({} rows)",
+ source_file.file_name(),
+ row_count
+ )));
+ }
+ let global =
+
file_offset.checked_add(position).ok_or_else(|| {
+ data_invalid("vector residual position
overflows u64")
+ })?;
+ live.insert(global);
}
- let global =
file_offset.checked_add(position).ok_or_else(|| {
- data_invalid("vector residual position
overflows u64")
- })?;
- live.insert(global);
}
}
}
@@ -841,6 +853,44 @@ mod tests {
assert_eq!(live.iter().collect::<Vec<u64>>(), vec![0]);
}
+ #[test]
+ fn test_whole_file_allow_list_matches_having_no_residual_at_all() {
+ // An adapter spells "unrestricted" out as an explicit whole-file
allow-list.
+ // That has to land on the same live set the no-residual path
produces, since
+ // it is the same statement said two ways.
+ let files = vec![
+ PkVectorSourceFile::new("f0".into(), 3).unwrap(),
+ PkVectorSourceFile::new("f1".into(), 2).unwrap(),
+ ];
+ let active = active_set(&["f0", "f1"]);
+ let mut residual = HashMap::new();
+ residual.insert("f0".to_string(), treemap(&[0, 1, 2]));
+ residual.insert("f1".to_string(), treemap(&[0, 1]));
+
+ let spelled_out = build_live_row_ids(&files, &active, &HashMap::new(),
Some(&residual))
+ .unwrap()
+ .unwrap();
+ assert_eq!(
+ spelled_out.iter().collect::<Vec<u64>>(),
+ vec![0, 1, 2, 3, 4]
+ );
+ }
+
+ #[test]
+ fn test_whole_file_allow_list_still_applies_the_deletion_vector() {
+ // The whole-file shortcut must not skip deletion vectors: f0 allows
every
+ // row, but position 1 is deleted and has to stay out.
+ let files = vec![PkVectorSourceFile::new("f0".into(), 3).unwrap()];
+ let mut dvs = HashMap::new();
+ dvs.insert("f0".to_string(), dv(&[1]));
+ let mut residual = HashMap::new();
+ residual.insert("f0".to_string(), treemap(&[0, 1, 2]));
+ let live = build_live_row_ids(&files, &active_set(&["f0"]), &dvs,
Some(&residual))
+ .unwrap()
+ .unwrap();
+ assert_eq!(live.iter().collect::<Vec<u64>>(), vec![0, 2]);
+ }
+
#[test]
fn test_build_live_row_ids_residual_maps_positions_across_file_offsets() {
// f0 rows global 0,1,2; f1 rows global 3,4. residual allows f0={2},
f1={1}.
