From f0faa6bfe710fab64736edf0da5eb2a9d1789ce1 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Sat, 10 Oct 2026 03:11:04 +0300 Subject: [PATCH 1/2] Keep semantic timestamp casts in physical comparisons and pruning --- datafusion/expr-common/src/casts.rs | 26 +++++++ .../src/simplify_expressions/unwrap_cast.rs | 12 +--- .../src/simplifier/unwrap_cast.rs | 67 +++++++++++++++---- datafusion/pruning/src/pruning_predicate.rs | 44 +++++++++++- 4 files changed, 126 insertions(+), 23 deletions(-) diff --git a/datafusion/expr-common/src/casts.rs b/datafusion/expr-common/src/casts.rs index 3518c02772672..926a2b131ee92 100644 --- a/datafusion/expr-common/src/casts.rs +++ b/datafusion/expr-common/src/casts.rs @@ -84,6 +84,19 @@ pub fn is_supported_type(data_type: &DataType) -> bool { || is_supported_binary_type(data_type) } +/// Returns true when moving a timestamp cast across a comparison may change its result. +/// +/// Narrowing can truncate values, widening can overflow the integer storage, +/// and changing the timezone can shift timestamp values. An identity cast does +/// none of these and remains eligible for unwrapping. +pub fn changes_timestamp_semantics(from_type: &DataType, to_type: &DataType) -> bool { + matches!( + (from_type, to_type), + (DataType::Timestamp(from_unit, from_tz), DataType::Timestamp(to_unit, to_tz)) + if from_unit != to_unit || from_tz != to_tz + ) +} + fn is_date_type(data_type: &DataType) -> bool { matches!(data_type, DataType::Date32 | DataType::Date64) } @@ -1015,6 +1028,19 @@ mod tests { )); } + #[test] + fn test_timestamp_cast_semantics_change() { + let ts_us = DataType::Timestamp(TimeUnit::Microsecond, None); + let ts_ns = DataType::Timestamp(TimeUnit::Nanosecond, None); + let ts_tz = DataType::Timestamp(TimeUnit::Microsecond, Some("+01:00".into())); + assert!(changes_timestamp_semantics(&ts_us, &ts_ns)); + assert!(changes_timestamp_semantics(&ts_ns, &ts_us)); + assert!(changes_timestamp_semantics(&ts_us, &ts_tz)); + assert!(changes_timestamp_semantics(&ts_tz, &ts_us)); + assert!(!changes_timestamp_semantics(&ts_us, &ts_us)); + assert!(!changes_timestamp_semantics(&DataType::Int64, &ts_us)); + } + #[test] fn test_is_date_narrowing_cast() { // Only Date64 -> Date32 narrows (ms -> days, many-to-one). diff --git a/datafusion/optimizer/src/simplify_expressions/unwrap_cast.rs b/datafusion/optimizer/src/simplify_expressions/unwrap_cast.rs index 67747ecf525ac..4ce587a9392f4 100644 --- a/datafusion/optimizer/src/simplify_expressions/unwrap_cast.rs +++ b/datafusion/optimizer/src/simplify_expressions/unwrap_cast.rs @@ -60,7 +60,8 @@ use datafusion_common::{internal_err, tree_node::Transformed}; use datafusion_expr::{BinaryExpr, lit}; use datafusion_expr::{Cast, Expr, Operator, TryCast, simplify::SimplifyContext}; use datafusion_expr_common::casts::{ - is_date_narrowing_cast, is_supported_type, try_cast_literal_to_type, + changes_timestamp_semantics, is_date_narrowing_cast, is_supported_type, + try_cast_literal_to_type, }; pub(super) fn unwrap_cast_in_comparison_for_binary( @@ -204,15 +205,6 @@ pub(super) fn is_cast_expr_and_support_unwrap_cast_in_comparison_for_inlist( true } -fn changes_timestamp_semantics(from_type: &DataType, to_type: &DataType) -> bool { - // Changing units can truncate or overflow; attaching a timezone can shift values. - matches!( - (from_type, to_type), - (DataType::Timestamp(from_unit, from_tz), DataType::Timestamp(to_unit, to_tz)) - if from_unit != to_unit || from_tz != to_tz - ) -} - ///// Tries to move a cast from an expression (such as column) to the literal other side of a comparison operator./ /// /// Specifically, rewrites diff --git a/datafusion/physical-expr/src/simplifier/unwrap_cast.rs b/datafusion/physical-expr/src/simplifier/unwrap_cast.rs index 3e67fc8291a4e..8db913cf564ca 100644 --- a/datafusion/physical-expr/src/simplifier/unwrap_cast.rs +++ b/datafusion/physical-expr/src/simplifier/unwrap_cast.rs @@ -37,8 +37,7 @@ use arrow::datatypes::{DataType, Schema}; use datafusion_common::{Result, ScalarValue, tree_node::Transformed}; use datafusion_expr::Operator; use datafusion_expr_common::casts::{ - is_date_narrowing_cast, is_timestamp_precision_narrowing_cast, - try_cast_literal_to_type, + changes_timestamp_semantics, is_date_narrowing_cast, try_cast_literal_to_type, }; use crate::PhysicalExpr; @@ -130,7 +129,7 @@ fn try_unwrap_cast_comparison( // Get the data type of the inner expression let inner_type = inner_expr.data_type(schema)?; - if is_timestamp_precision_narrowing_cast(&inner_type, cast_type) + if changes_timestamp_semantics(&inner_type, cast_type) || is_date_narrowing_cast(&inner_type, cast_type) { return Ok(None); @@ -251,6 +250,56 @@ mod tests { assert!(!result.transformed); } + #[test] + fn test_no_unwrap_timestamp_timezone_or_unit_change() { + let source_type = DataType::Timestamp(TimeUnit::Microsecond, None); + let schema = Schema::new(vec![Field::new("ts", source_type.clone(), false)]); + for target_type in [ + DataType::Timestamp(TimeUnit::Microsecond, Some("+01:00".into())), + DataType::Timestamp(TimeUnit::Nanosecond, None), + ] { + let target_literal = ScalarValue::TimestampMicrosecond(Some(0), None) + .cast_to(&target_type) + .unwrap(); + let column = col("ts", &schema).unwrap(); + for cast in [ + Arc::new(CastExpr::new( + Arc::clone(&column), + target_type.clone(), + None, + )) as Arc, + Arc::new(TryCastExpr::new(Arc::clone(&column), target_type.clone())), + ] { + let literal = lit(target_literal.clone()); + for comparison in [ + Arc::new(BinaryExpr::new( + Arc::clone(&cast), + Operator::Eq, + Arc::clone(&literal), + )), + Arc::new(BinaryExpr::new(Arc::clone(&literal), Operator::Eq, cast)), + ] { + let result = unwrap_cast_in_comparison(comparison, &schema).unwrap(); + assert!(!result.transformed, "cast to {target_type:?} was unwrapped"); + } + } + } + } + + #[test] + fn test_unwrap_identity_timestamp_cast() { + let ts_type = DataType::Timestamp(TimeUnit::Microsecond, None); + let schema = Schema::new(vec![Field::new("ts", ts_type.clone(), false)]); + let column = col("ts", &schema).unwrap(); + let comparison = Arc::new(BinaryExpr::new( + Arc::new(CastExpr::new(column, ts_type, None)), + Operator::Eq, + lit(ScalarValue::TimestampMicrosecond(Some(0), None)), + )); + let result = unwrap_cast_in_comparison(comparison, &schema).unwrap(); + assert!(result.transformed); + } + #[test] fn test_no_unwrap_when_types_unsupported() { let schema = Schema::new(vec![Field::new("f1", DataType::Float32, false)]); @@ -601,7 +650,7 @@ mod tests { } #[test] - fn test_unwrap_timestamp_precision_widening() { + fn test_no_unwrap_timestamp_precision_widening() { let schema = Schema::new(vec![Field::new( "ts", DataType::Timestamp(TimeUnit::Millisecond, None), @@ -620,14 +669,8 @@ mod tests { let result = unwrap_cast_in_comparison(binary_expr, &schema).unwrap(); - assert!(result.transformed); - let optimized_binary = result.data.downcast_ref::().unwrap(); - assert!(!is_cast_expr(optimized_binary.left())); - let right_literal = optimized_binary.right().downcast_ref::().unwrap(); - assert_eq!( - right_literal.value(), - &ScalarValue::TimestampMillisecond(Some(1), None) - ); + // The literal fits, but other column values may overflow when widened. + assert!(!result.transformed); } #[test] diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index 3a63451495e4c..32f9ce59e147e 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -2185,7 +2185,7 @@ mod tests { use datafusion_physical_expr::utils::collect_columns; use insta::assert_snapshot; - use arrow::array::Decimal128Array; + use arrow::array::{Decimal128Array, TimestampMicrosecondArray}; use arrow::{ array::{BinaryArray, Int32Array, Int64Array, StringArray, UInt64Array}, datatypes::TimeUnit, @@ -3618,6 +3618,48 @@ mod tests { Ok(()) } + #[test] + fn row_group_predicate_keeps_timezone_cast_match() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "ts", + DataType::Timestamp(TimeUnit::Microsecond, None), + false, + )])); + let target_type = + DataType::Timestamp(TimeUnit::Microsecond, Some("+01:00".into())); + let matching_literal = + ScalarValue::TimestampMicrosecond(Some(0), None).cast_to(&target_type)?; + let expr = try_cast(col("ts"), target_type).eq(lit(matching_literal)); + let physical = logical2physical(&expr, &schema); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(TimestampMicrosecondArray::from(vec![Some(0)]))], + )?; + let evaluated = physical.evaluate(&batch)?.into_array(1)?; + assert!(evaluated.as_boolean().value(0), "test row must match"); + + let statistics = TestStatistics::new() + .with( + "ts", + ContainerStats::default() + .with_min(Arc::new(TimestampMicrosecondArray::from(vec![ + Some(0), + Some(1), + ]))) + .with_max(Arc::new(TimestampMicrosecondArray::from(vec![ + Some(0), + Some(1), + ]))), + ) + .with_null_counts("ts", [Some(0), Some(0)]) + .with_row_counts("ts", [Some(1), Some(1)]); + let pruning = PruningPredicateBuilder::new() + .with_file_schema(schema) + .try_build(physical)?; + assert_eq!(pruning.prune(&statistics)?, vec![true, false]); + Ok(()) + } + #[test] fn row_group_predicate_cast_string_string() -> Result<()> { let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]); From 1ab1c1b645d47289e789b6c726ad7e4b9c4c82c8 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Sat, 10 Oct 2026 03:14:50 +0300 Subject: [PATCH 2/2] Preserve timestamp cast unwrap for safe timezone changes --- datafusion/expr-common/src/casts.rs | 13 ++++-- .../src/simplifier/unwrap_cast.rs | 45 +++++++++++++++++++ 2 files changed, 54 insertions(+), 4 deletions(-) diff --git a/datafusion/expr-common/src/casts.rs b/datafusion/expr-common/src/casts.rs index 926a2b131ee92..6b21487375a54 100644 --- a/datafusion/expr-common/src/casts.rs +++ b/datafusion/expr-common/src/casts.rs @@ -87,13 +87,14 @@ pub fn is_supported_type(data_type: &DataType) -> bool { /// Returns true when moving a timestamp cast across a comparison may change its result. /// /// Narrowing can truncate values, widening can overflow the integer storage, -/// and changing the timezone can shift timestamp values. An identity cast does -/// none of these and remains eligible for unwrapping. +/// and attaching a timezone to a timezone-less timestamp shifts its value. +/// Removing or changing an existing timezone does not change the stored value +/// when the unit stays the same, so those casts remain eligible for unwrapping. pub fn changes_timestamp_semantics(from_type: &DataType, to_type: &DataType) -> bool { matches!( (from_type, to_type), (DataType::Timestamp(from_unit, from_tz), DataType::Timestamp(to_unit, to_tz)) - if from_unit != to_unit || from_tz != to_tz + if from_unit != to_unit || (from_tz.is_none() && to_tz.is_some()) ) } @@ -1036,7 +1037,11 @@ mod tests { assert!(changes_timestamp_semantics(&ts_us, &ts_ns)); assert!(changes_timestamp_semantics(&ts_ns, &ts_us)); assert!(changes_timestamp_semantics(&ts_us, &ts_tz)); - assert!(changes_timestamp_semantics(&ts_tz, &ts_us)); + assert!(!changes_timestamp_semantics(&ts_tz, &ts_us)); + assert!(!changes_timestamp_semantics( + &ts_tz, + &DataType::Timestamp(TimeUnit::Microsecond, Some("+02:00".into())) + )); assert!(!changes_timestamp_semantics(&ts_us, &ts_us)); assert!(!changes_timestamp_semantics(&DataType::Int64, &ts_us)); } diff --git a/datafusion/physical-expr/src/simplifier/unwrap_cast.rs b/datafusion/physical-expr/src/simplifier/unwrap_cast.rs index 8db913cf564ca..0c596e0818c32 100644 --- a/datafusion/physical-expr/src/simplifier/unwrap_cast.rs +++ b/datafusion/physical-expr/src/simplifier/unwrap_cast.rs @@ -300,6 +300,51 @@ mod tests { assert!(result.transformed); } + #[test] + fn test_unwrap_timestamp_timezone_removal_or_replacement() { + let source_type = + DataType::Timestamp(TimeUnit::Microsecond, Some("+01:00".into())); + let schema = Schema::new(vec![Field::new("ts", source_type, false)]); + for target_type in [ + DataType::Timestamp(TimeUnit::Microsecond, None), + DataType::Timestamp(TimeUnit::Microsecond, Some("+02:00".into())), + ] { + let column = col("ts", &schema).unwrap(); + for cast in [ + Arc::new(CastExpr::new( + Arc::clone(&column), + target_type.clone(), + None, + )) as Arc, + Arc::new(TryCastExpr::new(Arc::clone(&column), target_type.clone())), + ] { + let comparison = Arc::new(BinaryExpr::new( + cast, + Operator::Eq, + lit(ScalarValue::TimestampMicrosecond( + Some(42), + match &target_type { + DataType::Timestamp(_, tz) => tz.clone(), + _ => unreachable!(), + }, + )), + )); + let result = unwrap_cast_in_comparison(comparison, &schema).unwrap(); + assert!( + result.transformed, + "cast to {target_type:?} was not unwrapped" + ); + let optimized = result.data; + let binary = optimized.downcast_ref::().unwrap(); + assert!(!is_cast_expr(binary.left())); + assert_eq!( + binary.right().downcast_ref::().unwrap().value(), + &ScalarValue::TimestampMicrosecond(Some(42), Some("+01:00".into())) + ); + } + } + } + #[test] fn test_no_unwrap_when_types_unsupported() { let schema = Schema::new(vec![Field::new("f1", DataType::Float32, false)]);