This is an automated email from the ASF dual-hosted git repository.
zzcclp pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new af7589eddc [GLUTEN-12652][CORE] Support merging two-phase aggregates
with FILTER clause (#12653)
af7589eddc is described below
commit af7589eddc519e23a18b058de2a86371a8d37df7
Author: Kaifei Yi <[email protected]>
AuthorDate: Mon Aug 3 13:09:29 2026 +0800
[GLUTEN-12652][CORE] Support merging two-phase aggregates with FILTER
clause (#12653)
Previously MergeTwoPhasesHashBaseAggregate skipped merging any two-phase
aggregate whose expressions carry a FILTER clause. This PR removes that
restriction.
The physical Final aggregate has its FILTER stripped by Spark's
AggUtils.mayRemoveAggFilters (FILTER is only kept in Partial/Complete modes).
So when merging into Complete mode, the FILTER is now restored from the partial
aggregate — whose expressions align 1:1 with the final ones — via a new
toCompleteAggregateExpressions helper (guarded by a length check). This covers
the hash, object-hash, and sort aggregate merge paths.
Fix #12652.
How was this patch tested?
Extended MergeTwoPhasesHashBaseAggregateSuite to assert both merge (single
aggregate node) and result correctness against vanilla Spark for filtered
hash/object/sort aggregates, plus a mixed filtered + non-filtered case that
verifies the FILTER is restored on the correct expression by positional
alignment.
---
.../columnar/MergeTwoPhasesHashBaseAggregate.scala | 43 +++++++++++---
.../MergeTwoPhasesHashBaseAggregateSuite.scala | 66 ++++++++++++++--------
2 files changed, 78 insertions(+), 31 deletions(-)
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/MergeTwoPhasesHashBaseAggregate.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/MergeTwoPhasesHashBaseAggregate.scala
index 16d4c0ebe6..ba0e2275b5 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/MergeTwoPhasesHashBaseAggregate.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/MergeTwoPhasesHashBaseAggregate.scala
@@ -20,7 +20,7 @@ import org.apache.gluten.config.GlutenConfig
import org.apache.spark.internal.Logging
import org.apache.spark.sql.SparkSession
-import org.apache.spark.sql.catalyst.expressions.aggregate.{Complete, Final,
Partial}
+import
org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression,
Complete, Final, Partial}
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.execution.SparkPlan
import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec,
HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec}
@@ -45,10 +45,19 @@ case class MergeTwoPhasesHashBaseAggregate(session:
SparkSession)
val mergeTwoPhasesAggEnabled: Boolean =
GlutenConfig.get.mergeTwoPhasesAggEnabled
private def isPartialAgg(partialAgg: BaseAggregateExec, finalAgg:
BaseAggregateExec): Boolean = {
- // TODO: now it can not support to merge agg which there are the filters
in the aggregate exprs.
+ // Aggregates with a FILTER clause can be merged as long as the FILTER
predicate is carried
+ // over to the Complete mode aggregate. Note the physical final aggregate
has its FILTER
+ // stripped (Spark's AggUtils.mayRemoveAggFilters only keeps FILTER in
Partial/Complete modes),
+ // so the FILTER must be restored from the partial aggregate when merging.
Spark's aggregate
+ // planning produces the partial and final aggregate expression lists
together and keeps them
+ // positionally aligned (the final phase reads the partial buffer by
position), so a
+ // partial/final pair can be matched by position. We cannot match by
`resultId`: for a single
+ // distinct aggregate, `AggUtils.planAggregateWithOneDistinct` builds the
partial and final
+ // distinct expressions with fresh `AggregateExpression` instances, so
their `resultId`s differ.
if (
- partialAgg.aggregateExpressions.forall(x => x.mode == Partial &&
x.filter.isEmpty) &&
- finalAgg.aggregateExpressions.forall(x => x.mode == Final &&
x.filter.isEmpty)
+ partialAgg.aggregateExpressions.forall(x => x.mode == Partial) &&
+ finalAgg.aggregateExpressions.forall(x => x.mode == Final) &&
+ partialAgg.aggregateExpressions.size ==
finalAgg.aggregateExpressions.size
) {
(finalAgg.logicalLink, partialAgg.logicalLink) match {
case (Some(agg1), Some(agg2)) => agg1.sameResult(agg2)
@@ -59,6 +68,23 @@ case class MergeTwoPhasesHashBaseAggregate(session:
SparkSession)
}
}
+ /**
+ * Builds Complete mode aggregate expressions from the final aggregate. The
physical final
+ * aggregate no longer carries the FILTER predicate (see `isPartialAgg`), so
the FILTER is
+ * restored from the partial aggregate. A partial/final pair is matched by
position: Spark's
+ * aggregate planning emits the two expression lists together and keeps them
positionally aligned
+ * (the final phase reads the partial buffer by position), and
`isPartialAgg` has already checked
+ * that the two lists have the same size.
+ */
+ private def toCompleteAggregateExpressions(
+ partialAgg: BaseAggregateExec,
+ finalAggExpressions: Seq[AggregateExpression]): Seq[AggregateExpression]
= {
+ finalAggExpressions.zip(partialAgg.aggregateExpressions).map {
+ case (finalExpr, partialExpr) =>
+ finalExpr.copy(mode = Complete, filter = partialExpr.filter)
+ }
+ }
+
override def apply(plan: SparkPlan): SparkPlan = {
if (!mergeTwoPhasesAggEnabled || !enableColumnarHashAgg) {
plan
@@ -75,7 +101,8 @@ case class MergeTwoPhasesHashBaseAggregate(session:
SparkSession)
resultExpressions,
child: HashAggregateExec) if !isStreaming && isPartialAgg(child,
hashAgg) =>
// convert to complete mode aggregate expressions
- val completeAggregateExpressions =
aggregateExpressions.map(_.copy(mode = Complete))
+ val completeAggregateExpressions =
+ toCompleteAggregateExpressions(child, aggregateExpressions)
hashAgg.copy(
groupingExpressions = child.groupingExpressions,
aggregateExpressions = completeAggregateExpressions,
@@ -94,7 +121,8 @@ case class MergeTwoPhasesHashBaseAggregate(session:
SparkSession)
child: ObjectHashAggregateExec)
if !isStreaming && isPartialAgg(child, objectHashAgg) =>
// convert to complete mode aggregate expressions
- val completeAggregateExpressions =
aggregateExpressions.map(_.copy(mode = Complete))
+ val completeAggregateExpressions =
+ toCompleteAggregateExpressions(child, aggregateExpressions)
objectHashAgg.copy(
requiredChildDistributionExpressions = None,
groupingExpressions = child.groupingExpressions,
@@ -114,7 +142,8 @@ case class MergeTwoPhasesHashBaseAggregate(session:
SparkSession)
child: SortAggregateExec)
if replaceSortAggWithHashAgg && !isStreaming &&
isPartialAgg(child, sortAgg) =>
// convert to complete mode aggregate expressions
- val completeAggregateExpressions =
aggregateExpressions.map(_.copy(mode = Complete))
+ val completeAggregateExpressions =
+ toCompleteAggregateExpressions(child, aggregateExpressions)
sortAgg.copy(
requiredChildDistributionExpressions = None,
groupingExpressions = child.groupingExpressions,
diff --git
a/gluten-ut/test/src/test/scala/org/apache/gluten/execution/MergeTwoPhasesHashBaseAggregateSuite.scala
b/gluten-ut/test/src/test/scala/org/apache/gluten/execution/MergeTwoPhasesHashBaseAggregateSuite.scala
index d3ca2acdd6..a5e6356d03 100644
---
a/gluten-ut/test/src/test/scala/org/apache/gluten/execution/MergeTwoPhasesHashBaseAggregateSuite.scala
+++
b/gluten-ut/test/src/test/scala/org/apache/gluten/execution/MergeTwoPhasesHashBaseAggregateSuite.scala
@@ -90,14 +90,28 @@ abstract class BaseMergeTwoPhasesHashBaseAggregateSuite
extends WholeStageTransf
1
)
- // with filter hash aggregate
- checkHashAggregateCount(
- spark.sql("""
- |SELECT key, count(key) FILTER (WHERE key LIKE '%1%') AS
pc2
- |FROM v1
- |GROUP BY key
- |""".stripMargin),
- 2
+ // with filter hash aggregate: it is merged into one complete-mode
aggregate, and the
+ // FILTER predicate must be preserved so the result still matches
vanilla Spark.
+ compareResultsAgainstVanillaSpark(
+ """
+ |SELECT key, count(key) FILTER (WHERE key LIKE '%1%') AS pc2
+ |FROM v1
+ |GROUP BY key
+ |""".stripMargin,
+ compareResult = true,
+ df => checkHashAggregateCount(df, 1)
+ )
+
+ // mix of filtered and non-filtered aggregates: verifies the FILTER is
restored on the
+ // right aggregate expression (positional alignment) after merging to
complete mode.
+ compareResultsAgainstVanillaSpark(
+ """
+ |SELECT key, count(key) AS c, count(key) FILTER (WHERE key LIKE
'%1%') AS pc2
+ |FROM v1
+ |GROUP BY key
+ |""".stripMargin,
+ compareResult = true,
+ df => checkHashAggregateCount(df, 1)
)
}
@@ -130,14 +144,16 @@ abstract class BaseMergeTwoPhasesHashBaseAggregateSuite
extends WholeStageTransf
1
)
- // with filter object aggregate
- checkObjectAggregateCount(
- spark.sql("""
- |SELECT key, collect_list(key) FILTER (WHERE key LIKE
'%1%') AS pc2
- |FROM v1
- |GROUP BY key
- |""".stripMargin),
- 2
+ // with filter object aggregate: it is merged into one complete-mode
aggregate, and the
+ // FILTER predicate must be preserved so the result still matches
vanilla Spark.
+ compareResultsAgainstVanillaSpark(
+ """
+ |SELECT key, collect_list(key) FILTER (WHERE key LIKE '%1%') AS pc2
+ |FROM v1
+ |GROUP BY key
+ |""".stripMargin,
+ compareResult = true,
+ df => checkObjectAggregateCount(df, 1)
)
}
@@ -171,14 +187,16 @@ abstract class BaseMergeTwoPhasesHashBaseAggregateSuite
extends WholeStageTransf
1
)
- // with filter sort aggregate
- checkSortAggregateCount(
- spark.sql("""
- |SELECT key, sum(if(key<0,0,key)) FILTER (WHERE key LIKE
'%1%') AS pc2
- |FROM v1
- |GROUP BY key
- |""".stripMargin),
- 2
+ // with filter sort aggregate: it is merged into one complete-mode
aggregate, and the
+ // FILTER predicate must be preserved so the result still matches
vanilla Spark.
+ compareResultsAgainstVanillaSpark(
+ """
+ |SELECT key, sum(if(key<0,0,key)) FILTER (WHERE key LIKE '%1%') AS
pc2
+ |FROM v1
+ |GROUP BY key
+ |""".stripMargin,
+ compareResult = true,
+ df => checkSortAggregateCount(df, 1)
)
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
