From bd68ed074d5ac6681c07241cf2697ddbd3717787 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Thu, 27 Aug 2026 15:49:28 +0800 Subject: [PATCH 1/7] refactor(table): separate PK-vector planning from searching MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `plan_and_search_pk_candidates_batch` resolved the query parameters, read the index manifest into a plan, and searched that plan in one body, so a caller that already holds a plan could not reuse the search path. Split it into three pieces with no behavior change: - `resolve_pk_vector_search_params` — the query-level parameters and the pre-filter guard: everything resolvable from the schema, the options and the queries alone, before planning. - `search_pk_raw_candidates_batch_with_plan` — search a supplied plan and return each query's raw indexed and exact candidate lists. Plan-dependent concurrency (segment count, batch-index parallelism, range-read bound) is derived from the plan actually being searched, so a narrowed plan can never be searched under limits computed for a wider one. - `search_pk_candidates_batch_with_plan` — the raw layer plus the optional exact rerank of the approximate candidates and the merge into one best-first list per query. `plan_and_search_pk_candidates_batch` keeps its signature and becomes a wrapper over the three. The empty-plan short circuit moves into the raw layer, still ahead of backend resolution, so a table with no searchable data does not error on an unrecognized index type. --- .../paimon/src/table/vector_search_builder.rs | 164 ++++++++++++++++-- 1 file changed, 153 insertions(+), 11 deletions(-) diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index aa1fdc6dc..e2db99708 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -734,7 +734,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, filter: Option<&Predicate>, @@ -742,11 +761,7 @@ async fn plan_and_search_pk_candidates_batch( pk_col: &str, queries: &[&[f32]], limit: usize, -) -> crate::Result<( - Vec>, - PkVectorScanPlan, - VectorSearchMetric, -)> { +) -> crate::Result { // 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 @@ -848,13 +863,56 @@ 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. +#[allow(clippy::too_many_arguments)] +async fn search_pk_raw_candidates_batch_with_plan( + table: &Table, + query_options: &HashMap, + filter: Option<&Predicate>, + core: &CoreOptions<'_>, + pk_col: &str, + queries: &[&[f32]], + limit: usize, + plan: &PkVectorScanPlan, + params: &PkVectorSearchParams, +) -> crate::Result> { + // 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 @@ -1122,6 +1180,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, + filter: Option<&Predicate>, + core: &CoreOptions<'_>, + pk_col: &str, + queries: &[&[f32]], + limit: usize, + plan: &PkVectorScanPlan, + params: &PkVectorSearchParams, +) -> crate::Result>> { + 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 @@ -1158,7 +1251,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, + filter: Option<&Predicate>, + core: &CoreOptions<'_>, + pk_col: &str, + queries: &[&[f32]], + limit: usize, +) -> crate::Result<( + Vec>, + 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> { From 6ce4e00abbd57ac89bbf67c7d6f1130522c3b222 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Thu, 27 Aug 2026 16:35:12 +0800 Subject: [PATCH 2/7] feat(table): plan a PK-vector search from a decoded bucket split A `BucketVectorSearchSplit` already carries everything a search needs for one bucket: the payload files, the rows each data file allows, and the snapshot the whole plan is pinned to. Planning could only be driven the other way round, by reading this table's index manifest, so a search could not be run over splits an engine planned elsewhere. `PkVectorScan::plan_for_bucket_vector_splits` builds a plan from such splits instead. The splits are authoritative -- no manifest is read -- and only the partition conjuncts of the scan's filter are re-applied, since a caller may narrow the query further than the planner that produced the splits. Bucket grouping, current-segment selection and exact-fallback eligibility reuse the manifest route's `plan_from_inputs`, so both routes pick segments the same way. A payload's index file is resolved where Java put it: its `_EXTERNAL_PATH` when it records one, and otherwise the bucket directory the split serialized when the table sets `index-file-in-data-file-dir`, or the table `index/` directory. Java records an external path only for an index stored outside the table -- `PkVectorAnnSegmentFile` writes `null` unless its path factory is external -- so an ordinary bucket-local payload carries none, and looking for it under `index/` would not find it. The manifest route keeps its existing `index/` assumption; the option is read only for splits an engine supplied. Four inputs are rejected rather than planned around: - No splits at all, which pins no snapshot to report, and the plan's snapshot id has to stay authoritative even when nothing is searchable. - Splits pinning different snapshots, checked before partition pruning so an inconsistent input cannot hide behind an empty plan. - Two splits for one bucket, which would search its rows twice. Java emits one split per bucket, but independently decoded buffers cannot enforce that. - A nested data split carrying its own row ranges, which would be a second authority over which physical rows are readable, free to disagree with the per-file ranges the bucket form carries. Row ranges become a per-split allow-list of physical positions on the plan, and the search intersects it with the residual predicate's allow-list: both sides list what is permitted, so a position needs to survive both. The normalization is where the two formats disagree -- Java records ranges only for the files its own pre-filter narrowed and omits the rest, while the search kernel reads a missing entry as "no rows allowed" -- so an omitted file is turned into an explicit full-file range. An empty list stays empty and excludes its file. A payload's `deletion_vectors_ranges` is ignored on purpose. Java reserves that field for deletion-vector index files, builds vector payloads through the overload that leaves it null, and takes a read's deletion vectors from the bucket's data split, so a value there describes something the payload is not. Planning from the Java golden fixture is covered end to end: the external payload path wins over both directory layouts, the five-billion-byte size survives, and a six-row file listed as rows 0-1 and 4-5 plans to exactly those positions. --- crates/paimon/src/spec/core_options.rs | 11 + .../src/table/pk_vector_bucket_split.rs | 89 +++ crates/paimon/src/table/pk_vector_scan.rs | 674 +++++++++++++++++- .../paimon/src/table/vector_search_builder.rs | 129 ++++ 4 files changed, 901 insertions(+), 2 deletions(-) diff --git a/crates/paimon/src/spec/core_options.rs b/crates/paimon/src/spec/core_options.rs index b17d31685..7afff6f67 100644 --- a/crates/paimon/src/spec/core_options.rs +++ b/crates/paimon/src/spec/core_options.rs @@ -32,6 +32,7 @@ const GLOBAL_INDEX_ROW_COUNT_PER_SHARD_OPTION: &str = "global-index.row-count-pe const GLOBAL_INDEX_THREAD_NUM_OPTION: &str = "global-index.thread-num"; const GLOBAL_INDEX_VINDEX_READ_THREAD_NUM_OPTION: &str = "global-index.vindex.read-thread-num"; const GLOBAL_INDEX_COLUMN_UPDATE_ACTION_OPTION: &str = "global-index.column-update-action"; +pub(crate) const INDEX_FILE_IN_DATA_FILE_DIR_OPTION: &str = "index-file-in-data-file-dir"; const SORTED_INDEX_RECORDS_PER_RANGE_OPTION: &str = "sorted-index.records-per-range"; const BTREE_INDEX_RECORDS_PER_RANGE_OPTION: &str = "btree-index.records-per-range"; const BTREE_INDEX_FALLBACK_SCAN_MAX_SIZE_OPTION: &str = "btree-index.fallback-scan-max-size"; @@ -686,6 +687,16 @@ impl<'a> CoreOptions<'a> { .unwrap_or(true) } + /// Whether index files are stored in the bucket data-file directory rather + /// than the table `index/` directory (option `index-file-in-data-file-dir`, + /// default false). + pub fn index_file_in_data_file_dir(&self) -> bool { + self.options + .get(INDEX_FILE_IN_DATA_FILE_DIR_OPTION) + .map(|value| value.eq_ignore_ascii_case("true")) + .unwrap_or(false) + } + pub fn global_index_search_mode(&self) -> crate::Result { self.index_search_mode(GLOBAL_INDEX_SEARCH_MODE_OPTION) } diff --git a/crates/paimon/src/table/pk_vector_bucket_split.rs b/crates/paimon/src/table/pk_vector_bucket_split.rs index 0c50b675b..ef469236e 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 the source-metadata + /// blob and the index metadata move instead of being cloned. + /// + /// 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, + 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, + IndexMap>, + ) { + (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 { 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, + row_ranges_by_file: IndexMap>, + ) -> 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>, + external_path: Option, + 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_scan.rs b/crates/paimon/src/table/pk_vector_scan.rs index cfc528b25..1f5fa4454 100644 --- a/crates/paimon/src/table/pk_vector_scan.rs +++ b/crates/paimon/src/table/pk_vector_scan.rs @@ -20,19 +20,56 @@ //! 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, GlobalIndexMeta, IndexManifest, Predicate, PrimaryKeyIndexSourceFile, PrimaryKeyIndexSourceMeta, }; +use crate::table::bucket_filter::split_partition_and_data_predicates; +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::{DataSplit, DataSplitBuilder, DeletionFile, RowRange}; use crate::table::Table; use crate::vindex::pkvector::bucket::{BucketActiveFile, BucketAnnSegment}; const INDEX_DIR: &str = "index"; +/// A bucket's identity across planning inputs: the partition's serialized bytes +/// (`BinaryRow` is not hashable) paired with the bucket number. +type BucketKey = (Vec, i32); + +/// Expand inclusive row ranges into the positions they allow. +#[cfg_attr(not(test), allow(dead_code))] +fn positions_in_ranges(ranges: &[RowRange]) -> crate::Result { + 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) +} + +/// Every position in a file, for a file the message left unrestricted. +#[cfg_attr(not(test), allow(dead_code))] +fn positions_in_whole_file(row_count: i64) -> crate::Result { + let rows = u64::try_from(row_count) + .map_err(|_| data_invalid("data file row count must not be negative"))?; + let mut positions = RoaringTreemap::new(); + if rows > 0 { + positions.insert_range(0..=rows - 1); + } + Ok(positions) +} + fn data_invalid(message: impl Into) -> crate::Error { crate::Error::DataInvalid { message: message.into(), @@ -197,6 +234,14 @@ pub(crate) struct PkVectorScanPlan { // at all (never written), which also yields empty `splits`. pub snapshot_id: i64, pub splits: Vec, + // Per-split allow-list of physical row positions, indexed parallel to `splits`: + // only the positions listed for a data file may produce candidates from it. + // 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>>, } pub(crate) struct PkVectorScan<'a> { @@ -262,6 +307,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?; @@ -313,8 +359,226 @@ 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. + /// 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, + ) -> crate::Result { + // 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, +) -> crate::Result { + // 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 = None; + let mut seen_buckets: HashSet = HashSet::new(); + let mut data_splits: Vec = Vec::with_capacity(splits.len()); + let mut index_entries: Vec<(BinaryRow, i32, GlobalIndexMeta, String, u64, String)> = Vec::new(); + let mut listed_ranges: HashMap>> = 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"))?; + // Java records an external path only for an index stored outside the + // table, so an ordinary bucket-local payload carries none and is + // resolved against the bucket directory the engine serialized when the + // table keeps index files there, and the table `index/` directory + // otherwise. Mirrors `IndexInDataFileDirPathFactory.toPath`. + let path = match parts.external_path { + Some(external) => external, + None if index_file_in_data_file_dir => { + format!("{}/{}", data_split.bucket_path(), parts.file_name) + } + None => format!("{table_path}/{INDEX_DIR}/{}", parts.file_name), + }; + index_entries.push(( + data_split.partition().clone(), + data_split.bucket(), + parts.global_index_meta, + path, + file_size, + parts.file_name, + )); + } + + 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)?; + + // 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)) { + Some(ranges) => positions_in_ranges(ranges)?, + None => positions_in_whole_file(file.row_count)?, + }; + Ok((file.file_name.clone(), allowed)) + }) + .collect::>>() + }) + .collect::>>()?; + + 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 @@ -390,6 +654,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) -> DataFileMeta { @@ -900,4 +1166,408 @@ 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, + payloads: Vec, + 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, + ) + } + + fn allowed(map: &HashMap, file: &str) -> Vec { + map.get(file) + .map(|positions| positions.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 manifest_planning_leaves_positions_unrestricted() { + // The manifest route must keep reporting `None`, which means "no positional + // restriction" -- not an empty allow-list, which would permit nothing. + let entries = vec![( + BinaryRow::new(0), + 0, + gim(2, 5, &[("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); + } + + #[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 e2db99708..3168f1cd7 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -882,6 +882,62 @@ fn resolve_pk_vector_search_params( /// 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 +/// positions 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. +fn intersect_row_allow_lists( + physical: Option<&[HashMap]>, + residual: Option>>, + split_count: usize, +) -> crate::Result>>> { + 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.to_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_map(|(file, allowed)| { + residual + .get(file) + .map(|kept| (file.clone(), allowed & kept)) + }) + .collect() + }) + .collect(), + )) + } + } +} + #[allow(clippy::too_many_arguments)] async fn search_pk_raw_candidates_batch_with_plan( table: &Table, @@ -1119,6 +1175,15 @@ async fn search_pk_raw_candidates_batch_with_plan( } 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 @@ -7384,4 +7449,68 @@ mod residual_positions_tests { }]; (reader, split, active) } + + // ---- combining the plan's positional restriction with the residual ---- + + fn allow_list(entries: &[(&str, &[u64])]) -> HashMap { + entries + .iter() + .map(|(file, positions)| ((*file).to_string(), positions.iter().copied().collect())) + .collect() + } + + fn listed(map: &HashMap, file: &str) -> Vec { + 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![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![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![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}"); + } } From 36485f174499441f65916e6cc3adb0edfa79d8ec Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Fri, 28 Aug 2026 12:51:18 +0800 Subject: [PATCH 3/7] perf(vindex): insert a whole-file allow-list as one range A producer that restricts the readable rows of only some data files leaves the rest unrestricted, and an adapter has to say so explicitly, as an allow-list covering the whole file. Building live row ids then walked that list one position at a time, costing an insert per row of the file, where the same statement made by omitting the residual entirely takes a single range insert. Recognize the whole-file shape and insert one range instead. A list whose length equals the file's row count and whose maximum is the last position can only be the full set, so the check also subsumes the per-position bound check it replaces. Deletion vectors still apply: the shortcut only replaces how positions enter the live set, not what happens to them afterwards. --- crates/paimon/src/vindex/pkvector/ann.rs | 72 ++++++++++++++++++++---- 1 file changed, 61 insertions(+), 11 deletions(-) diff --git a/crates/paimon/src/vindex/pkvector/ann.rs b/crates/paimon/src/vindex/pkvector/ann.rs index 76cf6aa49..e2be9f870 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![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![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![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}. From 90433b91e4014b5b992bf2bae2b0bbcff1a09e99 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Mon, 31 Aug 2026 17:50:24 +0800 Subject: [PATCH 4/7] docs(table): fix bucket-split planning comments and drop a vacuous test - `plan_for_bucket_vector_splits` carried its whole doc block twice. - `into_parts` claimed the index metadata moves rather than being cloned, which `plan_from_inputs` does not honor: it clones `_INDEX_META` out of the `GlobalIndexMeta` it is handed. Narrowed the claim to what the method itself does. - `manifest_planning_leaves_positions_unrestricted` could not observe what its name claims: it drives `plan_from_inputs`, which returns splits and not a `PkVectorScanPlan`, so `physical_row_ranges_by_split` is out of its reach and the only assertion left was a split count that `builds_one_split_per_bucket_with_data` already makes. --- .../src/table/pk_vector_bucket_split.rs | 4 +- crates/paimon/src/table/pk_vector_scan.rs | 37 ------------------- 2 files changed, 2 insertions(+), 39 deletions(-) diff --git a/crates/paimon/src/table/pk_vector_bucket_split.rs b/crates/paimon/src/table/pk_vector_bucket_split.rs index ef469236e..c82826757 100644 --- a/crates/paimon/src/table/pk_vector_bucket_split.rs +++ b/crates/paimon/src/table/pk_vector_bucket_split.rs @@ -127,8 +127,8 @@ impl BucketVectorPayload { .expect("a decoded payload always carries source metadata") } - /// Consume the payload into the pieces a planner needs, so the source-metadata - /// blob and the index metadata move instead of being cloned. + /// 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. diff --git a/crates/paimon/src/table/pk_vector_scan.rs b/crates/paimon/src/table/pk_vector_scan.rs index 1f5fa4454..cf1a67de2 100644 --- a/crates/paimon/src/table/pk_vector_scan.rs +++ b/crates/paimon/src/table/pk_vector_scan.rs @@ -363,18 +363,6 @@ impl<'a> PkVectorScan<'a> { }) } - /// 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. /// Build a plan from bucket splits an engine planned elsewhere, instead of from /// this table's index manifest. /// @@ -1268,31 +1256,6 @@ mod tests { assert_eq!(allowed(&ranges[0], "data-1.orc"), vec![0, 1, 4, 5]); } - #[test] - fn manifest_planning_leaves_positions_unrestricted() { - // The manifest route must keep reporting `None`, which means "no positional - // restriction" -- not an empty allow-list, which would permit nothing. - let entries = vec![( - BinaryRow::new(0), - 0, - gim(2, 5, &[("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); - } - #[test] fn rejects_empty_bucket_split_input() { let error = plan_from_bucket_splits("ivf-pq", 2, None, "/tbl", false, Vec::new()) From 122e0b9256072675db376d0e1cc51b1b74c5e9f3 Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Tue, 1 Sep 2026 13:32:59 +0800 Subject: [PATCH 5/7] perf(table): evaluate the residual over the rows the split allows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An engine-supplied bucket split can restrict a large data file to a handful of rows, but the residual predicate was evaluated by reading every physical row of every active file and discarding what fell outside the split's ranges afterwards. A ten-row split over a billion-row file therefore cost a full scan on every residual-filtered query, which is the opposite of what the split is for. Java evaluates the residual through an `IndexedSplit` built from the same candidate ranges. The plan now carries the selection as normalized `Vec` rather than materialized positions, since that is what a read is limited by; expanding a whole-file range of a large file into positions costs memory no reader needs. `DataFileReader` grew a ranges entry point beside the positions one — it coalesced positions into ranges internally anyway — and the residual read goes through it. Positions are recovered by walking the selection in step with the emitted rows instead of counting from zero, and the two are checked against each other: with no pushdown predicate and no deletion vector the read emits exactly what was selected, so a mismatch means the assumption broke. A file the plan lists no rows for is registered empty without being read at all. The intersection against the residual stays: it cannot remove anything once the residual was evaluated over the same ranges, which is the invariant it now states. --- crates/paimon/src/table/data_file_reader.rs | 32 ++- crates/paimon/src/table/pk_vector_scan.rs | 53 +++-- .../paimon/src/table/vector_search_builder.rs | 183 +++++++++++++----- 3 files changed, 193 insertions(+), 75 deletions(-) diff --git a/crates/paimon/src/table/data_file_reader.rs b/crates/paimon/src/table/data_file_reader.rs index 25e205496..cd72a1c12 100644 --- a/crates/paimon/src/table/data_file_reader.rs +++ b/crates/paimon/src/table/data_file_reader.rs @@ -599,6 +599,28 @@ impl DataFileReader { data_fields: Option>, dv: Option>, local_positions: Vec, + ) -> crate::Result { + 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. + pub(super) fn read_single_file_stream_local_ranges( + &self, + split: &DataSplit, + file_meta: DataFileMeta, + data_fields: Option>, + dv: Option>, + local_ranges: Vec, ) -> crate::Result { // Local-position selection is only sound against a predicate-free reader: a // row-filtering predicate drops arbitrary selected rows and desyncs the @@ -666,12 +688,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_scan.rs b/crates/paimon/src/table/pk_vector_scan.rs index cf1a67de2..932b8e12b 100644 --- a/crates/paimon/src/table/pk_vector_scan.rs +++ b/crates/paimon/src/table/pk_vector_scan.rs @@ -34,7 +34,7 @@ use crate::table::bucket_filter::split_partition_and_data_predicates; 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, RowRange}; +use crate::table::source::{merge_row_ranges, DataSplit, DataSplitBuilder, DeletionFile, RowRange}; use crate::table::Table; use crate::vindex::pkvector::bucket::{BucketActiveFile, BucketAnnSegment}; @@ -44,9 +44,9 @@ const INDEX_DIR: &str = "index"; /// (`BinaryRow` is not hashable) paired with the bucket number. type BucketKey = (Vec, i32); -/// Expand inclusive row ranges into the positions they allow. -#[cfg_attr(not(test), allow(dead_code))] -fn positions_in_ranges(ranges: &[RowRange]) -> crate::Result { +/// 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 { let mut positions = RoaringTreemap::new(); for range in ranges { let from = u64::try_from(range.from()) @@ -58,16 +58,15 @@ fn positions_in_ranges(ranges: &[RowRange]) -> crate::Result { Ok(positions) } -/// Every position in a file, for a file the message left unrestricted. -#[cfg_attr(not(test), allow(dead_code))] -fn positions_in_whole_file(row_count: i64) -> crate::Result { - let rows = u64::try_from(row_count) - .map_err(|_| data_invalid("data file row count must not be negative"))?; - let mut positions = RoaringTreemap::new(); - if rows > 0 { - positions.insert_range(0..=rows - 1); +/// 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". +fn whole_file_range(row_count: i64) -> Vec { + if row_count > 0 { + vec![RowRange::new(0, row_count - 1)] + } else { + Vec::new() } - Ok(positions) } fn data_invalid(message: impl Into) -> crate::Error { @@ -234,14 +233,17 @@ pub(crate) struct PkVectorScanPlan { // at all (never written), which also yields empty `splits`. pub snapshot_id: i64, pub splits: Vec, - // Per-split allow-list of physical row positions, indexed parallel to `splits`: - // only the positions listed for a data file may produce candidates from it. + // 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>>, + pub physical_row_ranges_by_split: Option>>>, } pub(crate) struct PkVectorScan<'a> { @@ -553,12 +555,14 @@ fn plan_from_bucket_splits( .iter() .map(|file| { let allowed = match listed.and_then(|ranges| ranges.get(&file.file_name)) { - Some(ranges) => positions_in_ranges(ranges)?, - None => positions_in_whole_file(file.row_count)?, + // 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::>>() + .collect::>>>() }) .collect::>>()?; @@ -1223,9 +1227,16 @@ mod tests { ) } - fn allowed(map: &HashMap, file: &str) -> Vec { + /// The positions a file's normalized ranges allow, for assertions that read + /// better as a row list than as ranges. + fn allowed(map: &HashMap>, file: &str) -> Vec { map.get(file) - .map(|positions| positions.iter().collect()) + .map(|ranges| { + positions_in_ranges(ranges) + .expect("planned ranges are in range") + .iter() + .collect() + }) .unwrap_or_default() } diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index 3168f1cd7..b398d9001 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::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::source::DataSplit; use crate::table::{ @@ -883,14 +883,19 @@ fn resolve_pk_vector_search_params( /// 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 -/// positions an engine-supplied plan restricts each file to, and the positions a -/// residual data predicate leaves behind. +/// 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]>, + physical: Option<&[HashMap>]>, residual: Option>>, split_count: usize, ) -> crate::Result>>> { @@ -907,7 +912,17 @@ fn intersect_row_allow_lists( } match (physical, residual) { (None, residual) => Ok(residual), - (Some(physical), None) => Ok(Some(physical.to_vec())), + (Some(physical), None) => Ok(Some( + physical + .iter() + .map(|per_file| { + per_file + .iter() + .map(|(file, ranges)| Ok((file.clone(), positions_in_ranges(ranges)?))) + .collect::>>() + }) + .collect::>>()?, + )), (Some(physical), Some(residual)) => { if residual.len() != split_count { return Err(crate::Error::DataInvalid { @@ -925,14 +940,15 @@ fn intersect_row_allow_lists( .map(|(physical, residual)| { physical .iter() - .filter_map(|(file, allowed)| { - residual - .get(file) - .map(|kept| (file.clone(), allowed & kept)) + .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() + .collect::>>() }) - .collect(), + .collect::>>()?, )) } } @@ -1159,13 +1175,21 @@ async fn search_pk_raw_candidates_batch_with_plan( 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?, ); @@ -2208,11 +2232,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 @@ -2229,6 +2258,7 @@ async fn residual_positions_by_file( split: &DataSplit, active_files: &[BucketActiveFile], residual: &FilePredicates, + allowed_rows: Option<&HashMap>>, ) -> crate::Result> { let scan_fields = reader.read_type().to_vec(); let active_names: HashSet<&str> = active_files.iter().map(|f| f.file_name.as_str()).collect(); @@ -2239,16 +2269,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 + 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( @@ -2257,24 +2319,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) @@ -7306,7 +7378,7 @@ 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]); @@ -7318,7 +7390,7 @@ mod residual_positions_tests { 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")); @@ -7329,7 +7401,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]); @@ -7347,7 +7419,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]); @@ -7390,7 +7462,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]); @@ -7406,7 +7478,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]); @@ -7459,6 +7531,21 @@ mod residual_positions_tests { .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> { + 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, file: &str) -> Vec { map.get(file) .map(|positions| positions.iter().collect()) @@ -7472,7 +7559,7 @@ mod residual_positions_tests { #[test] fn one_side_alone_passes_through() { - let physical = vec![allow_list(&[("d0", &[1, 2])])]; + 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"); @@ -7490,7 +7577,7 @@ mod residual_positions_tests { // `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![allow_list(&[("d0", &[1, 2, 3]), ("d1", &[0, 1])])]; + 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() @@ -7501,7 +7588,7 @@ mod residual_positions_tests { #[test] fn rejects_allow_lists_that_do_not_cover_every_split() { - let physical = vec![allow_list(&[("d0", &[1])])]; + 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"); From fcb9103c0d8a1f44a651d40d7936591179a735db Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Tue, 1 Sep 2026 13:55:00 +0800 Subject: [PATCH 6/7] perf(table): limit the exact fallback to the rows the split allows too MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The exact fallback had the same defect as the residual scan, and did not even need a data predicate to hit it: it opened every active file in full and let `is_excluded` reject the rows outside the split's ranges afterwards. Scoring is not the expensive part of that — reading the file is. `search_file` now takes the plan's selection for the file and reads through the ranges entry point, recovering each row's physical position from the selection the same way the residual scan does. `is_excluded` still applies on top, since it also folds in a residual predicate and the deletion vector, but it can no longer be the only thing standing between a ten-row split and a full file read. An empty selection returns without opening the file at all. Three tests, one per layer, because a result-only assertion cannot tell range-limited reading from full-read-then-filter: - the mosaic reader requests strictly fewer byte ranges, and fewer bytes, when the selection sits inside one row group. That is the granularity of the win: a row group is skipped before its column data is touched, so a narrow range does not seek within a row group. - the residual scan over a selection returns only positions inside it, and the rows it excludes would otherwise have matched the predicate. - the exact fallback returns the allowed positions rather than the nearest rows, which are outside the selection. The first two would fail loudly rather than quietly under a full read: positions come from walking the selection, so an over-long read exhausts it and reports that. --- crates/paimon/src/arrow/format/mosaic.rs | 58 +++++ .../src/table/pk_vector_data_file_reader.rs | 218 +++++++++++++++--- .../paimon/src/table/vector_search_builder.rs | 76 +++++- 3 files changed, 316 insertions(+), 36 deletions(-) diff --git a/crates/paimon/src/arrow/format/mosaic.rs b/crates/paimon/src/arrow/format/mosaic.rs index 1fdbbace7..b10c756c0 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>, + ) -> crate::Result>> { + let file_size = data.len() as u64; + let calls = Arc::new(Mutex::new(Vec::new())); + let _: Vec = 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/pk_vector_data_file_reader.rs b/crates/paimon/src/table/pk_vector_data_file_reader.rs index c3087beb4..2c406f18c 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,15 @@ 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 drained row count is checked against what was read + /// (both truncation and overrun fail loud). + /// + /// `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 +111,7 @@ impl DataFilePkVectorReaderFactory { metric: VectorSearchMetric, exact_limit: usize, is_excluded: &(dyn Fn(i64) -> bool + Sync), + allowed_rows: Option<&[RowRange]>, ) -> crate::Result>> { if exact_limit == 0 { return Err(data_invalid("vector search limit must be positive")); @@ -129,21 +137,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 + Send> = match allowed_rows { + Some(ranges) => { + let ranges: Vec = 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> = (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>> = Vec::new(); - let mut position: i64 = 0; while let Some(batch) = stream.try_next().await? { batch_vectors.clear(); append_batch_vectors( @@ -153,13 +188,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 +210,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 +460,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 = 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 +549,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 +600,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 +658,7 @@ mod integration_tests { VectorSearchMetric::L2, 2, &is_excluded, + None, ) .await .unwrap(); @@ -556,11 +666,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 +712,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 +740,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 +755,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 +774,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/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index b398d9001..f9da300aa 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -1217,8 +1217,12 @@ async fn search_pk_raw_candidates_batch_with_plan( // 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]], @@ -1234,11 +1238,24 @@ async fn search_pk_raw_candidates_batch_with_plan( row_count: file.row_count, }; let owned_queries: Vec> = 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 }) }, @@ -7384,6 +7401,61 @@ mod residual_positions_tests { 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. From 3b9ba6f21d1627e61351c0e91a06ba108cb3718d Mon Sep 17 00:00:00 2001 From: JunRuiLee Date: Tue, 1 Sep 2026 14:41:16 +0800 Subject: [PATCH 7/7] fix(table): fail on an unknown row count, and say what the range read saves Two corrections to the range pushdown. An unlisted file was read as "the whole file" through a helper that mapped any non-positive row count to an empty selection, so `DataFileMeta::ROW_COUNT_UNKNOWN` (-1) silently dropped the file from the search. Before the pushdown that conversion failed loudly, and it has to keep failing: the protocol decoder only validates the row count of a file it carries ranges for, which leaves the unlisted ones to this helper. What a range read saves is also format-dependent, and the reader's contract now says so rather than implying it is universal. Mosaic skips a row group before touching its column data, parquet skips pages through the offset index, `.row` prunes blocks. Avro is the exception: it loads the whole file and deserializes every record before applying the selection, so there a narrow selection saves only what comes after decoding. Java allows `file.format=avro` for a vector column, so that case is reachable rather than hypothetical; making Avro prune physically needs a block-aware reader and is not part of this change. Also narrows the exact fallback's comment about what its row-count check proves: a selected read vouches for the ranges it asked for and cannot notice a file truncated elsewhere, which a full read could. --- crates/paimon/src/table/data_file_reader.rs | 8 ++++ .../src/table/pk_vector_data_file_reader.rs | 6 ++- crates/paimon/src/table/pk_vector_scan.rs | 38 ++++++++++++++----- 3 files changed, 41 insertions(+), 11 deletions(-) diff --git a/crates/paimon/src/table/data_file_reader.rs b/crates/paimon/src/table/data_file_reader.rs index cd72a1c12..8958b3827 100644 --- a/crates/paimon/src/table/data_file_reader.rs +++ b/crates/paimon/src/table/data_file_reader.rs @@ -614,6 +614,14 @@ impl DataFileReader { /// 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, 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 2c406f18c..dc42f1728 100644 --- a/crates/paimon/src/table/pk_vector_data_file_reader.rs +++ b/crates/paimon/src/table/pk_vector_data_file_reader.rs @@ -95,8 +95,10 @@ 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 what was read - /// (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 diff --git a/crates/paimon/src/table/pk_vector_scan.rs b/crates/paimon/src/table/pk_vector_scan.rs index 932b8e12b..61cdfcea8 100644 --- a/crates/paimon/src/table/pk_vector_scan.rs +++ b/crates/paimon/src/table/pk_vector_scan.rs @@ -58,14 +58,20 @@ pub(super) fn positions_in_ranges(ranges: &[RowRange]) -> crate::Result Vec { - if row_count > 0 { - vec![RowRange::new(0, row_count - 1)] - } else { - Vec::new() +/// 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> { + 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}" + ))), } } @@ -558,7 +564,7 @@ fn plan_from_bucket_splits( // 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), + None => whole_file_range(file.row_count)?, }; Ok((file.file_name.clone(), allowed)) }) @@ -1267,6 +1273,20 @@ mod tests { 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())