[SPARK-34581][SQL] Don't optimize out grouping expressions from aggregate expressions

Author: peter-toth

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/grouping.scala

@@ -18,7 +18,7 @@
 package org.apache.spark.sql.catalyst.expressions
 
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
+import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodegenFallback, ExprCode}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
 
@@ -155,3 +155,20 @@ object GroupingID {
     if (SQLConf.get.integerGroupingIdEnabled) IntegerType else LongType
   }
 }
+
+/**
+ * Wrapper expression to avoid further optizations of child
+ */
+case class GroupingExpression(child: Expression) extends UnaryExpression {
+  override def eval(input: InternalRow): Any = {
+    child.eval(input)
+  }
+
+  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    child.genCode(ctx)
+  }
+
+  override def dataType: DataType = {
+    child.dataType
+  }
+}

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

@@ -870,8 +870,19 @@ object CollapseProject extends Rule[LogicalPlan] with AliasHelper {
       if (haveCommonNonDeterministicOutput(p.projectList, agg.aggregateExpressions)) {
         p
       } else {
-        agg.copy(aggregateExpressions = buildCleanedProjectList(
-          p.projectList, agg.aggregateExpressions))
+        val complexGroupingExpressions =
+          ExpressionSet(agg.groupingExpressions.filter(_.children.nonEmpty))
+
+        def wrapGroupingExpression(e: Expression): Expression = e match {
+          case _: AggregateExpression => e
+          case _ if complexGroupingExpressions.contains(e) => GroupingExpression(e)
+          case _ => e.mapChildren(wrapGroupingExpression)
+        }
+
+        val wrappedAggregateExpressions =
+          agg.aggregateExpressions.map(wrapGroupingExpression(_).asInstanceOf[NamedExpression])
+        agg.copy(aggregateExpressions =
+          buildCleanedProjectList(p.projectList, wrappedAggregateExpressions))
       }
     case Project(l1, g @ GlobalLimit(_, limit @ LocalLimit(_, p2 @ Project(l2, _))))
         if isRenaming(l1, l2) =>

sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala

@@ -4116,6 +4116,23 @@ class SQLQuerySuite extends QueryTest with SharedSparkSession with AdaptiveSpark
       }
     }
   }
+
+  test("SPARK-34581: Don't optimize out grouping expressions from aggregate expressions") {
+    withTempView("t") {
+      Seq[Integer](null, 1, 2, 3, null).toDF("id").createOrReplaceTempView("t")
+
+      val df = spark.sql(
+        """
+          |SELECT not(id), c
+          |FROM (
+          |  SELECT t.id IS NULL AS id, count(*) AS c
+          |  FROM t
+          |  GROUP BY t.id IS NULL
+          |) t
+          |""".stripMargin)
+      checkAnswer(df, Row(true, 3) :: Row(false, 2) :: Nil)
+    }
+  }
 }
 
 case class Foo(bar: Option[String])