Support AT TIME ZONE clause - #9647

…estamp (apache#25165)

## Which issue does this PR close?

- Closes apache#12218

> [!NOTE]
> **Earlier limitation, now fixed.** A first version chose the lowering
in the SQL planner from `Expr::get_type`, which runs before type
coercion. So a `CASE` or `UNION` whose first arm is naive but whose
coerced type is aware kept the old result. `AT TIME ZONE` now plans as a
`timezone(zone, expr)` function that chooses its behaviour from the
coerced type. `CASE WHEN b THEN naive ELSE aware END AT TIME ZONE
'America/Denver'` and the same `CASE` with the arms swapped both return
the naive `2024-01-01T05:00:00`, as in PostgreSQL 17. The cases are
tested in `datetime/timestamps.slt`.

## Rationale for this change

### Start here: what a user gets today

Two terms run through this description:

- an **aware timestamp** carries a timezone, so it names an instant.
Arrow spells it `Timestamp(unit, Some(tz))`. SQL spells it
`timestamptz`.
- a **naive timestamp** carries no timezone, so it names only a wall
clock. Arrow spells it `Timestamp(unit, None)`. SQL spells it
`timestamp`.

This is the case from the issue. Noon UTC is 05:00 in Denver.

```sql
SET datafusion.execution.time_zone = 'UTC';

SELECT ('2024-01-01T12:00:00Z'::timestamptz AT TIME ZONE 'America/Denver')::timestamp;
```

| Engine | Result |
| --- | --- |
| DataFusion, today | `2024-01-01T12:00:00` |
| DataFusion, this PR | `2024-01-01T05:00:00` |
| PostgreSQL 17.11 | `2024-01-01 05:00:00` |
| DuckDB 1.5.2 | `2024-01-01 05:00:00` |

DataFusion returns the UTC wall clock, so it drops the timezone the user
asks for. Drop the final cast and the cause appears:

```sql
SELECT
  '2024-01-01T12:00:00Z'::timestamptz AT TIME ZONE 'America/Denver' AS value,
  arrow_typeof('2024-01-01T12:00:00Z'::timestamptz AT TIME ZONE 'America/Denver') AS type;
```

| Engine | Value | Type |
| --- | --- | --- |
| DataFusion, today | `2024-01-01T05:00:00-07:00` | `Timestamp(ns,
"America/Denver")` |
| DataFusion, this PR | `2024-01-01T05:00:00` | `Timestamp(ns)` |
| PostgreSQL 17.11 | `2024-01-01 05:00:00` | `timestamp without time
zone` |
| DuckDB 1.5.2 | `2024-01-01 05:00:00` | `TIMESTAMP` |

The displayed value looks right today, but the type is wrong. The result
stays aware, so the next operation on it reads the UTC wall clock again.

### Overview

**This is a breaking semantic change.** It changes the result type of
`AT TIME ZONE`, and so the result value of anything downstream, when the
input is an aware timestamp.

`AT TIME ZONE` is asymmetric in PostgreSQL. It always returns the other
kind of timestamp:

- a **naive** input becomes an **aware** timestamp. The engine reads the
value as a wall clock in the target timezone and returns that instant.
- an **aware** input becomes a **naive** timestamp. The engine reads off
the wall clock that the instant has in the target timezone.

DataFusion implements the first half only. This PR adds the second half.
The first half does not change.

### Why the current lowering is wrong

DataFusion lowers both halves to `CAST(expr AS Timestamp(Nanosecond,
Some(tz)))`.

- For a naive input that cast is correct. Arrow reads the naive value as
local time in `tz`.
- For an aware input that cast only relabels the display timezone. It
keeps the instant, so the result stays aware.

## What changes are included in this PR?

`AT TIME ZONE` is now the new scalar function `timezone(zone,
expression)`, the function form that PostgreSQL and DuckDB also have.
The function is in `datafusion/functions/src/datetime/timezone.rs`. It
chooses its behaviour from the type of `expression` after type coercion:

- `Timestamp(unit, None)`, or any type that is not a timestamp: the same
result as before, `CAST(expr AS Timestamp(unit, Some(tz)))`.
- `Timestamp(unit, Some(_))`: the same cast, then `to_local_time`, which
drops the timezone and keeps the wall clock.
- A dictionary-encoded timestamp follows its value type.

The function reads the type every time: in `return_field_from_args` for
the plan, and in `invoke_with_args` for the data. It is deliberately
never rewritten into the `CAST` form during planning, because the type
at that point is not always final:

- the SQL planner runs before type coercion, so the first `CASE` arm or
`UNION` branch can have a different type from the coerced result,
- `PREPARE` runs the optimizer without the analyzer,
- an untyped placeholder has no type until `EXECUTE`.

`zone` must be a constant, because it decides the result type. A
non-constant zone gives an error.

### The planner hook

`datafusion-sql` must not depend on `datafusion-functions`, so `AT TIME
ZONE` goes through a new trait method:

- `ExprPlanner::plan_at_time_zone` in `datafusion/expr/src/planner.rs`.
It receives `[Utf8 literal tz, expr]`, the argument order of
`timezone(zone, expr)`.
- implemented by `DatetimeFunctionPlanner` in
`datafusion/functions/src/datetime/planner.rs`, which returns the
`timezone` call,
- with a default `PlannerResult::Original` body, so out-of-tree
implementations still compile.

`plan_extract` lowers `EXTRACT` to `date_part` through the same seam, so
this follows an established pattern. The SQL planner
(`sql_at_time_zone_to_expr` in `datafusion/sql/src/expr/mod.rs`) no
longer reads the input type at all. The trait doc states that an
implementation must return an expression that decides after type
coercion.

A session that registers no such planner now gets a plan error for every
`AT TIME ZONE` input, as for `EXTRACT`.

### The `TimeUnit` change

`AT TIME ZONE` no longer forces `Nanosecond`. It keeps the input's
`TimeUnit` when the input is a timestamp.

| Input | Type today | Type with this PR |
| --- | --- | --- |
| `Timestamp(µs, "UTC")` | `Timestamp(ns, "America/Denver")` |
`Timestamp(µs)` |
| `Timestamp(µs, None)` | `Timestamp(ns, "America/Denver")` |
`Timestamp(µs, "America/Denver")` |
| `Timestamp(s, "UTC")` | `Timestamp(ns, "America/Denver")` |
`Timestamp(s)` |
| `Utf8`, `Date32`, and other non-timestamps | `Timestamp(ns,
"America/Denver")` | `Timestamp(ns, "America/Denver")` |

PostgreSQL keeps the precision of its input too. `timestamptz
'2024-01-01 12:00:00.123456Z' AT TIME ZONE 'America/Denver'` returns
`2024-01-01 05:00:00.123456`.

### Not changed

- **Spark.** `SparkFunctionPlanner` implements `plan_extract` and
`plan_substring` only, and `with_spark_features` inserts it at index 0
of the list that `with_default_features` builds. So `plan_at_time_zone`
falls through to `DatetimeFunctionPlanner`, and no Spark-specific
function or planner behaviour changes. A Spark-flavoured session that
uses DataFusion's SQL planner does get the new `AT TIME ZONE` semantics,
but Spark SQL has no `AT TIME ZONE` syntax. It uses `from_utc_timestamp`
and `to_utc_timestamp`. Nothing under `datafusion/spark/` implements,
tests or documents the operator.
- **`AT LOCAL`.** PostgreSQL 17 supports it. sqlparser 0.62.0 has no
`AtLocal` AST node and rejects the syntax, so there is nothing to route.
- **Fixed-offset strings.** See the field research below. Only the
result type changes for these.
- **String-literal inputs.** `'2000-12-01T04:04:12-05:00' AT TIME ZONE
'x'` still takes the naive path, because `Utf8` is not a timestamp type.
That keeps today's behaviour: DataFusion honours the offset inside the
string. PostgreSQL resolves the unknown-typed literal to `timestamp` and
drops the offset. This gap is older than this PR and is separate from
apache#12218.

## What is the testing strategy for this PR?

### Field research

I measured two reference engines rather than trust the documentation:

- PostgreSQL 17.11, `docker run --rm -e POSTGRES_PASSWORD=pw -p
55452:5432 postgres:17`, with `SET TimeZone='UTC'`,
- DuckDB 1.5.2 CLI, with `SET TimeZone='UTC'`.

The DataFusion columns come from a `datafusion-cli` build of `main` and
of this branch, with `SET datafusion.execution.time_zone = 'UTC'`.

#### The four core cases

| Case | DataFusion, today | DataFusion, this PR | PostgreSQL 17.11 |
DuckDB 1.5.2 |
| --- | --- | --- | --- | --- |
| naive input:<br>`'2024-01-01 12:00:00'::timestamp AT TIME ZONE
'America/Denver'` | `2024-01-01T12:00:00-07:00`<br>`Timestamp(ns,
"America/Denver")` | same as today | `2024-01-01
19:00:00+00`<br>`timestamp with time zone` | `2024-01-01
19:00:00+00`<br>`TIMESTAMP WITH TIME ZONE` |
| aware input:<br>`'2024-01-01T12:00:00Z'::timestamptz AT TIME ZONE
'America/Denver'` | `2024-01-01T05:00:00-07:00`<br>`Timestamp(ns,
"America/Denver")` | `2024-01-01T05:00:00`<br>`Timestamp(ns)` |
`2024-01-01 05:00:00`<br>`timestamp without time zone` | `2024-01-01
05:00:00`<br>`TIMESTAMP` |
| the issue reproducer:<br>`(<aware> AT TIME ZONE
'America/Denver')::timestamp` | `2024-01-01T12:00:00` |
`2024-01-01T05:00:00` | `2024-01-01 05:00:00` | `2024-01-01 05:00:00` |
| chained:<br>`<aware> AT TIME ZONE 'America/Denver' AT TIME ZONE
'Europe/Brussels'` | `2024-01-01T13:00:00+01:00`<br>= 12:00 UTC |
`2024-01-01T05:00:00+01:00`<br>= 04:00 UTC | `2024-01-01 04:00:00+00` |
`2024-01-01 04:00:00+00` |

The naive row agrees across all four columns. `12:00` in Denver is
`19:00` UTC. Only the display convention differs.

The chained row is the one that changes value, not just type. Today
DataFusion relabels twice and the instant never moves. With this PR the
Denver wall clock (05:00) becomes a Brussels wall clock, which is 04:00
UTC. PostgreSQL and DuckDB both return 04:00 UTC.

#### Both 2024 DST transitions in `America/Denver`, aware input

| UTC instant | DataFusion, today | DataFusion, this PR | PostgreSQL
17.11 | DuckDB 1.5.2 |
| --- | --- | --- | --- | --- |
| `2024-01-01T12:00:00Z` | `2024-01-01T05:00:00-07:00` |
`2024-01-01T05:00:00` | `2024-01-01 05:00:00` | `2024-01-01 05:00:00` |
| `2024-03-10T08:59:00Z` | `2024-03-10T01:59:00-07:00` |
`2024-03-10T01:59:00` | `2024-03-10 01:59:00` | `2024-03-10 01:59:00` |
| `2024-03-10T09:00:00Z` | `2024-03-10T03:00:00-06:00` |
`2024-03-10T03:00:00` | `2024-03-10 03:00:00` | `2024-03-10 03:00:00` |
| `2024-07-01T12:00:00Z` | `2024-07-01T06:00:00-06:00` |
`2024-07-01T06:00:00` | `2024-07-01 06:00:00` | `2024-07-01 06:00:00` |
| `2024-11-03T07:59:00Z` | `2024-11-03T01:59:00-06:00` |
`2024-11-03T01:59:00` | `2024-11-03 01:59:00` | `2024-11-03 01:59:00` |
| `2024-11-03T08:00:00Z` | `2024-11-03T01:00:00-07:00` |
`2024-11-03T01:00:00` | `2024-11-03 01:00:00` | `2024-11-03 01:00:00` |

The rows cover the spring-forward gap (`01:59` jumps to `03:00`) and the
fall-back overlap (`01:59` MDT, then `01:00` MST). All four columns
agree on the wall clock. Only the DataFusion type changes.

#### The DST gap and overlap from the other side, naive input

This is the branch that this PR does not touch. It is here so the table
stays honest.

| Case | DataFusion, today and with this PR | PostgreSQL 17.11 | DuckDB
1.5.2 |
| --- | --- | --- | --- |
| gap: `'2024-03-10 02:30:00'::timestamp AT TIME ZONE 'America/Denver'`
| `Arrow error: Cast error: Cannot cast timezone to different timezone`
| `2024-03-10 09:30:00+00` | `2024-03-10 09:30:00+00` |
| overlap: `'2024-11-03 01:30:00'::timestamp AT TIME ZONE
'America/Denver'` | same error | `2024-11-03 08:30:00+00` | `2024-11-03
08:30:00+00` |

A bare `arrow_cast` to `Timestamp(Nanosecond, Some("America/Denver"))`
reproduces the error, so it comes from the arrow cast and not from the
SQL planner. This PR does not change it.

#### The case the two engines cannot arbitrate: a fixed-offset string

| Case | DataFusion, today | DataFusion, this PR | PostgreSQL 17.11 |
DuckDB 1.5.2 |
| --- | --- | --- | --- | --- |
| `<aware> AT TIME ZONE '+05:30'` |
`2024-01-01T17:30:00+05:30`<br>`Timestamp(ns, "+05:30")` |
`2024-01-01T17:30:00`<br>`Timestamp(ns)` | `2024-01-01
06:30:00`<br>`timestamp without time zone` | `Not implemented Error:
Unknown TimeZone '+05:30'!` |

DataFusion follows arrow and ISO 8601, where `+05:30` is east of UTC.
PostgreSQL applies the POSIX convention to an offset spelled as a string
and reads `+05:30` as west of UTC. DuckDB rejects the string and accepts
IANA names only. So DataFusion's reading matches neither engine.

PostgreSQL does agree with DataFusion when the offset carries a type:

```sql
postgres=# SELECT timestamptz '2024-01-01 12:00:00Z' AT TIME ZONE interval '+05:30';
 2024-01-01 17:30:00
```

This divergence already exists on main. It is separate from apache#12218 and
out of scope here. It is tracked in
apache#25170. This PR changes only
the result type for this case, not the value.

### Tests

New `sqllogictest` coverage in
`datafusion/sqllogictest/test_files/datetime/timestamps.slt`. The first
commit adds it as a characterization of today's behaviour. The second
commit flips it, so the diff shows exactly what changes. Every case
asserts `arrow_typeof` next to the value, and I checked every expected
value against PostgreSQL 17 by hand. I did not regenerate them with
`--complete`. The cases are:

- a naive column and an aware column,
- the exact `::timestamp` reproducer from apache#12218,
- two chained `AT TIME ZONE` applications,
- a fixed-offset zone,
- microsecond precision in both directions,
- `now()`,
- typed `NULL::timestamp` and `NULL::timestamptz` inputs,
- a multi-row aware column that spans both 2024 DST transitions in
`America/Denver`,
- `CASE`, `UNION ALL`, a CTE, a view and an aggregate whose first arm is
naive but whose coerced type is aware, including the `TimeUnit` of the
coerced type,
- a dictionary-encoded timestamp,
- `PREPARE` / `EXECUTE` with an untyped placeholder and with a `CASE`,
- the `EXPLAIN` plan shape and the column name,
- the `timezone` function form, and the error for a non-constant zone.

A unit test for `TimezoneFunc::invoke_with_args` in both directions.

Two planner tests in `datafusion/sql/tests/sql_integration.rs`: the plan
shape with `DatetimeFunctionPlanner` (a `timezone` call for every input
type), and the error when no `ExprPlanner` handles `AT TIME ZONE`.

Two existing expectations change:

- the `to_local_time` invalid-timezone error no longer comes from the
`simplify_expressions` rule, so its prefix changes,
- `t_europe WHERE column1 = '2024-01-31T16:00:01' AT TIME ZONE
'America/Los_Angeles'` loses a row that never matched. Before, `AT TIME
ZONE` was a `CAST`, and `unwrap_cast_in_comparison` relabelled the
literal instead of converting it
(apache#25095). DuckDB also returns
no row.

Green locally:

- the full `sqllogictest` suite,
- `cargo test -p datafusion-sql`,
- `cargo test -p datafusion-functions timezone` and `to_local_time`,
- `cargo fmt --all` and `cargo clippy --all-targets -- -D warnings`,
- the repo's extended workspace suite (`--features
avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption`).

One caveat. `AT TIME ZONE` appeared in exactly one test file on main
before apache#25175, and every use
there had a `Utf8` or naive inner expression. So the corpus had no
coverage of the broken case at all, which is why the bug survived from
apache#9647 in DataFusion 37. Few changed expectations is a weak safety signal
here, not a strong one. The field research above carries the argument
instead.

## Are there any user-facing changes?

**Yes. This is a breaking change to `AT TIME ZONE` semantics.** The `api
change` label is applied.

The changes, in order of impact:

1. An **aware** input now returns a **naive** timestamp, per the tables
above. Any expression downstream of it changes value.
2. `AT TIME ZONE` keeps the input's `TimeUnit` instead of a forced
`Nanosecond`.
3. An unaliased projection changes its column name for every input.
`SELECT aware AT TIME ZONE 'America/Denver' FROM t` now names the column
`t.aware AT TIME ZONE 'America/Denver'`.
4. The logical and physical plans show `timezone(zone, expr)` instead of
a `CAST`. So `unwrap_cast_in_comparison` no longer applies to it, which
fixes a wrong match (apache#25095)
for a literal compared with `naive AT TIME ZONE 'zone'`.
5. A session that does not register the datetime planner from
`datafusion-functions` now gets a plan error for `AT TIME ZONE`, for
every input type.
6. `ExprPlanner` gains a `plan_at_time_zone` method. It has a default
body, so out-of-tree implementations still compile.
7. A new scalar function `timezone(zone, expression)`.
8. New `AT TIME ZONE` section in
`docs/source/user-guide/sql/operators.md`, and a `timezone` entry in
`scalar_functions.md`. The `to_local_time` description says how they
relate.

A **naive** input does not change. Its value and its timezone label stay
the same. Only its column name and plan shape change.

**Maintainer input wanted:** does this land as a straight behaviour fix,
or behind a config flag with a deprecation period? I deliberately added
no flag. The current behaviour has no defensible justification, and a
flag means two timestamp typings through the planner for a long time.
But this changes results for anyone who relies on the old shape, so I
would rather you decide than assume.

**Interaction with apache#25175
that PR has merged. Its suite
`test_files/datetime/timestamps_timezone.slt` pinned the old behaviour
in SECTION 5b. This PR updates SECTION 5b to the naive results and
removes the `DIVERGES FROM POSTGRESQL (type, not instant)` comment,
because this PR removes that divergence. The `KNOWN WRONG` comment that
follows now refers only to
apache#25166.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Opus 5 <[email protected]>