From f1ad82afcf947990f598300baf74dcd246e9a983 Mon Sep 17 00:00:00 2001 From: QuakeWang Date: Sun, 9 Aug 2026 22:07:35 +0800 Subject: [PATCH] fix(scan): prune primary-key files with key stats Rust key-value files store primary-key statistics in key_stats while value_stats may be empty, so scan planning cannot prune them with primary-key predicates. Evaluate file-safe key predicates against key_stats, including schema-evolution mapping, while preserving fail-open behavior. Signed-off-by: QuakeWang --- crates/paimon/src/table/pk_vector_scan.rs | 43 +++--- crates/paimon/src/table/stats_filter.rs | 165 ++++++++++++++------ crates/paimon/src/table/table_scan.rs | 180 +++++++++++++++++++++- 3 files changed, 312 insertions(+), 76 deletions(-) diff --git a/crates/paimon/src/table/pk_vector_scan.rs b/crates/paimon/src/table/pk_vector_scan.rs index d9681d1cf..00bec032b 100644 --- a/crates/paimon/src/table/pk_vector_scan.rs +++ b/crates/paimon/src/table/pk_vector_scan.rs @@ -752,17 +752,11 @@ mod tests { ) } - /// Build a real single-file primary-key table via the public write path, in a - /// fresh temp dir. Persists the schema and writes one data batch, then commits - /// the written data file with real `value_stats` for the `id`/`score` columns. + /// Build a real single-file primary-key table via the public write path. /// - /// The stats injection mirrors the meta-modification the baseline fixture uses - /// for `level`/`file_source`: the Rust key-value (primary-key) writer records - /// column stats in `key_stats` and leaves `value_stats` empty, but scan-time - /// file pruning reads `value_stats`. Java primary-key writers populate value - /// stats, so committing them here makes the file prunable exactly as it would be - /// in a table written by the Java engine. Returns the temp dir (kept alive by - /// the caller) and the opened table. + /// The Rust key-value writer records primary-key stats in `key_stats`. For the + /// deletion-vector test, also populate `value_stats` so its non-key predicate + /// has the same metadata a Java primary-key writer produces. async fn build_pruning_test_table( with_deletion_vectors: bool, ) -> (tempfile::TempDir, Table) { @@ -793,19 +787,24 @@ mod tests { let base_meta = written.new_files[0].clone(); let bucket = written.bucket; let partition = written.partition.clone(); - - // Real value stats over the `id` (col 0) and `score` (col 1) columns, so a - // predicate outside the written [0, PRUNE_ROWS) range can prune the file. - let int = DataType::Int(IntType::new()); - let value_stats: BinaryTableStats = - compute_column_stats(&batch, &[0, 1], &[int.clone(), int]).unwrap(); - let indexed_meta = DataFileMeta { - value_stats, - value_stats_cols: Some(vec!["id".to_string(), "score".to_string()]), - ..base_meta + assert_eq!(base_meta.key_stats.null_counts(), &vec![Some(0)]); + assert!(base_meta.value_stats.null_counts().is_empty()); + assert_eq!(base_meta.value_stats_cols, Some(vec![])); + + let file_meta = if with_deletion_vectors { + let int = DataType::Int(IntType::new()); + let value_stats: BinaryTableStats = + compute_column_stats(&batch, &[0, 1], &[int.clone(), int]).unwrap(); + DataFileMeta { + value_stats, + value_stats_cols: Some(vec!["id".to_string(), "score".to_string()]), + ..base_meta + } + } else { + base_meta }; - let message = CommitMessage::new(partition, bucket, vec![indexed_meta]); + let message = CommitMessage::new(partition, bucket, vec![file_meta]); TableCommit::new(table.clone(), "pkvector-prune".to_string()) .commit(vec![message]) .await @@ -831,7 +830,7 @@ mod tests { } #[tokio::test] - async fn plan_prunes_file_when_pk_predicate_excludes_it() { + async fn plan_prunes_rust_pk_file_using_key_stats() { // Real PK table, one data file with id in [0, PRUNE_ROWS). A predicate // `id = OUT_OF_RANGE` cannot match the file's id stats, so the scan drops // the file and plan() returns no splits. Control (no filter) returns one. diff --git a/crates/paimon/src/table/stats_filter.rs b/crates/paimon/src/table/stats_filter.rs index 7de875c5c..8e2f06b05 100644 --- a/crates/paimon/src/table/stats_filter.rs +++ b/crates/paimon/src/table/stats_filter.rs @@ -37,11 +37,26 @@ pub(super) struct FileStatsRows { null_counts: Vec>, supports_in_min_max_pruning: bool, /// Maps schema field index → stats index. `None` means identity mapping - /// (stats cover all schema fields in order). `Some` is used when - /// `value_stats_cols` or `write_cols` is present (dense mode). + /// (stats cover all schema fields in order). `Some` maps dense stats such + /// as `value_stats_cols`, `write_cols`, or trimmed primary-key stats. stats_col_mapping: Option>>, } +fn dense_stats_col_mapping( + schema_fields: &[DataField], + stats_columns: &[String], +) -> Vec> { + let col_index: HashMap<&str, usize> = stats_columns + .iter() + .enumerate() + .map(|(i, c)| (c.as_str(), i)) + .collect(); + schema_fields + .iter() + .map(|field| col_index.get(field.name()).copied()) + .collect() +} + impl FileStatsRows { /// Build a `FileStatsRows` for manifest-level partition stats (no column mapping). pub(super) fn for_manifest_partition( @@ -66,31 +81,11 @@ impl FileStatsRows { /// When `value_stats_cols` is `Some`, stats are in dense mode — only covering those /// columns, and the mapping from schema field index to stats index is built by name. pub(super) fn from_data_file(file: &DataFileMeta, schema_fields: &[DataField]) -> Self { - let stats_col_mapping = if let Some(cols) = &file.value_stats_cols { - let col_index: HashMap<&str, usize> = cols - .iter() - .enumerate() - .map(|(i, c)| (c.as_str(), i)) - .collect(); - let mapping: Vec> = schema_fields - .iter() - .map(|field| col_index.get(field.name()).copied()) - .collect(); - Some(mapping) - } else if let Some(cols) = &file.write_cols { - let col_index: HashMap<&str, usize> = cols - .iter() - .enumerate() - .map(|(i, c)| (c.as_str(), i)) - .collect(); - let mapping: Vec> = schema_fields - .iter() - .map(|field| col_index.get(field.name()).copied()) - .collect(); - Some(mapping) - } else { - None - }; + let stats_col_mapping = file + .value_stats_cols + .as_ref() + .or(file.write_cols.as_ref()) + .map(|cols| dense_stats_col_mapping(schema_fields, cols)); Self { row_count: file.row_count, @@ -102,6 +97,25 @@ impl FileStatsRows { } } + /// Build file stats from `_KEY_STATS`. + /// + /// Key stats are dense and follow the table's trimmed primary-key order + /// (primary-key columns with partition columns removed). + pub(super) fn from_key_stats( + file: &DataFileMeta, + schema_fields: &[DataField], + key_fields: &[String], + ) -> Self { + Self { + row_count: file.row_count, + min_values: BinaryRow::from_serialized_bytes(file.key_stats.min_values()).ok(), + max_values: BinaryRow::from_serialized_bytes(file.key_stats.max_values()).ok(), + null_counts: file.key_stats.null_counts().clone(), + supports_in_min_max_pruning: true, + stats_col_mapping: Some(dense_stats_col_mapping(schema_fields, key_fields)), + } + } + /// Resolve a schema field index to the corresponding stats index. fn stats_index(&self, schema_index: usize) -> Option { match &self.stats_col_mapping { @@ -148,6 +162,7 @@ impl StatsAccessor for FileStatsRows { pub(super) struct ResolvedStatsSchema { file_fields: Vec, field_mapping: Vec>, + key_fields: Vec, } fn identity_field_mapping(num_fields: usize) -> Vec> { @@ -165,41 +180,89 @@ fn normalize_field_mapping(mapping: Option>, num_fields: usize) -> Vec< .unwrap_or_else(|| identity_field_mapping(num_fields)) } +fn has_always_false(predicates: &[Predicate], key_predicates: &[Predicate]) -> bool { + predicates + .iter() + .chain(key_predicates) + .any(|p| matches!(p, Predicate::AlwaysFalse)) +} + +fn matches_file_stats( + file: &DataFileMeta, + predicates: &[Predicate], + key_predicates: &[Predicate], + field_mapping: &[Option], + schema_fields: &[DataField], + key_fields: &[String], +) -> bool { + let value_stats = FileStatsRows::from_data_file(file, schema_fields); + if !predicates_may_match_with_schema(predicates, &value_stats, field_mapping, schema_fields) { + return false; + } + + if key_predicates.is_empty() { + return true; + } + + let key_stats = FileStatsRows::from_key_stats(file, schema_fields, key_fields); + predicates_may_match_with_schema(key_predicates, &key_stats, field_mapping, schema_fields) +} + /// Check whether a data file *may* contain rows matching all `predicates`. /// +/// Value predicates are evaluated against `_VALUE_STATS`; primary-key +/// predicates are additionally evaluated against `_KEY_STATS`. Both sources +/// must report that the file may match. +/// /// Pruning is evaluated per file and fails open when stats cannot be /// interpreted safely, including schema mismatches, incompatible stats arity, /// and missing or corrupted stats. -pub(super) fn data_file_matches_predicates( +pub(super) fn data_file_matches_predicates_with_key_stats( file: &DataFileMeta, predicates: &[Predicate], + key_predicates: &[Predicate], current_schema_id: i64, schema_fields: &[DataField], + key_fields: &[String], ) -> bool { - if predicates.is_empty() { + if predicates.is_empty() && key_predicates.is_empty() { return true; } - if predicates - .iter() - .any(|p| matches!(p, Predicate::AlwaysFalse)) - { + if has_always_false(predicates, key_predicates) { return false; } - if predicates - .iter() - .all(|p| matches!(p, Predicate::AlwaysTrue)) - { - return true; - } if file.schema_id != current_schema_id { return true; } - let stats = FileStatsRows::from_data_file(file, schema_fields); let field_mapping = identity_field_mapping(schema_fields.len()); - predicates_may_match_with_schema(predicates, &stats, &field_mapping, schema_fields) + matches_file_stats( + file, + predicates, + key_predicates, + &field_mapping, + schema_fields, + key_fields, + ) +} + +#[cfg(test)] +pub(super) fn data_file_matches_predicates( + file: &DataFileMeta, + predicates: &[Predicate], + current_schema_id: i64, + schema_fields: &[DataField], +) -> bool { + data_file_matches_predicates_with_key_stats( + file, + predicates, + &[], + current_schema_id, + schema_fields, + &[], + ) } async fn resolve_stats_schema( @@ -217,6 +280,7 @@ async fn resolve_stats_schema( Some(Arc::new(ResolvedStatsSchema { file_fields: current_fields.to_vec(), field_mapping: identity_field_mapping(current_fields.len()), + key_fields: table_schema.trimmed_primary_keys(), })) } else { let file_schema = table.schema_manager().schema(file_schema_id).await.ok()?; @@ -227,6 +291,7 @@ async fn resolve_stats_schema( current_fields.len(), ), file_fields, + key_fields: file_schema.trimmed_primary_keys(), })) }; @@ -238,18 +303,25 @@ pub(super) async fn data_file_matches_predicates_for_table( table: &Table, file: &DataFileMeta, predicates: &[Predicate], + key_predicates: &[Predicate], schema_cache: &mut HashMap>>, ) -> bool { - if predicates.is_empty() { + if predicates.is_empty() && key_predicates.is_empty() { return true; } + if has_always_false(predicates, key_predicates) { + return false; + } + if file.schema_id == table.schema().id() { - return data_file_matches_predicates( + return data_file_matches_predicates_with_key_stats( file, predicates, + key_predicates, table.schema().id(), table.schema().fields(), + &table.schema().trimmed_primary_keys(), ); } @@ -257,12 +329,13 @@ pub(super) async fn data_file_matches_predicates_for_table( return true; }; - let stats = FileStatsRows::from_data_file(file, &resolved.file_fields); - predicates_may_match_with_schema( + matches_file_stats( + file, predicates, - &stats, + key_predicates, &resolved.field_mapping, &resolved.file_fields, + &resolved.key_fields, ) } diff --git a/crates/paimon/src/table/table_scan.rs b/crates/paimon/src/table/table_scan.rs index 471606a15..6dcdff817 100644 --- a/crates/paimon/src/table/table_scan.rs +++ b/crates/paimon/src/table/table_scan.rs @@ -27,8 +27,8 @@ use super::global_index_types::normalize_sorted_global_index_type; use super::kv_file_reader::retain_primary_key_conjuncts; use super::partition_filter::PartitionFilter; use super::stats_filter::{ - data_evolution_group_matches_predicates, data_file_matches_predicates, - data_file_matches_predicates_for_table, group_by_overlapping_row_id, FileStatsRows, + data_evolution_group_matches_predicates, data_file_matches_predicates_for_table, + data_file_matches_predicates_with_key_stats, group_by_overlapping_row_id, FileStatsRows, ResolvedStatsSchema, }; use super::Table; @@ -123,8 +123,10 @@ async fn read_all_manifest_entries( partition_filter: Option<&PartitionFilter>, partition_fields: &[DataField], data_predicates: &[Predicate], + key_predicates: &[Predicate], current_schema_id: i64, schema_fields: &[DataField], + key_fields: &[String], bucket_predicate: Option<&Predicate>, bucket_key_fields: &[DataField], bucket_function_type: BucketFunctionType, @@ -239,12 +241,14 @@ async fn read_all_manifest_entries( counters.pruned_by_level += 1; continue; } - if !data_predicates.is_empty() - && !data_file_matches_predicates( + if (!data_predicates.is_empty() || !key_predicates.is_empty()) + && !data_file_matches_predicates_with_key_stats( entry.file(), data_predicates, + key_predicates, current_schema_id, schema_fields, + key_fields, ) { counters.pruned_by_data_stats += 1; @@ -1144,6 +1148,8 @@ impl<'a> PaimonTableScan<'a> { } else { self.stats_pruning_predicates() }; + let key_fields = self.table.schema().trimmed_primary_keys(); + let pushdown_key_predicates = self.key_stats_predicates(&pushdown_data_predicates); let bucket_key_fields: Vec = if self.bucket_predicate.is_none() { Vec::new() @@ -1179,8 +1185,10 @@ impl<'a> PaimonTableScan<'a> { self.partition_filter.as_ref(), &partition_fields, &pushdown_data_predicates, + &pushdown_key_predicates, self.table.schema().id(), self.table.schema().fields(), + &key_fields, self.bucket_predicate.as_ref(), &bucket_key_fields, bucket_function_type, @@ -1350,6 +1358,16 @@ impl<'a> PaimonTableScan<'a> { } } + /// Project file-safe predicates onto trimmed primary-key columns while + /// preserving table-schema field indices. + fn key_stats_predicates(&self, predicates: &[Predicate]) -> Vec { + let key_fields = self.table.schema().trimmed_primary_keys(); + if key_fields.is_empty() { + return Vec::new(); + } + retain_primary_key_conjuncts(predicates, self.table.schema().fields(), &key_fields) + } + /// Plan data splits from a snapshot's delta manifest list (APPEND deltas). /// /// Reuses the same split-building path as a full snapshot plan, but only @@ -1724,6 +1742,7 @@ impl<'a> PaimonTableScan<'a> { // For non-data-evolution tables, cross-schema files were kept (fail-open) // by the pushdown. Apply the full schema-aware filter for those files. let stats_pruning_predicates = self.stats_pruning_predicates(); + let key_stats_predicates = self.key_stats_predicates(&stats_pruning_predicates); let entries = if stats_pruning_predicates.is_empty() || data_evolution_enabled { entries } else { @@ -1747,6 +1766,7 @@ impl<'a> PaimonTableScan<'a> { self.table, entry.file(), &stats_pruning_predicates, + &key_stats_predicates, &mut schema_cache, ) .await @@ -2042,16 +2062,17 @@ mod tests { use crate::io::FileIOBuilder; use crate::spec::{ stats::BinaryTableStats, ArrayType, BinaryRow, BinaryRowBuilder, BucketFunctionType, - CommitKind, DataField, DataFileMeta, DataType, Datum, DeletionVectorMeta, FileKind, - IndexFileMeta, IndexManifestEntry, IntType, ManifestEntry, ManifestFileMeta, Predicate, - PredicateBuilder, PredicateOperator, Schema as PaimonSchema, Snapshot, TableSchema, - VarCharType, + ColumnMove, CommitKind, DataField, DataFileMeta, DataType, Datum, DeletionVectorMeta, + FileKind, IndexFileMeta, IndexManifestEntry, IntType, ManifestEntry, ManifestFileMeta, + Predicate, PredicateBuilder, PredicateOperator, Schema as PaimonSchema, SchemaChange, + Snapshot, TableSchema, VarCharType, }; use crate::table::bucket_filter::{compute_target_buckets, extract_predicate_for_keys}; use crate::table::partition_filter::PartitionFilter; use crate::table::source::{DataSplit, DataSplitBuilder, DeletionFile, RowRange}; use crate::table::stats_filter::{ data_evolution_group_matches_predicates, data_file_matches_predicates, + data_file_matches_predicates_for_table, data_file_matches_predicates_with_key_stats, group_by_overlapping_row_id, }; use crate::table::{CommitMessage, Table, TableCommit}; @@ -3198,6 +3219,149 @@ mod tests { )); } + #[test] + fn test_data_file_matches_composite_key_stats_in_primary_key_order() { + let fields = vec![ + DataField::new(0, "key_a".to_string(), DataType::Int(IntType::new())), + DataField::new(1, "payload".to_string(), DataType::Int(IntType::new())), + DataField::new(2, "key_b".to_string(), DataType::Int(IntType::new())), + ]; + let empty_stats = BinaryTableStats::empty(); + let mut file = test_data_file_meta( + empty_stats.min_values().to_vec(), + empty_stats.max_values().to_vec(), + Vec::new(), + 5, + ); + file.value_stats_cols = Some(Vec::new()); + // key_stats order is [key_b, key_a], independent of table field order. + file.key_stats = BinaryTableStats::new( + two_int_stats_row(Some(100), Some(10)), + two_int_stats_row(Some(200), Some(20)), + vec![Some(0), Some(0)], + ); + let key_fields = vec!["key_b".to_string(), "key_a".to_string()]; + let pb = PredicateBuilder::new(&fields); + + let key_a_out_of_range = pb.equal("key_a", Datum::Int(150)).unwrap(); + assert!(!data_file_matches_predicates_with_key_stats( + &file, + std::slice::from_ref(&key_a_out_of_range), + std::slice::from_ref(&key_a_out_of_range), + TEST_SCHEMA_ID, + &fields, + &key_fields, + )); + + let key_b_in_range = pb.equal("key_b", Datum::Int(150)).unwrap(); + assert!(data_file_matches_predicates_with_key_stats( + &file, + std::slice::from_ref(&key_b_in_range), + std::slice::from_ref(&key_b_in_range), + TEST_SCHEMA_ID, + &fields, + &key_fields, + )); + } + + #[tokio::test] + async fn test_data_file_matches_key_stats_across_renamed_reordered_schema() { + let old_schema = TableSchema::new( + 0, + &PaimonSchema::builder() + .column("payload", DataType::Int(IntType::new())) + .column("old_b", DataType::Int(IntType::new())) + .column("old_a", DataType::Int(IntType::new())) + .primary_key(["old_b", "old_a"]) + .build() + .unwrap(), + ); + let current_schema = old_schema + .apply_changes(vec![ + SchemaChange::rename_column("old_b".to_string(), "new_b".to_string()), + SchemaChange::rename_column("old_a".to_string(), "new_a".to_string()), + SchemaChange::update_column_position(ColumnMove::move_first("new_a".to_string())), + ]) + .unwrap(); + let table = Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("test_db", "cross_schema_key_stats"), + "memory:/cross_schema_key_stats".to_string(), + current_schema, + None, + ); + write_schema_file(&table, &old_schema).await; + + let empty_stats = BinaryTableStats::empty(); + let mut file = test_data_file_meta_with_schema( + empty_stats.min_values().to_vec(), + empty_stats.max_values().to_vec(), + Vec::new(), + 5, + old_schema.id(), + ); + file.value_stats_cols = Some(Vec::new()); + // Old key_stats order is [old_b, old_a]. Current fields are reordered to + // [new_a, payload, new_b], while field IDs remain [2, 0, 1]. + file.key_stats = BinaryTableStats::new( + two_int_stats_row(Some(100), Some(10)), + two_int_stats_row(Some(200), Some(20)), + vec![Some(0), Some(0)], + ); + + let pb = PredicateBuilder::new(table.schema().fields()); + let new_a_out_of_range = pb.equal("new_a", Datum::Int(150)).unwrap(); + let new_b_in_range = pb.equal("new_b", Datum::Int(150)).unwrap(); + let mut schema_cache = HashMap::new(); + + assert!( + !data_file_matches_predicates_for_table( + &table, + &file, + std::slice::from_ref(&new_a_out_of_range), + std::slice::from_ref(&new_a_out_of_range), + &mut schema_cache, + ) + .await + ); + assert!( + data_file_matches_predicates_for_table( + &table, + &file, + std::slice::from_ref(&new_b_in_range), + std::slice::from_ref(&new_b_in_range), + &mut schema_cache, + ) + .await + ); + } + + #[test] + fn test_data_file_matches_corrupt_key_stats_fails_open() { + let fields = int_field(); + let empty_stats = BinaryTableStats::empty(); + let mut file = test_data_file_meta( + empty_stats.min_values().to_vec(), + empty_stats.max_values().to_vec(), + Vec::new(), + 5, + ); + file.value_stats_cols = Some(Vec::new()); + file.key_stats = BinaryTableStats::new(vec![0], vec![0], vec![Some(0)]); + let predicate = PredicateBuilder::new(&fields) + .equal("id", Datum::Int(30)) + .unwrap(); + + assert!(data_file_matches_predicates_with_key_stats( + &file, + std::slice::from_ref(&predicate), + std::slice::from_ref(&predicate), + TEST_SCHEMA_ID, + &fields, + &["id".to_string()], + )); + } + #[tokio::test] async fn test_plan_with_trace_records_between_data_stats_pruning() { let table_path = "memory:/test_plan_with_trace_records_between_data_stats_pruning";