[SPARK-58927][SQL] Support hash-based aggregation for collated groupi…

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

@@ -367,7 +367,8 @@ abstract class Optimizer(catalogManager: CatalogManager)
       RewriteNearestByJoin,
       EvalInlineTables,
       ReplaceTranspose,
-      RewriteCollationJoin
+      RewriteCollationJoin,
+      RewriteCollationAggregate
     )
 
     override def apply(plan: LogicalPlan): LogicalPlan = {

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteCollationAggregate.scala

@@ -0,0 +1,108 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.optimizer
+
+import scala.collection.mutable
+
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate._
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreePattern.AGGREGATE
+import org.apache.spark.sql.catalyst.util.UnsafeRowUtils
+
+/**
+ * This rule rewrites Aggregate grouping expressions to ensure that non-binary collated strings
+ * are converted to their binary-stable collation keys (via [[CollationKey]]).
+ *
+ * This allows hash-based aggregation (e.g., ObjectHashAggregateExec) to work properly on data
+ * with non-binary collations, avoiding full sorting and spilling.
+ *
+ * Any original grouping expression referenced in the aggregate expressions (output) is preserved
+ * by wrapping it in `First(expr, ignoreNulls = false)`, an arbitrary representative of each
+ * collation-equal group.
+ */
+object RewriteCollationAggregate extends Rule[LogicalPlan] {
+  def apply(plan: LogicalPlan): LogicalPlan = {
+    if (!conf.collationHashAggregationEnabled) {
+      plan
+    } else {
+      plan.transformWithPruning(_.containsPattern(AGGREGATE)) {
+        case a @ Aggregate(groupingExpressions, aggregateExpressions, child, _)
+            if a.resolved &&
+              groupingExpressions.exists(e => !UnsafeRowUtils.isBinaryStable(e.dataType)) =>
+          val keyMapping = mutable.LinkedHashMap.empty[Expression, Expression]
+          val newGroupingExpressions = groupingExpressions.map { ge =>
+            val processed = CollationKey.injectCollationKey(ge)
+            if (!processed.fastEquals(ge)) {
+              keyMapping.put(ge.canonicalized, ge)
+              processed
+            } else {
+              ge
+            }
+          }
+
+          if (keyMapping.nonEmpty) {
+            def replaceGroupingKeyReferences(e: Expression): Expression = {
+              e match {
+                case _: AggregateExpression => e
+                case _ if e.foldable => e
+                case _ if keyMapping.contains(e.canonicalized) =>
+                  val origExpr = keyMapping(e.canonicalized)
+                  First(origExpr, ignoreNulls = false).toAggregateExpression()
+                case _ =>
+                  e.mapChildren(replaceGroupingKeyReferences)
+              }
+            }
+
+            val newAggregateExpressions = aggregateExpressions.map {
+              case a @ Alias(child, name) =>
+                val newChild = replaceGroupingKeyReferences(child)
+                if (!newChild.fastEquals(child)) {
+                  Alias(newChild, name)(
+                    exprId = a.exprId,
+                    qualifier = a.qualifier,
+                    explicitMetadata = a.explicitMetadata)
+                } else {
+                  a
+                }
+              case attr: Attribute =>
+                val newChild = replaceGroupingKeyReferences(attr)
+                if (!newChild.fastEquals(attr)) {
+                  Alias(newChild, attr.name)(exprId = attr.exprId, qualifier = attr.qualifier)
+                } else {
+                  attr
+                }
+              case other =>
+                val newOther = replaceGroupingKeyReferences(other)
+                newOther match {
+                  case ne: NamedExpression => ne
+                  case expr => Alias(expr, expr.prettyName)()
+                }
+            }
+
+            a.copy(
+              groupingExpressions = newGroupingExpressions,
+              aggregateExpressions = newAggregateExpressions)
+          } else {
+            a
+          }
+      }
+    }
+  }
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala

@@ -1337,7 +1337,9 @@ object Aggregate {
       return false
     }
 
-    aggregateExpressions.map(_.aggregateFunction).exists {
+    val schema = DataTypeUtils.fromAttributes(
+      aggregateExpressions.flatMap(_.aggregateFunction.aggBufferAttributes))
+    !isAggregateBufferMutable(schema) || aggregateExpressions.map(_.aggregateFunction).exists {
       case _: TypedImperativeAggregate[_] => true
       case _ => false
     }

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -1387,6 +1387,14 @@ object SQLConf {
       .booleanConf
       .createWithDefault(false)
 
+  val COLLATION_HASH_AGGREGATION_ENABLED =
+    buildConf("spark.sql.collation.hashAggregation.enabled")
+      .doc("When true, allows hash-based aggregation for non-binary collated strings by " +
+        "rewriting grouping keys to collation keys.")
+      .version("4.1.0")
+      .booleanConf
+      .createWithDefault(true)
+
   val OBJECT_LEVEL_COLLATIONS_ENABLED =
     buildConf("spark.sql.collation.objectLevel.enabled")
       .internal()
@@ -8844,6 +8852,8 @@ class SQLConf extends Serializable with Logging with SqlApiConf {
 
   def allowCollationsInMapKeys: Boolean = getConf(ALLOW_COLLATIONS_IN_MAP_KEYS)
 
+  def collationHashAggregationEnabled: Boolean = getConf(COLLATION_HASH_AGGREGATION_ENABLED)
+
   def objectLevelCollationsEnabled: Boolean = getConf(OBJECT_LEVEL_COLLATIONS_ENABLED)
 
   def schemaLevelCollationsEnabled: Boolean = getConf(SCHEMA_LEVEL_COLLATIONS_ENABLED)

sql/core/src/test/scala/org/apache/spark/sql/collation/CollationAggregationSuite.scala

@@ -20,36 +20,64 @@ package org.apache.spark.sql.collation
 import org.apache.spark.sql.Row
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
 import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
 
 class CollationAggregationSuite
   extends SharedSparkSession
   with AdaptiveSparkPlanHelper {
 
-  test("group by collated column doesn't work with obj hash aggregate") {
+  test("hash aggregate on collated grouping key with RewriteCollationAggregate") {
     val tblName = "grp_by_tbl"
     withTable(tblName) {
       sql(s"CREATE TABLE $tblName (c1 STRING COLLATE UTF8_LCASE, c2 INT) USING PARQUET")
       sql(s"INSERT INTO $tblName VALUES ('hello', 1), ('HELLO', 2), ('HeLlO', 3)")
 
-      // Result is correct without forcing object hash aggregate.
-      checkAnswer(
-        sql(s"SELECT COUNT(*) FROM $tblName GROUP BY c1"),
-        Seq(Row(3)))
+      val df = sql(s"SELECT c1, COUNT(*), SUM(c2) FROM $tblName GROUP BY c1")
+      val executedPlan = df.queryExecution.executedPlan
 
-      withSQLConf("spark.sql.test.forceApplyObjectHashAggregate" -> true.toString) {
-        checkAnswer(
-          sql(s"SELECT COUNT(*) FROM $tblName GROUP BY c1"),
-          Seq(Row(1), Row(1), Row(1)))
+      assert(collectFirst(executedPlan) {
+        case _: ObjectHashAggregateExec => true
+      }.nonEmpty)
+      assert(collectFirst(executedPlan) {
+        case _: SortAggregateExec => true
+      }.isEmpty)
 
-        checkAnswer(
-          sql(s"SELECT COLLECT_LIST(c2) AS c3 FROM $tblName GROUP BY c1 ORDER BY c3"),
-          Seq(Row(Seq(1)), Row(Seq(2)), Row(Seq(3))))
+      val res = df.collect()
+      assert(res.length == 1)
+      assert(res(0).getString(0).toLowerCase() == "hello")
+      assert(res(0).getLong(1) == 3L)
+      assert(res(0).getLong(2) == 6L)
+    }
+  }
+
+  test("disable hash aggregate on collated column via SQLConf") {
+    val tblName = "grp_by_disabled_tbl"
+    withTable(tblName) {
+      sql(s"CREATE TABLE $tblName (c1 STRING COLLATE UTF8_LCASE, c2 INT) USING PARQUET")
+      sql(s"INSERT INTO $tblName VALUES ('hello', 1), ('HELLO', 2), ('HeLlO', 3)")
+
+      withSQLConf(SQLConf.COLLATION_HASH_AGGREGATION_ENABLED.key -> "false") {
+        val df = sql(s"SELECT c1, COUNT(*) FROM $tblName GROUP BY c1")
+        val executedPlan = df.queryExecution.executedPlan
+
+        assert(collectFirst(executedPlan) {
+          case _: SortAggregateExec => true
+        }.nonEmpty)
+        assert(collectFirst(executedPlan) {
+          case _: ObjectHashAggregateExec => true
+          case _: HashAggregateExec => true
+        }.isEmpty)
+
+        val res = df.collect()
+        assert(res.length == 1)
+        assert(res(0).getString(0).toLowerCase() == "hello")
+        assert(res(0).getLong(1) == 3L)
       }
     }
   }
 
-  test("imperative aggregate fn does not use objectHashAggregate when group by collated column") {
+  test("imperative aggregate fn uses objectHashAggregate when group by collated column") {
     val tblName = "imp_agg"
     Seq(true, false).foreach { useObjHashAgg =>
       withTable(tblName) {
@@ -66,17 +94,16 @@ class CollationAggregationSuite
           val df = sql(s"SELECT COLLECT_LIST(c2) as list FROM $tblName GROUP BY c1")
           val executedPlan = df.queryExecution.executedPlan
 
-          // Plan should not have any hash aggregate nodes.
-          collectFirst(executedPlan) {
-            case _: ObjectHashAggregateExec => fail("ObjectHashAggregateExec should not be used.")
-            case _: HashAggregateExec => fail("HashAggregateExec should not be used.")
+          if (useObjHashAgg) {
+            assert(collectFirst(executedPlan) {
+              case _: ObjectHashAggregateExec => true
+            }.nonEmpty)
+          } else {
+            assert(collectFirst(executedPlan) {
+              case _: SortAggregateExec => true
+            }.nonEmpty)
           }
 
-          // Plan should have a [[SortAggregateExec]] node.
-          assert(collectFirst(executedPlan) {
-            case _: SortAggregateExec => true
-          }.nonEmpty)
-
           checkAnswer(
             // Sort the values to get deterministic output.
             df.selectExpr("array_sort(list)"),