From 2db4ace0abe82250179dc66b51d83e6f0ff65148 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Sat, 10 Oct 2026 01:57:17 +0300 Subject: [PATCH 1/3] Prune Iceberg data files by partition values within manifests --- datafusion_iceberg/src/pruning_statistics.rs | 165 ++++++++++++++++++- datafusion_iceberg/src/table/mod.rs | 94 +++++++---- 2 files changed, 222 insertions(+), 37 deletions(-) diff --git a/datafusion_iceberg/src/pruning_statistics.rs b/datafusion_iceberg/src/pruning_statistics.rs index 66204a45..4c82a7fc 100644 --- a/datafusion_iceberg/src/pruning_statistics.rs +++ b/datafusion_iceberg/src/pruning_statistics.rs @@ -12,7 +12,7 @@ * For the second level the trait PruningStatistics is implemented for the Manifest */ -use std::{any::Any, sync::Arc}; +use std::{any::Any, collections::HashSet, sync::Arc}; use crate::error::Error as DatafusionIcebergError; use datafusion::{ @@ -34,7 +34,7 @@ use iceberg_rust::{ error::Error, spec::{ decimal::{decimal_mantissa, decimal_scale, Decimal}, - manifest::ManifestEntry, + manifest::{Content, ManifestEntry}, manifest_list::ManifestListEntry, partition::{BoundPartitionField, Transform}, schema::Schema, @@ -154,6 +154,80 @@ pub(crate) struct PruneDataFiles<'table, 'manifests> { files: &'manifests [(ManifestPath, ManifestEntry)], } +pub(crate) struct PrunePartitionedDataFiles<'table, 'manifests> { + partition_fields: &'table [BoundPartitionField<'table>], + compatible_manifests: &'manifests HashSet<&'manifests str>, + files: &'manifests [(ManifestPath, ManifestEntry)], +} + +impl<'table, 'manifests> PrunePartitionedDataFiles<'table, 'manifests> { + pub(crate) fn new( + partition_fields: &'table [BoundPartitionField<'table>], + compatible_manifests: &'manifests HashSet<&'manifests str>, + files: &'manifests [(ManifestPath, ManifestEntry)], + ) -> Self { + Self { + partition_fields, + compatible_manifests, + files, + } + } + + fn partition_values(&self, column: &Column) -> Option { + let field = self + .partition_fields + .iter() + .find(|field| field.name() == column.name())?; + let field_type = field.field_type().tranform(field.transform()).ok()?; + let data_type = (&field_type).try_into().ok()?; + let values = self.files.iter().map(|(path, entry)| { + (entry.data_file().content() == &Content::Data + && self.compatible_manifests.contains(path.as_str())) + .then(|| entry.data_file().partition().get(column.name())) + .flatten() + .and_then(Option::as_ref) + .and_then(|value| value.clone().cast(&field_type).ok()) + .map(Value::into_any) + }); + any_iter_to_array(values, &data_type).ok() + } +} + +impl PruningStatistics for PrunePartitionedDataFiles<'_, '_> { + fn min_values(&self, column: &Column) -> Option { + self.partition_values(column) + } + + fn max_values(&self, column: &Column) -> Option { + self.partition_values(column) + } + + fn num_containers(&self) -> usize { + self.files.len() + } + + fn null_counts(&self, _column: &Column) -> Option { + None + } + + fn contained( + &self, + _column: &Column, + _values: &std::collections::HashSet, + ) -> Option { + None + } + + fn row_counts(&self) -> Option { + ScalarValue::iter_to_array( + self.files + .iter() + .map(|(_, entry)| ScalarValue::Int64(Some(*entry.data_file().record_count()))), + ) + .ok() + } +} + impl<'table, 'manifests> PruneDataFiles<'table, 'manifests> { pub(crate) fn new( schema: &'table Schema, @@ -714,6 +788,93 @@ mod tests { assert!(null_counts.is_null(0)); } + #[test] + fn partitioned_file_pruning_preserves_unknown_specs_and_delete_files() { + use datafusion::arrow::array::Int32Array; + + let source = StructField::new( + 1, + "ts", + false, + Type::Primitive(PrimitiveType::Timestamp), + None, + ); + let partition = PartitionField::new(1, 1000, "ts_year", Transform::Year); + let fields = [BoundPartitionField::new(&partition, &source)]; + let entry = |content: Content, year: Option| { + let file = DataFile::builder() + .with_content(content) + .with_file_path("/data.parquet".into()) + .with_file_format(FileFormat::Parquet) + .with_partition(Struct::from_iter(vec![( + "ts_year".to_owned(), + year.map(Value::Int), + )])) + .with_record_count(1) + .with_file_size_in_bytes(1) + .with_column_sizes(None) + .with_value_counts(None) + .with_null_value_counts(None) + .with_nan_value_counts(None) + .with_distinct_counts(None) + .with_lower_bounds(None) + .with_upper_bounds(None) + .build() + .unwrap(); + ManifestEntry::builder() + .with_format_version(FormatVersion::V2) + .with_status(Status::Added) + .with_data_file(file) + .build() + .unwrap() + }; + let files = vec![ + ("current".into(), entry(Content::Data, Some(54))), + ("current".into(), entry(Content::Data, Some(55))), + ("old_spec".into(), entry(Content::Data, Some(54))), + ("current".into(), entry(Content::Data, None)), + ("current".into(), entry(Content::PositionDeletes, Some(54))), + ]; + let compatible = HashSet::from(["current"]); + let pruning = PrunePartitionedDataFiles::new(&fields, &compatible, &files); + let values = pruning.min_values(&Column::from_name("ts_year")).unwrap(); + let values = values.as_any().downcast_ref::().unwrap(); + assert_eq!(values.len(), 5); + assert_eq!(values.value(0), 54); + assert_eq!(values.value(1), 55); + assert!(values.is_null(2)); + assert!(values.is_null(3)); + assert!(values.is_null(4)); + assert_eq!(pruning.num_containers(), 5); + + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "ts_year", + DataType::Int32, + true, + )])); + let session = SessionContext::new(); + let state = session.state(); + let physical = create_physical_expr( + &partition_comparison( + Expr::Column(Column::from_name("ts_year")), + Operator::Eq, + Expr::Literal(ScalarValue::Int32(Some(55)), None), + ), + &schema.clone().try_into().unwrap(), + state.execution_props(), + &PhysicalPlanningContext::default(), + ) + .unwrap(); + let predicate = PruningPredicateBuilder::new() + .with_file_schema(schema) + .try_build(physical) + .unwrap(); + assert_eq!( + predicate.prune(&pruning).unwrap(), + vec![false, true, true, true, true] + ); + } + #[test] fn month_projection_keeps_legacy_and_spec_manifests() { let source = StructField::new(1, "d", false, Type::Primitive(PrimitiveType::Date), None); diff --git a/datafusion_iceberg/src/table/mod.rs b/datafusion_iceberg/src/table/mod.rs index b9e80e63..d400799c 100644 --- a/datafusion_iceberg/src/table/mod.rs +++ b/datafusion_iceberg/src/table/mod.rs @@ -50,7 +50,9 @@ use crate::statistics::statistics_from_datafiles; use crate::variant_schema_adapter::IcebergPhysicalExprAdapterFactory; use crate::{ error::Error as DataFusionIcebergError, - pruning_statistics::{transform_predicate, PruneDataFiles, PruneManifests}, + pruning_statistics::{ + transform_predicate, PruneDataFiles, PruneManifests, PrunePartitionedDataFiles, + }, statistics::manifest_statistics, }; use datafusion::arrow::compute::SortOptions; @@ -796,44 +798,42 @@ async fn table_scan( .filter_map(|expr| transform_predicate(expr, partition_fields)), ); - // If there is a filter expression on the partition column, the manifest files to read are pruned. - let data_files: Vec<(ManifestPath, ManifestEntry)> = - if let Some(predicate) = partition_predicates { - let physical_partition_predicate = create_physical_expr( - &predicate, - &partition_schema.clone().try_into()?, - session.execution_props(), - &PhysicalPlanningContext::default(), - )?; - let pruning_predicate = PruningPredicateBuilder::new() + let partition_pruning_predicate = if let Some(predicate) = partition_predicates { + let physical_partition_predicate = create_physical_expr( + &predicate, + &partition_schema.clone().try_into()?, + session.execution_props(), + &PhysicalPlanningContext::default(), + )?; + Some( + PruningPredicateBuilder::new() .with_file_schema(partition_schema.clone()) - .try_build(physical_partition_predicate)?; - let manifests_to_prune = pruning_predicate.prune(&PruneManifests::new( + .try_build(physical_partition_predicate)?, + ) + } else { + None + }; + let manifests_to_prune = partition_pruning_predicate + .as_ref() + .map(|predicate| { + predicate.prune(&PruneManifests::new( partition_fields, table.metadata().default_spec_id, &manifests, - ))?; - - table - .datafiles( - &manifests, - Some(manifests_to_prune), - manifest_entry_sequence_range, - ) - .await - .map_err(DataFusionIcebergError::from)? - .try_collect() - .await - .map_err(DataFusionIcebergError::from)? - } else { - table - .datafiles(&manifests, None, manifest_entry_sequence_range) - .await - .map_err(DataFusionIcebergError::from)? - .try_collect() - .await - .map_err(DataFusionIcebergError::from)? - }; + )) + }) + .transpose()?; + let data_files: Vec<(ManifestPath, ManifestEntry)> = table + .datafiles( + &manifests, + manifests_to_prune, + manifest_entry_sequence_range, + ) + .await + .map_err(DataFusionIcebergError::from)? + .try_collect() + .await + .map_err(DataFusionIcebergError::from)?; let mut data_files = prepare_changed_data_files(data_files, scan_changed_data_files); if let Some(excluded) = excluded_data_file_paths.filter(|paths| !paths.is_empty()) { data_files.retain(|(_, entry)| { @@ -842,6 +842,30 @@ async fn table_scan( }); } + if let Some(predicate) = partition_pruning_predicate { + let compatible_manifests: HashSet<&str> = manifests + .iter() + .filter(|manifest| manifest.partition_spec_id == table.metadata().default_spec_id) + .map(|manifest| manifest.manifest_path.as_str()) + .collect(); + let has_prunable_files = data_files.iter().any(|(path, entry)| { + entry.data_file().content() == &Content::Data + && compatible_manifests.contains(path.as_str()) + }); + if has_prunable_files { + let files_to_keep = predicate.prune(&PrunePartitionedDataFiles::new( + partition_fields, + &compatible_manifests, + &data_files, + ))?; + data_files = data_files + .into_iter() + .zip(files_to_keep) + .filter_map(|(file, keep)| keep.then_some(file)) + .collect(); + } + } + let pruning_predicate = PruningPredicateBuilder::new() .with_file_schema(arrow_schema.clone()) .try_build(physical_predicate)?; From 0eb6d170f746867ad6fd4f3effe682fb9579cedc Mon Sep 17 00:00:00 2001 From: osipovartem Date: Sat, 10 Oct 2026 02:10:28 +0300 Subject: [PATCH 2/3] Project exact cast literals onto Iceberg time partitions --- datafusion_iceberg/src/pruning_statistics.rs | 81 +++++++++++++++++++- 1 file changed, 79 insertions(+), 2 deletions(-) diff --git a/datafusion_iceberg/src/pruning_statistics.rs b/datafusion_iceberg/src/pruning_statistics.rs index 4c82a7fc..23cc2f76 100644 --- a/datafusion_iceberg/src/pruning_statistics.rs +++ b/datafusion_iceberg/src/pruning_statistics.rs @@ -429,8 +429,8 @@ pub(crate) fn transform_predicate( return None; }; let (column, literal, column_on_left) = match (*left, *right) { - (Expr::Column(column), literal @ Expr::Literal(..)) => (column, literal, true), - (literal @ Expr::Literal(..), Expr::Column(column)) => (column, literal, false), + (Expr::Column(column), literal) => (column, fold_cast_literal(literal)?, true), + (literal, Expr::Column(column)) => (column, fold_cast_literal(literal)?, false), _ => return None, }; let field = partition_fields @@ -505,6 +505,27 @@ pub(crate) fn transform_predicate( Some(partition_comparison(column, op, literal)) } +fn fold_cast_literal(expr: Expr) -> Option { + match expr { + literal @ Expr::Literal(..) => Some(literal), + Expr::Cast(cast) => { + let Expr::Literal(value, _) = *cast.expr else { + return None; + }; + let value = value.cast_to(cast.field.data_type()).ok()?; + let value = match value { + ScalarValue::TimestampNanosecond(Some(nanos), None) if nanos % 1_000 == 0 => { + ScalarValue::TimestampMicrosecond(Some(nanos / 1_000), None) + } + ScalarValue::TimestampNanosecond(..) => return None, + value => value, + }; + Some(Expr::Literal(value, None)) + } + _ => None, + } +} + fn partition_comparison(column: Expr, op: Operator, literal: Expr) -> Expr { Expr::BinaryExpr(BinaryExpr::new(Box::new(column), op, Box::new(literal))) } @@ -1220,6 +1241,62 @@ mod tests { assert_eq!(result, input); } + #[test] + fn year_projection_folds_only_exact_timestamp_literal_casts() { + use datafusion_expr::expr::Cast; + + let source = StructField::new( + 1, + "ts", + false, + Type::Primitive(PrimitiveType::Timestamp), + None, + ); + let partition = PartitionField::new(1, 1000, "ts_year", Transform::Year); + let fields = [BoundPartitionField::new(&partition, &source)]; + let timestamp = Expr::Cast(Cast::new( + Box::new(Expr::Literal( + ScalarValue::Utf8(Some("2024-06-15".to_owned())), + None, + )), + DataType::Timestamp(TimeUnit::Nanosecond, None), + )); + assert!(matches!( + fold_cast_literal(timestamp.clone()), + Some(Expr::Literal( + ScalarValue::TimestampMicrosecond(Some(_), None), + _ + )) + )); + assert!(transform_predicate( + partition_comparison( + Expr::Column(Column::from_name("ts")), + Operator::Eq, + timestamp, + ), + &fields, + ) + .is_some()); + + let sub_microsecond = Expr::Cast(Cast::new( + Box::new(Expr::Literal( + ScalarValue::TimestampNanosecond(Some(-1), None), + None, + )), + DataType::Timestamp(TimeUnit::Nanosecond, None), + )); + assert!(fold_cast_literal(sub_microsecond.clone()).is_none()); + assert!(transform_predicate( + partition_comparison( + Expr::Column(Column::from_name("ts")), + Operator::Eq, + sub_microsecond, + ), + &fields, + ) + .is_none()); + } + fn project_id_predicate( transform: Transform, source_type: PrimitiveType, From a6f8047728f3979995d0158dcefce0803d248b8b Mon Sep 17 00:00:00 2001 From: osipovartem Date: Sat, 10 Oct 2026 02:27:07 +0300 Subject: [PATCH 3/3] Exercise cast projection through file-pruning mask --- datafusion_iceberg/src/pruning_statistics.rs | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/datafusion_iceberg/src/pruning_statistics.rs b/datafusion_iceberg/src/pruning_statistics.rs index 23cc2f76..9b45fff6 100644 --- a/datafusion_iceberg/src/pruning_statistics.rs +++ b/datafusion_iceberg/src/pruning_statistics.rs @@ -812,6 +812,7 @@ mod tests { #[test] fn partitioned_file_pruning_preserves_unknown_specs_and_delete_files() { use datafusion::arrow::array::Int32Array; + use datafusion_expr::expr::Cast; let source = StructField::new( 1, @@ -875,12 +876,20 @@ mod tests { )])); let session = SessionContext::new(); let state = session.state(); + let source_predicate = partition_comparison( + Expr::Column(Column::from_name("ts")), + Operator::Eq, + Expr::Cast(Cast::new( + Box::new(Expr::Literal( + ScalarValue::Utf8(Some("2025-06-15".to_owned())), + None, + )), + DataType::Timestamp(TimeUnit::Nanosecond, None), + )), + ); + let projected = transform_predicate(source_predicate, &fields).unwrap(); let physical = create_physical_expr( - &partition_comparison( - Expr::Column(Column::from_name("ts_year")), - Operator::Eq, - Expr::Literal(ScalarValue::Int32(Some(55)), None), - ), + &projected, &schema.clone().try_into().unwrap(), state.execution_props(), &PhysicalPlanningContext::default(),