kosiew commented on code in PR #25165: URL: https://github.com/apache/datafusion/pull/25165#discussion_r4081961452
########## 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; Review Comment: Small test-coverage suggestion: could we add SQL-level cases for both `NULL::timestamp AT TIME ZONE ...` and `NULL::timestamptz AT TIME ZONE ...`, checking both `arrow_typeof` and the NULL value? The current behavior looks correct: the naive form should produce `Timestamp(ns, "America/Denver")` with NULL, and the aware form should produce `Timestamp(ns)` with NULL. This would mainly protect the interaction between the type-based dispatch and `to_local_time`'s typed-NULL handling. ########## datafusion/sql/src/expr/mod.rs: ########## @@ -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 Review Comment: I think this needs to be addressed before merge. The aware vs naive branch is chosen here using `Expr::get_type`, but this happens before type coercion. That creates cases where the lowering depends on the pre-coercion expression shape rather than the final coerced type. `CASE` is one example. `Expr::Case::get_type` can reflect the first non-NULL `THEN` arm instead of the final coerced CASE type. The SQLLogicTest later in this PR demonstrates the result: two CASE expressions that both end up as `Timestamp(ns, "UTC")` produce different `AT TIME ZONE` result types and values just because the arms are reversed. `UNION` has the same issue and is not currently covered. With `datafusion.execution.time_zone = 'UTC'`, both subqueries below produce a column whose `arrow_typeof(x)` is `Timestamp(ns, "UTC")`, but `AT TIME ZONE` takes different paths depending on which input appears first: ```sql -- Takes the naive path, even though x is coerced to Timestamp(ns, "UTC") SELECT arrow_typeof(x AT TIME ZONE 'America/Denver'), x AT TIME ZONE 'America/Denver' FROM ( SELECT arrow_cast('2024-01-01T12:00:00', 'Timestamp(Nanosecond, None)') x UNION ALL SELECT '2024-01-01T12:00:00Z'::timestamptz ); -- Takes the aware path SELECT arrow_typeof(x AT TIME ZONE 'America/Denver'), x AT TIME ZONE 'America/Denver' FROM ( SELECT '2024-01-01T12:00:00Z'::timestamptz x UNION ALL SELECT arrow_cast('2024-01-01T12:00:00', 'Timestamp(Nanosecond, None)') ); ``` The first returns the old aware shape, while the second returns the intended naive wall clock. `coalesce` does not appear to have this problem because its arguments are coerced before this type check. Compared with `main`, these cases are not regressions introduced by this PR. They keep the old behavior. The problem is that this PR documents type-based PostgreSQL semantics where a timezone-aware input must return a timezone-naive wall clock, and common query shapes still violate that contract. Could we make this dispatch use the coerced input type instead? I would also replace the known-limitation CASE test with regression coverage for both CASE arm orders, and add the UNION case in both input orders. One possible direction, although I have not prototyped it, would be to lower `AT TIME ZONE` to a scalar function such as `at_time_zone(expr, tz)` and let its return type depend on the coerced argument type. Since type coercion recomputes schemas after rewriting, that may give the function the final type and could also avoid adding the new public `ExprPlanner::plan_at_time_zone` hook. I would treat that only as a possible implementation direction, not a requirement. The important part is that the CASE and UNION regression tests pass regardless of the chosen approach. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
