You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/03 17:41:47 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-6547-f980d6fb59f51ae20c0a6916eec9771ace8e50d1 created (now db5a40e902)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-6547-f980d6fb59f51ae20c0a6916eec9771ace8e50d1
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at db5a40e902 perf: read the shuffle natively in operators AQE reuses from the initial plan (#6547)
This branch includes the following new commits:
new db5a40e902 perf: read the shuffle natively in operators AQE reuses from the initial plan (#6547)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: perf: read the shuffle natively in operators AQE reuses from the initial plan (#6547)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6547-f980d6fb59f51ae20c0a6916eec9771ace8e50d1
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit db5a40e9020e94334a8aa4149e1ea9443ec5db19
Author: dustin <dw...@users.noreply.github.com>
AuthorDate: Sat Oct 3 17:41:23 2026 +0000
perf: read the shuffle natively in operators AQE reuses from the initial plan (#6547)
* perf: read the shuffle natively in operators AQE reuses from the initial plan
Under AQE, an operator that shares its logical node with a shuffle stage,
such as the final aggregate of a two-phase aggregate, comes back from
LogicalQueryStageStrategy as the same physical node. Its native plan was
serialized in the initial plan, when the input was a bare exchange and
was written as a plain Scan, so the reduce side decoded every block on
the JVM and imported it over FFI even with direct read on.
CometExecRule now refreshes such a node: when a child has become a
ShuffleScan sink and the parent's plan still reads that slot as a Scan
with the same field types, the leaf is replaced and the block is
serialized again. Nested native children propagate their leaves, write
operators are skipped, and any mismatch leaves the node as it was.
One AQE DPP test now runs with direct read off, since its first replan is
no longer unchanged; a sibling test covers the default and checks that
the broadcast roots keep their logical links.
* perf: flag a declined shuffle-read refresh and share the native leaf walk
When the refresh finds a stale Scan but the native plan's leaves do not
match its inputs, the node now carries an explain message, so a serde
change that reorders children shows up in extended explain instead of
only as a slower reduce stage.
The rule, findShuffleScanIndices and the tests now share one
nativeLeaves walk, the leaves are rebuilt with a recursive mapLeaves,
and the writers that build their own Scan share one predicate with the
firstNativeOp reset. The shuffle tests are table-driven and pair leaves
with the runtime input order, and the DPP replan tests share a helper.
---
docs/source/contributor-guide/native_shuffle.md | 5 +
.../org/apache/comet/rules/CometExecRule.scala | 96 ++++++-
.../org/apache/spark/sql/comet/operators.scala | 58 +++--
.../comet/exec/CometNativeShuffleSuite.scala | 204 ++++++++++++++-
.../apache/comet/rules/CometExecRuleSuite.scala | 288 ++++++++++++++++-----
5 files changed, 550 insertions(+), 101 deletions(-)
diff --git a/docs/source/contributor-guide/native_shuffle.md b/docs/source/contributor-guide/native_shuffle.md
index b882acd25a..f852dd497a 100644
--- a/docs/source/contributor-guide/native_shuffle.md
+++ b/docs/source/contributor-guide/native_shuffle.md
@@ -222,6 +222,11 @@ decides during plan serialization. When direct read is enabled and the sink's in
exchange, `convertToShuffleScan` emits a `ShuffleScan` operator. When either is false the sink falls
through to the base `CometSink.convert`, which emits the usual `Scan`.
+Under AQE, an operator that shares its logical node with a shuffle stage, such as the final aggregate
+of a two-phase aggregate, comes back from re-planning as the node already planned, whose input was
+serialized as a `Scan` before the stage existed. `CometExecRule` refreshes such a node: once its input
+is a sink that emits a `ShuffleScan`, that `ShuffleScan` replaces the stale `Scan` leaf.
+
The two are not alternatives on failure. If any output type fails `supportedSinkDataType`,
`convertToShuffleScan` records the fallback reason `Unsupported data type for shuffle direct read`
and returns `None`. It does not retry as a regular `Scan`, and retrying would not help, because
diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
index e5e250f75d..623c603af8 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
@@ -20,6 +20,7 @@
package org.apache.comet.rules
import scala.collection.mutable.ListBuffer
+import scala.jdk.CollectionConverters._
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, NamedExpression, Remainder}
@@ -57,6 +58,7 @@ import org.apache.comet.CometConf.{COMET_SPARK_TO_ARROW_ENABLED, COMET_SPARK_TO_
import org.apache.comet.CometSparkSessionExtensions._
import org.apache.comet.rules.CometExecRule.allExecs
import org.apache.comet.serde._
+import org.apache.comet.serde.OperatorOuterClass.Operator
import org.apache.comet.serde.operator._
import org.apache.comet.shims.{CometTypeShim, ShimCometStreaming, ShimCometWindowGroupLimit, ShimSubqueryBroadcast}
@@ -88,6 +90,13 @@ object CometExecRule {
val COMET_UNSAFE_PARTIAL: TreeNodeTag[String] =
TreeNodeTag[String]("comet.unsafePartialAgg")
+ /**
+ * Info message for a native operator that keeps a `Scan` over an input that now reads its
+ * shuffle directly, because its native plan's leaves do not line up with its inputs.
+ */
+ val STALE_SCAN_KEPT: String =
+ "Shuffle direct read not applied: the native plan's leaves do not match its inputs"
+
/**
* Fully native operators.
*/
@@ -619,7 +628,7 @@ case class CometExecRule(session: SparkSession)
}
plan.transformUp { case op =>
- val converted = convertNode(op)
+ val converted = convertNode(refreshStaleShuffleScans(op))
// Replace SubqueryBroadcastExec with CometSubqueryBroadcastExec in DPP expressions
// when the broadcast child has a Comet plan underneath. This enables exchange reuse
// between the DPP subquery and the join's CometBroadcastExchangeExec because both
@@ -628,6 +637,88 @@ case class CometExecRule(session: SparkSession)
}
}
+ /**
+ * AQE re-plans around a materialized stage by reusing the physical node linked to it, so a
+ * native operator that shares its logical node with a shuffle stage (the final aggregate of a
+ * two-phase aggregate) keeps the native plan it got while that input was a bare exchange, read
+ * through a plain `Scan`. Once the input is a sink that reads the shuffle directly, its
+ * `ShuffleScan` takes the place of the stale leaf. The leaf is patched in place because
+ * converting the node again from `originalPlan` would drop the stage's logical link that AQE
+ * relies on and re-run serde on a node that is already planned.
+ */
+ private def refreshStaleShuffleScans(op: SparkPlan): SparkPlan = op match {
+ case _ if scansChildAsSeparateBlock(op) => op
+ case native: CometNativeExec if native.children.nonEmpty =>
+ refreshedNativeOp(native) match {
+ case Some(newOp) =>
+ val refreshed = native.withRefreshedNativeOp(newOp)
+ // An operator that does not hold its native plan as a field cannot take a new one.
+ if (refreshed.nativeOp eq newOp) refreshed else op
+ case None => op
+ }
+ case _ => op
+ }
+
+ /**
+ * The native plan of `native` with each `Scan` leaf whose input is now a `ShuffleScan` of the
+ * same field types replaced by that `ShuffleScan`, or None if there is no such leaf or the plan
+ * children cannot be matched to the leaves. In the second case `native` gets an info message
+ * for extended explain, since the block still reads its shuffle through the JVM.
+ */
+ private def refreshedNativeOp(native: CometNativeExec): Option[Operator] = {
+ val children = native.children.collect { case child: CometNativeExec => child }
+ // Only a sink that reads a shuffle directly, or a native child that may hold one, can feed
+ // a `ShuffleScan`.
+ val mayFeedShuffleScan = children.exists {
+ case sink: CometSinkPlaceHolder => sink.nativeOp.hasShuffleScan
+ case _ => true
+ }
+ if (children.length != native.children.length || !mayFeedShuffleScan) return None
+ val leaves = CometExec.nativeLeaves(native.nativeOp)
+ if (!leaves.exists(_.hasScan)) return None
+
+ // Each plan child feeds a run of leaves, in order: a sink feeds one, and a native child
+ // feeds the leaves of its own native plan.
+ val current = children.flatMap {
+ case sink: CometSinkPlaceHolder => Seq(sink.nativeOp)
+ case child => CometExec.nativeLeaves(child.nativeOp)
+ }
+ def isStale(leaf: Operator, input: Operator): Boolean = leaf.hasScan && input.hasShuffleScan
+ val stale = leaves.zip(current).filter { case (leaf, input) => isStale(leaf, input) }
+ if (stale.isEmpty) return None
+ val isRefreshable = current.length == leaves.length &&
+ stale.forall { case (leaf, input) =>
+ leaf.getScan.getFieldsList == input.getShuffleScan.getFieldsList
+ }
+ if (!isRefreshable) {
+ withInfo(native, CometExecRule.STALE_SCAN_KEPT)
+ return None
+ }
+ val inputs = current.iterator
+ Some(mapLeaves(native.nativeOp) { leaf =>
+ val input = inputs.next()
+ if (isStale(leaf, input)) input else leaf
+ })
+ }
+
+ /** `op` with `f` applied to each childless operator, in `CometExec.nativeLeaves` order. */
+ private def mapLeaves(op: Operator)(f: Operator => Operator): Operator =
+ if (op.getChildrenCount == 0) {
+ f(op)
+ } else {
+ val children = op.getChildrenList.asScala.map(mapLeaves(_)(f))
+ op.toBuilder.clearChildren().addAllChildren(children.asJava).build()
+ }
+
+ /**
+ * Whether `op` is a writer whose native plan builds its own `Scan` over its child, so that the
+ * child runs as a separate native block instead of inside that plan.
+ */
+ private def scansChildAsSeparateBlock(op: SparkPlan): Boolean = op match {
+ case _: CometNativeWriteExec | _: CometIcebergWriteExec | _: CometWriteFilesExec => true
+ case _ => false
+ }
+
/**
* Replace SubqueryBroadcastExec with CometSubqueryBroadcastExec in a node's expressions
* (non-AQE DPP), and wrap SubqueryAdaptiveBroadcastExec in CometSubqueryAdaptiveBroadcastExec
@@ -943,8 +1034,7 @@ case class CometExecRule(session: SparkSession)
// its child (e.g., CometNativeScanExec, or a CometProject over an AQEShuffleRead)
// needs its own serialization. Reset the flag so children can start their own native
// execution blocks.
- if (op.isInstanceOf[CometNativeWriteExec] || op.isInstanceOf[CometIcebergWriteExec] ||
- op.isInstanceOf[CometWriteFilesExec]) {
+ if (scansChildAsSeparateBlock(op)) {
firstNativeOp = true
}
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
index 2a70c7de8e..91881c2c14 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala
@@ -728,6 +728,25 @@ object CometExec {
bytes
}
+ /**
+ * The childless operators of a native plan, depth first. Its `Scan` and `ShuffleScan` leaves
+ * read the block's inputs in this order.
+ */
+ def nativeLeaves(op: Operator): Seq[Operator] =
+ if (op.getChildrenCount == 0) Seq(op)
+ else op.getChildrenList.asScala.toSeq.flatMap(nativeLeaves)
+
+ /**
+ * The input indices of a native plan that its `ShuffleScan` leaves read. Each `Scan` or
+ * `ShuffleScan` leaf reads one input, in [[nativeLeaves]] order.
+ */
+ def findShuffleScanIndices(plan: Operator): Set[Int] =
+ nativeLeaves(plan)
+ .filter(leaf => leaf.hasScan || leaf.hasShuffleScan)
+ .zipWithIndex
+ .collect { case (leaf, index) if leaf.hasShuffleScan => index }
+ .toSet
+
def getCometIterator(
inputObjects: Array[Object],
numOutputCols: Int,
@@ -995,7 +1014,7 @@ abstract class CometNativeExec extends CometExec {
// (`ShuffleQueryStageExec`), so a bare non-AQE `CometShuffleExchangeExec` always serializes
// as a regular Scan regardless of `COMET_SHUFFLE_DIRECT_READ_ENABLED`. Driving the JVM
// dispatch from `shuffleScanIndices` instead of the conf keeps the two aligned.
- val shuffleScanIndices = findShuffleScanIndices(nativeOp)
+ val shuffleScanIndices = CometExec.findShuffleScanIndices(nativeOp)
def isBroadcastInput(plan: SparkPlan): Boolean = plan match {
case _: CometBroadcastExchangeExec => true
@@ -1176,27 +1195,6 @@ abstract class CometNativeExec extends CometExec {
}
}
- /**
- * Walk the protobuf operator tree depth-first to find which input indices correspond to
- * ShuffleScan vs Scan leaf nodes. Each Scan or ShuffleScan leaf consumes one input in order.
- */
- private def findShuffleScanIndices(plan: OperatorOuterClass.Operator): Set[Int] = {
- var scanIndex = 0
- val indices = mutable.Set.empty[Int]
- def walk(op: OperatorOuterClass.Operator): Unit = {
- if (op.hasShuffleScan) {
- indices += scanIndex
- scanIndex += 1
- } else if (op.hasScan) {
- scanIndex += 1
- } else {
- op.getChildrenList.asScala.foreach(walk)
- }
- }
- walk(plan)
- indices.toSet
- }
-
/**
* Converts this native Comet operator and its children into a native block which can be
* executed as a whole (i.e., in a single JNI call) from the native side.
@@ -1216,6 +1214,22 @@ abstract class CometNativeExec extends CometExec {
makeCopy(newArgs).asInstanceOf[CometNativeExec]
}
+ /**
+ * Copies this operator with `newOp` as its native plan and no serialized plan, so that
+ * `convertBlock` serializes the block again.
+ */
+ def withRefreshedNativeOp(newOp: Operator): CometNativeExec = {
+ def transform(arg: Any): AnyRef = arg match {
+ case op: Operator if op eq nativeOp => newOp
+ case _: SerializedPlan => SerializedPlan(None)
+ case other: AnyRef => other
+ case null => null
+ }
+
+ val newArgs = mapProductIterator(transform)
+ makeCopy(newArgs).asInstanceOf[CometNativeExec]
+ }
+
/**
* Cleans the serialized plan from this native Comet operator. Used to canonicalize the plan.
*/
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
index 03a57c18b4..b817ed7a50 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala
@@ -25,6 +25,7 @@ import java.nio.{ByteBuffer, ByteOrder}
import scala.collection.mutable
import scala.collection.mutable.ListBuffer
import scala.concurrent.duration.DurationInt
+import scala.jdk.CollectionConverters._
import scala.util.Random
import org.scalactic.source.Position
@@ -37,20 +38,22 @@ import org.apache.hadoop.fs.Path
import org.apache.spark.{SparkEnv, TaskContext}
import org.apache.spark.sql.{CometTestBase, DataFrame, Dataset, Row}
import org.apache.spark.sql.catalyst.expressions.AttributeReference
+import org.apache.spark.sql.catalyst.expressions.aggregate.Final
import org.apache.spark.sql.catalyst.plans.logical.LocalRelation
import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning
-import org.apache.spark.sql.comet.{CometExec, CometLocalTableScanExec, CometMetricNode, CometScanWrapper, CometSparkToColumnarExec, CometTakeOrderedAndProjectExec}
+import org.apache.spark.sql.comet.{CometExec, CometHashAggregateExec, CometLocalTableScanExec, CometMetricNode, CometNativeExec, CometScanWrapper, CometSparkToColumnarExec, CometTakeOrderedAndProjectExec}
import org.apache.spark.sql.comet.execution.arrow.CometArrowStream
import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec}
-import org.apache.spark.sql.execution.LocalTableScanExec
-import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
-import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
-import org.apache.spark.sql.functions.{col, count, sum}
+import org.apache.spark.sql.execution.{LocalTableScanExec, SparkPlan}
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, AQEShuffleReadExec, ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec}
+import org.apache.spark.sql.functions.{broadcast, col, count, countDistinct, sum}
+import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{ArrayType, DataType, LongType, MapType, StructField, StructType}
import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
import org.apache.comet.{CometConf, CometExecIterator, CometExplainInfo, CometShuffleBlockIterator, CometShuffleSizeLimitException, Native}
-import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus
+import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus}
import org.apache.comet.serde.{OperatorOuterClass, PartitioningOuterClass}
import org.apache.comet.shuffle.ShufflePartitionPusher
@@ -1415,6 +1418,195 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper
}
}
+ test("findShuffleScanIndices numbers the Scan and ShuffleScan leaves of a native plan") {
+ import OperatorOuterClass.{Filter, NativeScan, Operator, Projection, Scan, ShuffleScan}
+ val scan = Operator.newBuilder().setScan(Scan.getDefaultInstance).build()
+ val shuffleScan = Operator.newBuilder().setShuffleScan(ShuffleScan.getDefaultInstance).build()
+ val nativeScan = Operator.newBuilder().setNativeScan(NativeScan.getDefaultInstance).build()
+ val filterLeaf = Operator.newBuilder().setFilter(Filter.getDefaultInstance).build()
+ val emptyLeaf = Operator.getDefaultInstance
+ def node(children: Operator*): Operator =
+ Operator
+ .newBuilder()
+ .setProjection(Projection.getDefaultInstance)
+ .addAllChildren(children.asJava)
+ .build()
+
+ // Other leaves read no input, so they take no index.
+ val cases = Seq(
+ scan -> Set.empty[Int],
+ shuffleScan -> Set(0),
+ nativeScan -> Set.empty[Int],
+ node(scan, shuffleScan) -> Set(1),
+ node(shuffleScan, nativeScan, scan, shuffleScan) -> Set(0, 2),
+ node(nativeScan, filterLeaf, emptyLeaf) -> Set.empty[Int],
+ node(node(scan, nativeScan), node(filterLeaf, node(shuffleScan)), shuffleScan) -> Set(1, 2),
+ node(node(node(shuffleScan)), node(emptyLeaf, node(scan), shuffleScan), scan) -> Set(0, 2))
+ cases.foreach { case (plan, expected) =>
+ assert(CometExec.findShuffleScanIndices(plan) == expected, plan)
+ }
+ }
+
+ /**
+ * For each native block in `plan`, query stages included, that reads a shuffle: each input of
+ * the block, in the order execution collects them, paired with the kind of the `Scan` or
+ * `ShuffleScan` leaf that reads it.
+ */
+ private def shuffleReadingBlocks(plan: SparkPlan): Seq[Seq[(String, String)]] = {
+ def inputs(root: CometNativeExec): Seq[SparkPlan] = {
+ val plans = mutable.ArrayBuffer.empty[SparkPlan]
+ root.foreachUntilCometInput(root)(plans += _)
+ plans.filterNot(_.isInstanceOf[CometNativeExec]).toSeq
+ }
+ def readsShuffle(input: SparkPlan): Boolean = input match {
+ case _: ShuffleQueryStageExec | _: AQEShuffleReadExec | _: CometShuffleExchangeExec => true
+ case _ => false
+ }
+ // A node AQE reuses can keep the serialized plan of a block it used to be the root of.
+ val natives = collect(plan) { case native: CometNativeExec => native }
+ val nested = natives.flatMap(_.children.collect { case child: CometNativeExec => child })
+ natives
+ .filter(root => root.serializedPlanOpt.isDefined && !nested.exists(_ eq root))
+ .filter(inputs(_).exists(readsShuffle))
+ .map { root =>
+ val nativePlan = OperatorOuterClass.Operator.parseFrom(root.serializedPlanOpt.plan.get)
+ val blockInputs = inputs(root).map(_.getClass.getSimpleName)
+ val leafKinds = CometExec
+ .nativeLeaves(nativePlan)
+ .filter(leaf => leaf.hasScan || leaf.hasShuffleScan)
+ .map(_.getOpStructCase.name)
+ assert(blockInputs.length == leafKinds.length, s"$blockInputs vs $leafKinds")
+ blockInputs.zip(leafKinds)
+ }
+ }
+
+ private val shuffleStageRead = Seq("ShuffleQueryStageExec" -> "SHUFFLE_SCAN")
+
+ private val aqeWithoutCoalescing = Seq(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false")
+
+ /** A two-phase aggregate over many map partitions. */
+ private def groupedSum: DataFrame =
+ spark
+ .range(0, 10000, 1, 50)
+ .groupBy((col("id") % 997).as("k"))
+ .agg(sum("id"), count("id"))
+
+ test("final aggregate reads its shuffle through ShuffleScan only under AQE") {
+ // The coalesced case takes the stage's ShuffleScan before AQE puts the coalesced read
+ // between them, and the ShuffleScan then reads the partitions that the read specifies.
+ // (confs, whether the plan is adaptive, the block that reads the shuffle)
+ val cases = Seq(
+ (aqeWithoutCoalescing, true, shuffleStageRead),
+ (
+ Seq(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true"),
+ true,
+ Seq("AQEShuffleReadExec" -> "SHUFFLE_SCAN")),
+ (
+ Seq(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false"),
+ false,
+ Seq("CometShuffleExchangeExec" -> "SCAN")))
+ cases.foreach { case (confs, isAdaptive, expected) =>
+ withClue(s"$confs: ") {
+ withSQLConf(confs: _*) {
+ val (_, plan) = checkSparkAnswerAndOperator(groupedSum)
+ assert(plan.isInstanceOf[AdaptiveSparkPlanExec] == isAdaptive, plan)
+ if (isAdaptive) checkCometOperatorsInFinalPlan(plan)
+ assert(shuffleReadingBlocks(plan) == Seq(expected), plan)
+ }
+ }
+ }
+ }
+
+ test("AQE native aggregate chain over a shuffle stage reads it through ShuffleScan") {
+ withSQLConf(aqeWithoutCoalescing: _*) {
+ // A distinct aggregate plans two shuffles, with two native aggregates between them.
+ val df = spark
+ .range(0, 10000, 1, 20)
+ .groupBy((col("id") % 97).as("k"))
+ .agg(countDistinct(col("id") % 13), sum("id"))
+ val (_, plan) = checkSparkAnswerAndOperator(df)
+ checkCometOperatorsInFinalPlan(plan)
+ assert(shuffleReadingBlocks(plan) == Seq(shuffleStageRead, shuffleStageRead), plan)
+ }
+ }
+
+ test("AQE join of final aggregates reads only its shuffle stages through ShuffleScan") {
+ val agg = spark
+ .range(0, 10000, 1, 20)
+ .groupBy((col("id") % 97).as("k"))
+ .agg(sum("id").as("s"))
+ withSQLConf(
+ aqeWithoutCoalescing ++ Seq(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1"): _*) {
+ // The second side reuses the first side's shuffle.
+ val df = agg.join(agg.select(col("k"), (col("s") + 1).as("s2")), "k")
+ val (_, plan) = checkSparkAnswerAndOperator(df)
+ checkCometOperatorsInFinalPlan(plan, classOf[ReusedExchangeExec])
+ assert(collect(plan) { case r: ReusedExchangeExec => r }.nonEmpty, plan)
+ assert(shuffleReadingBlocks(plan) == Seq(shuffleStageRead ++ shuffleStageRead), plan)
+ }
+ withSQLConf(aqeWithoutCoalescing: _*) {
+ // The broadcast side stays a plain Scan.
+ val dim = spark.range(0, 50).select(col("id").as("k"), (col("id") * 2).as("w"))
+ val df = agg.join(broadcast(dim), "k")
+ val (_, plan) = checkSparkAnswerAndOperator(df)
+ checkCometOperatorsInFinalPlan(plan)
+ val joinBlock = shuffleStageRead :+ ("BroadcastQueryStageExec" -> "SCAN")
+ assert(shuffleReadingBlocks(plan) == Seq(joinBlock), plan)
+ }
+ }
+
+ test("AQE keeps exchange reuse above equivalent aggregates that read through ShuffleScan") {
+ withSQLConf(aqeWithoutCoalescing: _*) {
+ // Each side's aggregate reads its own shuffle stage, so their ShuffleScan sources differ.
+ val agg = spark
+ .range(0, 10000, 1, 20)
+ .groupBy((col("id") % 97).as("k"))
+ .agg(sum("id").as("s"))
+ .repartition(7, col("s"))
+ val (_, plan) = checkSparkAnswerAndOperator(agg.union(agg))
+ checkCometOperatorsInFinalPlan(plan, classOf[ReusedExchangeExec])
+ assertExchangeReuseOver(plan, "Expected the exchange above the aggregate to be reused") {
+ case a: CometHashAggregateExec if a.modes.contains(Final) => a
+ }
+ val blocks = shuffleReadingBlocks(plan)
+ val readsShuffleScan = blocks.forall(_.forall { case (_, leaf) => leaf == "SHUFFLE_SCAN" })
+ assert(blocks.nonEmpty && readsShuffleScan, plan)
+ }
+ }
+
+ test("AQE DPP query with a broadcast side reads its shuffle stage through ShuffleScan") {
+ assume(isSpark35Plus, "Native AQE DPP requires Spark 3.5+")
+ withTempDir { dir =>
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark
+ .range(0, 1000, 1, 10)
+ .selectExpr("id", "id % 10 AS p")
+ .write
+ .partitionBy("p")
+ .parquet(s"$dir/fact")
+ spark.range(0, 10).selectExpr("id AS p", "id % 3 AS x").write.parquet(s"$dir/dim")
+ }
+ withTempView("fact", "dim") {
+ spark.read.parquet(s"$dir/fact").createOrReplaceTempView("fact")
+ spark.read.parquet(s"$dir/dim").createOrReplaceTempView("dim")
+ withSQLConf(
+ aqeWithoutCoalescing :+
+ (SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true"): _*) {
+ val df = sql("""SELECT f.p, sum(f.id), count(*) FROM fact f JOIN dim d ON f.p = d.p
+ |WHERE d.x = 1 GROUP BY f.p""".stripMargin)
+ val (_, plan) = checkSparkAnswer(df)
+ assert(shuffleReadingBlocks(plan) == Seq(shuffleStageRead), plan)
+ }
+ }
+ }
+ }
+
// Regression test for https://github.com/apache/datafusion-comet/issues/3846
test("repartition count") {
withTempPath { dir =>
diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
index a48ba2e5e1..0024719ec4 100644
--- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala
@@ -36,6 +36,7 @@ import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, RoundRobinPartitioning}
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.comet._
+import org.apache.spark.sql.comet.CometExec.nativeLeaves
import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec}
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffleReadExec, BroadcastQueryStageExec, LogicalQueryStage, QueryStageExec, ShuffleQueryStageExec, SimpleCost, SimpleCostEvaluator}
@@ -335,8 +336,90 @@ class CometExecRuleSuite extends CometTestBase {
}
}
- test("AQE DPP broadcast roots retain temporary logical links after an unchanged replan") {
- assume(isSpark35Plus, "Native AQE DPP requires Spark 3.5+")
+ private def serializedLeafKinds(op: CometNativeExec): Seq[String] =
+ nativeLeaves(OperatorOuterClass.Operator.parseFrom(op.serializedPlanOpt.plan.get))
+ .map(_.getOpStructCase.name)
+
+ private def withAdaptiveAggregateConf(f: => Unit): Unit =
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false",
+ CometConf.COMET_SPARK_TO_ARROW_SUPPORTED_OPERATOR_LIST.key -> "Range")(f)
+
+ test("CometExecRule points an AQE-reused aggregate's stale Scan at its ShuffleScan stage") {
+ withAdaptiveAggregateConf {
+ val reused = createAdaptiveAggregate()
+ assert(nativeLeaves(reused.nativeOp).map(_.getOpStructCase.name) == Seq("SCAN"))
+
+ val refreshed = applyCometExecRule(reused).asInstanceOf[CometHashAggregateExec]
+ assert(nativeLeaves(refreshed.nativeOp).map(_.getOpStructCase.name) == Seq("SHUFFLE_SCAN"))
+ assert(serializedLeafKinds(refreshed) == Seq("SHUFFLE_SCAN"))
+ assert(refreshed.child.isInstanceOf[ShuffleQueryStageExec])
+ assert(refreshed.getTagValue(CometExplainInfo.EXTENSION_INFO).isEmpty)
+
+ // A second pass finds nothing stale and keeps the node as it is.
+ val again = applyCometExecRule(refreshed).asInstanceOf[CometHashAggregateExec]
+ assert(again.nativeOp eq refreshed.nativeOp)
+ assert(again.serializedPlanOpt.plan.get sameElements refreshed.serializedPlanOpt.plan.get)
+ }
+ }
+
+ test("CometExecRule keeps an AQE-reused aggregate's Scan with shuffle direct read off") {
+ withAdaptiveAggregateConf {
+ withSQLConf(CometConf.COMET_SHUFFLE_DIRECT_READ_ENABLED.key -> "false") {
+ val reused = createAdaptiveAggregate()
+ val result = applyCometExecRule(reused).asInstanceOf[CometHashAggregateExec]
+ assert(result.nativeOp eq reused.nativeOp)
+ assert(serializedLeafKinds(result) == Seq("SCAN"))
+ assert(result.getTagValue(CometExplainInfo.EXTENSION_INFO).isEmpty)
+ }
+ }
+ }
+
+ test("CometExecRule keeps an AQE-reused aggregate's Scan and says so when its leaves differ") {
+ withAdaptiveAggregateConf {
+ val stringType = QueryPlanSerde.serializeDataType(DataTypes.StringType).get
+ def retype(op: OperatorOuterClass.Operator): OperatorOuterClass.Operator =
+ if (op.hasScan) {
+ val fields = Seq.fill(op.getScan.getFieldsCount)(stringType)
+ op.toBuilder
+ .setScan(op.getScan.toBuilder.clearFields().addAllFields(fields.asJava))
+ .build()
+ } else {
+ op.toBuilder
+ .clearChildren()
+ .addAllChildren(op.getChildrenList.asScala.map(retype).asJava)
+ .build()
+ }
+ // A second Scan child gives the plan more leaves than the aggregate has inputs.
+ def addScanLeaf(op: OperatorOuterClass.Operator): OperatorOuterClass.Operator =
+ op.toBuilder.addChildren(nativeLeaves(op).head).build()
+
+ Seq("field types" -> (retype _), "leaf count" -> (addScanLeaf _)).foreach {
+ case (mismatch, change) =>
+ withClue(s"$mismatch: ") {
+ val reused = createAdaptiveAggregate()
+ val mismatched = reused.copy(nativeOp = change(reused.nativeOp))
+
+ val result = applyCometExecRule(mismatched).asInstanceOf[CometHashAggregateExec]
+ assert(result.nativeOp eq mismatched.nativeOp)
+ assert(nativeLeaves(result.nativeOp).forall(_.hasScan))
+ // The node stays native, with an info message rather than a fallback reason.
+ val info = CometExecRule.STALE_SCAN_KEPT
+ assert(result.getTagValue(CometExplainInfo.EXTENSION_INFO).contains(Set(info)))
+ assert(!CometSparkSessionExtensions.hasFallbackReason(result))
+ val explain = new ExtendedExplainInfo().generateVerboseInfo(result)
+ assert(explain.contains(s"CometHashAggregate [COMET-INFO: $info]"), explain)
+ }
+ }
+ }
+ }
+
+ /**
+ * Runs `f` with AQE DPP on and the fact and dim views of the DPP broadcast link tests, with
+ * shuffle direct read set to `directRead`.
+ */
+ private def withDppLinkViews(directRead: Boolean)(f: => Unit): Unit = {
withSQLConf(
SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
@@ -344,7 +427,8 @@ class CometExecRuleSuite extends CometTestBase {
SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true",
SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true",
SQLConf.SHUFFLE_PARTITIONS.key -> "2",
- CometConf.COMET_SHUFFLE_MODE.key -> "native") {
+ CometConf.COMET_SHUFFLE_MODE.key -> "native",
+ CometConf.COMET_SHUFFLE_DIRECT_READ_ENABLED.key -> directRead.toString) {
withTempDir { dir =>
withTempView("dpp_link_fact", "dpp_link_dim") {
withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
@@ -368,75 +452,139 @@ class CometExecRuleSuite extends CometTestBase {
assert(
spark.sessionState.conf.getConf(SQLConf.ADAPTIVE_CUSTOM_COST_EVALUATOR_CLASS).isEmpty)
- type Replan = (SparkPlan, LogicalPlan)
- val pending = new ThreadLocal[List[Option[Replan]]] {
- override def initialValue(): List[Option[Replan]] = Nil
- }
- val observed = new ConcurrentLinkedQueue[(CometBroadcastExchangeExec, LogicalPlan)]()
- val tempTag = AdaptiveSparkPlanExec.TEMP_LOGICAL_PLAN_TAG
- val costEvaluator = SimpleCostEvaluator(forceOptimizeSkewedJoin = false)
- beforeCometPreparation = plan => {
- val replan = plan match {
- case broadcast: CometBroadcastExchangeExec =>
- broadcast.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).collect {
- case stage: LogicalQueryStage =>
- assert(stage.physicalPlan eq broadcast)
- assert(broadcast.getTagValue(tempTag).exists(_ eq stage.logicalPlan))
- (broadcast.clone(), stage.logicalPlan)
- }
- case _ => None
- }
- pending.set(replan :: pending.get())
- }
- afterCometPreparation = plan => {
- val replan = pending.get().head
- val remaining = pending.get().tail
- if (remaining.isEmpty) pending.remove() else pending.set(remaining)
- replan.foreach { case (previous, logicalPlan) =>
- val broadcast = plan.asInstanceOf[CometBroadcastExchangeExec]
- assert(broadcast.logicalLink.exists(_ eq logicalPlan))
- assert(broadcast.getTagValue(tempTag).exists(_ eq logicalPlan))
- // Spark rejects an equal-cost candidate when its physical tree is unchanged.
- // Pin both inputs to that decision, including Comet's retained temporary link.
- assert(previous == broadcast)
- assert(costEvaluator.evaluateCost(previous) == SimpleCost(0))
- assert(costEvaluator.evaluateCost(broadcast) == SimpleCost(0))
- observed.add(
- (broadcast.clone().asInstanceOf[CometBroadcastExchangeExec], logicalPlan))
- }
+ f
+ }
+ }
+ }
+ }
+
+ /**
+ * Runs the DPP broadcast link query, checks its answer and that the join and the pruning
+ * broadcast are native, and returns the executed plan.
+ */
+ private def checkDppLinkQuery(): SparkPlan = {
+ val df = sql("""
+ |SELECT /*+ BROADCAST(d) */ f.k, f.total, d.total
+ |FROM (SELECT k, SUM(v) AS total FROM dpp_link_fact GROUP BY k) f
+ |JOIN (SELECT k, SUM(v) AS total FROM dpp_link_dim
+ | WHERE country = 'DE' GROUP BY k) d ON f.k = d.k
+ |""".stripMargin)
+ QueryTest.checkAnswer(
+ df,
+ (0 until 8 by 2).map(k => Row(k, 224L + 8L * k, 48L + 4L * k)),
+ checkToRDD = false)
+ val plan = df.queryExecution.executedPlan
+ assert(collect(plan) { case b: CometBroadcastHashJoinExec => b }.nonEmpty)
+ assert(collectWithSubqueries(plan) { case s: CometSubqueryBroadcastExec =>
+ s
+ }.nonEmpty)
+ plan
+ }
+
+ /**
+ * Runs `f` while observing each AQE replan of a broadcast root linked to a logical query stage.
+ * Before Comet prepares the root, `capture` takes the root and its stage. Afterwards `check`
+ * gets that captured value, the prepared root and the stage's logical plan.
+ */
+ private def withReplanObserver[T](
+ capture: (CometBroadcastExchangeExec, LogicalQueryStage) => T)(
+ check: (T, CometBroadcastExchangeExec, LogicalPlan) => Unit)(f: => Unit): Unit = {
+ val pending = new ThreadLocal[List[Option[(T, LogicalPlan)]]] {
+ override def initialValue(): List[Option[(T, LogicalPlan)]] = Nil
+ }
+ beforeCometPreparation = plan => {
+ val replan = plan match {
+ case broadcast: CometBroadcastExchangeExec =>
+ broadcast.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).collect {
+ case stage: LogicalQueryStage => (capture(broadcast, stage), stage.logicalPlan)
}
- try {
- val df = sql("""
- |SELECT /*+ BROADCAST(d) */ f.k, f.total, d.total
- |FROM (SELECT k, SUM(v) AS total FROM dpp_link_fact GROUP BY k) f
- |JOIN (SELECT k, SUM(v) AS total FROM dpp_link_dim
- | WHERE country = 'DE' GROUP BY k) d ON f.k = d.k
- |""".stripMargin)
- QueryTest.checkAnswer(
- df,
- (0 until 8 by 2).map(k => Row(k, 224L + 8L * k, 48L + 4L * k)),
- checkToRDD = false)
- val plan = df.queryExecution.executedPlan
- assert(collect(plan) { case b: CometBroadcastHashJoinExec => b }.nonEmpty)
- assert(collectWithSubqueries(plan) { case s: CometSubqueryBroadcastExec =>
- s
- }.nonEmpty)
- assert(!observed.isEmpty, "Expected a DPP broadcast root with a direct logical stage")
- observed.iterator().asScala.foreach { case (broadcast, logicalPlan) =>
- // Give the isolated snapshot conflicting links to pin Spark's TEMP-over-direct
- // precedence, which would otherwise be invisible after Comet repairs both.
- broadcast.setLogicalLink(LogicalQueryStage(logicalPlan, broadcast))
- val stage = BroadcastQueryStageExec(0, broadcast, broadcast.canonicalized)
- val setStageLink = PrivateMethod[Unit](Symbol("setLogicalLinkForNewQueryStage"))
- plan
- .asInstanceOf[AdaptiveSparkPlanExec]
- .invokePrivate(setStageLink(stage, broadcast))
- assert(stage.logicalLink.exists(_ eq logicalPlan))
- }
- } finally {
- beforeCometPreparation = _ => ()
- afterCometPreparation = _ => ()
+ case _ => None
+ }
+ pending.set(replan :: pending.get())
+ }
+ afterCometPreparation = plan => {
+ val replan = pending.get().head
+ val remaining = pending.get().tail
+ if (remaining.isEmpty) pending.remove() else pending.set(remaining)
+ replan.foreach { case (captured, logicalPlan) =>
+ check(captured, plan.asInstanceOf[CometBroadcastExchangeExec], logicalPlan)
+ }
+ }
+ try {
+ f
+ } finally {
+ beforeCometPreparation = _ => ()
+ afterCometPreparation = _ => ()
+ }
+ }
+
+ test("AQE DPP broadcast roots retain temporary logical links after an unchanged replan") {
+ assume(isSpark35Plus, "Native AQE DPP requires Spark 3.5+")
+ // With direct read, the replan that first sees the dim side's shuffle stage moves the
+ // aggregate under the broadcast onto a ShuffleScan, so that replan is not unchanged.
+ withDppLinkViews(directRead = false) {
+ val observed = new ConcurrentLinkedQueue[(CometBroadcastExchangeExec, LogicalPlan)]()
+ val tempTag = AdaptiveSparkPlanExec.TEMP_LOGICAL_PLAN_TAG
+ val costEvaluator = SimpleCostEvaluator(forceOptimizeSkewedJoin = false)
+ withReplanObserver { (broadcast, stage) =>
+ assert(stage.physicalPlan eq broadcast)
+ assert(broadcast.getTagValue(tempTag).exists(_ eq stage.logicalPlan))
+ broadcast.clone()
+ } { (previous, broadcast, logicalPlan) =>
+ assert(broadcast.logicalLink.exists(_ eq logicalPlan))
+ assert(broadcast.getTagValue(tempTag).exists(_ eq logicalPlan))
+ // Spark rejects an equal-cost candidate when its physical tree is unchanged.
+ // Pin both inputs to that decision, including Comet's retained temporary link.
+ assert(previous == broadcast)
+ assert(costEvaluator.evaluateCost(previous) == SimpleCost(0))
+ assert(costEvaluator.evaluateCost(broadcast) == SimpleCost(0))
+ observed.add((broadcast.clone().asInstanceOf[CometBroadcastExchangeExec], logicalPlan))
+ } {
+ val plan = checkDppLinkQuery()
+ assert(!observed.isEmpty, "Expected a DPP broadcast root with a direct logical stage")
+ observed.iterator().asScala.foreach { case (broadcast, logicalPlan) =>
+ // Give the isolated snapshot conflicting links to pin Spark's TEMP-over-direct
+ // precedence, which would otherwise be invisible after Comet repairs both.
+ broadcast.setLogicalLink(LogicalQueryStage(logicalPlan, broadcast))
+ val stage = BroadcastQueryStageExec(0, broadcast, broadcast.canonicalized)
+ val setStageLink = PrivateMethod[Unit](Symbol("setLogicalLinkForNewQueryStage"))
+ plan
+ .asInstanceOf[AdaptiveSparkPlanExec]
+ .invokePrivate(setStageLink(stage, broadcast))
+ assert(stage.logicalLink.exists(_ eq logicalPlan))
+ }
+ }
+ }
+ }
+
+ test("AQE DPP broadcast roots keep their logical links through a ShuffleScan replan") {
+ assume(isSpark35Plus, "Native AQE DPP requires Spark 3.5+")
+ withDppLinkViews(directRead = true) {
+ // For each replanned broadcast root: whether Comet changed it, whether it kept the
+ // stage's direct and temporary links, and the leaf kinds of the native plan under it.
+ val observed = new ConcurrentLinkedQueue[(Boolean, Boolean, Boolean, Seq[String])]()
+ val tempTag = AdaptiveSparkPlanExec.TEMP_LOGICAL_PLAN_TAG
+ withReplanObserver((broadcast, _) => broadcast.clone()) {
+ (previous, broadcast, logicalPlan) =>
+ val leafKinds = broadcast.child match {
+ case native: CometNativeExec =>
+ nativeLeaves(native.nativeOp).map(_.getOpStructCase.name)
+ case other => Seq(other.nodeName)
}
+ observed.add(
+ (
+ previous != broadcast,
+ broadcast.logicalLink.exists(_ eq logicalPlan),
+ broadcast.getTagValue(tempTag).exists(_ eq logicalPlan),
+ leafKinds))
+ } {
+ checkDppLinkQuery()
+ val replans = observed.iterator().asScala.toSeq
+ assert(replans.exists(_._1), s"Expected a replan that changes the broadcast: $replans")
+ replans.foreach { case (_, hasLink, hasTempLink, leafKinds) =>
+ assert(hasLink, s"Lost the stage's logical link: $replans")
+ assert(hasTempLink, s"Lost the stage's temporary logical link: $replans")
+ assert(leafKinds == Seq("SHUFFLE_SCAN"), s"Expected a ShuffleScan read: $replans")
}
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org