Repository: apache/incubator-gluten
Pull request #9280: [GLUTEN-9279] Not pulling out expression from PartialMerge aggregate function to avoid invalid reference binding in ProjectExecTransformer
Author: login: Z1Wu
Labels: name: CORE
description: works for Gluten Core
Created: 2025-04-09T16:46:19Z | Merged: 2025-04-21T01:39:38Z


## What changes were proposed in this pull request?

Pulling out expression from `PartialMerge` aggregate function results in invalid reference binding in `ProjectExecTransformer ` when there are `merge_sum` and `partial_count` aggreagate function with non-attribute expression in same `HashAggregate`.

Fixes issue: https://github.com/apache/incubator-gluten/issues/9279

## How was this patch tested?

UnitTest





Diff:
diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/utils/PullOutProjectHelper.scala b/gluten-substrait/src/main/scala/org/apache/gluten/utils/PullOutProjectHelper.scala
index e4fc11441031..4bb09bfcab6a 100644
--- a/gluten-substrait/src/main/scala/org/apache/gluten/utils/PullOutProjectHelper.scala
+++ b/gluten-substrait/src/main/scala/org/apache/gluten/utils/PullOutProjectHelper.scala
@@ -20,7 +20,7 @@ import org.apache.gluten.backendsapi.BackendsApiManager
 import org.apache.gluten.exception.{GlutenException, GlutenNotSupportException}
 
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateFunction}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateFunction, Complete, Partial}
 import org.apache.spark.sql.execution.aggregate._
 import org.apache.spark.sql.execution.window.WindowExec
 import org.apache.spark.sql.types.{ByteType, DateType, IntegerType, LongType, ShortType}
@@ -137,16 +137,20 @@ trait PullOutProjectHelper {
   protected def rewriteAggregateExpression(
       ae: AggregateExpression,
       expressionMap: mutable.HashMap[Expression, NamedExpression]): AggregateExpression = {
-    val newAggFuncChildren = ae.aggregateFunction.children.map {
-      case literal: Literal => literal
-      case other => replaceExpressionWithAttribute(other, expressionMap)
+    ae.mode match {
+      case Partial | Complete =>
+        val newAggFuncChildren = ae.aggregateFunction.children.map {
+          case literal: Literal => literal
+          case other => replaceExpressionWithAttribute(other, expressionMap)
+        }
+        val newAggFunc = ae.aggregateFunction
+          .withNewChildren(newAggFuncChildren)
+          .asInstanceOf[AggregateFunction]
+        val newFilter =
+          ae.filter.map(replaceExpressionWithAttribute(_, expressionMap))
+        ae.copy(aggregateFunction = newAggFunc, filter = newFilter)
+      case _ => ae
     }
-    val newAggFunc = ae.aggregateFunction
-      .withNewChildren(newAggFuncChildren)
-      .asInstanceOf[AggregateFunction]
-    val newFilter =
-      ae.filter.map(replaceExpressionWithAttribute(_, expressionMap))
-    ae.copy(aggregateFunction = newAggFunc, filter = newFilter)
   }
 
   private def needPreComputeRangeFrameBoundary(bound: Expression): Boolean = {
diff --git a/gluten-substrait/src/test/scala/org/apache/spark/sql/GlutenQueryTest.scala b/gluten-substrait/src/test/scala/org/apache/spark/sql/GlutenQueryTest.scala
index 1d4a63bdbc28..8852fdb137b7 100644
--- a/gluten-substrait/src/test/scala/org/apache/spark/sql/GlutenQueryTest.scala
+++ b/gluten-substrait/src/test/scala/org/apache/spark/sql/GlutenQueryTest.scala
@@ -397,6 +397,28 @@ abstract class GlutenQueryTest extends PlanTest {
     val executedPlan = getExecutedPlan(df)
     assert(executedPlan.exists(plan => tag.runtimeClass.isInstance(plan)))
   }
+
+  /**
+   * Check whether the executed plan of a dataframe contains expected number of expected plans.
+   *
+   * @param df:
+   *   the input dataframe.
+   * @param count:
+   *   expected number of expected plan.
+   * @param tag:
+   *   class of the expected plan.
+   * @tparam T:
+   *   type of the expected plan.
+   */
+  def checkGlutenOperatorCount[T <: GlutenPlan](df: DataFrame, count: Int)(implicit
+      tag: ClassTag[T]): Unit = {
+    val executedPlan = getExecutedPlan(df)
+    assert(
+      executedPlan.count(plan => tag.runtimeClass.isInstance(plan)) == count,
+      s"Expect $count ${tag.runtimeClass.getSimpleName} " +
+        s"in executedPlan:\n ${executedPlan.last}"
+    )
+  }
 }
 
 object GlutenQueryTest extends Assertions {
diff --git a/gluten-ut/test/src/test/scala/org/apache/gluten/extension/GlutenExtensionRewriteRuleSuite.scala b/gluten-ut/test/src/test/scala/org/apache/gluten/extension/GlutenExtensionRewriteRuleSuite.scala
index a29550815041..837d37236b65 100644
--- a/gluten-ut/test/src/test/scala/org/apache/gluten/extension/GlutenExtensionRewriteRuleSuite.scala
+++ b/gluten-ut/test/src/test/scala/org/apache/gluten/extension/GlutenExtensionRewriteRuleSuite.scala
@@ -16,7 +16,7 @@
  */
 package org.apache.gluten.extension
 
-import org.apache.gluten.execution.{ProjectExecTransformer, WholeStageTransformerSuite}
+import org.apache.gluten.execution.{HashAggregateExecBaseTransformer, ProjectExecTransformer, WholeStageTransformerSuite}
 import org.apache.gluten.utils.BackendTestUtils
 
 import org.apache.spark.SparkConf
@@ -62,4 +62,31 @@ class GlutenExtensionRewriteRuleSuite extends WholeStageTransformerSuite {
       }
     )
   }
+
+  test("GLUTEN-9279 - Not Pull out expression to avoid invalid reference binding") {
+    withTable("t") {
+      sql("CREATE TABLE t(f1 String, f2 String, f3 String, f4 String) USING PARQUET")
+      sql("INSERT INTO t values ('1', '2', '3', '4'), ('11' ,'22', '33', '4')")
+      var expectedProjectCount = 3
+      var noFallback = false
+      if (BackendTestUtils.isCHBackendLoaded()) {
+        // The `RewriteMultiChildrenCount` rule in the Velox-backend is the root cause of the
+        // additional ProjectExecTransformer, which leads to the invalid reference binding issue.
+        // We still conduct tests on the CH-backend here to ensure that the introduced modification
+        // in `PullOutPreProject` has no side effect on the CH-backend.
+        expectedProjectCount = 2
+        noFallback = true
+      }
+      runQueryAndCompare(
+        """
+          |SELECT SUM(f1) / COUNT(DISTINCT f2, f3) FROM t GROUP BY f4;
+          |""".stripMargin,
+        noFallBack = noFallback
+      )(
+        df => {
+          checkGlutenOperatorCount[ProjectExecTransformer](df, expectedProjectCount)
+          checkGlutenOperatorCount[HashAggregateExecBaseTransformer](df, 4)
+        })
+    }
+  }
 }


Reviews:
user: login: zml1206
body: One nit, thank you, cc @liujiayi771 
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T02:14:04Z
user: login: Z1Wu
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T14:55:59Z
user: login: liujiayi771
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T15:45:56Z
user: login: liujiayi771
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T15:46:31Z
user: login: Z1Wu
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T17:00:24Z
user: login: liujiayi771
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-18T01:33:18Z

Review comments:
diff_hunk: @@ -62,4 +62,15 @@ class GlutenExtensionRewriteRuleSuite extends WholeStageTransformerSuite {
       }
     )
   }
+
+  test("GLUTEN-9279 - Not Pull out expression to avoid invalid reference binding") {
+    withTable("t") {
+      sql("CREATE TABLE t(f1 String, f2 String, f3 String, f4 String) USING CSV")
+      sql("INSERT INTO t values ('1', '2', '3', '4'), ('11' ,'22', '33', '4')")
+      val df = sql("""
+                     |SELECT SUM(f1) / COUNT(DISTINCT f2, f3) FROM t GROUP BY f4;
+                     |""".stripMargin)
+      assert(df.collect()(0).getDouble(0) == 6)
path: gluten-ut/test/src/test/scala/org/apache/gluten/extension/GlutenExtensionRewriteRuleSuite.scala
user: login: zml1206
body: You can use runQueryAndCompare to add ProjectExecTransformer and HashAggregateExecBaseTransformer checks to ensure offload.
created_at: 2025-04-11T02:08:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -62,4 +62,15 @@ class GlutenExtensionRewriteRuleSuite extends WholeStageTransformerSuite {
       }
     )
   }
+
+  test("GLUTEN-9279 - Not Pull out expression to avoid invalid reference binding") {
+    withTable("t") {
+      sql("CREATE TABLE t(f1 String, f2 String, f3 String, f4 String) USING CSV")
+      sql("INSERT INTO t values ('1', '2', '3', '4'), ('11' ,'22', '33', '4')")
+      val df = sql("""
+                     |SELECT SUM(f1) / COUNT(DISTINCT f2, f3) FROM t GROUP BY f4;
+                     |""".stripMargin)
+      assert(df.collect()(0).getDouble(0) == 6)
path: gluten-ut/test/src/test/scala/org/apache/gluten/extension/GlutenExtensionRewriteRuleSuite.scala
user: login: Z1Wu
body: fixed
created_at: 2025-04-11T14:55:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -62,4 +62,21 @@ class GlutenExtensionRewriteRuleSuite extends WholeStageTransformerSuite {
       }
     )
   }
+
+  test("GLUTEN-9279 - Not Pull out expression to avoid invalid reference binding") {
+    withTable("t") {
+      sql("CREATE TABLE t(f1 String, f2 String, f3 String, f4 String) USING CSV")
+      sql("INSERT INTO t values ('1', '2', '3', '4'), ('11' ,'22', '33', '4')")
+      runQueryAndCompare(
+        """
+          |SELECT SUM(f1) / COUNT(DISTINCT f2, f3) FROM t GROUP BY f4;
+          |""".stripMargin,
+        noFallBack = false
+      )(
+        df => {
+          checkGlutenOperatorMatch[ProjectExecTransformer](df)
path: gluten-ut/test/src/test/scala/org/apache/gluten/extension/GlutenExtensionRewriteRuleSuite.scala
user: login: liujiayi771
body: Also check the number of ProjectExecTransformer and HashAggregateExecBaseTransformer.
created_at: 2025-04-11T15:45:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -62,4 +62,21 @@ class GlutenExtensionRewriteRuleSuite extends WholeStageTransformerSuite {
       }
     )
   }
+
+  test("GLUTEN-9279 - Not Pull out expression to avoid invalid reference binding") {
+    withTable("t") {
+      sql("CREATE TABLE t(f1 String, f2 String, f3 String, f4 String) USING CSV")
path: gluten-ut/test/src/test/scala/org/apache/gluten/extension/GlutenExtensionRewriteRuleSuite.scala
user: login: liujiayi771
body: No need to add USING CSV.
created_at: 2025-04-11T15:46:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -62,4 +62,21 @@ class GlutenExtensionRewriteRuleSuite extends WholeStageTransformerSuite {
       }
     )
   }
+
+  test("GLUTEN-9279 - Not Pull out expression to avoid invalid reference binding") {
+    withTable("t") {
+      sql("CREATE TABLE t(f1 String, f2 String, f3 String, f4 String) USING CSV")
path: gluten-ut/test/src/test/scala/org/apache/gluten/extension/GlutenExtensionRewriteRuleSuite.scala
user: login: Z1Wu
body: If we not add `USING CSV`, Spark will create `Hive serde table` by default, which requires some hive-related settings being configured in Spark Context. 

``` 
WARN org.apache.spark.sql.catalyst.analysis.ResolveSessionCatalog: A Hive serde table will be created as there is no table provider specified. You can set spark.sql.legacy.createHiveTableByDefault to false so that native data source table will be created instead.


[NOT_SUPPORTED_COMMAND_WITHOUT_HIVE_SUPPORT] CREATE Hive TABLE (AS SELECT) is not supported, if you want to enable it, please set "spark.sql.catalogImplementation" to "hive".;
'CreateTable `spark_catalog`.`default`.`t`, org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe, ErrorIfExists
```
created_at: 2025-04-11T17:00:24Z
author_association: CONTRIBUTOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-09T16:46:36Z
body: https://github.com/apache/incubator-gluten/issues/9279
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-04-09T16:46:53Z
body: Run Gluten ClickHouse CI on ARM
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-04-09T16:52:09Z
body: Run Gluten ClickHouse CI on ARM
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-04-09T17:10:18Z
body: Run Gluten ClickHouse CI on ARM
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-04-11T14:55:53Z
body: Run Gluten ClickHouse CI on ARM
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-04-11T15:07:01Z
body: Run Gluten ClickHouse CI on ARM
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-04-11T17:00:56Z
body: Run Gluten ClickHouse CI on ARM
author_association: NONE
user: login: liujiayi771
created_at: 2025-04-14T01:29:57Z
body: @Z1Wu CH UT `GLUTEN-9279 - Not Pull out expression to avoid invalid reference binding *** FAILED ***` is failed, can you also modify this UT.
author_association: CONTRIBUTOR
user: login: Z1Wu
created_at: 2025-04-14T15:10:05Z
body: ![image](https://github.com/user-attachments/assets/7651c0e0-b36e-4b16-8bef-358cb0565f44)

Unable to see the details of failed UT. Where could I get the username and password of jenkins. 
author_association: CONTRIBUTOR
user: login: zml1206
created_at: 2025-04-15T00:02:35Z
body: https://github.com/apache/incubator-gluten/blob/main/docs/get-started/ClickHouse.md#new-ci-system @Z1Wu 
author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-04-17T21:02:03Z
body: Run Gluten ClickHouse CI on ARM
author_association: NONE
user: login: liujiayi771
created_at: 2025-04-18T01:32:51Z
body: cc @baibaichen.
author_association: CONTRIBUTOR