From d6eb60a9baefca98484b934808d2f344c208359a Mon Sep 17 00:00:00 2001 From: osipovartem Date: Thu, 15 Oct 2026 21:13:07 +0300 Subject: [PATCH] Prune pre-epoch day and hour partitions conservatively --- datafusion_iceberg/src/pruning_statistics.rs | 190 +++++++++++++++++-- 1 file changed, 173 insertions(+), 17 deletions(-) diff --git a/datafusion_iceberg/src/pruning_statistics.rs b/datafusion_iceberg/src/pruning_statistics.rs index 24906645..adb06f28 100644 --- a/datafusion_iceberg/src/pruning_statistics.rs +++ b/datafusion_iceberg/src/pruning_statistics.rs @@ -505,19 +505,18 @@ pub(crate) fn transform_predicate( let field = partition_fields .iter() .find(|field| field.source_name() == column.name())?; - if field.transform() != &Transform::Identity { + let legacy_negative_day_hour = if field.transform() != &Transform::Identity { let Expr::Literal(value, _) = &literal else { return None; }; if !literal_matches_source(value, field) { return None; } - if matches!(field.transform(), Transform::Day | Transform::Hour) + matches!(field.transform(), Transform::Day | Transform::Hour) && matches!(value, ScalarValue::TimestampMicrosecond(Some(value), _) if *value < 0) - { - return None; - } - } + } else { + false + }; let op = if column_on_left { op } else { @@ -530,6 +529,18 @@ pub(crate) fn transform_predicate( _ => return None, } }; + if legacy_negative_day_hour && field.transform() == &Transform::Hour { + // Hour IDs wrap outside i32. Equality remains safe with both encodings, but a + // monotone range can drop matching rows whose partition ID has wrapped. + if op != Operator::Eq { + return None; + } + let Expr::Literal(ScalarValue::TimestampMicrosecond(Some(micros), _), _) = &literal else { + return None; + }; + let hour = i32::try_from(micros.div_euclid(3_600_000_000)).ok()?; + hour.checked_add(1)?; + } let op = match field.transform() { Transform::Identity => op, // Legacy numeric manifests and truncated String min/max statistics can @@ -550,7 +561,8 @@ pub(crate) fn transform_predicate( }; let literal = transform_literal(literal, field.transform())?; let column = Expr::Column(Column::new(column.relation, field.name().to_owned())); - if field.transform() == &Transform::Month { + // Legacy pre-epoch day/hour IDs may be truncated toward zero instead of floored. + if field.transform() == &Transform::Month || legacy_negative_day_hour { let legacy = Expr::BinaryExpr(BinaryExpr::new( Box::new(literal.clone()), Operator::Plus, @@ -638,6 +650,21 @@ fn transform_literal(expr: Expr, transform: &Transform) -> Option { return None; } match transform { + Transform::Day | Transform::Hour if matches!(value, ScalarValue::TimestampMicrosecond(Some(micros), _) if *micros < 0) => + { + let ScalarValue::TimestampMicrosecond(Some(micros), _) = value else { + return None; + }; + let unit = if transform == &Transform::Day { + 86_400_000_000 + } else { + 3_600_000_000 + }; + Some(Expr::Literal( + ScalarValue::Int32(Some(i32::try_from(micros.div_euclid(unit)).ok()?)), + None, + )) + } Transform::Year => Some(Expr::ScalarFunction(ScalarFunction::new_udf( Arc::new(ScalarUDF::new_from_impl(DateTransform::new())), vec![Expr::Literal(ScalarValue::new_utf8("year"), None), expr], @@ -1695,23 +1722,152 @@ mod tests { } #[test] - fn negative_timestamp_day_and_hour_skip_legacy_unsafe_pruning() { + fn negative_timestamp_day_and_hour_prune_both_partition_encodings() { + let manifests = (-2..=1) + .map(|partition_id| ManifestListEntry { + format_version: FormatVersion::V2, + manifest_path: format!("/{partition_id}.avro"), + manifest_length: 1, + partition_spec_id: 1, + content: ManifestContent::Data, + sequence_number: 1, + min_sequence_number: 1, + added_snapshot_id: 1, + added_files_count: Some(1), + existing_files_count: Some(0), + deleted_files_count: Some(0), + added_rows_count: Some(1), + existing_rows_count: Some(0), + deleted_rows_count: Some(0), + partitions: Some(vec![FieldSummary { + contains_null: false, + contains_nan: None, + lower_bound: Some(Value::Int(partition_id)), + upper_bound: Some(Value::Int(partition_id)), + }]), + key_metadata: None, + first_row_id: None, + }) + .collect::>(); + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "partition_id", + DataType::Int32, + true, + )])); + let session = SessionContext::new(); + let state = session.state(); for transform in [Transform::Day, Transform::Hour] { - assert!(project_id_predicate( - transform, - PrimitiveType::Timestamp, - Operator::Eq, - ScalarValue::TimestampMicrosecond(Some(-1), None), - ) - .is_none()); + let source = StructField::new( + 1, + "id", + false, + Type::Primitive(PrimitiveType::Timestamp), + None, + ); + let partition = PartitionField::new(1, 1000, "partition_id", transform.clone()); + let fields = [BoundPartitionField::new(&partition, &source)]; + let cases = if transform == Transform::Hour { + vec![(Operator::Eq, [false, true, true, false])] + } else { + vec![ + (Operator::Eq, [false, true, true, false]), + (Operator::Lt, [true, true, true, false]), + (Operator::Gt, [false, true, true, true]), + ] + }; + for (op, expected) in cases { + let expr = project_id_predicate( + transform.clone(), + PrimitiveType::Timestamp, + op, + ScalarValue::TimestampMicrosecond(Some(-1), None), + ) + .expect("pre-epoch timestamp should permit conservative pruning"); + let physical = create_physical_expr( + &expr, + &schema.clone().try_into().unwrap(), + state.execution_props(), + &PhysicalPlanningContext::default(), + ) + .unwrap(); + let predicate = PruningPredicateBuilder::new() + .with_file_schema(schema.clone()) + .try_build(physical) + .unwrap(); + assert_eq!( + predicate + .prune(&PruneManifests::new(&fields, 1, &manifests)) + .unwrap(), + expected, + "{transform:?} {op:?} must retain floor and legacy IDs" + ); + } + if transform == Transform::Hour { + for op in [Operator::Lt, Operator::LtEq, Operator::Gt, Operator::GtEq] { + assert!( + project_id_predicate( + Transform::Hour, + PrimitiveType::Timestamp, + op, + ScalarValue::TimestampMicrosecond(Some(-1), None), + ) + .is_none(), + "negative hour ranges cannot safely prune wrapped partition IDs" + ); + } + assert!( + project_id_predicate( + Transform::Hour, + PrimitiveType::Timestamp, + Operator::Eq, + ScalarValue::TimestampMicrosecond(Some(-7_730_941_136_400_000_000), None), + ) + .is_none(), + "out-of-range hour ID must not overflow its adjacent candidate" + ); + } } - assert!(project_id_predicate( + let extreme_day = project_id_predicate( + Transform::Day, + PrimitiveType::Timestamp, + Operator::Eq, + ScalarValue::TimestampMicrosecond(Some(i64::MIN), None), + ) + .expect("extreme pre-epoch day should project without Chrono conversion"); + let Expr::BinaryExpr(BinaryExpr { + left, + op: Operator::And, + .. + }) = extreme_day + else { + panic!("expected conservative day interval"); + }; + let Expr::BinaryExpr(BinaryExpr { right, .. }) = *left else { + panic!("expected lower day bound"); + }; + assert_eq!( + *right, + Expr::Literal( + ScalarValue::Int32(Some( + i32::try_from(i64::MIN.div_euclid(86_400_000_000)).unwrap() + )), + None + ) + ); + let date = project_id_predicate( Transform::Day, PrimitiveType::Date, Operator::Eq, ScalarValue::Date32(Some(-1)), ) - .is_some()); + .expect("Date32 does not need a legacy timestamp candidate"); + assert!(matches!( + date, + Expr::BinaryExpr(BinaryExpr { + op: Operator::Eq, + .. + }) + )); } #[test]