fix: AT TIME ZONE on a timezone-aware timestamp returns a naive timestamp

datafusion/expr/src/planner.rs

@@ -218,6 +218,26 @@ pub trait ExprPlanner: Debug + Send + Sync {
         Ok(PlannerResult::Original(args))
     }
 
+    /// Plan `<expr> AT TIME ZONE '<tz>'` when `<expr>` is **already
+    /// timezone-aware**.
+    ///
+    /// Following PostgreSQL, `<tz-aware timestamp> AT TIME ZONE zone` returns
+    /// the wall clock the instant has in `zone`, as a timezone-*naive*
+    /// timestamp. (The timezone-*naive* case is the plain
+    /// `CAST(<expr> AS Timestamp(unit, Some(tz)))` that the SQL planner
+    /// handles on its own and never reaches this method.)
+    ///
+    /// `args` holds a single expression: the input already cast to
+    /// `Timestamp(unit, Some(tz))`. That cast preserves the instant and only
+    /// relabels the timezone, so all an implementation has to do is drop the
+    /// timezone while keeping the displayed value — exactly what
+    /// `to_local_time` does.
+    ///
+    /// Returns original expression arguments if not possible
+    fn plan_at_time_zone(&self, args: Vec<Expr>) -> Result<PlannerResult<Vec<Expr>>> {
+        Ok(PlannerResult::Original(args))
+    }
+
     /// Plans a struct literal, such as  `{'field1' : expr1, 'field2' : expr2, ...}`
     ///
     /// This function takes a vector of expressions and a boolean flag

datafusion/functions/src/datetime/planner.rs

@@ -32,4 +32,20 @@ impl ExprPlanner for DatetimeFunctionPlanner {
             ScalarFunction::new_udf(crate::datetime::date_part(), args),
         )))
     }
+
+    /// `<tz-aware timestamp> AT TIME ZONE 'tz'` returns the wall clock in `tz`
+    /// as a timezone-naive timestamp, matching PostgreSQL.
+    ///
+    /// The SQL planner hands us the input already cast to
+    /// `Timestamp(unit, Some(tz))` — an instant-preserving relabel — so the
+    /// remaining step is to strip the timezone while keeping the displayed
+    /// value, which is `to_local_time`.
+    fn plan_at_time_zone(
+        &self,
+        args: Vec<Expr>,
+    ) -> datafusion_common::Result<PlannerResult<Vec<Expr>>> {
+        Ok(PlannerResult::Planned(Expr::ScalarFunction(
+            ScalarFunction::new_udf(crate::datetime::to_local_time(), args),
+        )))
+    }
 }

datafusion/functions/src/datetime/to_local_time.rs

@@ -44,7 +44,11 @@ use datafusion_macros::user_doc;
 /// while keep the display value of the timestamp the same.
 #[user_doc(
     doc_section(label = "Time and Date Functions"),
-    description = "Converts a timestamp with a timezone to a timestamp without a timezone (with no offset or timezone information). This function handles daylight saving time changes.",
+    description = r#"Converts a timestamp with a timezone to a timestamp without a timezone (with no offset or timezone information). This function handles daylight saving time changes.
+
+A timestamp that already has no timezone is returned unchanged.
+
+`AT TIME ZONE` applied to a timezone-*aware* timestamp is defined in terms of this function: it relabels the instant into the target timezone and then applies `to_local_time`. Wrapping such an expression in `to_local_time` is therefore redundant. The examples below instead apply `AT TIME ZONE` to timezone-*naive* values, which is the form that produces a timezone-aware timestamp for `to_local_time` to strip."#,
     syntax_example = "to_local_time(expression)",
     sql_example = r#"```sql
 > SELECT to_local_time('2024-04-01T00:00:20Z'::timestamp);

datafusion/sql/src/expr/mod.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 use std::ops::ControlFlow;
+use std::sync::Arc;
 
 use arrow::datatypes::{DataType, TimeUnit};
 use datafusion_expr::planner::{
@@ -659,24 +660,12 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
             SQLExpr::AtTimeZone {
                 timestamp,
                 time_zone,
-            } => Ok(Expr::Cast(Cast::new(
-                Box::new(self.sql_expr_to_logical_expr_internal(
-                    *timestamp,
-                    schema,
-                    planner_context,
-                )?),
-                match *time_zone {
-                    SQLExpr::Value(ValueWithSpan {
-                        value: Value::SingleQuotedString(s),
-                        span: _,
-                    }) => DataType::Timestamp(TimeUnit::Nanosecond, Some(s.into())),
-                    _ => {
-                        return not_impl_err!(
-                            "Unsupported ast node in sqltorel: {time_zone:?}"
-                        );
-                    }
-                },
-            ))),
+            } => self.sql_at_time_zone_to_expr(
+                *timestamp,
+                *time_zone,
+                schema,
+                planner_context,
+            ),
             SQLExpr::Dictionary(fields) => {
                 self.try_plan_dictionary_literal(fields, schema, planner_context)
             }
@@ -814,6 +803,86 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
         }
     }
 
+    /// Plan `<timestamp> AT TIME ZONE '<tz>'`.
+    ///
+    /// The meaning of `AT TIME ZONE` depends on whether its input carries a
+    /// timezone, and it always returns the *other* kind of timestamp. This
+    /// follows PostgreSQL (and DuckDB):
+    ///
+    /// * a timezone-**naive** input is read as a wall clock in `tz`, and the
+    ///   result is the corresponding timezone-**aware** instant. That is a
+    ///   plain `CAST(expr AS Timestamp(unit, Some(tz)))`, because arrow's
+    ///   `Timestamp(_, None) -> Timestamp(_, Some(tz))` cast interprets the
+    ///   naive value as local time in `tz`.
+    /// * a timezone-**aware** input is an instant, and the result is the wall
+    ///   clock that instant has in `tz`, as a timezone-**naive** timestamp.
+    ///   The same cast is still the first half of that (casting between two
+    ///   aware types preserves the instant and only relabels the zone); the
+    ///   second half — dropping the zone while keeping the displayed value —
+    ///   is delegated to [`ExprPlanner::plan_at_time_zone`], which
+    ///   `datafusion-functions` implements with `to_local_time`.
+    ///
+    /// Anything that is not a timestamp (a string literal, for instance) takes
+    /// the naive path, since a `CAST` to a timezone-aware timestamp is the
+    /// natural reading of `AT TIME ZONE` for it.
+    ///
+    /// [`ExprPlanner::plan_at_time_zone`]: datafusion_expr::planner::ExprPlanner::plan_at_time_zone
+    fn sql_at_time_zone_to_expr(
+        &self,
+        timestamp: SQLExpr,
+        time_zone: SQLExpr,
+        schema: &DFSchema,
+        planner_context: &mut PlannerContext,
+    ) -> Result<Expr> {
+        let tz: Arc<str> = match time_zone {
+            SQLExpr::Value(ValueWithSpan {
+                value: Value::SingleQuotedString(s),
+                span: _,
+            }) => s.into(),
+            _ => {
+                return not_impl_err!("Unsupported ast node in sqltorel: {time_zone:?}");
+            }
+        };
+
+        let expr =
+            self.sql_expr_to_logical_expr_internal(timestamp, schema, planner_context)?;
+
+        // `AT TIME ZONE` does not change the precision of its input, so keep
+        // the input's `TimeUnit` when it has one.
+        let (unit, input_is_tz_aware) = match expr.get_type(schema)? {
+            DataType::Timestamp(unit, tz) => (unit, tz.is_some()),
+            _ => (TimeUnit::Nanosecond, false),
+        };
+
+        // Instant-preserving relabel into `tz` for an aware input; local-time
+        // interpretation for a naive one.
+        let relabeled = Expr::Cast(Cast::new(
+            Box::new(expr),
+            DataType::Timestamp(unit, Some(tz)),
+        ));
+
+        if !input_is_tz_aware {
+            return Ok(relabeled);
+        }
+
+        let mut args = vec![relabeled];
+        for planner in self.context_provider.get_expr_planners() {
+            match planner.plan_at_time_zone(args)? {
+                PlannerResult::Planned(expr) => return Ok(expr),
+                PlannerResult::Original(original) => {
+                    args = original;
+                }
+            }
+        }
+
+        plan_err!(
+            "AT TIME ZONE on a timezone-aware timestamp is not supported by any \
+             ExprPlanner. It needs the `to_local_time` function; register \
+             `datafusion_functions::datetime` (or its `DatetimeFunctionPlanner`) \
+             with the session"
+        )
+    }
+
     fn sql_position_to_expr(
         &self,
         substr_expr: SQLExpr,

datafusion/sql/tests/sql_integration.rs

@@ -3805,6 +3805,39 @@ fn plan_merge_into_rejects_invalid_actions_and_structure(
     );
 }
 
+/// A timezone-naive input is read as a wall clock in the target timezone, so it
+/// lowers to a plain `CAST` and preserves the input's `TimeUnit`.
+#[test]
+fn plan_at_time_zone_on_naive_timestamp() {
+    let plan =
+        logical_plan("SELECT birth_date AT TIME ZONE 'America/Denver' FROM person")
+            .unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+    Projection: CAST(person.birth_date AS Timestamp(ns, "America/Denver"))
+      TableScan: person
+    "#
+    );
+}
+
+/// A timezone-aware input needs `to_local_time`, which is provided by
+/// `datafusion-functions`' `DatetimeFunctionPlanner`. The mock context provider
+/// used by these tests does not register it, so planning must fail with a clear
+/// message rather than silently falling back to the naive lowering.
+#[test]
+fn plan_at_time_zone_on_tz_aware_timestamp_without_planner() {
+    let err = logical_plan(
+        "SELECT (birth_date AT TIME ZONE 'UTC') AT TIME ZONE 'America/Denver' FROM person",
+    )
+    .unwrap_err();
+    assert!(
+        err.strip_backtrace()
+            .contains("AT TIME ZONE on a timezone-aware timestamp"),
+        "unexpected error: {err}"
+    );
+}
+
 fn logical_plan(sql: &str) -> Result<LogicalPlan> {
     logical_plan_with_options(sql, ParserOptions::default())
 }

datafusion/sqllogictest/test_files/datetime/timestamps.slt

@@ -4049,6 +4049,120 @@ SELECT '2000-12-01 04:04:12' AT TIME ZONE 'America/New York';
 statement error
 SELECT '2023-03-12 02:00:00' AT TIME ZONE 'EDT';
 
+##########
+## AT TIME ZONE applied to a timezone-*aware* timestamp
+##
+## https://github.com/apache/datafusion/issues/12218
+##
+## PostgreSQL (and DuckDB) semantics are asymmetric:
+##
+##   * `<tz-naive timestamp> AT TIME ZONE zone` reads the value as a wall clock
+##     in `zone` and returns the matching tz-*aware* instant, and
+##   * `<tz-aware timestamp> AT TIME ZONE zone` returns the wall clock that the
+##     instant has in `zone`, as a tz-*naive* timestamp.
+##
+## Every result below was checked against PostgreSQL 17.
+##########
+
+statement ok
+SET datafusion.execution.time_zone = 'UTC';
+
+statement ok
+CREATE TABLE at_tz_t AS
+SELECT
+  '2024-01-01T12:00:00Z'::timestamptz AS tstz,
+  '2024-01-01 12:00:00'::timestamp AS tsn;
+
+# The tz-naive column is unaffected by this issue: noon-in-Denver is the same
+# instant PostgreSQL reports (`2024-01-01 19:00:00+00`).
+query TP
+SELECT arrow_typeof(tsn AT TIME ZONE 'America/Denver'), tsn AT TIME ZONE 'America/Denver' FROM at_tz_t;
+----
+Timestamp(ns, "America/Denver") 2024-01-01T12:00:00-07:00
+
+# The tz-aware column: PostgreSQL returns `timestamp` (naive) `2024-01-01 05:00:00`.
+query TP
+SELECT arrow_typeof(tstz AT TIME ZONE 'America/Denver'), tstz AT TIME ZONE 'America/Denver' FROM at_tz_t;
+----
+Timestamp(ns) 2024-01-01T05:00:00
+
+# The exact reproducer from #12218. PostgreSQL returns `2024-01-01 05:00:00`.
+query P
+SELECT (tstz AT TIME ZONE 'America/Denver')::timestamp FROM at_tz_t;
+----
+2024-01-01T05:00:00
+
+# Chained. PostgreSQL: `timestamptz` `2024-01-01 04:00:00+00`, i.e. the Denver
+# wall clock (05:00) re-read as a Brussels wall clock.
+query TP
+SELECT
+  arrow_typeof(tstz AT TIME ZONE 'America/Denver' AT TIME ZONE 'Europe/Brussels'),
+  tstz AT TIME ZONE 'America/Denver' AT TIME ZONE 'Europe/Brussels'
+FROM at_tz_t;
+----
+Timestamp(ns, "Europe/Brussels") 2024-01-01T05:00:00+01:00
+
+# Fixed offset. Note DataFusion follows the arrow/ISO-8601 sign convention here
+# (`+05:30` is 5h30m *east* of UTC), while PostgreSQL applies the POSIX
+# convention to offsets spelled as strings. That difference is pre-existing and
+# is not what #12218 is about; only the result *type* changes below.
+query TP
+SELECT arrow_typeof(tstz AT TIME ZONE '+05:30'), tstz AT TIME ZONE '+05:30' FROM at_tz_t;
+----
+Timestamp(ns) 2024-01-01T17:30:00
+
+# `AT TIME ZONE` preserves the precision of its input.
+query T
+SELECT arrow_typeof(arrow_cast('2024-01-01T12:00:00Z', 'Timestamp(Microsecond, Some("UTC"))') AT TIME ZONE 'America/Denver');
+----
+Timestamp(µs)
+
+query T
+SELECT arrow_typeof(arrow_cast('2024-01-01T12:00:00', 'Timestamp(Microsecond, None)') AT TIME ZONE 'America/Denver');
+----
+Timestamp(µs, "America/Denver")
+
+# A real (multi-row) tz-aware column, spanning both DST transitions in
+# America/Denver.
+statement ok
+CREATE TABLE at_tz_dst(ts timestamptz) AS VALUES
+  (arrow_cast('2024-01-01T12:00:00Z', 'Timestamp(Nanosecond, Some("UTC"))')),
+  (arrow_cast('2024-03-10T08:59:00Z', 'Timestamp(Nanosecond, Some("UTC"))')),
+  (arrow_cast('2024-03-10T09:00:00Z', 'Timestamp(Nanosecond, Some("UTC"))')),
+  (arrow_cast('2024-07-01T12:00:00Z', 'Timestamp(Nanosecond, Some("UTC"))')),
+  (arrow_cast('2024-11-03T07:59:00Z', 'Timestamp(Nanosecond, Some("UTC"))')),
+  (arrow_cast('2024-11-03T08:00:00Z', 'Timestamp(Nanosecond, Some("UTC"))'));
+
+query PPT
+SELECT
+  ts,
+  ts AT TIME ZONE 'America/Denver' AS denver,
+  arrow_typeof(ts AT TIME ZONE 'America/Denver') AS denver_type
+FROM at_tz_dst ORDER BY ts;
+----
+2024-01-01T12:00:00Z 2024-01-01T05:00:00 Timestamp(ns)
+2024-03-10T08:59:00Z 2024-03-10T01:59:00 Timestamp(ns)
+2024-03-10T09:00:00Z 2024-03-10T03:00:00 Timestamp(ns)
+2024-07-01T12:00:00Z 2024-07-01T06:00:00 Timestamp(ns)
+2024-11-03T07:59:00Z 2024-11-03T01:59:00 Timestamp(ns)
+2024-11-03T08:00:00Z 2024-11-03T01:00:00 Timestamp(ns)
+
+# `now()` is tz-aware whenever `datafusion.execution.time_zone` is set.
+query T
+SELECT arrow_typeof(now() AT TIME ZONE 'America/Denver');
+----
+Timestamp(ns)
+
+statement ok
+DROP TABLE at_tz_dst;
+
+statement ok
+DROP TABLE at_tz_t;
+
+statement ok
+RESET datafusion.execution.time_zone;
+
+
 # Test current_time without parentheses
 query B
 select current_time = current_time;
@@ -5529,3 +5643,41 @@ query P
 SELECT date_bin(NULL, TIMESTAMP '2023-01-01 12:30:00', TIMESTAMP '2023-01-01 12:00:00')
 ----
 NULL
+
+# KNOWN LIMITATION: `AT TIME ZONE` chooses its behaviour from the type that
+# `Expr::get_type` reports in the SQL planner, which runs BEFORE the type
+# coercion analyzer. For a `CASE` expression, `get_type` reports the type of the
+# first non-null `THEN` arm and ignores coercion. So the two queries below have
+# the same coerced input type, `Timestamp(ns, "UTC")`, but give different
+# results: the first still takes the timezone-naive path and keeps the old
+# aware-in aware-out shape.
+#
+# PostgreSQL gives `2024-01-01 05:00:00` as a naive timestamp for both.
+# A correct fix has to dispatch after coercion. Tracked here so the behaviour is
+# visible rather than silent.
+statement ok
+CREATE TABLE at_tz_case AS
+SELECT arrow_cast('2024-01-01T12:00:00Z', 'Timestamp(Nanosecond, Some("UTC"))') AS aware,
+       arrow_cast('2024-01-01T12:00:00', 'Timestamp(Nanosecond, None)') AS naive,
+       true AS b
+
+# both CASE expressions have the same type
+query TT
+SELECT arrow_typeof(CASE WHEN b THEN naive ELSE aware END),
+       arrow_typeof(CASE WHEN b THEN aware ELSE naive END)
+FROM at_tz_case
+----
+Timestamp(ns, "UTC") Timestamp(ns, "UTC")
+
+# but AT TIME ZONE does not agree on them
+query TPTP
+SELECT arrow_typeof((CASE WHEN b THEN naive ELSE aware END) AT TIME ZONE 'America/Denver'),
+       (CASE WHEN b THEN naive ELSE aware END) AT TIME ZONE 'America/Denver',
+       arrow_typeof((CASE WHEN b THEN aware ELSE naive END) AT TIME ZONE 'America/Denver'),
+       (CASE WHEN b THEN aware ELSE naive END) AT TIME ZONE 'America/Denver'
+FROM at_tz_case
+----
+Timestamp(ns, "America/Denver") 2024-01-01T05:00:00-07:00 Timestamp(ns) 2024-01-01T05:00:00
+
+statement ok
+DROP TABLE at_tz_case