Copilot commented on code in PR #12691:
URL: https://github.com/apache/gluten/pull/12691#discussion_r3714057904
##########
backends-velox/src/test/scala/org/apache/gluten/execution/StageExecutionModeSuite.scala:
##########
@@ -113,19 +113,19 @@ class StageExecutionModeSuite extends
VeloxWholeStageTransformerSuite {
s"Expected GPU AQE shuffle reader, but got
${reader.executionMode}")
}
- val shuffleStages = plan.collect {
- case stage: ShuffleQueryStageExec => stage
+ val shuffleStages: Seq[ShuffleQueryStageExec] =
shuffleReaders.map(_.delegate).map {
+ case a: AQEShuffleReadExec =>
+ assert(a.child.isInstanceOf[ShuffleQueryStageExec])
+ a.child.asInstanceOf[ShuffleQueryStageExec]
+ case s: ShuffleQueryStageExec => s
+ case _ =>
+ throw new IllegalArgumentException("Unexpected child of
ColumnarAQEShuffleReadExec")
}
- val exchanges = shuffleStages.flatMap {
- _.plan.collect {
- case exchange: ColumnarShuffleExchangeExec => exchange
- }
- }
-
- assert(exchanges.nonEmpty)
-
- exchanges.foreach {
+ shuffleStages.foreach {
+ shuffleStage =>
+
assert(shuffleStage.shuffle.isInstanceOf[ColumnarShuffleExchangeExec])
+ val exchange =
shuffleStage.shuffle.asInstanceOf[ColumnarShuffleExchangeExec]
exchange =>
assert(
!exchange.mapperStageMode.contains(MockGPUStageMode),
Review Comment:
This `foreach` block contains an extra `exchange =>` lambda (likely leftover
from the previous `exchanges.foreach { exchange => ... }`). As written, this
should not compile. Remove the stray `exchange =>` and keep a single lambda
body that uses the `exchange` value you already derived on line 128.
##########
gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala:
##########
@@ -31,40 +31,56 @@ import org.apache.spark.sql.vectorized.ColumnarBatch
* ShuffleQueryStageExec if executionMode is set by the planner.
*
* @param delegate
- * The AQEShuffleReadExec or ShuffleQueryStageExec.
+ * AQEShuffleReadExec or ShuffleQueryStageExec. Or ShuffleExchange during
canonicalization.
* @param executionMode
* The execution mode of the current AQE stage.
*/
case class ColumnarAQEShuffleReadExec(
- delegate: Either[AQEShuffleReadExec, ShuffleQueryStageExec],
+ delegate: SparkPlan,
executionMode: StageExecutionMode) extends UnaryExecNode {
override def nodeName: String =
s"ColumnarAQEShuffleRead(${executionMode.name})"
- private val isAQEShuffleRead = delegate.isLeft
-
- private val aqeReader: AQEShuffleReadExec = {
- if (isAQEShuffleRead) {
- delegate.left.get
- } else {
- // Wrap ShuffleQueryStageExe with dummy PartitionSpecs.
- val queryStageExec = delegate.right.get
- // Create CoalescedPartitionSpec for each partition.
- val partitionSpecs =
- Array.tabulate(queryStageExec.shuffle.numPartitions)(i =>
CoalescedPartitionSpec(i, i + 1))
- AQEShuffleReadExec(queryStageExec, partitionSpecs)
- }
+ override def supportsColumnar: Boolean = true
+
+ override def child: SparkPlan = delegate match {
+ case AQEShuffleReadExec(c, _) => c
+ case _ => delegate
}
- override def supportsColumnar: Boolean = true
+ override def output: Seq[Attribute] = delegate.output
- override def child: SparkPlan = aqeReader.child
+ override lazy val outputPartitioning: Partitioning =
delegate.outputPartitioning
- override def output: Seq[Attribute] = aqeReader.child.output
+ override protected def stringArgs: Iterator[Any] = {
+ delegate match {
+ case a: AQEShuffleReadExec => a.stringArgs
+ case _ => super.stringArgs
+ }
+ }
- override lazy val outputPartitioning: Partitioning =
aqeReader.outputPartitioning
+ override protected def withNewChildInternal(newChild: SparkPlan):
ColumnarAQEShuffleReadExec = {
+ delegate match {
+ case a: AQEShuffleReadExec => copy(delegate =
a.withNewChildren(Seq(newChild)))
+ case _ => copy(delegate = newChild)
+ }
+ }
- override def stringArgs: Iterator[Any] = aqeReader.stringArgs
+ private lazy val aqeReader: AQEShuffleReadExec = {
+ delegate match {
+ case a: AQEShuffleReadExec => a
+ case s: ShuffleQueryStageExec =>
+ // Wrap ShuffleQueryStageExe with dummy PartitionSpecs by creating
CoalescedPartitionSpec
+ // for each partition.
+ val partitionSpecs =
+ Array.tabulate(s.shuffle.numPartitions)(i =>
CoalescedPartitionSpec(i, i + 1))
+ AQEShuffleReadExec(s, partitionSpecs)
+ case _ =>
+ // The child is Exchange during canonicalization.
+ throw new IllegalStateException(
+ s"Cannot get aqeReader from delegate class
${delegate.getClass.getSimpleName}.")
Review Comment:
This exception can be hard to diagnose in practice because it omits the
expected delegate types and any plan context. Consider including the allowed
delegate types (`AQEShuffleReadExec` / `ShuffleQueryStageExec`), plus
`delegate.nodeName` (and optionally `executionMode.name`) to make failures
actionable without needing extra logging.
##########
gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala:
##########
@@ -31,40 +31,56 @@ import org.apache.spark.sql.vectorized.ColumnarBatch
* ShuffleQueryStageExec if executionMode is set by the planner.
*
* @param delegate
- * The AQEShuffleReadExec or ShuffleQueryStageExec.
+ * AQEShuffleReadExec or ShuffleQueryStageExec. Or ShuffleExchange during
canonicalization.
Review Comment:
Scaladoc wording is a bit awkward/ambiguous. Consider rephrasing to clearly
list the valid delegate shapes, e.g. 'AQEShuffleReadExec,
ShuffleQueryStageExec, or (during canonicalization) ShuffleExchange', so it
reads as a single structured contract.
##########
backends-velox/src/main/scala/org/apache/spark/sql/execution/AdjustStageExecutionMode.scala:
##########
@@ -85,21 +85,19 @@ object AdjustStageExecutionMode extends Logging {
// TODO: support BroadcastQueryStageExec.
case aqeShuffleRead @ AQEShuffleReadExec(s @ ShuffleQueryStageExec(_, _,
_), _)
if s.shuffle.isInstanceOf[ColumnarShuffleExchangeExec] =>
- ColumnarAQEShuffleReadExec(
- Left(aqeShuffleRead),
- stageExecutionMode)
+ ColumnarAQEShuffleReadExec(aqeShuffleRead, stageExecutionMode)
case queryStageExec: ShuffleQueryStageExec
if queryStageExec.shuffle.isInstanceOf[ColumnarShuffleExchangeExec]
=>
- ColumnarAQEShuffleReadExec(
- Right(queryStageExec),
- stageExecutionMode)
+ ColumnarAQEShuffleReadExec(queryStageExec, stageExecutionMode)
case shuffle: ColumnarShuffleExchangeExec =>
shuffle
.copy(mapperStageMode = Some(stageExecutionMode))
.withNewChildren(Seq(adjustExecutionMode(shuffle.child,
stageExecutionMode)))
- case resizeBatches: VeloxResizeBatchesExec =>
+ case r: VeloxResizeBatchesExec
+ if r.child.isInstanceOf[ShuffleQueryStageExec] ||
+ r.child.isInstanceOf[AQEShuffleReadExec] =>
VeloxResizeBatchesExec(
- adjustExecutionMode(resizeBatches.child, stageExecutionMode),
+ adjustExecutionMode(r.child, stageExecutionMode),
Some(stageExecutionMode))
Review Comment:
Narrowing this case can lead to inconsistent propagation of
`Some(stageExecutionMode)` when `VeloxResizeBatchesExec` wraps a shuffle reader
that has already been rewritten (e.g., `ColumnarAQEShuffleReadExec`), because
it will fall through to the default branch and keep the old
`stageExecutionMode` (possibly `None`). Consider extending the guard to include
`ColumnarAQEShuffleReadExec` (and any other expected shuffle-reader wrappers),
or document why it is safe/desired to not set the stage mode in those cases.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
