marin-ma commented on code in PR #12691:
URL: https://github.com/apache/gluten/pull/12691#discussion_r3720205916
##########
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:
> When delegate is a ShuffleQueryStageExec, this can be incorrect because
the effective partitioning for a shuffle reader is determined by
AQEShuffleReadExec's partitionSpecs
It's not true. `delegate` is the wrapped original plan. When it's
`ShuffleQueryStageExec`, it means the original plan doesn't have a
`AQEShuffleReadExec`.
--
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]
