Re: [PR] [GLUTEN-11524][VL] Fix ColumnarAQEShuffleReadExec [gluten]

2026-08-05 Thread via GitHub


marin-ma merged PR #12691:
URL: https://github.com/apache/gluten/pull/12691


-- 
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]



Re: [PR] [GLUTEN-11524][VL] Fix ColumnarAQEShuffleReadExec [gluten]

2026-08-05 Thread via GitHub


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]



Re: [PR] [GLUTEN-11524][VL] Fix ColumnarAQEShuffleReadExec [gluten]

2026-08-05 Thread via GitHub


github-actions[bot] commented on PR #12691:
URL: https://github.com/apache/gluten/pull/12691#issuecomment-5191033266

   Run Gluten Clickhouse CI on x86


-- 
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]



Re: [PR] [GLUTEN-11524][VL] Fix ColumnarAQEShuffleReadExec [gluten]

2026-08-05 Thread via GitHub


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]



Re: [PR] [GLUTEN-11524][VL] Fix ColumnarAQEShuffleReadExec [gluten]

2026-08-05 Thread via GitHub


github-actions[bot] commented on PR #12691:
URL: https://github.com/apache/gluten/pull/12691#issuecomment-5190237797

   Run Gluten Clickhouse CI on x86


-- 
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]



Re: [PR] [GLUTEN-11524][VL] Fix ColumnarAQEShuffleReadExec [gluten]

2026-08-05 Thread via GitHub


github-actions[bot] commented on PR #12691:
URL: https://github.com/apache/gluten/pull/12691#issuecomment-5189798062

   Run Gluten Clickhouse CI on x86


-- 
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]



Re: [PR] [GLUTEN-11524][VL] Fix ColumnarAQEShuffleReadExec [gluten]

2026-08-04 Thread via GitHub


github-actions[bot] commented on PR #12691:
URL: https://github.com/apache/gluten/pull/12691#issuecomment-5182231837

   Run Gluten Clickhouse CI on x86


-- 
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]



Re: [PR] [GLUTEN-11524][VL] Fix ColumnarAQEShuffleReadExec [gluten]

2026-08-04 Thread via GitHub


github-actions[bot] commented on PR #12691:
URL: https://github.com/apache/gluten/pull/12691#issuecomment-5181647953

   Run Gluten Clickhouse CI on x86


-- 
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]



Re: [PR] [GLUTEN-11524][VL] Fix ColumnarAQEShuffleReadExec [gluten]

2026-08-04 Thread via GitHub


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