Copilot commented on code in PR #12691:
URL: https://github.com/apache/gluten/pull/12691#discussion_r3719955882
##########
gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala:
##########
@@ -31,49 +31,59 @@ import org.apache.spark.sql.vectorized.ColumnarBatch
* ShuffleQueryStageExec if executionMode is set by the planner.
*
* @param delegate
- * The AQEShuffleReadExec or ShuffleQueryStageExec.
+ * AQEShuffleReadExec, ShuffleQueryStageExec, or (during canonicalization)
ShuffleExchange.
* @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 = aqeReader.child
+ override def child: SparkPlan = delegate match {
+ case AQEShuffleReadExec(c, _) => c
+ case _ => delegate
+ }
- override def output: Seq[Attribute] = aqeReader.child.output
+ override def output: Seq[Attribute] = delegate.output
- override lazy val outputPartitioning: Partitioning =
aqeReader.outputPartitioning
+ override lazy val outputPartitioning: Partitioning =
delegate.outputPartitioning
Review Comment:
`outputPartitioning` now delegates directly to
`delegate.outputPartitioning`. When `delegate` is a `ShuffleQueryStageExec`,
this can be incorrect because the effective partitioning for a shuffle *reader*
is determined by `AQEShuffleReadExec`'s `partitionSpecs` (e.g., after
coalescing), not the query stage's original shuffle partitioning. Consider
returning `aqeReader.outputPartitioning` when `delegate` is
`AQEShuffleReadExec` or `ShuffleQueryStageExec`, and only falling back to
`delegate.outputPartitioning` for the canonicalization-only
`ShuffleExchangeExec` case.
--
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]
