[akupchinskiy]

Which issue does this PR close?

• Closes first_value/last_value ORDER BY a time zone aware timestamp gets slower with more partitions (Arc<str> refcount contention in FirstLastGroupsAccumulator) #27046.

Rationale for this change

A grouped first_value(x ORDER BY ts) / last_value(x ORDER BY ts) with ts a
Timestamp(_, Some(tz)) gets slower as target_partitions grows, and so does every
DISTINCT ON with a time zone aware key (it is planned as first_value per column). At 10
partitions it is 3.5-14x slower than the same query ordered by a naive timestamp. We found it with
a benchmark (digibench) where a DISTINCT ON (session_id, event_name, client_ts) dedup was 25x
slower than DuckDB.

The cause is in FirstLastGroupsAccumulator: it keeps each group's winning ordering values as
Vec<Vec<ScalarValue>>. A time zone aware ScalarValue::Timestamp* holds its time zone as an
Arc<str>, and that Arc comes from the column's DataType, so it is shared by all partitions.
Extracting a winner row, replacing a group's orderings (clear() + extend_from_slice()), and
growing the groups (resize(n, default_orderings.clone())) each clone or drop it, i.e. several
atomic operations on one cache line per updated group, from every thread at once.

The time zone does not take part in the comparison (the values are UTC instants), so the
accumulator does not need it.

What changes are included in this PR?

In FirstLastGroupsAccumulator (datafusion/functions-aggregate/src/first_last.rs):

• try_new records the time zone of each ordering column (ordering_timezones) and builds
  default_orderings without it.
• update_batch and merge_batch relabel the time zone aware ordering columns of their input as
  time-zone-free before anything else (PrimitiveArray::with_timezone_opt(None): no copy, no value
  change, unlike a cast). From there on the stored ScalarValues carry None and cloning them is a
  plain copy. When no ordering column has a time zone the input is passed through unchanged
  (Cow::Borrowed).
• state puts the original time zone back on the ordering columns it emits, so the intermediate
  state schema is unchanged.

Results, the state schema, plans and the non-timestamp paths are unchanged. Size accounting now
reports slightly less per stored ordering value, because a time-zone-free ScalarValue::size()
no longer counts the (shared) time zone string.

Not changed: the non-grouped FirstValueAccumulator / LastValueAccumulator extract their
ordering row once per accumulator and batch rather than once per updated group, so they are much
less exposed; nth_value keeps ScalarValue ordering rows too and may have a milder form of the
same contention; ordered array_agg already stores encoded rows. The general fix, keeping
per-group orderings in the arrow::row format or as typed columns instead of ScalarValues
(which would also remove the remaining per-group ScalarValue work for every ordering type, and
cover time zone aware timestamps nested in structs), is a larger refactor and is left as a
follow-up.

What is the testing strategy for this PR?

• Unit test test_group_acc_timezone_ordering (first and last): the stored orderings carry no time
  zone, size_of_orderings stays consistent, the state() ordering column has the original
  Timestamp(µs, "+02:00") type also when the state is emitted in parts (EmitTo::First, then
  EmitTo::All), and merging those states with another partial state gives the right values.
• aggregate.slt: first_value / last_value ordered by a Timestamp(µs, "+05:00") column over a
  UNION ALL of two tables (multi-partition, partial state merged in the final stage), with NULL
  ordering values (NULLS FIRST / NULLS LAST), ties broken by a second key, the time zone aware
  column as the second ordering key, DESC, the time zone aware value itself as the result
  (arrow_typeof keeps the time zone, displayed values are not shifted), DISTINCT ON with a
  time zone aware key, and second, millisecond and nanosecond timestamps with three different time
  zones as ordering keys, with and without IGNORE NULLS.
• Existing first_last unit tests and the aggregate, first_last_*, distinct_on and
  ordered_aggregate_* sqllogictests pass.

Benchmark (datafusion-cli --release, 2M rows in 1M groups of 2 rows whose order in ts is
scrambled, data generated in SQL; queries in the issue). Median of 5 hot runs after a warm-up, ms,
two interleaved rounds (main, branch, main, branch) shown as round 1 / round 2, Apple M2 Pro,
10 cores, otherwise idle:

query (ordered by Timestamp(µs, "UTC"))	target_partitions	main	this PR
first_value(v ORDER BY ts) GROUP BY k1, k2	1	193 / 194	189 / 194
	4	324 / 329	104 / 104
	10	491 / 505	97 / 99
last_value(v ORDER BY ts) GROUP BY k1, k2	1	191 / 191	192 / 191
	4	337 / 318	106 / 104
	10	485 / 502	98 / 96
DISTINCT ON (k1) ... ORDER BY k1, ts, o	1	266 / 267	267 / 266
	4	443 / 440	165 / 165
	10	558 / 573	157 / 153

With this PR the time zone aware queries run at the speed of the same queries ordered by a naive
timestamp (10 partitions: 97-99 ms for first_value / last_value, 153-161 ms for DISTINCT ON,
on both builds), i.e. about 5x faster for first_value / last_value and 3.6x for DISTINCT ON at
10 partitions and unchanged at 1. A latest-state-per-order query (10M events, 2M orders,
DISTINCT ON (order_id) ... ORDER BY order_id, ingest_ts DESC, default partitions) goes from
6.77-6.92 s to 0.47-0.52 s (about 14x), the same as with a naive ingest_ts (0.50-0.54 s). Query
results are identical between main and the branch (and, for the latest-state query, with DuckDB).

Are there any user-facing changes?

No, only performance.