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..c82826757 100644 --- a/crates/paimon/src/table/pk_vector_bucket_split.rs +++ b/crates/paimon/src/table/pk_vector_bucket_split.rs @@ -126,6 +126,39 @@ impl BucketVectorPayload { .as_deref() .expect("a decoded payload always carries source metadata") } + + /// Consume the payload into the pieces a planner needs, so its decoded metadata + /// moves out of the payload rather than being cloned out of it. + /// + /// Two decoded fields are deliberately left behind. `row_count` is the payload's + /// own row count, which the read path derives from the source metadata instead. + /// `deletion_vectors_ranges` belongs to deletion-vector index files -- Java + /// builds a vector payload through the overload that leaves it null, and a read + /// takes its deletion vectors from the bucket's data split -- so a value here + /// describes something this payload is not, and is ignored rather than applied. + #[cfg_attr(not(test), allow(dead_code))] + pub(crate) fn into_parts(self) -> BucketVectorPayloadParts { + BucketVectorPayloadParts { + index_type: self.index_type, + file_name: self.file_name, + file_size: self.file_size, + external_path: self.external_path, + global_index_meta: self.global_index_meta, + } + } +} + +/// The owned pieces of a [`BucketVectorPayload`], produced by +/// [`BucketVectorPayload::into_parts`]. +#[cfg_attr(not(test), allow(dead_code))] +pub(crate) struct BucketVectorPayloadParts { + pub(crate) index_type: String, + pub(crate) file_name: String, + /// As decoded: Java writes a signed length, so a negative value is possible on + /// the wire and is rejected where it is converted, not here. + pub(crate) file_size: i64, + pub(crate) external_path: Option, + 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..cf1a67de2 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,214 @@ impl<'a> PkVectorScan<'a> { Ok(PkVectorScanPlan { snapshot_id, splits, + physical_row_ranges_by_split: None, }) } + + /// Build a plan from bucket splits an engine planned elsewhere, instead of from + /// this table's index manifest. + /// + /// The splits are the planning input and are taken as authoritative: their + /// payload files, their per-file row ranges, and the snapshot they pin are used + /// as given, and no index manifest is read. Only the partition conjuncts of this + /// scan's filter are re-applied, because a caller may narrow the query further + /// than the planner that produced the splits. + /// + /// Mirrors what Java's `PrimaryKeyVectorRead` does with a + /// `BucketVectorSearchSplit`: search the payloads the split names, over the rows + /// the split allows. + // Entry point for engine-supplied splits; no in-tree caller reads a plan from + // them yet, and the tests drive `plan_from_bucket_splits` directly. + #[allow(dead_code)] + pub(crate) fn plan_for_bucket_vector_splits( + &self, + splits: Vec, + ) -> 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 +642,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 +1154,383 @@ 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 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 aa1fdc6dc..3168f1cd7 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,112 @@ async fn plan_and_search_pk_candidates_batch( } } - let plan = PkVectorScan::new(table, field_id, index_type.clone(), filter.cloned()) - .plan() - .await?; + Ok(PkVectorSearchParams { + metric, + concurrency, + index_type, + field_id, + vector_field, + skip_exact_fallback, + refine_factor, + indexed_limit, + }) +} + +/// Search an already-resolved plan across every query and return each query's raw +/// indexed and exact candidate lists, before any rerank or merge. +/// +/// Plan-dependent concurrency — the vindex segment count, batch-index parallelism +/// and the range-read bound — is derived here from the plan that is actually being +/// searched, so a narrowed plan can never be searched under limits computed for a +/// wider one. +/// Combine the two per-split row allow-lists a search can be handed: the physical +/// 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, + 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 @@ -1061,6 +1175,15 @@ async fn plan_and_search_pk_candidates_batch( } None => None, }; + // Fold the plan's own positional restriction into the same allow-list. A plan + // built from engine-supplied bucket splits carries the physical positions each + // file is limited to; a plan read from the index manifest carries none. Both + // sides list what is permitted, so combining them is an intersection. + let residual_by_split = intersect_row_allow_lists( + plan.physical_row_ranges_by_split.as_deref(), + residual_by_split, + plan.splits.len(), + )?; // Build the exact-fallback search on demand: the kernel calls this only for a // file it actually searches (uncovered by ANN, residual-allowed, and only when @@ -1122,6 +1245,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 +1316,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> { @@ -7242,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}"); + } } 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}.