fix: choose the AT TIME ZONE branch from the final input type (stacked on apache/datafusion#25165)

[adriangb]

Stacked on apache#25165 (base branch fix-at-time-zone-on-tz-aware-timestamps). This PR is for evaluation. If we like it, the two commits go into the upstream PR.

Which issue does this PR close?

• Removes the known limitation in fix: AT TIME ZONE on a timezone-aware timestamp returns a naive timestamp apache/datafusion#25165: AT TIME ZONE chooses its branch before type coercion.

Rationale for this change

Terms: an aware timestamp has a timezone (Timestamp(unit, Some(tz))). A naive timestamp has none (Timestamp(unit, None)).

The upstream PR reads the input type in the SQL planner, which runs before type coercion. That type is wrong for a CASE or a UNION whose first arm is naive while the coerced type is aware. It also sets the TimeUnit from that first arm.

SET datafusion.execution.time_zone = 'UTC';
CREATE TABLE t 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,
  arrow_cast('2024-01-01T12:00:00.123456789Z', 'Timestamp(Nanosecond, Some("UTC"))') AS aware_frac,
  arrow_cast('2024-01-01T12:00:00', 'Timestamp(Second, None)') AS naive_s,
  true AS b, false AS f;

Measured with datafusion-cli builds of main (33574a1), the upstream PR (bf445f0) and this PR:

Query	main	upstream PR	this PR	DuckDB 1.5.2
(CASE WHEN b THEN aware ELSE naive END) AT TIME ZONE 'America/Denver'	05:00:00-07:00 aware	05:00:00 naive	05:00:00 naive	05:00:00 naive
(CASE WHEN b THEN naive ELSE aware END) AT TIME ZONE 'America/Denver'	05:00:00-07:00 aware	05:00:00-07:00 aware ❌	05:00:00 naive	05:00:00 naive
x AT TIME ZONE 'America/Denver' over SELECT naive AS x UNION ALL SELECT aware	aware	aware ❌	naive	naive
the same over a CTE column CASE WHEN b THEN naive ELSE aware END	aware	aware ❌	naive	naive
(CASE WHEN f THEN naive_s ELSE aware_frac END) AT TIME ZONE 'UTC'	12:00:00.123456789Z	12:00:00Z ❌ (fraction lost)	12:00:00.123456789	n/a

The last row is a regression of the upstream PR against main: the planner casts to the unit of the first arm (seconds), so the other arm loses its fraction of a second.

A mix of naive and aware always coerces to aware, so the early type can only miss an aware input. It never reports a naive input as aware.

What changes are included in this PR?

Root cause

The result type of AT TIME ZONE depends on the input type. Every failure so far came from one mistake: code chose the branch from a type that was not yet final.

Where the branch was chosen	Why the type there is not final	Failing shapes
SQL planner (upstream PR)	runs before type coercion	CASE, UNION, CTE, view; TimeUnit of the first arm
optimizer simplify (first version of this PR)	PREPARE runs the optimizer without the analyzer; an untyped placeholder has no type until EXECUTE	PREPARE of a CASE; $1 AT TIME ZONE ...

The fix removes every early choice. AT TIME ZONE becomes the function timezone(zone, expr), and that function chooses the branch only from its argument type, at the two places where DataFusion reads that type:

• return_field_from_args, each time the plan is typed. The physical planner types the final, coerced plan.
• invoke_with_args, from the data type of the actual input.

Nothing rewrites the call into a CAST during planning, so no stale choice can stay in the plan.

Commits

1. feat: add timezone(zone, expression). A new scalar function, the function form of AT TIME ZONE. A dictionary-encoded timestamp follows its value type (item 4 of the upstream self-review). The column is named with the operator syntax: t.x AT TIME ZONE 'America/Denver'.
2. fix: choose the AT TIME ZONE branch after type coercion. The SQL planner no longer reads the input type. It builds timezone(zone, expr) through ExprPlanner::plan_at_time_zone, which now receives the raw operands [zone, expr], like plan_extract. That resolves item 3 of the upstream self-review.
3. fix: never fix the AT TIME ZONE branch before the input type is final. Removes the simplify lowering that the second commit still had, and the SQL planner's no-planner fallback that read the input type. Without any ExprPlanner for AT TIME ZONE, planning now fails for every input, as for EXTRACT.

Why the name timezone(zone, expression)

Engine	Function form	Result for an aware input
PostgreSQL	timezone(zone, timestamp), documented as the same as AT TIME ZONE	naive wall clock
DuckDB	timezone(text, timestamp), AT TIME ZONE is "syntactic sugar" for it	naive wall clock
Redshift	TIMEZONE('zone', timestamp)	naive wall clock
Trino	at_timezone(ts, zone)	aware, same instant, new label
ClickHouse	toTimeZone(ts, zone)	aware, same instant, new label
Snowflake	CONVERT_TIMEZONE(target, ts)	aware, same instant, new label
Spark, Databricks, MySQL	convert_timezone, from_utc_timestamp, CONVERT_TZ	naive input only

Only timezone has the same semantics and argument order in the engines that AT TIME ZONE already follows. The other names mean something different, so they are not good aliases.

Checked shapes

With the final build, each shape below gives the PostgreSQL / DuckDB result for an aware input: CASE in both arm orders, UNION ALL, a CTE, a view, a view whose body has AT TIME ZONE, a scalar subquery, GROUP BY / ORDER BY, a filter, a join key, a nested AT TIME ZONE, max(...), first_value(...) OVER (), coalesce, and PREPARE / EXECUTE with a CASE, an untyped placeholder, a placeholder inside a CASE, and a typed placeholder.

Are these changes tested?

• New timestamps.slt cases: CASE in both arm orders, UNION ALL, a CTE, a view, an aggregate, mixed time units in CASE and UNION, a dictionary input in both directions, PREPARE with an untyped placeholder (bound to an aware and to a naive value) and with a CASE, EXPLAIN, the function form, and the error for a non-constant zone. DuckDB 1.5.2 gives the same values for the CASE and UNION cases.
• A unit test for invoke_with_args.
• sql_integration.rs: the plan with DatetimeFunctionPlanner, and the error without it.
• In the whole sqllogictest suite, one existing expectation changes. See the last user-facing item.
• Green locally: the full sqllogictest suite (524 files), cargo test -p datafusion-sql -p datafusion-functions -p datafusion-expr, cargo clippy --all-targets --workspace --features avro,integration-tests,extended_tests -- -D warnings, cargo fmt, dev/update_function_docs.sh, and the doc prettier check.

Are there any user-facing changes?

In addition to the upstream PR:

• New SQL function timezone(zone, expression).
• AT TIME ZONE uses the coerced input type, for both the branch and the TimeUnit.
• An unaliased AT TIME ZONE column is named x AT TIME ZONE 'zone' for every input. On main a naive input gives x. The upstream PR gives to_local_time(x) for an aware input.
• ExprPlanner::plan_at_time_zone receives [zone, expr] instead of the relabelled CAST. The hook is new in the upstream PR, so no released API changes.
• Without any ExprPlanner for AT TIME ZONE (only when datafusion-sql is used without datafusion-functions), planning fails for every input. On main it gave a CAST.
• Plans show timezone(...) where main showed CAST(... AS Timestamp(..., Some(zone))). So the optimizer no longer sees a CAST for naive AT TIME ZONE 'zone':
  • unwrap_cast_in_comparison no longer relabels a literal compared with it. That relabel is wrong for this cast (unwrap_cast_in_comparison drops the timezone shift when unwrapping CAST(timestamp AS timestamptz) = literal apache/datafusion#25095). One timestamps.slt query returned a row whose instant is one hour away from the literal; it now returns no row, as in DuckDB.
  • A sort on naive AT TIME ZONE 'UTC' or a fixed offset no longer reuses the input ordering. For a named zone, the CAST never preserved ordering (daylight saving time). If this matters, timezone can implement output_ordering for fixed-offset zones.

🤖 Generated with Claude Code

[adriangb]

Folded into apache#25165.