feat(scan): read bucket metadata for manifest pruning

crates/paimon/src/spec/avro/manifest_file_meta_decode.rs

@@ -40,6 +40,7 @@ impl AvroRecordDecode for ManifestFileMeta {
         let mut max_level: Option<i32> = None;
         let mut min_row_id: Option<i64> = None;
         let mut max_row_id: Option<i64> = None;
+        let mut total_buckets: Option<i32> = None;
         let mut extra_files: Option<Vec<String>> = 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,
         ))
     }

crates/paimon/src/spec/manifest_file_meta.rs

@@ -104,6 +104,13 @@ pub struct ManifestFileMeta {
     )]
     max_row_id: Option<i64>,
 
+    /// 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<i32>,
+
     /// 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<i32> {
+        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<i32>) -> 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<i32>,
         min_row_id: Option<i64>,
         max_row_id: Option<i64>,
+        total_buckets: Option<i32>,
         extra_files: Option<Vec<String>>,
     ) -> ManifestFileMeta {
         Self {
@@ -298,6 +322,7 @@ impl ManifestFileMeta {
             max_level,
             min_row_id,
             max_row_id,
+            total_buckets,
             extra_files,
         }
     }

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::<ManifestFileMeta>(&bytes).unwrap();
+        assert_eq!(decoded[0].total_buckets(), Some(8));
+    }
+
     #[test]
     fn test_roundtrip_manifest_entry() {
         let original = vec![manifest_entry()];

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<crate::spec::ManifestFileMeta>,
+    only_real_buckets: bool,
+    bucket_predicate: Option<&Predicate>,
+    bucket_key_fields: &[DataField],
+    bucket_function_type: BucketFunctionType,
+) {
+    let mut target_cache: HashMap<i32, Option<HashSet<i32>>> = 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<_>>(),
+            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)]);