[SPARK-47690][SQL] Enable hash aggregation support for all collations (StringType) #46640
Conversation
uros-db
changed the title
| // This rewrite rule is used to enabled hash aggregation on collated string columns. However, | ||
| // hash aggregation is currently only supported for grouping aggregations - this means that no | ||
| // string type can be found in the aggregate expressions, so we avoid rewrite in this case. | ||
| !aggregate.aggregateExpressions.exists(e => e.dataType.isInstanceOf[StringType]) |
There was a problem hiding this comment.
I think that you should check if hash aggregation is supported at all, regardless of StringType.
If we are going to end up doing merge agg there is no need to insert collation_key.
There was a problem hiding this comment.
If I just try to use supportsHashAggregate here, I might find that the aggregate does not support hash aggregation before the rewrite, but will support it after the rewrite (as a result of this, the rewrite rule will never actually execute)
However, we perform this check before doing the plan rewrite, so the point of this check is to verify that the current Aggregate is only a grouping aggregate with respect to StringType (i.e. StringType is not found in aggregateExpressions).
Any ideas on how to make this better?
There was a problem hiding this comment.
Maybe you can call supportsHashAggregate by just passing agg keys and empty seq for group by?
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExprUtils.scala
Outdated
Show resolved
Hide resolved
sql/core/src/test/scala/org/apache/spark/sql/CollationSuite.scala
Outdated
Show resolved
Hide resolved
...catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteGroupByCollation.scala
Show resolved
Hide resolved
What changes were proposed in this pull request?
Enable collation support for hash aggregation on StringType, for aggregates where aggregate expressions don't include a non-binary collation expression. Note: support for complex types will be added separately.
CollationKeyCollationKeyis a unary expression that transformsStringTypetoBinaryTypeWhy are the changes needed?
Improve GROUP BY performance for collated strings.
Does this PR introduce any user-facing change?
No.
How was this patch tested?
RewriteGroupByCollationinCollationSuiteWas this patch authored or co-authored using generative AI tooling?
No.