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 ba0e2275b51..cb8100a6f4b 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 @@ -47,13 +47,11 @@ case class MergeTwoPhasesHashBaseAggregate(session: SparkSession) private def isPartialAgg(partialAgg: BaseAggregateExec, finalAgg: BaseAggregateExec): Boolean = { // 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. + // 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. if ( partialAgg.aggregateExpressions.forall(x => x.mode == Partial) && finalAgg.aggregateExpressions.forall(x => x.mode == Final) && 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 a5e6356d034..3b4d371ed85 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 @@ -113,6 +113,43 @@ abstract class BaseMergeTwoPhasesHashBaseAggregateSuite extends WholeStageTransf compareResult = true, df => checkHashAggregateCount(df, 1) ) + + // pure distinct + FILTER: Spark's planAggregateWithOneDistinct inserts an exchange + // between the partial and final stages, so the merge rule's direct parent-child pattern + // match does NOT fire. This records that boundary: distinct + FILTER is not merged into + // a single aggregate, and the FILTER is preserved through the unmerged stages. + compareResultsAgainstVanillaSpark( + """ + |SELECT count(DISTINCT key) FILTER (WHERE key > 50) AS pc + |FROM v1 + |""".stripMargin, + compareResult = true, + df => { + df.collect() + val plans = collect(df.queryExecution.executedPlan) { + case agg: HashAggregateExecBaseTransformer => agg + } + assert(plans.size > 1, "distinct + FILTER should not be merged into a single aggregate") + } + ) + + // mixed distinct + non-distinct + FILTER: same exchange boundary; the merge rule does + // not fire. Verifies FILTER is preserved on both sides through the unmerged stages. + compareResultsAgainstVanillaSpark( + """ + |SELECT count(DISTINCT key) FILTER (WHERE key > 50) AS dc, + | count(key) FILTER (WHERE key LIKE '%1%') AS pc + |FROM v1 + |""".stripMargin, + compareResult = true, + df => { + df.collect() + val plans = collect(df.queryExecution.executedPlan) { + case agg: HashAggregateExecBaseTransformer => agg + } + assert(plans.size > 1, "distinct + FILTER should not be merged into a single aggregate") + } + ) } // with exchange hash aggregate