fix: read timezone-naive timestamps in datafusion.execution.time_zone when implicitly converted to timezone-aware

datafusion/common/src/config.rs

@@ -940,7 +940,7 @@ config_namespace! {
 
         /// The default time zone
         ///
-        /// Some functions, e.g. `now` return timestamps in this time zone
+        /// Some functions, e.g. `now`, return timestamps in this time zone. It is also the time zone a timezone-naive timestamp is read in when it is implicitly converted to a timezone-aware one, for example in comparisons, arithmetic, `UNION`, `CASE`, function arguments, `VALUES` and `INSERT`. Explicit conversions (`AT TIME ZONE`, `arrow_cast`) are unaffected.
         pub time_zone: Option<String>, default = None
 
         /// Parquet options

datafusion/core/tests/expr_api/mod.rs

@@ -17,6 +17,7 @@
 
 use arrow::array::{
     ArrayRef, Int64Array, RecordBatch, StringArray, StructArray,
+    TimestampNanosecondArray,
     builder::{ListBuilder, StringBuilder},
 };
 use arrow::datatypes::{DataType, Field};
@@ -364,6 +365,43 @@ async fn test_create_physical_expr_coercion() {
     );
 }
 
+/// `SessionContext::create_physical_expr` coerces the way SQL planning does, so
+/// a timezone-naive timestamp is read in `datafusion.execution.time_zone`.
+#[test]
+fn test_create_physical_expr_timestamp_subtraction_uses_session_timezone() {
+    // 2024-11-01T04:00:00Z, which is 2024-11-01T00:00:00 in New York.
+    let timestamp_tz = Arc::new(
+        TimestampNanosecondArray::from(vec![1_730_433_600_000_000_000])
+            .with_timezone("America/New_York"),
+    ) as ArrayRef;
+    // The same wall clock reading, with no zone attached. Read in +08:00 it is
+    // 2024-10-31T16:00:00Z, twelve hours before `timestamp_tz`.
+    let timestamp = Arc::new(TimestampNanosecondArray::from(vec![
+        1_730_419_200_000_000_000,
+    ])) as ArrayRef;
+    let batch = RecordBatch::try_from_iter(vec![
+        ("timestamp_tz", timestamp_tz),
+        ("timestamp", timestamp),
+    ])
+    .unwrap();
+    let df_schema = DFSchema::try_from(batch.schema()).unwrap();
+
+    let mut config = SessionConfig::new();
+    config.options_mut().execution.time_zone = Some("+08:00".to_string());
+    let ctx = SessionContext::new_with_config(config);
+
+    let physical_expr = ctx
+        .create_physical_expr(col("timestamp_tz") - col("timestamp"), &df_schema)
+        .unwrap();
+    let result = physical_expr.evaluate(&batch).unwrap();
+    let result = ScalarValue::try_from_array(&result.into_array(1).unwrap(), 0).unwrap();
+
+    assert_eq!(
+        result,
+        ScalarValue::DurationNanosecond(Some(12 * 3_600 * 1_000_000_000))
+    );
+}
+
 /// Evaluates the specified expr as an aggregate and compares the result to the
 /// expected result.
 async fn evaluate_agg_test(expr: Expr, expected_lines: Vec<&str>) {

datafusion/expr-common/src/type_coercion/binary.rs

@@ -285,6 +285,34 @@ impl<'a> BinaryTypeCoercer<'a> {
             return Ok(Signature { lhs, rhs, ret });
         }
         Plus | Minus | Multiply | Divide | Modulo  =>  {
+            // `Timestamp(u, Some(tz)) - Timestamp(u, None)` at *equal* units is
+            // the one mixed timezone-aware/naive pair arrow can subtract
+            // directly, so the `self.get_result` probe below would accept it
+            // and no cast would be inserted: arrow then subtracts the raw
+            // values, i.e. reads the naive operand as UTC. Every other unit
+            // pairing falls through to `temporal_coercion_strict_timezone`
+            // below and reads the naive operand in the aware operand's zone,
+            // as comparisons already do. Coerce here so that all unit pairings
+            // agree.
+            if self.op == &Minus
+                && matches!(
+                    (lhs, rhs),
+                    (Timestamp(_, Some(_)), Timestamp(_, None))
+                        | (Timestamp(_, None), Timestamp(_, Some(_)))
+                )
+                && let Some(coerced) = temporal_coercion_strict_timezone(lhs, rhs)
+            {
+                let ret = self.get_result(&coerced, &coerced).map_err(|e| {
+                    plan_datafusion_err!(
+                        "Cannot get result type for temporal operation {coerced} {} {coerced}: {e}", self.op
+                    )
+                })?;
+                return Ok(Signature {
+                    lhs: coerced.clone(),
+                    rhs: coerced,
+                    ret,
+                });
+            }
             if let Ok(ret) = self.get_result(lhs, rhs) {
 
                 // Temporal arithmetic, e.g. Date32 + Interval

datafusion/expr-common/src/type_coercion/binary/tests/arithmetic.rs

@@ -475,3 +475,32 @@ fn test_decimal_precision_overflow_cross_variant() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_timestamp_minus_mixed_timezone_awareness() -> Result<()> {
+    let aware_tz: Arc<str> = Arc::from("America/New_York");
+    let aware_ms = DataType::Timestamp(Millisecond, Some(Arc::clone(&aware_tz)));
+    let aware_ns = DataType::Timestamp(Nanosecond, Some(Arc::clone(&aware_tz)));
+    let naive_ns = DataType::Timestamp(Nanosecond, None);
+
+    // Differing units already coerced to the aware zone; equal units did not.
+    // Both must now coerce both operands to the aware zone at the finer unit,
+    // in either operand order.
+    test_coercion_binary_rule!(aware_ms, naive_ns, Operator::Minus, aware_ns);
+    test_coercion_binary_rule!(naive_ns, aware_ms, Operator::Minus, aware_ns);
+    test_coercion_binary_rule!(aware_ns, naive_ns, Operator::Minus, aware_ns);
+    test_coercion_binary_rule!(naive_ns, aware_ns, Operator::Minus, aware_ns);
+
+    // Pairs that are both aware or both naive are untouched.
+    test_coercion_binary_rule!(aware_ns, aware_ns, Operator::Minus, aware_ns);
+    test_coercion_binary_rule!(naive_ns, naive_ns, Operator::Minus, naive_ns);
+    // Two different aware zones keep their own types (arrow subtracts the
+    // instants directly); unchanged by this rule.
+    let other_tz = DataType::Timestamp(Nanosecond, Some(Arc::from("+08:00")));
+    test_coercion_binary_rule!(aware_ns, other_tz, Operator::Minus, aware_ns, other_tz);
+
+    // For contrast, comparisons already coerced the mixed pair this way.
+    test_coercion_binary_rule!(aware_ms, naive_ns, Operator::Eq, aware_ns);
+
+    Ok(())
+}

datafusion/expr/src/expr_rewriter/mod.rs

@@ -24,6 +24,7 @@ use std::sync::Arc;
 
 use crate::expr::{Alias, Sort, Unnest};
 use crate::logical_plan::Projection;
+use crate::type_coercion::session_time_zone::cast_to_with_session_time_zone;
 use crate::{Expr, ExprSchemable, LogicalPlan, LogicalPlanBuilder};
 
 use datafusion_common::TableReference;
@@ -221,20 +222,46 @@ pub fn strip_outer_reference(expr: Expr) -> Expr {
 
 /// Returns plan with expressions coerced to types compatible with
 /// schema types
+///
+/// Timezone-naive timestamps are read in the zone of the type they are coerced
+/// to. Use [`coerce_plan_expr_for_schema_with_session_time_zone`] to read them
+/// in `datafusion.execution.time_zone` instead.
 pub fn coerce_plan_expr_for_schema(
     plan: LogicalPlan,
     schema: &DFSchema,
+) -> Result<LogicalPlan> {
+    coerce_plan_expr_for_schema_with_session_time_zone(plan, schema, None)
+}
+
+/// Returns plan with expressions coerced to types compatible with schema types,
+/// reading a timezone-naive timestamp in `session_time_zone` wherever a cast
+/// this function inserts makes one timezone-aware.
+///
+/// `session_time_zone` is `datafusion.execution.time_zone`; see
+/// [`cast_to_with_session_time_zone`] for exactly what changes. Passing `None`
+/// is [`coerce_plan_expr_for_schema`]: every plan is then byte-identical to the
+/// one DataFusion produced before the session time zone existed.
+///
+/// Only the casts inserted here are affected. An expression that already has
+/// the target type is returned untouched, so a cast the *user* wrote — `AT TIME
+/// ZONE`, `arrow_cast` — keeps arrow's semantics.
+pub fn coerce_plan_expr_for_schema_with_session_time_zone(
+    plan: LogicalPlan,
+    schema: &DFSchema,
+    session_time_zone: Option<&str>,
 ) -> Result<LogicalPlan> {
     match plan {
         // special case Projection to avoid adding multiple projections
         LogicalPlan::Projection(Projection { expr, input, .. }) => {
-            let new_exprs = coerce_exprs_for_schema(expr, input.schema(), schema)?;
+            let new_exprs =
+                coerce_exprs_for_schema(expr, input.schema(), schema, session_time_zone)?;
             let projection = Projection::try_new(new_exprs, input)?;
             Ok(LogicalPlan::Projection(projection))
         }
         _ => {
             let exprs: Vec<Expr> = plan.schema().iter().map(Expr::from).collect();
-            let new_exprs = coerce_exprs_for_schema(exprs, plan.schema(), schema)?;
+            let new_exprs =
+                coerce_exprs_for_schema(exprs, plan.schema(), schema, session_time_zone)?;
             let add_project = new_exprs.iter().any(|expr| expr.try_as_col().is_none());
             if add_project {
                 let projection = Projection::try_new(new_exprs, Arc::new(plan))?;
@@ -250,6 +277,7 @@ fn coerce_exprs_for_schema(
     exprs: Vec<Expr>,
     src_schema: &DFSchema,
     dst_schema: &DFSchema,
+    session_time_zone: Option<&str>,
 ) -> Result<Vec<Expr>> {
     exprs
         .into_iter()
@@ -259,7 +287,13 @@ fn coerce_exprs_for_schema(
             if new_type != &expr.get_type(src_schema)? {
                 match expr {
                     Expr::Alias(Alias { expr, name, .. }) => {
-                        Ok(expr.cast_to(new_type, src_schema)?.alias(name))
+                        Ok(cast_to_with_session_time_zone(
+                            *expr,
+                            new_type,
+                            src_schema,
+                            session_time_zone,
+                        )?
+                        .alias(name))
                     }
                     #[expect(deprecated)]
                     Expr::Wildcard { .. } => Ok(expr),
@@ -270,9 +304,20 @@ fn coerce_exprs_for_schema(
                             // (see: https://github.com/apache/datafusion/issues/18818)
                             Expr::Column(ref column) => {
                                 let name = column.name().to_owned();
-                                Ok(expr.cast_to(new_type, src_schema)?.alias(name))
+                                Ok(cast_to_with_session_time_zone(
+                                    expr,
+                                    new_type,
+                                    src_schema,
+                                    session_time_zone,
+                                )?
+                                .alias(name))
                             }
-                            _ => Ok(expr.cast_to(new_type, src_schema)?),
+                            _ => cast_to_with_session_time_zone(
+                                expr,
+                                new_type,
+                                src_schema,
+                                session_time_zone,
+                            ),
                         }
                     }
                 }

datafusion/expr/src/expr_schema.rs

@@ -26,6 +26,7 @@ use crate::expr::{FieldMetadata, LambdaVariable};
 use crate::higher_order_function::HigherOrderReturnFieldArgs;
 use crate::type_coercion::functions::value_fields_with_higher_order_udf_and_lambdas;
 use crate::type_coercion::functions::{UDFCoercionExt, fields_with_udf};
+use crate::type_coercion::session_time_zone::cast_to_with_session_time_zone;
 use crate::udf::ReturnFieldArgs;
 use crate::{
     LogicalPlan, Operator, Projection, Subquery, WindowFunctionDefinition, utils,
@@ -822,25 +823,55 @@ fn scalar_subquery_nullable(subquery: &Subquery) -> bool {
 ///    new projection with the casted expression.
 /// 2. **Non-projection plan**: If the subquery isn't a projection, it adds a projection to the plan
 ///    with the casted first column.
+///
+/// A timezone-naive timestamp is read in the zone of `cast_to_type`. Use
+/// [`cast_subquery_with_session_time_zone`] to read it in
+/// `datafusion.execution.time_zone` instead.
 pub fn cast_subquery(subquery: Subquery, cast_to_type: &DataType) -> Result<Subquery> {
+    cast_subquery_with_session_time_zone(subquery, cast_to_type, None)
+}
+
+/// [`cast_subquery`], reading a timezone-naive timestamp in
+/// `session_time_zone` if the cast it inserts makes one timezone-aware.
+///
+/// `session_time_zone` is `datafusion.execution.time_zone`; see
+/// [`cast_to_with_session_time_zone`] for exactly what changes. Passing `None`
+/// is [`cast_subquery`]: the subquery is then byte-identical to the one
+/// DataFusion produced before the session time zone existed.
+///
+/// Only the cast inserted here is affected. A subquery whose projected
+/// expression already has `cast_to_type` is returned untouched, so a cast the
+/// *user* wrote — `AT TIME ZONE`, `arrow_cast` — keeps arrow's semantics.
+pub fn cast_subquery_with_session_time_zone(
+    subquery: Subquery,
+    cast_to_type: &DataType,
+    session_time_zone: Option<&str>,
+) -> Result<Subquery> {
     if subquery.subquery.schema().field(0).data_type() == cast_to_type {
         return Ok(subquery);
     }
 
     let plan = subquery.subquery.as_ref();
     let new_plan = match plan {
         LogicalPlan::Projection(projection) => {
-            let cast_expr = projection.expr[0]
-                .clone()
-                .cast_to(cast_to_type, projection.input.schema())?;
+            let cast_expr = cast_to_with_session_time_zone(
+                projection.expr[0].clone(),
+                cast_to_type,
+                projection.input.schema().as_ref(),
+                session_time_zone,
+            )?;
             LogicalPlan::Projection(Projection::try_new(
                 vec![cast_expr],
                 Arc::clone(&projection.input),
             )?)
         }
         _ => {
-            let cast_expr = Expr::Column(Column::from(plan.schema().qualified_field(0)))
-                .cast_to(cast_to_type, subquery.subquery.schema())?;
+            let cast_expr = cast_to_with_session_time_zone(
+                Expr::Column(Column::from(plan.schema().qualified_field(0))),
+                cast_to_type,
+                subquery.subquery.schema().as_ref(),
+                session_time_zone,
+            )?;
             LogicalPlan::Projection(Projection::try_new(
                 vec![cast_expr],
                 subquery.subquery,