fix: use session timezone for timestamp subtraction

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() {
     );
 }
 
+#[test]
+fn test_create_physical_expr_timestamp_subtraction_uses_session_timezone() {
+    const NANOS_PER_SECOND: i64 = 1_000_000_000;
+    // 2024-11-01 00:00:00 in America/New_York (04:00 UTC).
+    let timestamp_tz: ArrayRef = Arc::new(
+        TimestampNanosecondArray::from(vec![1_730_433_600 * NANOS_PER_SECOND])
+            .with_timezone("America/New_York"),
+    );
+    // The same wall-clock time without a timezone.
+    let timestamp: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![
+        1_730_419_200 * NANOS_PER_SECOND,
+    ]));
+    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".into());
+    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()
+        .into_array(batch.num_rows())
+        .unwrap();
+
+    assert_eq!(
+        ScalarValue::try_from_array(&result, 0).unwrap(),
+        ScalarValue::DurationNanosecond(Some(12 * 60 * 60 * NANOS_PER_SECOND)),
+    );
+}
+
 /// 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/optimizer/src/analyzer/type_coercion.rs

@@ -102,7 +102,13 @@ impl AnalyzerRule for TypeCoercion {
 
         // recurse
         let transformed_plan = plan
-            .transform_up_with_subqueries(|plan| analyze_internal(&EMPTY_SCHEMA, plan))?
+            .transform_up_with_subqueries(|plan| {
+                analyze_internal(
+                    &EMPTY_SCHEMA,
+                    plan,
+                    config.execution.time_zone.as_deref(),
+                )
+            })?
             .data;
 
         // finish
@@ -116,6 +122,7 @@ impl AnalyzerRule for TypeCoercion {
 fn analyze_internal(
     external_schema: &DFSchema,
     plan: LogicalPlan,
+    session_time_zone: Option<&str>,
 ) -> Result<Transformed<LogicalPlan>> {
     // get schema representing all available input fields. This is used for data type
     // resolution only, so order does not matter here
@@ -157,7 +164,8 @@ fn analyze_internal(
         plan
     };
 
-    let mut expr_rewrite = TypeCoercionRewriter::new(&schema);
+    let mut expr_rewrite =
+        TypeCoercionRewriter::new(&schema).with_session_time_zone(session_time_zone);
 
     let name_preserver = NamePreserver::new(&plan);
     // apply coercion rewrite all expressions in the plan individually
@@ -175,13 +183,25 @@ fn analyze_internal(
 /// Rewrite expressions to apply type coercion.
 pub struct TypeCoercionRewriter<'a> {
     pub(crate) schema: &'a DFSchema,
+    session_time_zone: Option<&'a str>,
 }
 
 impl<'a> TypeCoercionRewriter<'a> {
     /// Create a new [`TypeCoercionRewriter`] with a provided schema
     /// representing both the inputs and output of the [`LogicalPlan`] node.
     pub fn new(schema: &'a DFSchema) -> Self {
-        Self { schema }
+        Self {
+            schema,
+            session_time_zone: None,
+        }
+    }
+
+    pub(crate) fn with_session_time_zone(
+        mut self,
+        session_time_zone: Option<&'a str>,
+    ) -> Self {
+        self.session_time_zone = session_time_zone;
+        self
     }
 
     /// Coerce the [`LogicalPlan`].
@@ -398,9 +418,14 @@ impl<'a> TypeCoercionRewriter<'a> {
     ) -> Result<(Expr, Expr)> {
         let left_data_type = left.get_type(left_schema)?;
         let right_data_type = right.get_type(right_schema)?;
-        let (left_type, right_type) =
+        let (left_type, right_type) = if let Some(types) =
+            self.timestamp_subtraction_input_types(&left_data_type, &op, &right_data_type)
+        {
+            types
+        } else {
             BinaryTypeCoercer::new(&left_data_type, &op, &right_data_type)
-                .get_input_types()?;
+                .get_input_types()?
+        };
         let left_cast_ok = can_cast_types(&left_data_type, &left_type);
         let right_cast_ok = can_cast_types(&right_data_type, &right_type);
 
@@ -434,6 +459,47 @@ impl<'a> TypeCoercionRewriter<'a> {
         Ok((left_expr, right_expr))
     }
 
+    /// Coerces the timezone-naive side of timestamp subtraction using the session
+    /// timezone, matching PostgreSQL and DuckDB. The timezone-aware side keeps its
+    /// timezone while both operands are widened to the same precision.
+    fn timestamp_subtraction_input_types(
+        &self,
+        left_type: &DataType,
+        op: &Operator,
+        right_type: &DataType,
+    ) -> Option<(DataType, DataType)> {
+        if op != &Operator::Minus {
+            return None;
+        }
+        let session_time_zone = self.session_time_zone?;
+        let (left_time_zone, right_time_zone) = match (left_type, right_type) {
+            (
+                DataType::Timestamp(_, Some(left_time_zone)),
+                DataType::Timestamp(_, None),
+            ) => (
+                Some(Arc::clone(left_time_zone)),
+                Some(Arc::from(session_time_zone)),
+            ),
+            (
+                DataType::Timestamp(_, None),
+                DataType::Timestamp(_, Some(right_time_zone)),
+            ) => (
+                Some(Arc::from(session_time_zone)),
+                Some(Arc::clone(right_time_zone)),
+            ),
+            _ => return None,
+        };
+        let DataType::Timestamp(unit, _) = comparison_coercion(left_type, right_type)?
+        else {
+            return None;
+        };
+
+        Some((
+            DataType::Timestamp(unit, left_time_zone),
+            DataType::Timestamp(unit, right_time_zone),
+        ))
+    }
+
     fn coerce_date_time_math_op(
         expr: Expr,
         op: &Operator,
@@ -588,8 +654,12 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> {
                 outer_ref_columns,
                 spans,
             }) => {
-                let new_plan =
-                    analyze_internal(self.schema, Arc::unwrap_or_clone(subquery))?.data;
+                let new_plan = analyze_internal(
+                    self.schema,
+                    Arc::unwrap_or_clone(subquery),
+                    self.session_time_zone,
+                )?
+                .data;
                 Ok(Transformed::yes(Expr::ScalarSubquery(Subquery {
                     subquery: Arc::new(new_plan),
                     outer_ref_columns,
@@ -600,6 +670,7 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> {
                 let new_plan = analyze_internal(
                     self.schema,
                     Arc::unwrap_or_clone(subquery.subquery),
+                    self.session_time_zone,
                 )?
                 .data;
                 Ok(Transformed::yes(Expr::Exists(Exists {
@@ -619,6 +690,7 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> {
                 let new_plan = analyze_internal(
                     self.schema,
                     Arc::unwrap_or_clone(subquery.subquery),
+                    self.session_time_zone,
                 )?
                 .data;
                 let expr_type = expr.get_type(self.schema)?;
@@ -648,6 +720,7 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> {
                 let new_plan = analyze_internal(
                     self.schema,
                     Arc::unwrap_or_clone(subquery.subquery),
+                    self.session_time_zone,
                 )?
                 .data;
                 let expr_type = expr.get_type(self.schema)?;
@@ -2721,7 +2794,7 @@ mod test {
             vec![Field::new("a", DataType::Int64, true)].into(),
             std::collections::HashMap::new(),
         )?);
-        let mut rewriter = TypeCoercionRewriter { schema: &schema };
+        let mut rewriter = TypeCoercionRewriter::new(&schema);
         let expr = is_true(lit(12i32).gt(lit(13i64)));
         let expected = is_true(cast(lit(12i32), DataType::Int64).gt(lit(13i64)));
         let result = expr.rewrite(&mut rewriter).data()?;
@@ -2732,7 +2805,7 @@ mod test {
             vec![Field::new("a", DataType::Int64, true)].into(),
             std::collections::HashMap::new(),
         )?);
-        let mut rewriter = TypeCoercionRewriter { schema: &schema };
+        let mut rewriter = TypeCoercionRewriter::new(&schema);
         let expr = is_true(lit(12i32).eq(lit(13i64)));
         let expected = is_true(cast(lit(12i32), DataType::Int64).eq(lit(13i64)));
         let result = expr.rewrite(&mut rewriter).data()?;
@@ -2743,7 +2816,7 @@ mod test {
             vec![Field::new("a", DataType::Int64, true)].into(),
             std::collections::HashMap::new(),
         )?);
-        let mut rewriter = TypeCoercionRewriter { schema: &schema };
+        let mut rewriter = TypeCoercionRewriter::new(&schema);
         let expr = is_true(lit(12i32).lt(lit(13i64)));
         let expected = is_true(cast(lit(12i32), DataType::Int64).lt(lit(13i64)));
         let result = expr.rewrite(&mut rewriter).data()?;
@@ -3201,6 +3274,53 @@ mod test {
         )
     }
 
+    #[test]
+    fn timestamp_subtraction_uses_session_timezone() -> Result<()> {
+        let schema = DFSchema::from_unqualified_fields(
+            vec![
+                Field::new(
+                    "tstz",
+                    DataType::Timestamp(
+                        TimeUnit::Second,
+                        Some("America/New_York".into()),
+                    ),
+                    true,
+                ),
+                Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), true),
+            ]
+            .into(),
+            std::collections::HashMap::new(),
+        )?;
+        let input = Arc::new(LogicalPlan::EmptyRelation(EmptyRelation {
+            produce_one_row: false,
+            schema: Arc::new(schema),
+        }));
+        let subtract = |left, right| {
+            Expr::BinaryExpr(BinaryExpr::new(
+                Box::new(col(left)),
+                Operator::Minus,
+                Box::new(col(right)),
+            ))
+        };
+        let plan = LogicalPlan::Projection(Projection::try_new(
+            vec![subtract("tstz", "ts"), subtract("ts", "tstz")],
+            input,
+        )?);
+
+        let mut options = ConfigOptions::default();
+        options.execution.time_zone = Some("+08:00".to_string());
+        let rule = Arc::new(TypeCoercion::new());
+        assert_analyzed_plan_with_config_eq_snapshot!(
+            options,
+            rule,
+            plan,
+            @r#"
+        Projection: CAST(tstz AS Timestamp(ns, "America/New_York")) - CAST(ts AS Timestamp(ns, "+08:00")), CAST(ts AS Timestamp(ns, "+08:00")) - CAST(tstz AS Timestamp(ns, "America/New_York"))
+          EmptyRelation: rows=0
+        "#,
+        )
+    }
+
     #[test]
     fn in_subquery_cast_subquery() -> Result<()> {
         let empty_int32 = empty_with_type(DataType::Int32);

datafusion/optimizer/src/scalar_subquery_to_join.rs

@@ -398,9 +398,7 @@ fn build_join(
     // itself be NULL) otherwise.
     let mut compensation_exprs = HashMap::new();
     if let Some(expr_map) = collected_count_expr_map {
-        let mut expr_rewrite = TypeCoercionRewriter {
-            schema: new_plan.schema(),
-        };
+        let mut expr_rewrite = TypeCoercionRewriter::new(new_plan.schema());
         let having_arm = pull_up
             .pull_up_having_expr
             .as_ref()

datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs

@@ -243,7 +243,9 @@ impl ExprSimplifier {
     /// See the [type coercion module](datafusion_expr::type_coercion)
     /// documentation for more details on type coercion
     pub fn coerce(&self, expr: Expr, schema: &DFSchema) -> Result<Expr> {
-        let mut expr_rewrite = TypeCoercionRewriter { schema };
+        let mut expr_rewrite = TypeCoercionRewriter::new(schema).with_session_time_zone(
+            self.info.config_options().execution.time_zone.as_deref(),
+        );
         expr.rewrite(&mut expr_rewrite).data()
     }
 

datafusion/optimizer/src/utils.rs

@@ -244,7 +244,7 @@ fn evaluate_expr_with_null_column<'a>(
 }
 
 fn coerce(expr: Expr, schema: &DFSchema) -> Result<Expr> {
-    let mut expr_rewrite = TypeCoercionRewriter { schema };
+    let mut expr_rewrite = TypeCoercionRewriter::new(schema);
     expr.rewrite(&mut expr_rewrite).data()
 }
 

datafusion/sqllogictest/test_files/datetime/timestamps.slt

@@ -1980,6 +1980,60 @@ SELECT '2000-01-01T00:00:00'::timestamp - '2010-01-01T00:00:00'::timestamp;
 ----
 -3653 days 0 hours 0 mins 0.000000000 secs
 
+# Regression test for https://github.com/apache/datafusion/issues/13212
+statement ok
+SET TIME ZONE = '+08'
+
+query ???
+WITH timestamps(ts_tz, ts) AS (
+  VALUES ('2024-11-01T00:00:00+00:00'::timestamptz, '2024-11-01T00:00:00'::timestamp)
+)
+SELECT
+  '2024-11-01T00:00:00+00:00'::timestamptz - '2024-11-01T00:00:00'::timestamp,
+  ts_tz - ts,
+  ts - ts_tz
+FROM timestamps;
+----
+0 days 8 hours 0 mins 0.000000000 secs 0 days 8 hours 0 mins 0.000000000 secs 0 days -8 hours 0 mins 0.000000000 secs
+
+# The session timezone is propagated when coercing a scalar subquery.
+query ?
+SELECT (
+  SELECT ts_tz - ts
+  FROM (
+    VALUES ('2024-11-01T00:00:00+00:00'::timestamptz, '2024-11-01T00:00:00'::timestamp)
+  ) AS timestamps(ts_tz, ts)
+);
+----
+0 days 8 hours 0 mins 0.000000000 secs
+
+# The session timezone, not the aware operand's timezone, controls the cast.
+query ??
+WITH timestamps(ts_tz, ts) AS (
+  VALUES (
+    arrow_cast('2024-11-01T00:00:00-04:00', 'Timestamp(Nanosecond, Some("America/New_York"))'),
+    '2024-11-01T00:00:00'::timestamp
+  )
+)
+SELECT ts_tz - ts, ts - ts_tz
+FROM timestamps;
+----
+0 days 12 hours 0 mins 0.000000000 secs 0 days -12 hours 0 mins 0.000000000 secs
+
+statement ok
+SET TIME ZONE = 'America/New_York'
+
+# The timezone-naive operand is interpreted in EST and EDT, respectively.
+query ??
+SELECT
+  '2024-01-15T00:00:00+00:00'::timestamptz - '2024-01-15T00:00:00'::timestamp,
+  '2024-07-15T00:00:00+00:00'::timestamptz - '2024-07-15T00:00:00'::timestamp;
+----
+0 days -5 hours 0 mins 0.000000000 secs 0 days -4 hours 0 mins 0.000000000 secs
+
+statement ok
+RESET datafusion.execution.time_zone
+
 # Interval - Timestamp => error
 query error Cannot coerce arithmetic expression Interval\(MonthDayNano\) - Timestamp\(ns\) to valid types
 SELECT i - ts1 from FOO;