diff --git a/crates/paimon/src/spec/avro/manifest_file_meta_decode.rs b/crates/paimon/src/spec/avro/manifest_file_meta_decode.rs index 6beff6166..6ba791cc2 100644 --- a/crates/paimon/src/spec/avro/manifest_file_meta_decode.rs +++ b/crates/paimon/src/spec/avro/manifest_file_meta_decode.rs @@ -40,6 +40,7 @@ impl AvroRecordDecode for ManifestFileMeta { let mut max_level: Option = None; let mut min_row_id: Option = None; let mut max_row_id: Option = None; + let mut total_buckets: Option = None; let mut extra_files: Option> = None; for field in &writer_schema.fields { @@ -64,6 +65,7 @@ impl AvroRecordDecode for ManifestFileMeta { "_MAX_LEVEL" => max_level = read_optional_int(cursor, field.nullable)?, "_MIN_ROW_ID" => min_row_id = read_optional_long(cursor, field.nullable)?, "_MAX_ROW_ID" => max_row_id = read_optional_long(cursor, field.nullable)?, + "_TOTAL_BUCKETS" => total_buckets = read_optional_int(cursor, field.nullable)?, "_EXTRA_FILES" => { extra_files = decode_nullable_string_array(cursor, field.nullable)? } @@ -85,6 +87,7 @@ impl AvroRecordDecode for ManifestFileMeta { max_level, min_row_id, max_row_id, + total_buckets, extra_files, )) } diff --git a/crates/paimon/src/spec/manifest_file_meta.rs b/crates/paimon/src/spec/manifest_file_meta.rs index adcc95c6f..0d37b7145 100644 --- a/crates/paimon/src/spec/manifest_file_meta.rs +++ b/crates/paimon/src/spec/manifest_file_meta.rs @@ -104,6 +104,13 @@ pub struct ManifestFileMeta { )] max_row_id: Option, + /// Common positive bucket count recorded by an external manifest writer. + /// + /// Rust consumes this field for manifest pruning but intentionally does not + /// serialize it into manifest lists. + #[serde(rename = "_TOTAL_BUCKETS", default, skip_serializing)] + total_buckets: Option, + /// Files owned by this manifest and sharing its lifecycle. /// /// `None` preserves the distinction between legacy manifest lists (where the @@ -195,6 +202,12 @@ impl ManifestFileMeta { self.max_row_id } + /// Get the common positive bucket count for entries in this manifest. + #[inline] + pub fn total_buckets(&self) -> Option { + self.total_buckets + } + /// Get files owned by this manifest, if the metadata was recorded. #[inline] pub fn extra_files(&self) -> Option<&[String]> { @@ -231,6 +244,15 @@ impl ManifestFileMeta { self } + /// Attach external manifest metadata in read-path tests. + #[cfg(test)] + #[inline] + #[must_use] + pub(crate) fn with_total_buckets(mut self, total_buckets: Option) -> Self { + self.total_buckets = total_buckets.filter(|value| *value > 0); + self + } + /// Attach files whose lifecycle is owned by this manifest. #[inline] #[must_use] @@ -262,6 +284,7 @@ impl ManifestFileMeta { max_level: None, min_row_id: None, max_row_id: None, + total_buckets: None, extra_files: None, } } @@ -282,6 +305,7 @@ impl ManifestFileMeta { max_level: Option, min_row_id: Option, max_row_id: Option, + total_buckets: Option, extra_files: Option>, ) -> ManifestFileMeta { Self { @@ -298,6 +322,7 @@ impl ManifestFileMeta { max_level, min_row_id, max_row_id, + total_buckets, extra_files, } } diff --git a/crates/paimon/src/spec/objects_file.rs b/crates/paimon/src/spec/objects_file.rs index 78421b0a0..8f5bfaad8 100644 --- a/crates/paimon/src/spec/objects_file.rs +++ b/crates/paimon/src/spec/objects_file.rs @@ -229,6 +229,7 @@ mod tests { None, Some(100), Some(199), + None, Some(vec!["manifest-row-tracking-0.idx".to_string()]), )]; let bytes = to_avro_bytes(MANIFEST_FILE_META_SCHEMA, &original).unwrap(); @@ -242,6 +243,62 @@ mod tests { ); } + #[test] + fn test_read_manifest_file_meta_total_buckets_without_writing_it() { + assert!(!MANIFEST_FILE_META_SCHEMA.contains("_TOTAL_BUCKETS")); + + let original = vec![ManifestFileMeta::new( + "manifest-java-0".to_string(), + 1024, + 5, + 0, + BinaryTableStats::empty(), + 0, + ) + .with_bucket_level_stats(Some(2), Some(2), Some(0), Some(0))]; + let bytes = to_avro_bytes(MANIFEST_FILE_META_SCHEMA, &original).unwrap(); + let mut value = Reader::new(bytes.as_slice()) + .unwrap() + .next() + .unwrap() + .unwrap(); + let fields = match &mut value { + Value::Union(_, record) => match record.as_mut() { + Value::Record(fields) => fields, + other => panic!("Expected an Avro record, got {other:?}"), + }, + other => panic!("Expected an Avro union, got {other:?}"), + }; + let extra_files_index = fields + .iter() + .position(|(name, _)| name == "_EXTRA_FILES") + .unwrap(); + fields.insert( + extra_files_index, + ( + "_TOTAL_BUCKETS".to_string(), + Value::Union(1, Box::new(Value::Int(8))), + ), + ); + + let java_schema = MANIFEST_FILE_META_SCHEMA.replacen( + r#"{"name": "_EXTRA_FILES", "type": ["null", {"type": "array", "items": "string"}], "default": null}"#, + concat!( + r#"{"name": "_TOTAL_BUCKETS", "type": ["null", "int"], "default": null},"#, + "\n ", + r#"{"name": "_EXTRA_FILES", "type": ["null", {"type": "array", "items": "string"}], "default": null}"# + ), + 1, + ); + let schema = Schema::parse_str(&java_schema).unwrap(); + let mut writer = Writer::new(&schema, Vec::new()); + writer.append(value).unwrap(); + let bytes = writer.into_inner().unwrap(); + + let decoded = from_avro_bytes_fast::(&bytes).unwrap(); + assert_eq!(decoded[0].total_buckets(), Some(8)); + } + #[test] fn test_roundtrip_manifest_entry() { let original = vec![manifest_entry()]; diff --git a/crates/paimon/src/table/table_scan.rs b/crates/paimon/src/table/table_scan.rs index a83880c80..49b999ce1 100644 --- a/crates/paimon/src/table/table_scan.rs +++ b/crates/paimon/src/table/table_scan.rs @@ -212,6 +212,14 @@ async fn read_all_manifest_entries( trace.manifest_files_after_partition_pruning = manifest_files.len(); } + retain_manifest_buckets( + &mut manifest_files, + has_primary_keys && !scan_all_files, + bucket_predicate, + bucket_key_fields, + bucket_function_type, + ); + if let Some(index) = row_range_index { let before = manifest_files.len(); retain_manifest_row_ranges(&mut manifest_files, index); @@ -364,6 +372,58 @@ fn retain_manifest_row_ranges( manifests.retain(|manifest| manifest_file_overlaps_row_range_index(manifest, row_range_index)); } +/// Conservatively prune manifests using their bucket envelope and a common +/// positive `_TOTAL_BUCKETS` value. Missing, invalid, mixed, or legacy metadata +/// always fails open. Tight envelopes become especially effective when +/// manifests are bucket-sorted. +fn retain_manifest_buckets( + manifests: &mut Vec, + only_real_buckets: bool, + bucket_predicate: Option<&Predicate>, + bucket_key_fields: &[DataField], + bucket_function_type: BucketFunctionType, +) { + let mut target_cache: HashMap>> = HashMap::new(); + manifests.retain(|manifest| { + let (Some(min_bucket), Some(max_bucket)) = (manifest.min_bucket(), manifest.max_bucket()) + else { + return true; + }; + if min_bucket > max_bucket { + return true; + } + if only_real_buckets && max_bucket < 0 { + return false; + } + + let Some(predicate) = bucket_predicate else { + return true; + }; + let Some(total_buckets) = manifest.total_buckets().filter(|value| *value > 0) else { + return true; + }; + // A mixed range containing the unassigned bucket is not a safe + // representation of the real-bucket subset. + if min_bucket < 0 || max_bucket >= total_buckets { + return true; + } + + let targets = target_cache.entry(total_buckets).or_insert_with(|| { + compute_target_buckets( + predicate, + bucket_key_fields, + bucket_function_type, + total_buckets, + ) + }); + targets.as_ref().is_none_or(|targets| { + targets + .iter() + .any(|bucket| *bucket >= min_bucket && *bucket <= max_bucket) + }) + }); +} + fn data_file_overlaps_row_range_index( file: &DataFileMeta, row_range_index: &RowRangeIndex, @@ -1830,6 +1890,14 @@ impl<'a> PaimonTableScan<'a> { }; let bucket_function_type = core_options.bucket_function_type()?; + retain_manifest_buckets( + &mut manifest_metas, + has_primary_keys && !self.scan_all_files, + self.bucket_predicate.as_ref(), + &bucket_key_fields, + bucket_function_type, + ); + let base_path = format!("{}/{}", table_path.trim_end_matches('/'), MANIFEST_DIR); let shared_cache = SharedSchemaCache::new(); let partition_filter = self.partition_filter.as_ref(); @@ -2394,10 +2462,10 @@ mod tests { data_evolution_row_range_groups, data_file_overlaps_row_range_index, group_data_files_by_partition_bucket, manifest_file_overlaps_row_range_index, prune_data_evolution_group_by_read_fields, retain_index_manifest_entry, - retain_index_manifest_entry_for_scan, retain_manifest_entry_row_ranges, - retain_manifest_row_ranges, scan_predicate_field_ids, should_skip_level_zero_for_scan, - split_row_ranges_for_files, LimitPushdownAccumulator, PaimonTableScan, RowRangeIndex, - TableScan, + retain_index_manifest_entry_for_scan, retain_manifest_buckets, + retain_manifest_entry_row_ranges, retain_manifest_row_ranges, scan_predicate_field_ids, + should_skip_level_zero_for_scan, split_row_ranges_for_files, LimitPushdownAccumulator, + PaimonTableScan, RowRangeIndex, TableScan, }; use crate::catalog::Identifier; use crate::io::FileIOBuilder; @@ -2500,6 +2568,50 @@ mod tests { )); } + #[test] + fn test_manifest_bucket_pruning_is_exact_and_fails_open() { + let fields = int_field(); + let predicate = PredicateBuilder::new(&fields) + .equal("id", Datum::Int(42)) + .unwrap(); + let target = *compute_target_buckets(&predicate, &fields, BucketFunctionType::Default, 8) + .unwrap() + .iter() + .next() + .unwrap(); + let other = (target + 1) % 8; + let stats = BinaryTableStats::new(Vec::new(), Vec::new(), Vec::new()); + let manifest = |name: &str, bucket: i32, total_buckets| { + ManifestFileMeta::new(name.to_string(), 1, 1, 0, stats.clone(), 0) + .with_bucket_level_stats(Some(bucket), Some(bucket), Some(0), Some(0)) + .with_total_buckets(total_buckets) + }; + let mut manifests = vec![ + manifest("target", target, Some(8)), + manifest("other", other, Some(8)), + manifest("legacy", other, None), + manifest("unassigned", -1, Some(8)), + manifest("at-total", 8, Some(8)), + manifest("above-total", 9, Some(8)), + ]; + + retain_manifest_buckets( + &mut manifests, + true, + Some(&predicate), + &fields, + BucketFunctionType::Default, + ); + + assert_eq!( + manifests + .iter() + .map(ManifestFileMeta::file_name) + .collect::>(), + vec!["target", "legacy", "at-total", "above-total"] + ); + } + #[test] fn test_manifest_row_range_pruning_uses_each_envelope() { let index = RowRangeIndex::create(vec![RowRange::new(2, 2)]);