You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@gluten.apache.org by ja...@apache.org on 2025/03/24 01:13:11 UTC
(incubator-gluten) branch main updated: [GLUTEN-8966][VL] Propagate HashAggregate's ignoreNullKeys when possible (#8967)
This is an automated email from the ASF dual-hosted git repository.
jackylee pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-gluten.git
The following commit(s) were added to refs/heads/main by this push:
new e368a32ac7 [GLUTEN-8966][VL] Propagate HashAggregate's ignoreNullKeys when possible (#8967)
e368a32ac7 is described below
commit e368a32ac7a7b832def212596363361e2c245a90
Author: WangGuangxin <wa...@bytedance.com>
AuthorDate: Sun Mar 23 18:13:06 2025 -0700
[GLUTEN-8966][VL] Propagate HashAggregate's ignoreNullKeys when possible (#8967)
---
.../gluten/backendsapi/velox/VeloxRuleApi.scala | 2 +
.../org/apache/gluten/config/VeloxConfig.scala | 12 +++
.../execution/HashAggregateExecTransformer.scala | 24 ++++--
.../HashAggregateIgnoreNullKeysRule.scala | 89 ++++++++++++++++++++++
.../execution/VeloxAggregateFunctionsSuite.scala | 24 +++++-
cpp/velox/substrait/SubstraitToVeloxPlan.cc | 5 ++
.../HashAggregateExecBaseTransformer.scala | 7 +-
7 files changed, 153 insertions(+), 10 deletions(-)
diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala
index 566cffd1bf..82a20b8cc5 100644
--- a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala
+++ b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala
@@ -99,6 +99,7 @@ object VeloxRuleApi {
injector.injectPostTransform(_ => PullOutDuplicateProject)
injector.injectPostTransform(_ => CollapseProjectExecTransformer)
injector.injectPostTransform(c => FlushableHashAggregateRule.apply(c.session))
+ injector.injectPostTransform(c => HashAggregateIgnoreNullKeysRule.apply(c.session))
injector.injectPostTransform(c => InsertTransitions.create(c.outputsColumnar, VeloxBatch))
// Gluten columnar: Fallback policies.
@@ -183,6 +184,7 @@ object VeloxRuleApi {
injector.injectPostTransform(_ => PullOutDuplicateProject)
injector.injectPostTransform(_ => CollapseProjectExecTransformer)
injector.injectPostTransform(c => FlushableHashAggregateRule.apply(c.session))
+ injector.injectPostTransform(c => HashAggregateIgnoreNullKeysRule.apply(c.session))
injector.injectPostTransform(c => InsertTransitions.create(c.outputsColumnar, VeloxBatch))
injector.injectPostTransform(c => RemoveTopmostColumnarToRow(c.session, c.caller.isAqe()))
SparkShimLoader.getSparkShims
diff --git a/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala b/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala
index 9fd8064d6e..e872c9b279 100644
--- a/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala
+++ b/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala
@@ -60,6 +60,9 @@ class VeloxConfig(conf: SQLConf) extends GlutenConfig(conf) {
def veloxOrcScanEnabled: Boolean =
getConf(VELOX_ORC_SCAN_ENABLED)
+
+ def enablePropagateIgnoreNullKeys: Boolean =
+ getConf(VELOX_PROPAGATE_IGNORE_NULL_KEYS_ENABLED)
}
object VeloxConfig {
@@ -520,4 +523,13 @@ object VeloxConfig {
.internal()
.stringConf
.createWithDefault("")
+
+ val VELOX_PROPAGATE_IGNORE_NULL_KEYS_ENABLED =
+ buildConf("spark.gluten.sql.columnar.backend.velox.propagateIgnoreNullKeys")
+ .doc(
+ "If enabled, we will identify aggregation followed by an inner join " +
+ "on the grouping keys, and mark the ignoreNullKeys flag to true to " +
+ "avoid unnecessary aggregation on null keys.")
+ .booleanConf
+ .createWithDefault(true)
}
diff --git a/backends-velox/src/main/scala/org/apache/gluten/execution/HashAggregateExecTransformer.scala b/backends-velox/src/main/scala/org/apache/gluten/execution/HashAggregateExecTransformer.scala
index 197bc60248..70a81e1cdb 100644
--- a/backends-velox/src/main/scala/org/apache/gluten/execution/HashAggregateExecTransformer.scala
+++ b/backends-velox/src/main/scala/org/apache/gluten/execution/HashAggregateExecTransformer.scala
@@ -50,7 +50,8 @@ abstract class HashAggregateExecTransformer(
aggregateAttributes: Seq[Attribute],
initialInputBufferOffset: Int,
resultExpressions: Seq[NamedExpression],
- child: SparkPlan)
+ child: SparkPlan,
+ ignoreNullKeys: Boolean)
extends HashAggregateExecBaseTransformer(
requiredChildDistributionExpressions,
groupingExpressions,
@@ -58,7 +59,9 @@ abstract class HashAggregateExecTransformer(
aggregateAttributes,
initialInputBufferOffset,
resultExpressions,
- child) {
+ child,
+ ignoreNullKeys
+ ) {
override def output: Seq[Attribute] = {
// TODO: We should have a check to make sure the returned schema actually matches the output
@@ -192,7 +195,8 @@ abstract class HashAggregateExecTransformer(
private def formatExtOptimizationString(isStreaming: Boolean): String = {
val isStreamingStr = if (isStreaming) "1" else "0"
val allowFlushStr = if (allowFlush) "1" else "0"
- s"isStreaming=$isStreamingStr\nallowFlush=$allowFlushStr\n"
+ val ignoreNullKeysStr = if (ignoreNullKeys) "1" else "0"
+ s"isStreaming=$isStreamingStr\nallowFlush=$allowFlushStr\nignoreNullKeys=$ignoreNullKeysStr\n"
}
// Create aggregate function node and add to list.
@@ -705,7 +709,8 @@ case class RegularHashAggregateExecTransformer(
aggregateAttributes: Seq[Attribute],
initialInputBufferOffset: Int,
resultExpressions: Seq[NamedExpression],
- child: SparkPlan)
+ child: SparkPlan,
+ ignoreNullKeys: Boolean = false)
extends HashAggregateExecTransformer(
requiredChildDistributionExpressions,
groupingExpressions,
@@ -713,7 +718,9 @@ case class RegularHashAggregateExecTransformer(
aggregateAttributes,
initialInputBufferOffset,
resultExpressions,
- child) {
+ child,
+ ignoreNullKeys
+ ) {
override protected def allowFlush: Boolean = false
@@ -737,7 +744,8 @@ case class FlushableHashAggregateExecTransformer(
aggregateAttributes: Seq[Attribute],
initialInputBufferOffset: Int,
resultExpressions: Seq[NamedExpression],
- child: SparkPlan)
+ child: SparkPlan,
+ ignoreNullKeys: Boolean = false)
extends HashAggregateExecTransformer(
requiredChildDistributionExpressions,
groupingExpressions,
@@ -745,7 +753,9 @@ case class FlushableHashAggregateExecTransformer(
aggregateAttributes,
initialInputBufferOffset,
resultExpressions,
- child) {
+ child,
+ ignoreNullKeys
+ ) {
override protected def allowFlush: Boolean = true
diff --git a/backends-velox/src/main/scala/org/apache/gluten/extension/HashAggregateIgnoreNullKeysRule.scala b/backends-velox/src/main/scala/org/apache/gluten/extension/HashAggregateIgnoreNullKeysRule.scala
new file mode 100644
index 0000000000..1f57d93de0
--- /dev/null
+++ b/backends-velox/src/main/scala/org/apache/gluten/extension/HashAggregateIgnoreNullKeysRule.scala
@@ -0,0 +1,89 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.gluten.extension
+
+import org.apache.gluten.config.VeloxConfig
+import org.apache.gluten.execution._
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.catalyst.plans.Inner
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.adaptive.ShuffleQueryStageExec
+import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
+import org.apache.spark.sql.execution.joins.BaseJoinExec
+
+/**
+ * To identify aggregates that the groupby key is used as inner join keys. In this case, we can set
+ * ignoreNullKeys to true when convert to velox's AggregateNode.
+ */
+case class HashAggregateIgnoreNullKeysRule(session: SparkSession) extends Rule[SparkPlan] {
+ override def apply(plan: SparkPlan): SparkPlan = {
+ if (!VeloxConfig.get.enablePropagateIgnoreNullKeys) {
+ return plan
+ }
+ plan.transformUp {
+ case join: BaseJoinExec if join.joinType == Inner =>
+ val newLeftChild = setIgnoreKeysIfAggregateOnJoinKeys(join.left, join.leftKeys)
+ val newRightChild = setIgnoreKeysIfAggregateOnJoinKeys(join.right, join.rightKeys)
+ if (newLeftChild.fastEquals(join.left) && newRightChild.fastEquals(join.right)) {
+ join
+ } else {
+ join.withNewChildren(Seq(newLeftChild, newRightChild))
+ }
+ case p => p
+ }
+ }
+
+ private def setIgnoreKeysIfAggregateOnJoinKeys(
+ plan: SparkPlan,
+ joinKeys: Seq[Expression]): SparkPlan = plan match {
+ case agg: HashAggregateExecTransformer =>
+ val newChild = setIgnoreKeysIfAggregateOnJoinKeys(agg.child, joinKeys)
+ val canIgnoreNullKeysRule = semanticEquals(agg.groupingExpressions, joinKeys)
+ agg match {
+ case f: FlushableHashAggregateExecTransformer =>
+ f.copy(ignoreNullKeys = canIgnoreNullKeysRule, child = newChild)
+ case r: RegularHashAggregateExecTransformer =>
+ r.copy(ignoreNullKeys = canIgnoreNullKeysRule, child = newChild)
+ case _ => agg
+ }
+ case s: ShuffleQueryStageExec =>
+ s.copy(plan = setIgnoreKeysIfAggregateOnJoinKeys(s.plan, joinKeys))
+ case p if !canPropagate(p) => p
+ case other =>
+ other.withNewChildren(
+ other.children.map(c => setIgnoreKeysIfAggregateOnJoinKeys(c, joinKeys)))
+ }
+
+ private def canPropagate(plan: SparkPlan): Boolean = plan match {
+ case _: ProjectExecTransformer => true
+ case _: WholeStageTransformer => true
+ case _: VeloxResizeBatchesExec => true
+ case _: ShuffleExchangeLike => true
+ case _: VeloxColumnarToRowExec => true
+ case _: SortExecTransformer => true
+ case _ => false
+ }
+
+ private def semanticEquals(aggExpression: Seq[Expression], joinKeys: Seq[Expression]): Boolean = {
+ aggExpression.size == joinKeys.size && aggExpression.zip(joinKeys).forall {
+ case (e1: Expression, e2: Expression) => e1.semanticEquals(e2)
+ }
+ }
+}
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 8a61379df1..a2117a78aa 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
@@ -16,7 +16,7 @@
*/
package org.apache.gluten.execution
-import org.apache.gluten.config.VeloxConfig
+import org.apache.gluten.config.{GlutenConfig, VeloxConfig}
import org.apache.gluten.extension.columnar.validator.FallbackInjects
import org.apache.spark.SparkConf
@@ -1193,6 +1193,28 @@ class VeloxAggregateFunctionsDefaultSuite extends VeloxAggregateFunctionsSuite {
}
}
}
+
+ test("aggregate on join keys can set ignoreNullKeys") {
+ val s =
+ """
+ |select count(1) from
+ | (select l_orderkey, max(l_partkey) from lineitem group by l_orderkey) a
+ |inner join
+ | (select l_orderkey from lineitem) b
+ |on a.l_orderkey = b.l_orderkey
+ |""".stripMargin
+ withSQLConf(GlutenConfig.COLUMNAR_FORCE_SHUFFLED_HASH_JOIN_ENABLED.key -> "true") {
+ runQueryAndCompare(s) {
+ df =>
+ val executedPlan = getExecutedPlan(df)
+ assert(executedPlan.exists {
+ case a: RegularHashAggregateExecTransformer if a.ignoreNullKeys => true
+ case a: FlushableHashAggregateExecTransformer if a.ignoreNullKeys => true
+ case _ => false
+ })
+ }
+ }
+ }
}
class VeloxAggregateFunctionsFlushSuite extends VeloxAggregateFunctionsSuite {
diff --git a/cpp/velox/substrait/SubstraitToVeloxPlan.cc b/cpp/velox/substrait/SubstraitToVeloxPlan.cc
index 0ef2cd6e03..b7fbe9a47f 100644
--- a/cpp/velox/substrait/SubstraitToVeloxPlan.cc
+++ b/cpp/velox/substrait/SubstraitToVeloxPlan.cc
@@ -418,6 +418,11 @@ core::PlanNodePtr SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::
preGroupingExprs.insert(preGroupingExprs.begin(), veloxGroupingExprs.begin(), veloxGroupingExprs.end());
}
+ if (aggRel.has_advanced_extension() &&
+ SubstraitParser::configSetInOptimization(aggRel.advanced_extension(), "ignoreNullKeys=")) {
+ ignoreNullKeys = true;
+ }
+
// Get the output names of Aggregation.
std::vector<std::string> aggOutNames;
aggOutNames.reserve(aggRel.measures().size());
diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/execution/HashAggregateExecBaseTransformer.scala b/gluten-substrait/src/main/scala/org/apache/gluten/execution/HashAggregateExecBaseTransformer.scala
index 35809be0be..a1572073fe 100644
--- a/gluten-substrait/src/main/scala/org/apache/gluten/execution/HashAggregateExecBaseTransformer.scala
+++ b/gluten-substrait/src/main/scala/org/apache/gluten/execution/HashAggregateExecBaseTransformer.scala
@@ -40,7 +40,8 @@ abstract class HashAggregateExecBaseTransformer(
aggregateAttributes: Seq[Attribute],
initialInputBufferOffset: Int,
resultExpressions: Seq[NamedExpression],
- child: SparkPlan)
+ child: SparkPlan,
+ ignoreNullKeys: Boolean = false)
extends BaseAggregateExec
with UnaryTransformSupport {
@@ -87,11 +88,13 @@ abstract class HashAggregateExecBaseTransformer(
s"HashAggregateTransformer(keys=$keyString, " +
s"functions=$functionString, " +
s"isStreamingAgg=$isCapableForStreamingAggregation, " +
+ s"ignoreNullKeys=$ignoreNullKeys, " +
s"output=$outputString)"
} else {
s"HashAggregateTransformer(keys=$keyString, " +
s"functions=$functionString, " +
- s"isStreamingAgg=$isCapableForStreamingAggregation)"
+ s"isStreamingAgg=$isCapableForStreamingAggregation, " +
+ s"ignoreNullKeys=$ignoreNullKeys)"
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@gluten.apache.org
For additional commands, e-mail: commits-help@gluten.apache.org