Shinoaki0145 opened a new pull request, #58196:
URL: https://github.com/apache/spark/pull/58196
### What changes were proposed in this pull request?
This PR mirrors the approach introduced in `RewriteCollationJoin`
(SPARK-48000) for hash joins and introduces hash-based aggregation support for
non-binary collated grouping keys:
1. **`RewriteCollationAggregate` Optimizer Rule**:
- Added a new Catalyst optimizer rule `RewriteCollationAggregate`
registered under the `Finish Analysis` batch.
- For `Aggregate` operators with non-binary-stable grouping keys (e.g.
`UTF8_LCASE`, structs/arrays containing non-binary collated strings), it
rewrites the grouping expressions using `CollationKey(...)` so grouping is
performed on binary-stable collation keys.
- Preserves original grouping key expressions referenced in the
output/aggregate expressions by wrapping them in `First(origExpr, ignoreNulls =
false).toAggregateExpression()`.
2. **ObjectHashAggregate Planner Update**:
- Updated `Aggregate.supportsObjectHashAggregate` to support aggregations
with non-mutable buffer schemas (such as those containing `First(string)` from
collation rewrites) when grouping expressions are binary-stable, allowing these
queries to plan as `ObjectHashAggregateExec` instead of falling back to
`SortAggregateExec`.
3. **Feature Flag**:
- Added a new configuration `spark.sql.collation.hashAggregation.enabled`
(default: `true`) to allow reverting to the previous sort-based behavior if
necessary.
### Why are the changes needed?
Currently, a `GROUP BY` (or any aggregation) on non-binary collated keys
(e.g., `UTF8_LCASE`) is planned as `SortAggregateExec` because
`UnsafeRowUtils.isBinaryStable` returns `false` for non-binary collations,
preventing hash aggregation.
This forces a full-input sort and can cause heavy disk spilling on large
datasets. While hash join already solves this by injecting `CollationKey` in
`RewriteCollationJoin`, aggregations were not given equivalent optimization.
Rewriting grouping keys to collation keys and routing to
`ObjectHashAggregateExec` avoids the mandatory sort and disk spill,
significantly improving query performance on collated data.
### Does this PR introduce _any_ user-facing change?
Yes. Aggregations on non-binary collated columns will now be executed using
`ObjectHashAggregateExec` instead of `SortAggregateExec` by default,
dramatically improving execution speed without requiring manual casts to binary.
Users can disable this optimization and restore the legacy sort-based
execution by setting:
```sql
SET spark.sql.collation.hashAggregation.enabled = false;
```
### How was this patch tested?
- Added/updated unit tests in
`org.apache.spark.sql.collation.CollationAggregationSuite`:
- Verified `ObjectHashAggregateExec` is selected and output results are
correct for `GROUP BY` on `UTF8_LCASE`.
- Verified `spark.sql.collation.hashAggregation.enabled = false` falls
back to `SortAggregateExec`.
- Verified compatibility with imperative aggregate functions (e.g.
`collect_list`).
### Was this patch authored or co-authored using generative AI tooling?
No.
--
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]
