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 1cf2c0d956 [VL] Eager aggregation: Refactor option partialMerge.mode
(#13328)
1cf2c0d956 is described below
commit 1cf2c0d956b555422785d673dfb76c9ca9046269
Author: Hongze Zhang <[email protected]>
AuthorDate: Fri Oct 16 16:38:54 2026 +0100
[VL] Eager aggregation: Refactor option partialMerge.mode (#13328)
---
.../execution/StarSchemaJoinAggregateSuite.scala | 51 ++++++--
docs/Configuration.md | 3 +-
.../org/apache/gluten/config/GlutenConfig.scala | 32 +++--
.../extension/joinagg/ImplementJoinAggregate.scala | 138 ++++++++++++++++-----
.../joinagg/PushAggregateThroughJoin.scala | 9 +-
.../execution/PushAggregateThroughJoinSuite.scala | 77 ++++++------
6 files changed, 215 insertions(+), 95 deletions(-)
diff --git
a/backends-velox/src/test/scala/org/apache/gluten/execution/StarSchemaJoinAggregateSuite.scala
b/backends-velox/src/test/scala/org/apache/gluten/execution/StarSchemaJoinAggregateSuite.scala
index 373ceedef3..0b8bb11e41 100644
---
a/backends-velox/src/test/scala/org/apache/gluten/execution/StarSchemaJoinAggregateSuite.scala
+++
b/backends-velox/src/test/scala/org/apache/gluten/execution/StarSchemaJoinAggregateSuite.scala
@@ -32,9 +32,12 @@ class VanillaJoinAggregateLogicalOnlyExtensions extends
(SparkSessionExtensions
}
}
-class VanillaStarSchemaJoinAggregateLogicalOnlySuite extends
StarSchemaJoinAggregateSuite {
+class StarSchemaJoinAggregateVanillaLogicalOnlySuite extends
StarSchemaJoinAggregateSuite {
override protected def sparkConf: SparkConf = {
super.sparkConf
+ .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_MAX_DEPTH.key,
Int.MaxValue.toString)
+ .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_PARTIAL_MERGE_ENABLED.key,
"false")
+ .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_TWO_STAGE_MODE.key, "none")
.set(GlutenConfig.GLUTEN_ENABLED.key, "false")
.set("spark.shuffle.manager", "sort")
.set(
@@ -52,9 +55,12 @@ class VanillaJoinAggregateExtensions extends
(SparkSessionExtensions => Unit) {
}
}
-class VanillaStarSchemaJoinAggregateSuite extends StarSchemaJoinAggregateSuite
{
+class StarSchemaJoinAggregateVanillaSuite extends StarSchemaJoinAggregateSuite
{
override protected def sparkConf: SparkConf = {
super.sparkConf
+ .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_MAX_DEPTH.key,
Int.MaxValue.toString)
+ .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_PARTIAL_MERGE_ENABLED.key,
"false")
+ .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_TWO_STAGE_MODE.key, "none")
.set(GlutenConfig.GLUTEN_ENABLED.key, "false")
.set("spark.shuffle.manager", "sort")
.set("spark.sql.extensions",
classOf[VanillaJoinAggregateExtensions].getCanonicalName)
@@ -63,28 +69,45 @@ class VanillaStarSchemaJoinAggregateSuite extends
StarSchemaJoinAggregateSuite {
override protected def checkPlan(df: DataFrame): Unit = {}
}
-class StarSchemaJoinAggregateSingleDepthSuite extends
StarSchemaJoinAggregateSuite {
+class StarSchemaJoinAggregateMaxDepthOneSuite extends
StarSchemaJoinAggregateSuite {
override protected def sparkConf: SparkConf = {
super.sparkConf
.set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_MAX_DEPTH.key, "1")
+ .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_PARTIAL_MERGE_ENABLED.key,
"false")
+ .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_TWO_STAGE_MODE.key, "none")
}
}
-class StarSchemaJoinAggregatePartialMergeEnabledSuite extends
StarSchemaJoinAggregateSuite {
+class StarSchemaJoinAggregatePartialMergeTwoStageSuite extends
StarSchemaJoinAggregateSuite {
override protected def sparkConf: SparkConf = {
super.sparkConf
- .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_PARTIAL_MERGE_MODE.key,
"all")
+ .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_MAX_DEPTH.key,
Int.MaxValue.toString)
+ .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_PARTIAL_MERGE_ENABLED.key,
"true")
+ .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_TWO_STAGE_MODE.key,
"final")
}
}
-class StarSchemaJoinAggregateTopPartialMergeSuite extends
StarSchemaJoinAggregateSuite {
+class StarSchemaJoinAggregateTwoStageFinalSuite extends
StarSchemaJoinAggregateSuite {
override protected def sparkConf: SparkConf = {
super.sparkConf
- .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_PARTIAL_MERGE_MODE.key,
"top")
+ .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_MAX_DEPTH.key,
Int.MaxValue.toString)
+ .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_PARTIAL_MERGE_ENABLED.key,
"false")
+ .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_TWO_STAGE_MODE.key,
"final")
}
}
-class StarSchemaJoinAggregateSuite extends VeloxTPCHTableSupport with
AdaptiveSparkPlanHelper {
+class StarSchemaJoinAggregateDefaultSuite extends StarSchemaJoinAggregateSuite
{
+ override protected def sparkConf: SparkConf = {
+ super.sparkConf
+ .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_MAX_DEPTH.key,
Int.MaxValue.toString)
+ .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_PARTIAL_MERGE_ENABLED.key,
"false")
+ .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_TWO_STAGE_MODE.key, "none")
+ }
+}
+
+abstract class StarSchemaJoinAggregateSuite
+ extends VeloxTPCHTableSupport
+ with AdaptiveSparkPlanHelper {
private val factMeasureColumnNames = Set(
"sales_price",
"return_amt",
@@ -108,8 +131,6 @@ class StarSchemaJoinAggregateSuite extends
VeloxTPCHTableSupport with AdaptiveSp
super.sparkConf
.set(GlutenConfig.COLUMNAR_FORCE_SHUFFLED_HASH_JOIN_ENABLED.key, "true")
.set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_ENABLED.key, "true")
- .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_MAX_DEPTH.key,
s"${Int.MaxValue}")
- .set(GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_PARTIAL_MERGE_MODE.key,
"none")
.set("spark.sql.adaptive.enabled", "false")
}
@@ -1074,4 +1095,14 @@ class StarSchemaJoinAggregateSuite extends
VeloxTPCHTableSupport with AdaptiveSp
runQueryAndCompare(query)(df => checkDf(df))
}
+
+ test("Support filtered aggregates alongside pushed aggregates") {
+ val query = """
+ |SELECT SUM(ss.ss_sales_price),
+ | MAX(d.d_moy) FILTER (WHERE d.d_year = 1999)
+ |FROM store_sales ss
+ |JOIN date_dim d ON ss.ss_sold_date_sk = d.d_date_sk
+ |""".stripMargin
+ runQueryAndCompare(query)(df => checkDf(df))
+ }
}
diff --git a/docs/Configuration.md b/docs/Configuration.md
index 117d5782b0..fa5dd41b7c 100644
--- a/docs/Configuration.md
+++ b/docs/Configuration.md
@@ -138,7 +138,8 @@ nav_order: 15
| spark.gluten.sql.orc.charType.scan.fallback.enabled | 🔄
Dynamic | true | Force fallback for orc char type scan.
[...]
| spark.gluten.sql.pushAggregateThroughJoin.enabled | 🔄
Dynamic | false | Enables the push-aggregate-through-join
optimization in Gluten. When enabled, aggregate operators may be pushed below
joins during logical optimization and corresponding physical plans may be
rewritten to execute the aggregation earlier.
[...]
| spark.gluten.sql.pushAggregateThroughJoin.maxDepth | 🔄
Dynamic | 2147483647 | Maximum join traversal depth when applying the
push-aggregate-through-join optimization. A value of 1 allows pushing an
aggregate through a single join; larger values allow the rule to traverse and
push through multiple consecutive joins.
[...]
-| spark.gluten.sql.pushAggregateThroughJoin.partialMerge.mode | 🔄
Dynamic | NONE | Controls where PartialMerge aggregates are
retained during aggregate pushdown: none retains none, top retains one at the
original partial aggregate position, and all retains one at every successful
join push.
[...]
+| spark.gluten.sql.pushAggregateThroughJoin.partialMerge.enabled | 🔄
Dynamic | false | Retains local partial merges between joins
crossed by aggregate pushdown.
[...]
+| spark.gluten.sql.pushAggregateThroughJoin.twoStage.mode | 🔄
Dynamic | NONE | Controls physical staging of the final join
aggregate: NONE uses a single final stage; FINAL adds a local partial merge
before the final aggregation exchange.
[...]
| spark.gluten.sql.removeNativeWriteFilesSortAndProject | 🔄
Dynamic | true | When true, Gluten will remove the vanilla
Spark V1Writes added sort and project for velox backend.
[...]
| spark.gluten.sql.rewrite.dateTimestampComparison | 🔄
Dynamic | true | Rewrite the comparision between date and
timestamp to timestamp comparison.For example `from_unixtime(ts) > date` will
be rewritten to `ts > to_unixtime(date)`
[...]
| spark.gluten.sql.scan.detailedMetrics.enabled | 🔄
Dynamic | true | When true (default), Velox backend scan
operators register all detailed SQL metrics. When false, only essential scan
metrics are registered to reduce driver memory usage. Also enabled
automatically when spark.gluten.sql.debug is true. Does not affect the
ClickHouse backend.
[...]
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/config/GlutenConfig.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/config/GlutenConfig.scala
index 6699c6eb22..1459e7fed0 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/config/GlutenConfig.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/config/GlutenConfig.scala
@@ -30,8 +30,10 @@ import java.util.Locale
import scala.collection.JavaConverters._
import scala.collection.mutable
-object PartialMergeMode extends Enumeration {
- val NONE, TOP, ALL = Value
+object JoinAggregateTwoStageMode extends Enumeration {
+ // TODO: Support ALL at existing shuffle boundaries; splitting below joins
can otherwise
+ // introduce unnecessary shuffles, especially for broadcast hash joins.
+ val NONE, FINAL = Value
}
trait ShuffleWriterType {
@@ -162,8 +164,11 @@ class GlutenConfig(conf: SQLConf) extends
GlutenCoreConfig(conf) {
def pushAggregateThroughJoinMaxDepth: Int =
getConf(PUSH_AGGREGATE_THROUGH_JOIN_MAX_DEPTH)
- def pushAggregateThroughJoinPartialMergeMode: PartialMergeMode.Value =
- getConf(PUSH_AGGREGATE_THROUGH_JOIN_PARTIAL_MERGE_MODE)
+ def pushAggregateThroughJoinPartialMergeEnabled: Boolean =
+ getConf(PUSH_AGGREGATE_THROUGH_JOIN_PARTIAL_MERGE_ENABLED)
+
+ def pushAggregateThroughJoinTwoStageMode: JoinAggregateTwoStageMode.Value =
+ getConf(PUSH_AGGREGATE_THROUGH_JOIN_TWO_STAGE_MODE)
def forceOrcCharTypeScanFallbackEnabled: Boolean =
getConf(VELOX_FORCE_ORC_CHAR_TYPE_SCAN_FALLBACK)
@@ -791,14 +796,19 @@ object GlutenConfig extends ConfigRegistry {
.checkValue(_ >= 1, "must be greater than or equal to 1.")
.createWithDefault(Int.MaxValue)
- val PUSH_AGGREGATE_THROUGH_JOIN_PARTIAL_MERGE_MODE =
- buildConf("spark.gluten.sql.pushAggregateThroughJoin.partialMerge.mode")
+ val PUSH_AGGREGATE_THROUGH_JOIN_PARTIAL_MERGE_ENABLED =
+ buildConf("spark.gluten.sql.pushAggregateThroughJoin.partialMerge.enabled")
+ .doc("Retains local partial merges between joins crossed by aggregate
pushdown.")
+ .booleanConf
+ .createWithDefault(false)
+
+ val PUSH_AGGREGATE_THROUGH_JOIN_TWO_STAGE_MODE =
+ buildConf("spark.gluten.sql.pushAggregateThroughJoin.twoStage.mode")
.doc(
- "Controls where PartialMerge aggregates are retained during aggregate
pushdown: " +
- "none retains none, top retains one at the original partial
aggregate position, " +
- "and all retains one at every successful join push.")
- .enumConf(PartialMergeMode)
- .createWithDefault(PartialMergeMode.NONE)
+ "Controls physical staging of the final join aggregate: NONE uses a
single final " +
+ "stage; FINAL adds a local partial merge before the final
aggregation exchange.")
+ .enumConf(JoinAggregateTwoStageMode)
+ .createWithDefault(JoinAggregateTwoStageMode.NONE)
val GLUTEN_SOFT_AFFINITY_ENABLED =
buildConf("spark.gluten.soft-affinity.enabled")
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/ImplementJoinAggregate.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/ImplementJoinAggregate.scala
index 5a9434e473..6a48a33ccb 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/ImplementJoinAggregate.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/ImplementJoinAggregate.scala
@@ -16,11 +16,11 @@
*/
package org.apache.gluten.extension.joinagg
-import org.apache.gluten.config.GlutenConfig
+import org.apache.gluten.config.{GlutenConfig, JoinAggregateTwoStageMode}
import org.apache.spark.sql.SparkSession
-import org.apache.spark.sql.catalyst.expressions.{Alias, GetStructField,
NamedExpression}
-import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.expressions.{Alias, Expression,
GetStructField, NamedExpression}
+import
org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression,
Complete, Final, Partial, PartialMerge}
import org.apache.spark.sql.catalyst.planning.PhysicalAggregation
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
import org.apache.spark.sql.execution.{ProjectExec, SparkPlan, SparkStrategy}
@@ -29,6 +29,11 @@ import
org.apache.spark.sql.execution.aggregate.HashAggregateExec
import scala.collection.mutable.ArrayBuffer
case class ImplementJoinAggregate(spark: SparkSession) extends SparkStrategy {
+ private object AggregationStage extends Enumeration {
+ // DISTRIBUTED requires grouped input but does not add a mapper aggregate.
+ val LOCAL, DISTRIBUTED, TWO_STAGE = Value
+ }
+
/*
* This strategy lowers the logical join-aggregate wrapper shape produced by
* `PushAggregateThroughJoin` into ordinary Spark physical hash aggregates.
@@ -117,7 +122,8 @@ case class ImplementJoinAggregate(spark: SparkSession)
extends SparkStrategy {
resultExpressions,
childPlan,
unpackInputBuffers = false,
- packOutputBuffers = true)
+ packOutputBuffers = true,
+ stage = AggregationStage.LOCAL)
}
private def planPartialMergePhase(
@@ -131,7 +137,31 @@ case class ImplementJoinAggregate(spark: SparkSession)
extends SparkStrategy {
resultExpressions,
childPlan,
unpackInputBuffers = true,
- packOutputBuffers = true)
+ packOutputBuffers = true,
+ stage = AggregationStage.LOCAL)
+ }
+
+ private def planFinalPhase(
+ grouping: Seq[NamedExpression],
+ aggregateExpressions: Seq[AggregateExpression],
+ resultExpressions: Seq[NamedExpression],
+ childPlan: SparkPlan): Option[SparkPlan] = {
+ val stage =
+ if (
+ GlutenConfig.get.pushAggregateThroughJoinTwoStageMode ==
JoinAggregateTwoStageMode.FINAL
+ ) {
+ AggregationStage.TWO_STAGE
+ } else {
+ AggregationStage.DISTRIBUTED
+ }
+ planPhase(
+ grouping,
+ aggregateExpressions,
+ resultExpressions,
+ childPlan,
+ unpackInputBuffers = true,
+ packOutputBuffers = false,
+ stage = stage)
}
private def planPhase(
@@ -140,7 +170,8 @@ case class ImplementJoinAggregate(spark: SparkSession)
extends SparkStrategy {
resultExpressions: Seq[NamedExpression],
childPlan: SparkPlan,
unpackInputBuffers: Boolean,
- packOutputBuffers: Boolean): Option[SparkPlan] = {
+ packOutputBuffers: Boolean,
+ stage: AggregationStage.Value): Option[SparkPlan] = {
val rewrittenAggExprs = rewriteWrapperAggregates(aggregateExpressions)
if (rewrittenAggExprs.isEmpty) {
return None
@@ -150,16 +181,75 @@ case class ImplementJoinAggregate(spark: SparkSession)
extends SparkStrategy {
} else {
childPlan
}
+ val stageChild = if (stage == AggregationStage.TWO_STAGE) {
+ buildMapperAggregate(grouping, preparedChild, rewrittenAggExprs)
+ } else {
+ preparedChild
+ }
+ val stageExpressions = if (stage == AggregationStage.TWO_STAGE) {
+ rewrittenAggExprs.map {
+ ae =>
+ val mode = ae.mode match {
+ case Partial => PartialMerge
+ case Complete => Final
+ case other => other
+ }
+ // Like Spark's AggUtils.mayRemoveAggFilters, remove filters from
buffer-merge stages.
+ ae.copy(mode = mode, filter = None)
+ }
+ } else {
+ rewrittenAggExprs
+ }
+ val distribution = stage match {
+ case AggregationStage.LOCAL => None
+ case AggregationStage.DISTRIBUTED | AggregationStage.TWO_STAGE =>
+ Some(grouping.map(_.toAttribute))
+ }
if (packOutputBuffers) {
- planBufferPhase(
+ buildBufferAggregate(
grouping,
aggregateExpressions,
resultExpressions,
- preparedChild,
- rewrittenAggExprs)
+ stageChild,
+ stageExpressions,
+ distribution)
} else {
- planFinalOutput(grouping, resultExpressions, preparedChild,
rewrittenAggExprs)
+ buildFinalAggregate(grouping, resultExpressions, stageChild,
stageExpressions, distribution)
+ }
+ }
+
+ private def buildMapperAggregate(
+ grouping: Seq[NamedExpression],
+ childPlan: SparkPlan,
+ aggregateExpressions: Seq[AggregateExpression]): SparkPlan = {
+ val mergeExpressions = aggregateExpressions.map {
+ ae =>
+ val mode = ae.mode match {
+ case Partial | Complete => Partial
+ case _ => PartialMerge
+ }
+ ae.copy(mode = mode)
}
+ // Preserve the reducer stage's input-buffer IDs when emitting mapper
buffers.
+ val bufferOutputs = mergeExpressions.flatMap {
+ ae =>
+ val function = ae.aggregateFunction
+
function.aggBufferAttributes.zip(function.inputAggBufferAttributes).map {
+ case (output, input) =>
+ Alias(output, input.name)(exprId = input.exprId, qualifier =
input.qualifier)
+ }
+ }
+ HashAggregateExec(
+ requiredChildDistributionExpressions = None,
+ isStreaming = false,
+ numShufflePartitions = None,
+ groupingExpressions = grouping,
+ aggregateExpressions = mergeExpressions,
+ aggregateAttributes =
mergeExpressions.flatMap(_.aggregateFunction.aggBufferAttributes),
+ initialInputBufferOffset = 0,
+ resultExpressions = grouping.map(_.toAttribute) ++ bufferOutputs,
+ child = childPlan
+ )
}
private def unpackInputBufferFields(
@@ -199,14 +289,15 @@ case class ImplementJoinAggregate(spark: SparkSession)
extends SparkStrategy {
}
}
- private def planBufferPhase(
+ private def buildBufferAggregate(
grouping: Seq[NamedExpression],
aggregateExpressions: Seq[AggregateExpression],
resultExpressions: Seq[NamedExpression],
childPlan: SparkPlan,
- rewrittenAggExprs: Seq[AggregateExpression]): Option[SparkPlan] = {
+ rewrittenAggExprs: Seq[AggregateExpression],
+ requiredChildDistributionExpressions: Option[Seq[Expression]]):
Option[SparkPlan] = {
val hashAgg = HashAggregateExec(
- requiredChildDistributionExpressions = None,
+ requiredChildDistributionExpressions =
requiredChildDistributionExpressions,
isStreaming = false,
numShufflePartitions = None,
groupingExpressions = grouping,
@@ -276,32 +367,19 @@ case class ImplementJoinAggregate(spark: SparkSession)
extends SparkStrategy {
Some(ProjectExec(rewrittenResultExpressions, hashAgg))
}
- private def planFinalPhase(
- grouping: Seq[NamedExpression],
- aggregateExpressions: Seq[AggregateExpression],
- resultExpressions: Seq[NamedExpression],
- childPlan: SparkPlan): Option[SparkPlan] = {
- planPhase(
- grouping,
- aggregateExpressions,
- resultExpressions,
- childPlan,
- unpackInputBuffers = true,
- packOutputBuffers = false)
- }
-
- private def planFinalOutput(
+ private def buildFinalAggregate(
grouping: Seq[NamedExpression],
resultExpressions: Seq[NamedExpression],
childPlan: SparkPlan,
- rewrittenAggExprs: Seq[AggregateExpression]): Option[SparkPlan] = {
+ rewrittenAggExprs: Seq[AggregateExpression],
+ requiredChildDistributionExpressions: Option[Seq[Expression]]):
Option[SparkPlan] = {
val aggregateAttrs = rewrittenAggExprs.map(_.resultAttribute)
val rewrittenResultExpressions =
rewriteResultAsAggregateAttributes(resultExpressions, rewrittenAggExprs)
Some(
HashAggregateExec(
- requiredChildDistributionExpressions =
Some(grouping.map(_.toAttribute)),
+ requiredChildDistributionExpressions =
requiredChildDistributionExpressions,
isStreaming = false,
numShufflePartitions = None,
groupingExpressions = grouping,
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/PushAggregateThroughJoin.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/PushAggregateThroughJoin.scala
index 4f318f10c8..a9a39290f7 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/PushAggregateThroughJoin.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/extension/joinagg/PushAggregateThroughJoin.scala
@@ -16,7 +16,7 @@
*/
package org.apache.gluten.extension.joinagg
-import org.apache.gluten.config.{GlutenConfig, PartialMergeMode}
+import org.apache.gluten.config.GlutenConfig
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute,
AttributeReference, AttributeSet, EqualTo, Expression, NamedExpression,
PredicateHelper}
@@ -174,7 +174,7 @@ case class PushAggregateThroughJoin(spark: SparkSession)
private def pushPartialWrapperAggregate(lowerPartialAgg: Aggregate):
LogicalPlan = {
// Push one join edge per iteration. This keeps the rewrite local and lets
`maxDepth` bound
// how far a pushed aggregate is allowed to travel through a multi-join
subtree.
- val partialMergeMode =
GlutenConfig.get.pushAggregateThroughJoinPartialMergeMode
+ val partialMergeEnabled =
GlutenConfig.get.pushAggregateThroughJoinPartialMergeEnabled
var current: LogicalPlan = lowerPartialAgg
var changed = true
var pushCount = 0
@@ -182,9 +182,8 @@ case class PushAggregateThroughJoin(spark: SparkSession)
changed = false
current = current.transformUp {
case partialAgg: Aggregate if
isPurePartialWrapperAggregate(partialAgg) =>
- // The first successful push leaves a merge at the original partial
aggregate position.
- val retainPartialMerge = partialMergeMode == PartialMergeMode.ALL ||
- (partialMergeMode == PartialMergeMode.TOP && pushCount == 0)
+ // The physical strategy owns the top merge; retain logical merges
only between joins.
+ val retainPartialMerge = partialMergeEnabled && pushCount > 0
val pushed = if (retainPartialMerge) {
val (partialMergeAgg, freshLowerPartialAgg) =
splitPartialAggregateForMerge(partialAgg)
pushOnce(freshLowerPartialAgg).map {
diff --git
a/gluten-substrait/src/test/scala/org/apache/gluten/execution/PushAggregateThroughJoinSuite.scala
b/gluten-substrait/src/test/scala/org/apache/gluten/execution/PushAggregateThroughJoinSuite.scala
index 9f204ba688..bdbaf07caf 100644
---
a/gluten-substrait/src/test/scala/org/apache/gluten/execution/PushAggregateThroughJoinSuite.scala
+++
b/gluten-substrait/src/test/scala/org/apache/gluten/execution/PushAggregateThroughJoinSuite.scala
@@ -38,7 +38,7 @@ class PushAggregateThroughJoinSuite extends PlanTest with
SharedSparkSession {
private case class PushdownCase(inputSql: String)
- private case class ExpectedAggCounts(none: Int, top: Int, all: Int)
+ private case class ExpectedAggCounts(partialMergeDisabled: Int,
partialMergeEnabled: Int)
override protected def sparkConf: SparkConf = {
// Avoid Janino projection codegen here because Spark 4's
QueryExecutionErrors
@@ -95,18 +95,19 @@ class PushAggregateThroughJoinSuite extends PlanTest with
SharedSparkSession {
maxDepth: Int,
expectedPushCount: Int,
expectedAggCounts: ExpectedAggCounts): Unit = {
- val modes = Seq(
- "none" -> expectedAggCounts.none,
- "top" -> expectedAggCounts.top,
- "all" -> expectedAggCounts.all)
- modes.foreach {
- case (partialMergeMode, expectedAggCount) =>
- withClue(s"partialMergeMode=$partialMergeMode, maxDepth=$maxDepth: ") {
+ Seq(false, true).foreach {
+ partialMergeEnabled =>
+ val expectedAggCount = if (partialMergeEnabled) {
+ expectedAggCounts.partialMergeEnabled
+ } else {
+ expectedAggCounts.partialMergeDisabled
+ }
+ withClue(s"partialMergeEnabled=$partialMergeEnabled,
maxDepth=$maxDepth: ") {
withSQLConf(
GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_ENABLED.key -> "true",
GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_MAX_DEPTH.key ->
maxDepth.toString,
- GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_PARTIAL_MERGE_MODE.key ->
- partialMergeMode
+ GlutenConfig.PUSH_AGGREGATE_THROUGH_JOIN_PARTIAL_MERGE_ENABLED.key
->
+ partialMergeEnabled.toString
) {
val (withoutRuleRows, withoutRuleLogicalPlan,
withoutRulePhysicalPlan) =
withExtraPlanning(Nil, Nil) {
@@ -160,25 +161,25 @@ class PushAggregateThroughJoinSuite extends PlanTest with
SharedSparkSession {
if (debugMode) {
// scalastyle:off println
- println("=== Optimized Plan Before (without
PushJoinAggregatePreAggregation) ===")
+ println("=== Optimized Plan Before (without
PushAggregateThroughJoin) ===")
println(withoutRuleLogicalPlan.treeString)
- println("=== Optimized Plan After (with
PushJoinAggregatePreAggregation) ===")
+ println("=== Optimized Plan After (with
PushAggregateThroughJoin) ===")
println(withRuleLogicalPlan.treeString)
println(
- "=== Optimized Plan After (with
PushJoinAggregatePreAggregation and strategy) ===")
+ "=== Optimized Plan After (with PushAggregateThroughJoin and
strategy) ===")
println(withRuleAndStrategyLogicalPlan.treeString)
- println("=== Physical Plan Before (without
PushJoinAggregatePreAggregation) ===")
+ println("=== Physical Plan Before (without
PushAggregateThroughJoin) ===")
println(withoutRulePhysicalPlan.treeString)
- println("=== Physical Plan After (with
PushJoinAggregatePreAggregation only) ===")
+ println("=== Physical Plan After (with PushAggregateThroughJoin
only) ===")
println(withRulePhysicalPlan.treeString)
println(
- "=== Physical Plan After (with PushJoinAggregatePreAggregation
and strategy) ===")
+ "=== Physical Plan After (with PushAggregateThroughJoin and
strategy) ===")
println(withRuleAndStrategyPhysicalPlan.treeString)
- println("=== Result Before (without
PushJoinAggregatePreAggregation) ===")
+ println("=== Result Before (without PushAggregateThroughJoin)
===")
println(withoutRuleRows.mkString("\n"))
- println("=== Result After (with PushJoinAggregatePreAggregation
only) ===")
+ println("=== Result After (with PushAggregateThroughJoin only)
===")
println(withRuleRows.mkString("\n"))
- println("=== Result After (with PushJoinAggregatePreAggregation
and strategy) ===")
+ println("=== Result After (with PushAggregateThroughJoin and
strategy) ===")
println(withRuleAndStrategyRows.mkString("\n"))
// scalastyle:on println
}
@@ -235,7 +236,7 @@ class PushAggregateThroughJoinSuite extends PlanTest with
SharedSparkSession {
pushdownCase,
maxDepth = Int.MaxValue,
expectedPushCount = 2,
- expectedAggCounts = ExpectedAggCounts(none = 2, top = 3, all = 4))
+ expectedAggCounts = ExpectedAggCounts(partialMergeDisabled = 2,
partialMergeEnabled = 3))
}
test("pre-aggregate store_sales for sum") {
@@ -253,7 +254,7 @@ class PushAggregateThroughJoinSuite extends PlanTest with
SharedSparkSession {
pushdownCase,
maxDepth = Int.MaxValue,
expectedPushCount = 1,
- expectedAggCounts = ExpectedAggCounts(none = 2, top = 3, all = 3))
+ expectedAggCounts = ExpectedAggCounts(partialMergeDisabled = 2,
partialMergeEnabled = 2))
}
test("pre-aggregate store_sales for avg") {
@@ -271,7 +272,7 @@ class PushAggregateThroughJoinSuite extends PlanTest with
SharedSparkSession {
pushdownCase,
maxDepth = Int.MaxValue,
expectedPushCount = 1,
- expectedAggCounts = ExpectedAggCounts(none = 2, top = 3, all = 3))
+ expectedAggCounts = ExpectedAggCounts(partialMergeDisabled = 2,
partialMergeEnabled = 2))
}
test("pre-aggregate store_sales for sum on fact table") {
@@ -289,7 +290,7 @@ class PushAggregateThroughJoinSuite extends PlanTest with
SharedSparkSession {
pushdownCase,
maxDepth = Int.MaxValue,
expectedPushCount = 1,
- expectedAggCounts = ExpectedAggCounts(none = 2, top = 3, all = 3))
+ expectedAggCounts = ExpectedAggCounts(partialMergeDisabled = 2,
partialMergeEnabled = 2))
}
test("pre-aggregate store_sales for avg on fact table") {
@@ -307,7 +308,7 @@ class PushAggregateThroughJoinSuite extends PlanTest with
SharedSparkSession {
pushdownCase,
maxDepth = Int.MaxValue,
expectedPushCount = 1,
- expectedAggCounts = ExpectedAggCounts(none = 2, top = 3, all = 3))
+ expectedAggCounts = ExpectedAggCounts(partialMergeDisabled = 2,
partialMergeEnabled = 2))
}
test("pre-aggregate store_sales for sum on three-way join") {
@@ -327,7 +328,7 @@ class PushAggregateThroughJoinSuite extends PlanTest with
SharedSparkSession {
pushdownCase,
maxDepth = Int.MaxValue,
expectedPushCount = 2,
- expectedAggCounts = ExpectedAggCounts(none = 2, top = 3, all = 4))
+ expectedAggCounts = ExpectedAggCounts(partialMergeDisabled = 2,
partialMergeEnabled = 3))
}
test("pre-aggregate store_sales for sum and avg on different fact columns on
three-way join") {
@@ -348,7 +349,7 @@ class PushAggregateThroughJoinSuite extends PlanTest with
SharedSparkSession {
pushdownCase,
maxDepth = Int.MaxValue,
expectedPushCount = 2,
- expectedAggCounts = ExpectedAggCounts(none = 2, top = 3, all = 4))
+ expectedAggCounts = ExpectedAggCounts(partialMergeDisabled = 2,
partialMergeEnabled = 3))
}
test("pre-aggregate store_sales for sum and avg on same fact column on
three-way join") {
@@ -369,7 +370,7 @@ class PushAggregateThroughJoinSuite extends PlanTest with
SharedSparkSession {
pushdownCase,
maxDepth = Int.MaxValue,
expectedPushCount = 2,
- expectedAggCounts = ExpectedAggCounts(none = 2, top = 3, all = 4))
+ expectedAggCounts = ExpectedAggCounts(partialMergeDisabled = 2,
partialMergeEnabled = 3))
}
test("pre-aggregate store_sales by i_item_desc") {
@@ -387,7 +388,7 @@ class PushAggregateThroughJoinSuite extends PlanTest with
SharedSparkSession {
pushdownCase,
maxDepth = Int.MaxValue,
expectedPushCount = 1,
- expectedAggCounts = ExpectedAggCounts(none = 2, top = 3, all = 3))
+ expectedAggCounts = ExpectedAggCounts(partialMergeDisabled = 2,
partialMergeEnabled = 2))
}
test("pre-aggregate store_sales by substr(i_item_desc, 3), 3 ways") {
@@ -407,7 +408,7 @@ class PushAggregateThroughJoinSuite extends PlanTest with
SharedSparkSession {
pushdownCase,
maxDepth = Int.MaxValue,
expectedPushCount = 2,
- expectedAggCounts = ExpectedAggCounts(none = 2, top = 3, all = 4))
+ expectedAggCounts = ExpectedAggCounts(partialMergeDisabled = 2,
partialMergeEnabled = 3))
}
test("pre-aggregate store_sales for sum with item filter") {
@@ -424,7 +425,7 @@ class PushAggregateThroughJoinSuite extends PlanTest with
SharedSparkSession {
pushdownCase,
maxDepth = Int.MaxValue,
expectedPushCount = 1,
- expectedAggCounts = ExpectedAggCounts(none = 2, top = 3, all = 3))
+ expectedAggCounts = ExpectedAggCounts(partialMergeDisabled = 2,
partialMergeEnabled = 2))
}
test("pre-aggregate three-way joins independently under union all") {
@@ -467,17 +468,17 @@ class PushAggregateThroughJoinSuite extends PlanTest with
SharedSparkSession {
pushdownCase,
maxDepth = 1,
expectedPushCount = 3,
- expectedAggCounts = ExpectedAggCounts(none = 6, top = 9, all = 9))
+ expectedAggCounts = ExpectedAggCounts(partialMergeDisabled = 6,
partialMergeEnabled = 6))
runCaseWithMaxDepth(
pushdownCase,
maxDepth = 2,
expectedPushCount = 6,
- expectedAggCounts = ExpectedAggCounts(none = 6, top = 9, all = 12))
+ expectedAggCounts = ExpectedAggCounts(partialMergeDisabled = 6,
partialMergeEnabled = 9))
runCaseWithMaxDepth(
pushdownCase,
maxDepth = Int.MaxValue,
expectedPushCount = 6,
- expectedAggCounts = ExpectedAggCounts(none = 6, top = 9, all = 12))
+ expectedAggCounts = ExpectedAggCounts(partialMergeDisabled = 6,
partialMergeEnabled = 9))
}
test("pre-aggregate store_sales for sum on three-way join with maxDepth=1 /
maxDepth=2") {
@@ -497,17 +498,17 @@ class PushAggregateThroughJoinSuite extends PlanTest with
SharedSparkSession {
pushdownCase,
maxDepth = 1,
expectedPushCount = 1,
- expectedAggCounts = ExpectedAggCounts(none = 2, top = 3, all = 3))
+ expectedAggCounts = ExpectedAggCounts(partialMergeDisabled = 2,
partialMergeEnabled = 2))
runCaseWithMaxDepth(
pushdownCase,
maxDepth = 2,
expectedPushCount = 2,
- expectedAggCounts = ExpectedAggCounts(none = 2, top = 3, all = 4))
+ expectedAggCounts = ExpectedAggCounts(partialMergeDisabled = 2,
partialMergeEnabled = 3))
runCaseWithMaxDepth(
pushdownCase,
maxDepth = Int.MaxValue,
expectedPushCount = 2,
- expectedAggCounts = ExpectedAggCounts(none = 2, top = 3, all = 4))
+ expectedAggCounts = ExpectedAggCounts(partialMergeDisabled = 2,
partialMergeEnabled = 3))
}
test("partial merge is not retained when no side can be pushed") {
@@ -522,7 +523,7 @@ class PushAggregateThroughJoinSuite extends PlanTest with
SharedSparkSession {
pushdownCase,
maxDepth = Int.MaxValue,
expectedPushCount = 0,
- expectedAggCounts = ExpectedAggCounts(none = 2, top = 2, all = 2))
+ expectedAggCounts = ExpectedAggCounts(partialMergeDisabled = 2,
partialMergeEnabled = 2))
}
test("pre-aggregate with filter inside inner equi-join") {
@@ -540,6 +541,6 @@ class PushAggregateThroughJoinSuite extends PlanTest with
SharedSparkSession {
pushdownCase,
maxDepth = Int.MaxValue,
expectedPushCount = 1,
- expectedAggCounts = ExpectedAggCounts(none = 2, top = 3, all = 3))
+ expectedAggCounts = ExpectedAggCounts(partialMergeDisabled = 2,
partialMergeEnabled = 2))
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
