viirya opened a new issue, #26818:
URL: https://github.com/apache/datafusion/issues/26818
### Describe the bug
`AggregateExec` adds `Constraint::Unique(group keys)` to its equivalence
properties in every mode, including `Partial` and `PartialReduce`. Partial
state is not unique on the group keys:
- Operators that merge partitions, e.g. `RepartitionExec`,
`SortPreservingMergeExec` and `CoalescePartitionsExec`, keep constraints, so
after a merge the same key can come from several input partitions.
- Even within one partition, a partial stage can emit the same group more
than once (early emission under memory pressure, skipping partial aggregation,
clustered partial aggregation emitting incomplete groups under memory pressure).
`EquivalenceProperties::ordering_satisfy` uses `Unique` constraints: an
ordering on `[k]` plus `Unique(k)` satisfies `[k, x]` for any `x`. With a stale
`Unique(k)` this removes a sort that is needed.
### To Reproduce
Two input partitions, each sorted on non-null `k`: `(1, 9), (2, 5)` and `(1,
1), (2, 7)`. A `Partial` `sum(v)` grouped by `k`, under a parent that requires
`Hash(k)` and `[k, sum]` within each partition. Run `EnsureRequirements` with
`target_partitions = 2` and `prefer_existing_sort = true`:
```
RepartitionExec: partitioning=Hash([k@0], 2), input_partitions=2,
preserve_order=true, sort_exprs=k@0 ASC
AggregateExec: mode=Partial, gby=[k@0 as k], aggr=[sum_v],
group_clustering_mode=Full
DataSourceExec: partitions=2, partition_sizes=[1, 1],
output_ordering=k@0 ASC
```
The sort on `[k, sum]` is gone, and the partition for `k = 1` comes out as
`(1, 9), (1, 1)`. `SanityCheckPlan` accepts the plan because it uses the same
`ordering_satisfy`. Without `prefer_existing_sort` the optimizer keeps a
`SortExec(k, sum)` and the output is correct.
The same happens directly in `SortExec`, which skips sorting when its input
already satisfies the requested ordering. Each of the following under
`SortExec([k, sum])` returns rows with equal `k` out of order:
- `RepartitionExec(Hash(k)).with_preserve_order()` over the Partial aggregate
- `SortPreservingMergeExec([k])` over it
- `CoalescePartitionsExec` then `SortExec([k])` over it
### Expected behavior
Only aggregation modes that produce final values (`Final`,
`FinalPartitioned`, `Single`, `SinglePartitioned`) should declare `Unique` on
the group keys.
### Additional context
With the default SQL planner, the operator directly above the merge over a
Partial aggregate is always the Final aggregate. That aggregate's ordering
requirement covers only the group keys, and it declares its own (correct)
`Unique`. I couldn't produce wrong results from SQL: I checked plans for `ORDER
BY k, sum`, window functions ordered by `k, sum`, sort-merge joins of grouped
subqueries, `DISTINCT` and `LIMIT`, with `prefer_existing_sort` and
`repartition_sorts` on and off. The bug shows up in plans that put something
other than an aggregate above Partial output: hand-built plans, custom
optimizer rules, or planners that split stages between the Partial and Final
aggregates.
The constraint was added for all modes in #14111.
--
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]
