This is an automated email from the ASF dual-hosted git repository.
zhztheplayer 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 b9dda4b955 [GLUTEN-13349][VL] Fix flushable aggregation
double-counting distinct keys (#13350)
b9dda4b955 is described below
commit b9dda4b95536b2a4c72b3820234ca2fffc79c9a1
Author: YangJie <[email protected]>
AuthorDate: Mon Oct 19 07:31:16 2026 -0400
[GLUTEN-13349][VL] Fix flushable aggregation double-counting distinct keys
(#13350)
---
.../extension/FlushableHashAggregateRule.scala | 16 +++--
.../execution/VeloxAggregateFunctionsSuite.scala | 68 +++++++++++++++++++++-
2 files changed, 78 insertions(+), 6 deletions(-)
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/extension/FlushableHashAggregateRule.scala
b/backends-velox/src/main/scala/org/apache/gluten/extension/FlushableHashAggregateRule.scala
index d5c1292fd1..7aff5b73ea 100644
---
a/backends-velox/src/main/scala/org/apache/gluten/extension/FlushableHashAggregateRule.scala
+++
b/backends-velox/src/main/scala/org/apache/gluten/extension/FlushableHashAggregateRule.scala
@@ -22,7 +22,8 @@ import org.apache.gluten.execution._
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.expressions.aggregate._
import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.{SortExec, SparkPlan}
+import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
import org.apache.spark.sql.types.{DataType, DoubleType, FloatType}
/**
@@ -122,7 +123,9 @@ case class FlushableHashAggregateRule(session:
SparkSession) extends Rule[SparkP
private def collectProtectedOneDistinctPartialMergeAggIds(plan: SparkPlan):
Set[Int] = {
val protectedAggIds = Set.newBuilder[Int]
plan.foreach {
- case agg: HashAggregateExecTransformer =>
+ // The distinct-partial aggregate may run in vanilla Spark, e.g. when
one of its functions is
+ // not offloaded, while the PartialMerge aggregate below it is still
offloaded.
+ case agg: BaseAggregateExec =>
findProtectedPartialMergeAgg(agg).foreach {
protectedAgg => protectedAggIds += protectedAgg.id
}
@@ -131,9 +134,11 @@ case class FlushableHashAggregateRule(session:
SparkSession) extends Rule[SparkP
protectedAggIds.result()
}
- /** If this aggregate is the distinct-partial stage, return its child
PartialMerge aggregate. */
+ /**
+ * If this aggregate is the distinct-partial stage, return the PartialMerge
aggregate below it.
+ */
private def findProtectedPartialMergeAgg(
- distinctPartialAgg: HashAggregateExecTransformer):
Option[HashAggregateExecTransformer] = {
+ distinctPartialAgg: BaseAggregateExec):
Option[HashAggregateExecTransformer] = {
if (
!distinctPartialAgg.aggregateExpressions.exists(
expr => expr.isDistinct && expr.mode == Partial)
@@ -149,6 +154,9 @@ case class FlushableHashAggregateRule(session:
SparkSession) extends Rule[SparkP
private def asAggregate(plan: SparkPlan):
Option[HashAggregateExecTransformer] = plan match {
case agg: HashAggregateExecTransformer => Some(agg)
+ // A sort-based distinct-partial aggregate reads its child through a local
sort.
+ case sort: SortExecTransformer if !sort.global => asAggregate(sort.child)
+ case sort: SortExec if !sort.global => asAggregate(sort.child)
case _ => None
}
}
diff --git
a/backends-velox/src/test/scala/org/apache/gluten/execution/VeloxAggregateFunctionsSuite.scala
b/backends-velox/src/test/scala/org/apache/gluten/execution/VeloxAggregateFunctionsSuite.scala
index d9c8486154..67f5e79803 100644
---
a/backends-velox/src/test/scala/org/apache/gluten/execution/VeloxAggregateFunctionsSuite.scala
+++
b/backends-velox/src/test/scala/org/apache/gluten/execution/VeloxAggregateFunctionsSuite.scala
@@ -20,7 +20,7 @@ import org.apache.gluten.config.{GlutenConfig, VeloxConfig}
import org.apache.gluten.extension.columnar.validator.FallbackInjects
import org.apache.spark.SparkConf
-import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial,
PartialMerge}
import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
import org.apache.spark.sql.internal.SQLConf
@@ -1307,7 +1307,8 @@ class VeloxAggregateFunctionsFlushSuite extends
VeloxAggregateFunctionsSuite {
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
SQLConf.FILES_MAX_PARTITION_BYTES.key -> "1k") {
- runQueryAndCompare("select count(distinct l_partkey) from lineitem group
by l_orderkey") {
+ runQueryAndCompare(
+ "select count(distinct l_partkey), count(*) from lineitem group by
l_orderkey") {
df =>
val executedPlan = getExecutedPlan(df)
val regularAggCount = executedPlan.count {
@@ -1327,6 +1328,69 @@ class VeloxAggregateFunctionsFlushSuite extends
VeloxAggregateFunctionsSuite {
}
}
+ test("flushable aggregate rule - vanilla distinct-partial keeps partial
merge regular") {
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ withTable("flush_distinct_src") {
+ spark
+ .range(0, 30000, 1, 10)
+ .selectExpr("id % 28000 as z")
+ .selectExpr("z % 3 as k", "z")
+ .write
+ .format("parquet")
+ .saveAsTable("flush_distinct_src")
+ // histogram_numeric is not offloaded, so the distinct-partial
aggregate runs in vanilla
+ // Spark on top of the offloaded PartialMerge aggregate. The raw
histogram depends on input
+ // order, so only the sum of its counts is compared.
+ val sql =
+ """
+ |select k, count(distinct z), count(*),
+ | aggregate(histogram_numeric(distinct z, 5), 0D, (acc, e) -> acc
+ e.y)
+ |from flush_distinct_src group by k
+ |""".stripMargin
+ // With useObjectHashAggregateExec=false the vanilla aggregates are
SortAggregateExec, which
+ // reads its child through a local sort; with columnar sort off that
sort is a SortExec.
+ Seq(
+ Seq(SQLConf.USE_OBJECT_HASH_AGG.key -> "true"),
+ Seq(SQLConf.USE_OBJECT_HASH_AGG.key -> "false"),
+ Seq(
+ SQLConf.USE_OBJECT_HASH_AGG.key -> "false",
+ GlutenConfig.COLUMNAR_SORT_ENABLED.key -> "false")
+ ).foreach {
+ confs =>
+ withSQLConf(confs: _*) {
+ runQueryAndCompare(sql) {
+ df =>
+ val distinctPartialAggs = getExecutedPlan(df).collect {
+ case agg: BaseAggregateExec
+ if agg.aggregateExpressions.exists(
+ e => e.isDistinct && e.mode == Partial) =>
+ agg
+ }
+ assert(
+ distinctPartialAggs.nonEmpty &&
+
distinctPartialAggs.forall(!_.isInstanceOf[HashAggregateExecTransformer]),
+ "expected the distinct-partial aggregate to run in vanilla
Spark"
+ )
+ val partialMergeAggs =
distinctPartialAggs.flatMap(_.collectFirst {
+ case agg: HashAggregateExecTransformer
+ if agg.aggregateExpressions.forall(_.mode ==
PartialMerge) =>
+ agg
+ })
+ assert(
+ partialMergeAggs.nonEmpty,
+ "expected an offloaded PartialMerge aggregate below the
distinct-partial one")
+ assert(
+
partialMergeAggs.forall(!_.isInstanceOf[FlushableHashAggregateExecTransformer]),
+ "the PartialMerge aggregate below the distinct-partial
aggregate must not be " +
+ "flushable"
+ )
+ }
+ }
+ }
+ }
+ }
+ }
+
test("flushable aggregate rule - distinct feeding a join keeps final agg
regular") {
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
