akupchinskiy opened a new pull request, #27047:
URL: https://github.com/apache/datafusion/pull/27047
## Which issue does this PR close?
- Closes #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 `ScalarValue`s 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 `ScalarValue`s
(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.
--
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]
