You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "andygrove (via GitHub)" <gi...@apache.org> on 2026/09/30 14:53:08 UTC
[PR] fix: coalesce shuffle partitions under Comet unions the way Spark does [datafusion-comet]
andygrove opened a new pull request, #6459:
URL: https://github.com/apache/datafusion-comet/pull/6459
## Which issue does this PR close?
Closes #6454.
## Rationale for this change
With AQE on, a union that Comet converts to `CometUnionExec` got no shuffle partition coalescing when one of its branches has a leaf that is not a shuffle, such as a scan or a table-cache stage, so the shuffled branches kept `spark.sql.shuffle.partitions` partitions. The issue has the analysis and a reproduction that ends with 201 partitions where Spark has 2.
Spark's `CoalesceShufflePartitions` coalesces each child of a `UnionExec` as its own group, and from Spark 4.0 each child of a `CartesianProductExec`, `BroadcastHashJoinExec` or `BroadcastNestedLoopJoinExec` too, but it matches those classes. Comet replaces them while AQE prepares a stage, before the optimizer rules run, so Spark's rule sees the Comet operators and falls through to the case that coalesces only when every leaf below is an exchange stage.
## What changes are included in this PR?
A new query-stage optimizer rule, `CometCoalesceShufflePartitions`, registered after Comet's other two, so it runs after Spark's `CoalesceShufflePartitions`. It looks for Comet operators whose `originalPlan` is one of the operators above, and whose shuffle stages no AQE rule has put a read over. For each one it rebuilds the Spark operators over the Comet children, runs Spark's own `CoalesceShufflePartitions` over that, and swaps the Comet operators back in over the new children. Spark's code makes every coalescing decision, per Spark version.
The rule extends `AQEShuffleReadRule`, so AQE treats it like Spark's rule: it is skipped for the final stage when `spark.sql.adaptive.applyFinalStageShuffleOptimizations` is off, and its result is discarded if it breaks a distribution required above it.
It also covers a union whose branches are all shuffles with different partition counts. Spark's rule coalesces a Comet union's shuffles together, finds the counts differ, and coalesces nothing, where for a Spark union it would coalesce each branch on its own.
Two limits:
- When every leaf below a Comet union is an exchange stage and the counts agree, Spark's rule already coalesces its shuffles, together rather than branch by branch, and this rule leaves them as they are.
- Spark 3.4 has no hook for query-stage optimizer rules, so nothing changes there.
Once this lands, #5634 can drop the `IgnoreComet` it puts on Spark's `SPARK-42101` union test.
## How are these changes tested?
- A new test in `CometExecSuite` runs the issue's query, a union of a repartitioned DataFrame with a Parquet scan. It checks for a coalesced read under the `CometUnionExec`, and for the same number of result partitions as with Comet disabled.
- A port of Spark's `SPARK-42101: Coalesce shuffle partition with union even if exists TableCacheQueryStage` to `CometInMemoryCacheSuite`.
Both fail with the rule unregistered. They pass on Spark 3.5, 4.0, 4.1 and 4.2, and are canceled on 3.4.
Spark 4.1's `AdaptiveQueryExecSuite`, `CoalesceShufflePartitionsSuite`, `DataFrameSetOperationsSuite`, `ExchangeSuite` and `DataFrameSuite`, run from the test jar with Comet enabled, give the same results with and without the rule, except that `CoalesceShufflePartitionsSuite`'s `Union two datasets with different pre-shuffle partition number` fails without it and passes with it.
`CometExecSuite`, `CometJoinSuite`, `CometSetOpWithGroupBySuite`, `CometPlanEqualitySuite`, `CometInMemoryCacheSuite` and `CometNativeShuffleSuite` pass on 4.1. The strict-warnings compile and scalafix pass on 3.5.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: coalesce shuffle partitions under Comet unions the way Spark does [datafusion-comet]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #6459:
URL: https://github.com/apache/datafusion-comet/pull/6459#discussion_r4154984789
##########
spark/src/main/scala/org/apache/comet/rules/CometCoalesceShufflePartitions.scala:
##########
@@ -0,0 +1,128 @@
+/*
+ * 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.comet.rules
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.comet.CometExec
+import org.apache.spark.sql.execution.{SparkPlan, UnionExec}
+import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, AQEShuffleReadRule, CoalesceShufflePartitions, ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.exchange.ShuffleOrigin
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, CartesianProductExec}
+
+/**
+ * Coalesces the shuffle partitions below a Comet operator that Spark's CoalesceShufflePartitions
+ * coalesces child by child but does not recognize.
+ *
+ * Spark coalesces each child of a `UnionExec` as a group of its own, and from Spark 4.0 each
+ * child of a `CartesianProductExec`, `BroadcastHashJoinExec` or `BroadcastNestedLoopJoinExec`
+ * too. It matches those classes, and the Comet operators that replace them are other classes, so
+ * it falls through to the case that coalesces only when every leaf below the operator is an
+ * exchange stage. A union with a scan or a table-cache stage in one branch then keeps every
+ * partition of the shuffles in the others: `spark.sql.shuffle.partitions` tasks for a query that
+ * needs a few.
+ *
+ * Comet replaces these operators while AQE prepares a stage, before its optimizer rules run, and
+ * plans the operators above them against the Comet versions. So this runs after Spark's rule
+ * instead, on each such operator whose shuffle stages that rule left untouched. It rebuilds the
+ * Spark operator each Comet one replaced over the Comet children, has Spark's own rule coalesce
+ * that, and swaps the Comet operators back in. The partitions come out as Spark would have
+ * coalesced them, down to which operators count, since it is Spark's code deciding. The one
+ * difference is that Spark divides its minimum partition count among the coalesce groups of the
+ * whole plan, and this among those below the Comet operator, which are usually all of them.
+ *
+ * When every leaf below such an operator is an exchange stage, Spark's rule already coalesces its
+ * shuffles, together rather than child by child, and this leaves them as they are.
+ *
+ * Extending `AQEShuffleReadRule` gets this the same treatment from AQE as Spark's rule: it is
+ * skipped for the final stage when that stage's shuffle optimizations are off, and its result is
+ * discarded if it breaks a distribution required above it.
+ */
+case object CometCoalesceShufflePartitions extends AQEShuffleReadRule {
+
+ // The Comet operator that a stand-in Spark operator was rebuilt from.
+ private val COMET_OPERATOR = TreeNodeTag[SparkPlan]("cometCoalesceShufflePartitions")
+
+ // Required by the trait. Which shuffles are coalesced is decided by Spark's rule, which applies
+ // its own list.
+ override protected def supportedShuffleOrigins: Seq[ShuffleOrigin] =
+ CoalesceShufflePartitions(SparkSession.active).supportedShuffleOrigins
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ if (!conf.coalesceShufflePartitionsEnabled || !plan.exists(replaced(_).isDefined)) {
+ return plan
+ }
+ plan.transformDown {
+ case p if replaced(p).isDefined && untouched(p) => coalesceBelow(p)
+ }
+ }
+
+ // The Spark operator a Comet operator replaced, if Spark's rule coalesces its children one by
+ // one. The class match mirrors Spark's, and Spark's rule decides, for its version, which of
+ // these it actually treats that way.
+ private def replaced(plan: SparkPlan): Option[SparkPlan] = plan match {
+ case comet: CometExec =>
+ comet.originalPlan match {
+ case original @ (_: UnionExec | _: CartesianProductExec | _: BroadcastHashJoinExec |
+ _: BroadcastNestedLoopJoinExec)
+ if original.children.length == comet.children.length =>
+ Some(original)
+ case _ => None
+ }
+ case _ => None
+ }
+
+ // No AQE rule has put a read over any shuffle stage below `plan`: Spark's rule coalesced none
+ // of them, and none is a skew-split or local read that coalescing now could disturb.
+ private def untouched(plan: SparkPlan): Boolean =
+ plan.exists(_.isInstanceOf[ShuffleQueryStageExec]) &&
+ !plan.exists(_.isInstanceOf[AQEShuffleReadExec])
+
+ private def coalesceBelow(plan: SparkPlan): SparkPlan = {
+ val asSpark = plan.transformUp { case p =>
+ replaced(p) match {
+ case Some(original) =>
+ val standIn = original.withNewChildren(p.children)
+ // `withNewChildren` hands back the original itself when the children are the same ones,
+ // and the tag must not land on the operator that the Comet one keeps.
+ if (standIn eq original) {
+ p
+ } else {
+ standIn.setTagValue(COMET_OPERATOR, p)
+ standIn
+ }
+ case None => p
+ }
+ }
+ val coalesced = CoalesceShufflePartitions(SparkSession.active).apply(asSpark)
Review Comment:
[P2] Preserve ancestor-join context when invoking Spark's coalescer. On Spark 4.0+, a `CometUnionExec` containing a coalescible shuffle and a scan can sit below `CartesianProductExec`, for example after a `crossJoin` with broadcasting disabled. Calling Spark's rule only on the union subtree hides the Cartesian ancestor, so Spark loses `hasExplodingJoin` and uses the ordinary advisory size instead of its smaller join-specific target. With 64 partitions reporting 1 MiB each, `minPartitionNum=1`, a 64 MiB advisory size and a 1 MiB minimum size, Spark and pre-PR Comet retain 64 partitions, while this rule reduces them to one. This concentrates the shuffled branch's cross-product work into one partition instead of 64. Please retain the containing join context when delegating, or skip these subtrees until that context can be preserved.
Evidence: Reproduced on Spark 4.1.3 with the exact-head rule and extracted current `CometUnionExec`. The plan is `CartesianProductExec(CometUnionExec(ShuffleQueryStageExec, scan), scan)` with synthetic statistics of 64 × 1 MiB. `python3 /tmp/comet-6459-root-review/run.py` exits 0 and reports `EXPLODING_SPARK_COUNTS=Vector(64); EXPLODING_BEFORE_PR_COUNTS=Vector(64); EXPLODING_AFTER_PR_COUNTS=Vector(1)`. Source and output are in `CurrentPlanProbe.scala` and `probe.log` in that directory. Spark 4.0.4, 4.1.3 and 4.2.0 sources propagate `hasExplodingJoin` from ancestors and use the minimum partition size for these groups. This establishes a partition-parallelism regression without claiming a measured wall-clock slowdown.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: coalesce shuffle partitions under Comet unions the way Spark does [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove merged PR #6459:
URL: https://github.com/apache/datafusion-comet/pull/6459
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: coalesce shuffle partitions under Comet unions the way Spark does [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #6459:
URL: https://github.com/apache/datafusion-comet/pull/6459#issuecomment-6015273159
9d54ba50f covers the case that the review of #5634 raised. The rule no longer leaves a stage alone because it holds an `AQEShuffleReadExec`. Below a Cartesian product, Spark's rule coalesces a shuffle on the other side first, and the read it leaves there kept the union's shuffled branch at every partition. The rule now hides each read behind a leaf while Spark's rule runs over the whole stage, and puts it back afterwards, so Spark's rule leaves alone the shuffles it would coalesce together with a read one and coalesces the rest. The ancestor sizing and the `UnknownPartitioning` guard are unchanged.
`AQE coalesces the shuffle partitions of a union in a stage Spark coalesced in part` covers it. It fails against the previous revision, whose plan has one coalesced read where Spark's has two.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: coalesce shuffle partitions under Comet unions the way Spark does [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #6459:
URL: https://github.com/apache/datafusion-comet/pull/6459#discussion_r4146871073
##########
spark/src/main/scala/org/apache/comet/rules/CometCoalesceShufflePartitions.scala:
##########
@@ -0,0 +1,128 @@
+/*
+ * 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.comet.rules
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.comet.CometExec
+import org.apache.spark.sql.execution.{SparkPlan, UnionExec}
+import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, AQEShuffleReadRule, CoalesceShufflePartitions, ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.exchange.ShuffleOrigin
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, CartesianProductExec}
+
+/**
+ * Coalesces the shuffle partitions below a Comet operator that Spark's CoalesceShufflePartitions
+ * coalesces child by child but does not recognize.
+ *
+ * Spark coalesces each child of a `UnionExec` as a group of its own, and from Spark 4.0 each
+ * child of a `CartesianProductExec`, `BroadcastHashJoinExec` or `BroadcastNestedLoopJoinExec`
+ * too. It matches those classes, and the Comet operators that replace them are other classes, so
+ * it falls through to the case that coalesces only when every leaf below the operator is an
+ * exchange stage. A union with a scan or a table-cache stage in one branch then keeps every
+ * partition of the shuffles in the others: `spark.sql.shuffle.partitions` tasks for a query that
+ * needs a few.
+ *
+ * Comet replaces these operators while AQE prepares a stage, before its optimizer rules run, and
+ * plans the operators above them against the Comet versions. So this runs after Spark's rule
+ * instead, on each such operator whose shuffle stages that rule left untouched. It rebuilds the
+ * Spark operator each Comet one replaced over the Comet children, has Spark's own rule coalesce
+ * that, and swaps the Comet operators back in. The partitions come out as Spark would have
+ * coalesced them, down to which operators count, since it is Spark's code deciding. The one
+ * difference is that Spark divides its minimum partition count among the coalesce groups of the
+ * whole plan, and this among those below the Comet operator, which are usually all of them.
+ *
+ * When every leaf below such an operator is an exchange stage, Spark's rule already coalesces its
+ * shuffles, together rather than child by child, and this leaves them as they are.
+ *
+ * Extending `AQEShuffleReadRule` gets this the same treatment from AQE as Spark's rule: it is
+ * skipped for the final stage when that stage's shuffle optimizations are off, and its result is
+ * discarded if it breaks a distribution required above it.
+ */
+case object CometCoalesceShufflePartitions extends AQEShuffleReadRule {
+
+ // The Comet operator that a stand-in Spark operator was rebuilt from.
+ private val COMET_OPERATOR = TreeNodeTag[SparkPlan]("cometCoalesceShufflePartitions")
+
+ // Required by the trait. Which shuffles are coalesced is decided by Spark's rule, which applies
+ // its own list.
+ override protected def supportedShuffleOrigins: Seq[ShuffleOrigin] =
+ CoalesceShufflePartitions(SparkSession.active).supportedShuffleOrigins
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ if (!conf.coalesceShufflePartitionsEnabled || !plan.exists(replaced(_).isDefined)) {
+ return plan
+ }
+ plan.transformDown {
+ case p if replaced(p).isDefined && untouched(p) => coalesceBelow(p)
+ }
+ }
+
+ // The Spark operator a Comet operator replaced, if Spark's rule coalesces its children one by
+ // one. The class match mirrors Spark's, and Spark's rule decides, for its version, which of
+ // these it actually treats that way.
+ private def replaced(plan: SparkPlan): Option[SparkPlan] = plan match {
+ case comet: CometExec =>
+ comet.originalPlan match {
+ case original @ (_: UnionExec | _: CartesianProductExec | _: BroadcastHashJoinExec |
+ _: BroadcastNestedLoopJoinExec)
+ if original.children.length == comet.children.length =>
+ Some(original)
+ case _ => None
+ }
+ case _ => None
+ }
+
+ // No AQE rule has put a read over any shuffle stage below `plan`: Spark's rule coalesced none
+ // of them, and none is a skew-split or local read that coalescing now could disturb.
+ private def untouched(plan: SparkPlan): Boolean =
+ plan.exists(_.isInstanceOf[ShuffleQueryStageExec]) &&
+ !plan.exists(_.isInstanceOf[AQEShuffleReadExec])
+
+ private def coalesceBelow(plan: SparkPlan): SparkPlan = {
+ val asSpark = plan.transformUp { case p =>
+ replaced(p) match {
+ case Some(original) =>
+ val standIn = original.withNewChildren(p.children)
+ // `withNewChildren` hands back the original itself when the children are the same ones,
+ // and the tag must not land on the operator that the Comet one keeps.
+ if (standIn eq original) {
+ p
+ } else {
+ standIn.setTagValue(COMET_OPERATOR, p)
+ standIn
+ }
+ case None => p
+ }
+ }
+ val coalesced = CoalesceShufflePartitions(SparkSession.active).apply(asSpark)
+ if (coalesced eq asSpark) plan else restore(coalesced)
Review Comment:
On Spark 4.1+ a union advertises its children's partitioning when they match (`spark.sql.unionOutputPartitioning`), so a final aggregate above it can skip its shuffle. If only one branch is coalescible, for example the other is `repartition(200, $"k")`, coalescing it alone changes that partitioning, and I expect AQE's `ValidateRequirements` to accept it because `CometHashAggregateExec` doesn't declare `requiredChildDistribution`. I haven't run this, so it may not be reachable. Would it make sense to skip operators whose `outputPartitioning` isn't `UnknownPartitioning`, with a 4.1 test that groups by the key over such a union?
##########
spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala:
##########
@@ -107,6 +107,7 @@ class CometSparkSessionExtensions
}
injectQueryStageOptimizerRuleShim(extensions, CometPlanAdaptiveDynamicPruningFilters)
injectQueryStageOptimizerRuleShim(extensions, CometReuseSubquery)
+ injectQueryStageOptimizerRuleShim(extensions, CometCoalesceShufflePartitions)
Review Comment:
Nit: the class scaladoc above lists the AQE query-stage optimizer rules for each stage and notes which ones are not registered on Spark 3.4. Would it make sense to add this rule to both places?
##########
spark/src/main/scala/org/apache/comet/rules/CometCoalesceShufflePartitions.scala:
##########
@@ -0,0 +1,128 @@
+/*
+ * 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.comet.rules
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.comet.CometExec
+import org.apache.spark.sql.execution.{SparkPlan, UnionExec}
+import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, AQEShuffleReadRule, CoalesceShufflePartitions, ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.exchange.ShuffleOrigin
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, CartesianProductExec}
+
+/**
+ * Coalesces the shuffle partitions below a Comet operator that Spark's CoalesceShufflePartitions
+ * coalesces child by child but does not recognize.
+ *
+ * Spark coalesces each child of a `UnionExec` as a group of its own, and from Spark 4.0 each
+ * child of a `CartesianProductExec`, `BroadcastHashJoinExec` or `BroadcastNestedLoopJoinExec`
+ * too. It matches those classes, and the Comet operators that replace them are other classes, so
+ * it falls through to the case that coalesces only when every leaf below the operator is an
+ * exchange stage. A union with a scan or a table-cache stage in one branch then keeps every
+ * partition of the shuffles in the others: `spark.sql.shuffle.partitions` tasks for a query that
+ * needs a few.
+ *
+ * Comet replaces these operators while AQE prepares a stage, before its optimizer rules run, and
+ * plans the operators above them against the Comet versions. So this runs after Spark's rule
+ * instead, on each such operator whose shuffle stages that rule left untouched. It rebuilds the
+ * Spark operator each Comet one replaced over the Comet children, has Spark's own rule coalesce
+ * that, and swaps the Comet operators back in. The partitions come out as Spark would have
+ * coalesced them, down to which operators count, since it is Spark's code deciding. The one
+ * difference is that Spark divides its minimum partition count among the coalesce groups of the
+ * whole plan, and this among those below the Comet operator, which are usually all of them.
+ *
+ * When every leaf below such an operator is an exchange stage, Spark's rule already coalesces its
+ * shuffles, together rather than child by child, and this leaves them as they are.
+ *
+ * Extending `AQEShuffleReadRule` gets this the same treatment from AQE as Spark's rule: it is
+ * skipped for the final stage when that stage's shuffle optimizations are off, and its result is
+ * discarded if it breaks a distribution required above it.
+ */
+case object CometCoalesceShufflePartitions extends AQEShuffleReadRule {
+
+ // The Comet operator that a stand-in Spark operator was rebuilt from.
+ private val COMET_OPERATOR = TreeNodeTag[SparkPlan]("cometCoalesceShufflePartitions")
+
+ // Required by the trait. Which shuffles are coalesced is decided by Spark's rule, which applies
+ // its own list.
+ override protected def supportedShuffleOrigins: Seq[ShuffleOrigin] =
+ CoalesceShufflePartitions(SparkSession.active).supportedShuffleOrigins
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ if (!conf.coalesceShufflePartitionsEnabled || !plan.exists(replaced(_).isDefined)) {
+ return plan
+ }
+ plan.transformDown {
+ case p if replaced(p).isDefined && untouched(p) => coalesceBelow(p)
+ }
+ }
+
+ // The Spark operator a Comet operator replaced, if Spark's rule coalesces its children one by
+ // one. The class match mirrors Spark's, and Spark's rule decides, for its version, which of
+ // these it actually treats that way.
+ private def replaced(plan: SparkPlan): Option[SparkPlan] = plan match {
+ case comet: CometExec =>
+ comet.originalPlan match {
+ case original @ (_: UnionExec | _: CartesianProductExec | _: BroadcastHashJoinExec |
+ _: BroadcastNestedLoopJoinExec)
+ if original.children.length == comet.children.length =>
+ Some(original)
+ case _ => None
+ }
+ case _ => None
+ }
+
+ // No AQE rule has put a read over any shuffle stage below `plan`: Spark's rule coalesced none
+ // of them, and none is a skew-split or local read that coalescing now could disturb.
+ private def untouched(plan: SparkPlan): Boolean =
+ plan.exists(_.isInstanceOf[ShuffleQueryStageExec]) &&
+ !plan.exists(_.isInstanceOf[AQEShuffleReadExec])
+
+ private def coalesceBelow(plan: SparkPlan): SparkPlan = {
+ val asSpark = plan.transformUp { case p =>
+ replaced(p) match {
+ case Some(original) =>
+ val standIn = original.withNewChildren(p.children)
+ // `withNewChildren` hands back the original itself when the children are the same ones,
+ // and the tag must not land on the operator that the Comet one keeps.
+ if (standIn eq original) {
Review Comment:
If a Comet union is created during an AQE re-plan on top of stages that are already materialized, I think `originalPlan.children` equals `children`, so `withNewChildren` returns `original` and this branch returns `p` without coalescing. I haven't run it, but AQE adopts a re-planned plan when it differs at equal cost, so I'd expect this to happen for queries where a join changes strategy below the union. Would it make sense to force a fresh copy of `original` here (for example with `makeCopy`) so those unions are handled too, and to cover that shape in a test?
##########
spark/src/main/scala/org/apache/comet/rules/CometCoalesceShufflePartitions.scala:
##########
@@ -0,0 +1,128 @@
+/*
+ * 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.comet.rules
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.comet.CometExec
+import org.apache.spark.sql.execution.{SparkPlan, UnionExec}
+import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, AQEShuffleReadRule, CoalesceShufflePartitions, ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.exchange.ShuffleOrigin
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, CartesianProductExec}
+
+/**
+ * Coalesces the shuffle partitions below a Comet operator that Spark's CoalesceShufflePartitions
+ * coalesces child by child but does not recognize.
+ *
+ * Spark coalesces each child of a `UnionExec` as a group of its own, and from Spark 4.0 each
+ * child of a `CartesianProductExec`, `BroadcastHashJoinExec` or `BroadcastNestedLoopJoinExec`
+ * too. It matches those classes, and the Comet operators that replace them are other classes, so
+ * it falls through to the case that coalesces only when every leaf below the operator is an
+ * exchange stage. A union with a scan or a table-cache stage in one branch then keeps every
+ * partition of the shuffles in the others: `spark.sql.shuffle.partitions` tasks for a query that
+ * needs a few.
+ *
+ * Comet replaces these operators while AQE prepares a stage, before its optimizer rules run, and
+ * plans the operators above them against the Comet versions. So this runs after Spark's rule
+ * instead, on each such operator whose shuffle stages that rule left untouched. It rebuilds the
+ * Spark operator each Comet one replaced over the Comet children, has Spark's own rule coalesce
+ * that, and swaps the Comet operators back in. The partitions come out as Spark would have
+ * coalesced them, down to which operators count, since it is Spark's code deciding. The one
+ * difference is that Spark divides its minimum partition count among the coalesce groups of the
+ * whole plan, and this among those below the Comet operator, which are usually all of them.
+ *
+ * When every leaf below such an operator is an exchange stage, Spark's rule already coalesces its
+ * shuffles, together rather than child by child, and this leaves them as they are.
+ *
+ * Extending `AQEShuffleReadRule` gets this the same treatment from AQE as Spark's rule: it is
+ * skipped for the final stage when that stage's shuffle optimizations are off, and its result is
+ * discarded if it breaks a distribution required above it.
+ */
+case object CometCoalesceShufflePartitions extends AQEShuffleReadRule {
+
+ // The Comet operator that a stand-in Spark operator was rebuilt from.
+ private val COMET_OPERATOR = TreeNodeTag[SparkPlan]("cometCoalesceShufflePartitions")
+
+ // Required by the trait. Which shuffles are coalesced is decided by Spark's rule, which applies
+ // its own list.
+ override protected def supportedShuffleOrigins: Seq[ShuffleOrigin] =
+ CoalesceShufflePartitions(SparkSession.active).supportedShuffleOrigins
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ if (!conf.coalesceShufflePartitionsEnabled || !plan.exists(replaced(_).isDefined)) {
+ return plan
+ }
+ plan.transformDown {
+ case p if replaced(p).isDefined && untouched(p) => coalesceBelow(p)
+ }
+ }
+
+ // The Spark operator a Comet operator replaced, if Spark's rule coalesces its children one by
+ // one. The class match mirrors Spark's, and Spark's rule decides, for its version, which of
+ // these it actually treats that way.
+ private def replaced(plan: SparkPlan): Option[SparkPlan] = plan match {
+ case comet: CometExec =>
+ comet.originalPlan match {
+ case original @ (_: UnionExec | _: CartesianProductExec | _: BroadcastHashJoinExec |
Review Comment:
Nit: I don't see a Comet counterpart of `CartesianProductExec` under `spark/src/main`, so that arm looks unreachable today. The `BroadcastHashJoinExec` and `BroadcastNestedLoopJoinExec` arms have no test either, and I'm not sure which plan shape triggers them, since a Comet join with only exchange stages below it is already coalesced by Spark's rule. Would it make sense to narrow this to what a test exercises?
##########
spark/src/main/scala/org/apache/comet/rules/CometCoalesceShufflePartitions.scala:
##########
@@ -0,0 +1,128 @@
+/*
+ * 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.comet.rules
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.comet.CometExec
+import org.apache.spark.sql.execution.{SparkPlan, UnionExec}
+import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, AQEShuffleReadRule, CoalesceShufflePartitions, ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.exchange.ShuffleOrigin
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, CartesianProductExec}
+
+/**
+ * Coalesces the shuffle partitions below a Comet operator that Spark's CoalesceShufflePartitions
+ * coalesces child by child but does not recognize.
+ *
+ * Spark coalesces each child of a `UnionExec` as a group of its own, and from Spark 4.0 each
Review Comment:
Nit: I think this holds from Spark 3.5 rather than 4.0. In `v3.5.8`, `CoalesceShufflePartitions.collectCoalesceGroups` already has a `case` for each of `UnionExec`, `CartesianProductExec`, `BroadcastHashJoinExec` and `BroadcastNestedLoopJoinExec` (lines 150 to 157), and `v3.4.3` only has `UnionExec`. Might be worth adjusting here and in the PR description.
##########
spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala:
##########
@@ -3573,6 +3573,44 @@ class CometExecSuite extends CometTestBase {
}
}
+ // https://github.com/apache/datafusion-comet/issues/6454
+ test("AQE coalesces the shuffle partitions of a union whose other branch is a scan") {
Review Comment:
The description says the rule also covers a union whose shuffles Spark's rule gives up on together, but both new tests use a scan or a table cache stage as the other branch. It might be worth adding that shape, for example a hash-shuffled join unioned with a global aggregate as in Spark's `Union two datasets with different pre-shuffle partition number`. It differs from these tests because every leaf is an exchange stage, and I expect the `SinglePartition` shuffle of the aggregate to be what makes Spark's rule skip the whole Comet union.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: coalesce shuffle partitions under Comet unions the way Spark does [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #6459:
URL: https://github.com/apache/datafusion-comet/pull/6459#discussion_r4186236774
##########
spark/src/main/scala/org/apache/comet/rules/CometCoalesceShufflePartitions.scala:
##########
@@ -0,0 +1,128 @@
+/*
+ * 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.comet.rules
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.comet.CometExec
+import org.apache.spark.sql.execution.{SparkPlan, UnionExec}
+import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, AQEShuffleReadRule, CoalesceShufflePartitions, ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.exchange.ShuffleOrigin
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, CartesianProductExec}
+
+/**
+ * Coalesces the shuffle partitions below a Comet operator that Spark's CoalesceShufflePartitions
+ * coalesces child by child but does not recognize.
+ *
+ * Spark coalesces each child of a `UnionExec` as a group of its own, and from Spark 4.0 each
+ * child of a `CartesianProductExec`, `BroadcastHashJoinExec` or `BroadcastNestedLoopJoinExec`
+ * too. It matches those classes, and the Comet operators that replace them are other classes, so
+ * it falls through to the case that coalesces only when every leaf below the operator is an
+ * exchange stage. A union with a scan or a table-cache stage in one branch then keeps every
+ * partition of the shuffles in the others: `spark.sql.shuffle.partitions` tasks for a query that
+ * needs a few.
+ *
+ * Comet replaces these operators while AQE prepares a stage, before its optimizer rules run, and
+ * plans the operators above them against the Comet versions. So this runs after Spark's rule
+ * instead, on each such operator whose shuffle stages that rule left untouched. It rebuilds the
+ * Spark operator each Comet one replaced over the Comet children, has Spark's own rule coalesce
+ * that, and swaps the Comet operators back in. The partitions come out as Spark would have
+ * coalesced them, down to which operators count, since it is Spark's code deciding. The one
+ * difference is that Spark divides its minimum partition count among the coalesce groups of the
+ * whole plan, and this among those below the Comet operator, which are usually all of them.
+ *
+ * When every leaf below such an operator is an exchange stage, Spark's rule already coalesces its
+ * shuffles, together rather than child by child, and this leaves them as they are.
+ *
+ * Extending `AQEShuffleReadRule` gets this the same treatment from AQE as Spark's rule: it is
+ * skipped for the final stage when that stage's shuffle optimizations are off, and its result is
+ * discarded if it breaks a distribution required above it.
+ */
+case object CometCoalesceShufflePartitions extends AQEShuffleReadRule {
+
+ // The Comet operator that a stand-in Spark operator was rebuilt from.
+ private val COMET_OPERATOR = TreeNodeTag[SparkPlan]("cometCoalesceShufflePartitions")
+
+ // Required by the trait. Which shuffles are coalesced is decided by Spark's rule, which applies
+ // its own list.
+ override protected def supportedShuffleOrigins: Seq[ShuffleOrigin] =
+ CoalesceShufflePartitions(SparkSession.active).supportedShuffleOrigins
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ if (!conf.coalesceShufflePartitionsEnabled || !plan.exists(replaced(_).isDefined)) {
+ return plan
+ }
+ plan.transformDown {
+ case p if replaced(p).isDefined && untouched(p) => coalesceBelow(p)
+ }
+ }
+
+ // The Spark operator a Comet operator replaced, if Spark's rule coalesces its children one by
+ // one. The class match mirrors Spark's, and Spark's rule decides, for its version, which of
+ // these it actually treats that way.
+ private def replaced(plan: SparkPlan): Option[SparkPlan] = plan match {
+ case comet: CometExec =>
+ comet.originalPlan match {
+ case original @ (_: UnionExec | _: CartesianProductExec | _: BroadcastHashJoinExec |
+ _: BroadcastNestedLoopJoinExec)
+ if original.children.length == comet.children.length =>
+ Some(original)
+ case _ => None
+ }
+ case _ => None
+ }
+
+ // No AQE rule has put a read over any shuffle stage below `plan`: Spark's rule coalesced none
+ // of them, and none is a skew-split or local read that coalescing now could disturb.
+ private def untouched(plan: SparkPlan): Boolean =
+ plan.exists(_.isInstanceOf[ShuffleQueryStageExec]) &&
+ !plan.exists(_.isInstanceOf[AQEShuffleReadExec])
+
+ private def coalesceBelow(plan: SparkPlan): SparkPlan = {
+ val asSpark = plan.transformUp { case p =>
+ replaced(p) match {
+ case Some(original) =>
+ val standIn = original.withNewChildren(p.children)
+ // `withNewChildren` hands back the original itself when the children are the same ones,
+ // and the tag must not land on the operator that the Comet one keeps.
+ if (standIn eq original) {
+ p
+ } else {
+ standIn.setTagValue(COMET_OPERATOR, p)
+ standIn
+ }
+ case None => p
+ }
+ }
+ val coalesced = CoalesceShufflePartitions(SparkSession.active).apply(asSpark)
+ if (coalesced eq asSpark) plan else restore(coalesced)
Review Comment:
Reproduced. `keys.repartition($"k").union(keys.repartition(20, $"k")).groupBy("k").count()` plans the aggregate directly over the union on 4.1, the rule coalesced only the first branch, and each key came back twice. The rule now stands in only for operators whose output partitioning is `UnknownPartitioning`, since nothing in the stage can rely on that. `AQE leaves the shuffles of a union alone when an aggregate relies on its partitioning` runs that query.
##########
spark/src/main/scala/org/apache/comet/rules/CometCoalesceShufflePartitions.scala:
##########
@@ -0,0 +1,128 @@
+/*
+ * 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.comet.rules
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.comet.CometExec
+import org.apache.spark.sql.execution.{SparkPlan, UnionExec}
+import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, AQEShuffleReadRule, CoalesceShufflePartitions, ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.exchange.ShuffleOrigin
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, CartesianProductExec}
+
+/**
+ * Coalesces the shuffle partitions below a Comet operator that Spark's CoalesceShufflePartitions
+ * coalesces child by child but does not recognize.
+ *
+ * Spark coalesces each child of a `UnionExec` as a group of its own, and from Spark 4.0 each
+ * child of a `CartesianProductExec`, `BroadcastHashJoinExec` or `BroadcastNestedLoopJoinExec`
+ * too. It matches those classes, and the Comet operators that replace them are other classes, so
+ * it falls through to the case that coalesces only when every leaf below the operator is an
+ * exchange stage. A union with a scan or a table-cache stage in one branch then keeps every
+ * partition of the shuffles in the others: `spark.sql.shuffle.partitions` tasks for a query that
+ * needs a few.
+ *
+ * Comet replaces these operators while AQE prepares a stage, before its optimizer rules run, and
+ * plans the operators above them against the Comet versions. So this runs after Spark's rule
+ * instead, on each such operator whose shuffle stages that rule left untouched. It rebuilds the
+ * Spark operator each Comet one replaced over the Comet children, has Spark's own rule coalesce
+ * that, and swaps the Comet operators back in. The partitions come out as Spark would have
+ * coalesced them, down to which operators count, since it is Spark's code deciding. The one
+ * difference is that Spark divides its minimum partition count among the coalesce groups of the
+ * whole plan, and this among those below the Comet operator, which are usually all of them.
+ *
+ * When every leaf below such an operator is an exchange stage, Spark's rule already coalesces its
+ * shuffles, together rather than child by child, and this leaves them as they are.
+ *
+ * Extending `AQEShuffleReadRule` gets this the same treatment from AQE as Spark's rule: it is
+ * skipped for the final stage when that stage's shuffle optimizations are off, and its result is
+ * discarded if it breaks a distribution required above it.
+ */
+case object CometCoalesceShufflePartitions extends AQEShuffleReadRule {
+
+ // The Comet operator that a stand-in Spark operator was rebuilt from.
+ private val COMET_OPERATOR = TreeNodeTag[SparkPlan]("cometCoalesceShufflePartitions")
+
+ // Required by the trait. Which shuffles are coalesced is decided by Spark's rule, which applies
+ // its own list.
+ override protected def supportedShuffleOrigins: Seq[ShuffleOrigin] =
+ CoalesceShufflePartitions(SparkSession.active).supportedShuffleOrigins
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ if (!conf.coalesceShufflePartitionsEnabled || !plan.exists(replaced(_).isDefined)) {
+ return plan
+ }
+ plan.transformDown {
+ case p if replaced(p).isDefined && untouched(p) => coalesceBelow(p)
+ }
+ }
+
+ // The Spark operator a Comet operator replaced, if Spark's rule coalesces its children one by
+ // one. The class match mirrors Spark's, and Spark's rule decides, for its version, which of
+ // these it actually treats that way.
+ private def replaced(plan: SparkPlan): Option[SparkPlan] = plan match {
+ case comet: CometExec =>
+ comet.originalPlan match {
+ case original @ (_: UnionExec | _: CartesianProductExec | _: BroadcastHashJoinExec |
+ _: BroadcastNestedLoopJoinExec)
+ if original.children.length == comet.children.length =>
+ Some(original)
+ case _ => None
+ }
+ case _ => None
+ }
+
+ // No AQE rule has put a read over any shuffle stage below `plan`: Spark's rule coalesced none
+ // of them, and none is a skew-split or local read that coalescing now could disturb.
+ private def untouched(plan: SparkPlan): Boolean =
+ plan.exists(_.isInstanceOf[ShuffleQueryStageExec]) &&
+ !plan.exists(_.isInstanceOf[AQEShuffleReadExec])
+
+ private def coalesceBelow(plan: SparkPlan): SparkPlan = {
+ val asSpark = plan.transformUp { case p =>
+ replaced(p) match {
+ case Some(original) =>
+ val standIn = original.withNewChildren(p.children)
+ // `withNewChildren` hands back the original itself when the children are the same ones,
+ // and the tag must not land on the operator that the Comet one keeps.
+ if (standIn eq original) {
Review Comment:
Done. When `withNewChildren` hands back the original, the rule copies it with `makeCopy` and tags the copy. `AQE coalesces the shuffle partitions of a union planned again over its stages` builds such a union from a finished query's stages, and fails without the copy.
##########
spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala:
##########
@@ -3573,6 +3573,44 @@ class CometExecSuite extends CometTestBase {
}
}
+ // https://github.com/apache/datafusion-comet/issues/6454
+ test("AQE coalesces the shuffle partitions of a union whose other branch is a scan") {
Review Comment:
Added `AQE coalesces each branch of a union whose shuffles have different partition counts`, a port of that Spark test. Spark's rule gives up on the Comet union because of the aggregate's single-partition shuffle, and the rule coalesces the join's two shuffles.
##########
spark/src/main/scala/org/apache/comet/rules/CometCoalesceShufflePartitions.scala:
##########
@@ -0,0 +1,128 @@
+/*
+ * 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.comet.rules
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.comet.CometExec
+import org.apache.spark.sql.execution.{SparkPlan, UnionExec}
+import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, AQEShuffleReadRule, CoalesceShufflePartitions, ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.exchange.ShuffleOrigin
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, CartesianProductExec}
+
+/**
+ * Coalesces the shuffle partitions below a Comet operator that Spark's CoalesceShufflePartitions
+ * coalesces child by child but does not recognize.
+ *
+ * Spark coalesces each child of a `UnionExec` as a group of its own, and from Spark 4.0 each
+ * child of a `CartesianProductExec`, `BroadcastHashJoinExec` or `BroadcastNestedLoopJoinExec`
+ * too. It matches those classes, and the Comet operators that replace them are other classes, so
+ * it falls through to the case that coalesces only when every leaf below the operator is an
+ * exchange stage. A union with a scan or a table-cache stage in one branch then keeps every
+ * partition of the shuffles in the others: `spark.sql.shuffle.partitions` tasks for a query that
+ * needs a few.
+ *
+ * Comet replaces these operators while AQE prepares a stage, before its optimizer rules run, and
+ * plans the operators above them against the Comet versions. So this runs after Spark's rule
+ * instead, on each such operator whose shuffle stages that rule left untouched. It rebuilds the
+ * Spark operator each Comet one replaced over the Comet children, has Spark's own rule coalesce
+ * that, and swaps the Comet operators back in. The partitions come out as Spark would have
+ * coalesced them, down to which operators count, since it is Spark's code deciding. The one
+ * difference is that Spark divides its minimum partition count among the coalesce groups of the
+ * whole plan, and this among those below the Comet operator, which are usually all of them.
+ *
+ * When every leaf below such an operator is an exchange stage, Spark's rule already coalesces its
+ * shuffles, together rather than child by child, and this leaves them as they are.
+ *
+ * Extending `AQEShuffleReadRule` gets this the same treatment from AQE as Spark's rule: it is
+ * skipped for the final stage when that stage's shuffle optimizations are off, and its result is
+ * discarded if it breaks a distribution required above it.
+ */
+case object CometCoalesceShufflePartitions extends AQEShuffleReadRule {
+
+ // The Comet operator that a stand-in Spark operator was rebuilt from.
+ private val COMET_OPERATOR = TreeNodeTag[SparkPlan]("cometCoalesceShufflePartitions")
+
+ // Required by the trait. Which shuffles are coalesced is decided by Spark's rule, which applies
+ // its own list.
+ override protected def supportedShuffleOrigins: Seq[ShuffleOrigin] =
+ CoalesceShufflePartitions(SparkSession.active).supportedShuffleOrigins
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ if (!conf.coalesceShufflePartitionsEnabled || !plan.exists(replaced(_).isDefined)) {
+ return plan
+ }
+ plan.transformDown {
+ case p if replaced(p).isDefined && untouched(p) => coalesceBelow(p)
+ }
+ }
+
+ // The Spark operator a Comet operator replaced, if Spark's rule coalesces its children one by
+ // one. The class match mirrors Spark's, and Spark's rule decides, for its version, which of
+ // these it actually treats that way.
+ private def replaced(plan: SparkPlan): Option[SparkPlan] = plan match {
+ case comet: CometExec =>
+ comet.originalPlan match {
+ case original @ (_: UnionExec | _: CartesianProductExec | _: BroadcastHashJoinExec |
Review Comment:
Dropped `CartesianProductExec`. The two broadcast join arms matter more now: the rule hands Spark's rule the whole stage (see [the thread on ancestor joins](https://github.com/apache/datafusion-comet/pull/6459#discussion_r4154984789)), so a union below a Comet broadcast join is reached only through the join's stand-in, and a nested loop join's stand-in also brings Spark's smaller target size. `AQE coalesces the shuffle partitions of a union below a join as Spark does` covers both, and a Cartesian product above the union.
##########
spark/src/main/scala/org/apache/comet/rules/CometCoalesceShufflePartitions.scala:
##########
@@ -0,0 +1,128 @@
+/*
+ * 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.comet.rules
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.comet.CometExec
+import org.apache.spark.sql.execution.{SparkPlan, UnionExec}
+import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, AQEShuffleReadRule, CoalesceShufflePartitions, ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.exchange.ShuffleOrigin
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, CartesianProductExec}
+
+/**
+ * Coalesces the shuffle partitions below a Comet operator that Spark's CoalesceShufflePartitions
+ * coalesces child by child but does not recognize.
+ *
+ * Spark coalesces each child of a `UnionExec` as a group of its own, and from Spark 4.0 each
Review Comment:
Fixed in the scaladoc: 3.5 has all four cases and 3.4 only the union.
##########
spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala:
##########
@@ -107,6 +107,7 @@ class CometSparkSessionExtensions
}
injectQueryStageOptimizerRuleShim(extensions, CometPlanAdaptiveDynamicPruningFilters)
injectQueryStageOptimizerRuleShim(extensions, CometReuseSubquery)
+ injectQueryStageOptimizerRuleShim(extensions, CometCoalesceShufflePartitions)
Review Comment:
Added it to the per-stage list, after Spark's three shuffle-read rules, and to the Spark 3.4 note.
##########
spark/src/main/scala/org/apache/comet/rules/CometCoalesceShufflePartitions.scala:
##########
@@ -0,0 +1,128 @@
+/*
+ * 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.comet.rules
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.trees.TreeNodeTag
+import org.apache.spark.sql.comet.CometExec
+import org.apache.spark.sql.execution.{SparkPlan, UnionExec}
+import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, AQEShuffleReadRule, CoalesceShufflePartitions, ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.exchange.ShuffleOrigin
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, CartesianProductExec}
+
+/**
+ * Coalesces the shuffle partitions below a Comet operator that Spark's CoalesceShufflePartitions
+ * coalesces child by child but does not recognize.
+ *
+ * Spark coalesces each child of a `UnionExec` as a group of its own, and from Spark 4.0 each
+ * child of a `CartesianProductExec`, `BroadcastHashJoinExec` or `BroadcastNestedLoopJoinExec`
+ * too. It matches those classes, and the Comet operators that replace them are other classes, so
+ * it falls through to the case that coalesces only when every leaf below the operator is an
+ * exchange stage. A union with a scan or a table-cache stage in one branch then keeps every
+ * partition of the shuffles in the others: `spark.sql.shuffle.partitions` tasks for a query that
+ * needs a few.
+ *
+ * Comet replaces these operators while AQE prepares a stage, before its optimizer rules run, and
+ * plans the operators above them against the Comet versions. So this runs after Spark's rule
+ * instead, on each such operator whose shuffle stages that rule left untouched. It rebuilds the
+ * Spark operator each Comet one replaced over the Comet children, has Spark's own rule coalesce
+ * that, and swaps the Comet operators back in. The partitions come out as Spark would have
+ * coalesced them, down to which operators count, since it is Spark's code deciding. The one
+ * difference is that Spark divides its minimum partition count among the coalesce groups of the
+ * whole plan, and this among those below the Comet operator, which are usually all of them.
+ *
+ * When every leaf below such an operator is an exchange stage, Spark's rule already coalesces its
+ * shuffles, together rather than child by child, and this leaves them as they are.
+ *
+ * Extending `AQEShuffleReadRule` gets this the same treatment from AQE as Spark's rule: it is
+ * skipped for the final stage when that stage's shuffle optimizations are off, and its result is
+ * discarded if it breaks a distribution required above it.
+ */
+case object CometCoalesceShufflePartitions extends AQEShuffleReadRule {
+
+ // The Comet operator that a stand-in Spark operator was rebuilt from.
+ private val COMET_OPERATOR = TreeNodeTag[SparkPlan]("cometCoalesceShufflePartitions")
+
+ // Required by the trait. Which shuffles are coalesced is decided by Spark's rule, which applies
+ // its own list.
+ override protected def supportedShuffleOrigins: Seq[ShuffleOrigin] =
+ CoalesceShufflePartitions(SparkSession.active).supportedShuffleOrigins
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ if (!conf.coalesceShufflePartitionsEnabled || !plan.exists(replaced(_).isDefined)) {
+ return plan
+ }
+ plan.transformDown {
+ case p if replaced(p).isDefined && untouched(p) => coalesceBelow(p)
+ }
+ }
+
+ // The Spark operator a Comet operator replaced, if Spark's rule coalesces its children one by
+ // one. The class match mirrors Spark's, and Spark's rule decides, for its version, which of
+ // these it actually treats that way.
+ private def replaced(plan: SparkPlan): Option[SparkPlan] = plan match {
+ case comet: CometExec =>
+ comet.originalPlan match {
+ case original @ (_: UnionExec | _: CartesianProductExec | _: BroadcastHashJoinExec |
+ _: BroadcastNestedLoopJoinExec)
+ if original.children.length == comet.children.length =>
+ Some(original)
+ case _ => None
+ }
+ case _ => None
+ }
+
+ // No AQE rule has put a read over any shuffle stage below `plan`: Spark's rule coalesced none
+ // of them, and none is a skew-split or local read that coalescing now could disturb.
+ private def untouched(plan: SparkPlan): Boolean =
+ plan.exists(_.isInstanceOf[ShuffleQueryStageExec]) &&
+ !plan.exists(_.isInstanceOf[AQEShuffleReadExec])
+
+ private def coalesceBelow(plan: SparkPlan): SparkPlan = {
+ val asSpark = plan.transformUp { case p =>
+ replaced(p) match {
+ case Some(original) =>
+ val standIn = original.withNewChildren(p.children)
+ // `withNewChildren` hands back the original itself when the children are the same ones,
+ // and the tag must not land on the operator that the Comet one keeps.
+ if (standIn eq original) {
+ p
+ } else {
+ standIn.setTagValue(COMET_OPERATOR, p)
+ standIn
+ }
+ case None => p
+ }
+ }
+ val coalesced = CoalesceShufflePartitions(SparkSession.active).apply(asSpark)
Review Comment:
The rule now gives Spark's rule the whole stage rather than the subtree below each Comet operator, so Spark's walk sees a Cartesian product or nested loop join above the union and uses the minimum partition size, and it divides the minimum partition count over every group in the stage, as it would for Spark's operators. It runs only when no AQE rule has put a read over any shuffle in the stage, since Spark's rule would coalesce those reads a second time. In the new `AQE coalesces the shuffle partitions of a union below a join as Spark does`, with a `crossJoin` above the union, the shuffled branch keeps its 10 partitions on Spark 4.1, as it does without Comet, where the previous revision coalesced them to 1.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org