Merge pull request #130 from Embucket/fix-df55-physical-timestamp-cast

Keep semantic timestamp casts in physical comparison pruning

Author: osipovartem

datafusion/expr-common/src/casts.rs

@@ -84,6 +84,20 @@ 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 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.is_none() && to_tz.is_some())
+    )
+}
+
 fn is_date_type(data_type: &DataType) -> bool {
     matches!(data_type, DataType::Date32 | DataType::Date64)
 }
@@ -1015,6 +1029,23 @@ 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_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));
+    }
+
     #[test]
     fn test_is_date_narrowing_cast() {
         // Only Date64 -> Date32 narrows (ms -> days, many-to-one).

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

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,101 @@ 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<dyn PhysicalExpr>,
+                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_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<dyn PhysicalExpr>,
+                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::<BinaryExpr>().unwrap();
+                assert!(!is_cast_expr(binary.left()));
+                assert_eq!(
+                    binary.right().downcast_ref::<Literal>().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)]);
@@ -601,7 +695,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 +714,8 @@ mod tests {
 
         let result = unwrap_cast_in_comparison(binary_expr, &schema).unwrap();
 
-        assert!(result.transformed);
-        let optimized_binary = result.data.downcast_ref::<BinaryExpr>().unwrap();
-        assert!(!is_cast_expr(optimized_binary.left()));
-        let right_literal = optimized_binary.right().downcast_ref::<Literal>().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]

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)]);