[GLUTEN-11524][VL] Fix ColumnarAQEShuffleReadExec

[marin-ma]

ColumnarAQEShuffleReadExec is a wrapper for ShuffleQueryStageExec or a replacement for AQEShuffleReadExec. Its execution mode determines whether the shuffle reader outputs Velox RowVectors for the CPU pipeline or CudfVectors for the GPU pipeline.

However, during canonicalization, the delegate field can be set to ShuffleExchange, which currently causes a failure. This PR fixes the issue and adds UT.

This PR also fixes adding ResizeBatchesExec before shuffle write when cudf is enabled.

Related issue: #11524

[Copilot]

Pull request overview

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Fixes ColumnarAQEShuffleReadExec canonicalization by allowing its delegate to be a generic SparkPlan (including ShuffleExchange during canonicalization), while preserving AQE shuffle reader behavior for runtime execution.

Changes:

• Refactors ColumnarAQEShuffleReadExec to store delegate: SparkPlan and adjusts child/output/outputPartitioning/withNewChildInternal accordingly.
• Updates stage execution mode adjustment to construct ColumnarAQEShuffleReadExec with the new delegate type and tweaks resize-batches handling.
• Updates Velox tests to accommodate the new delegate representation.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.

File	Description
gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala	Refactors delegate handling to avoid canonicalization failures and updates plan plumbing accordingly.
backends-velox/src/test/scala/org/apache/gluten/execution/StageExecutionModeSuite.scala	Updates assertions to derive shuffle stages via the new delegate: SparkPlan shape.
backends-velox/src/main/scala/org/apache/spark/sql/execution/AdjustStageExecutionMode.scala	Constructs ColumnarAQEShuffleReadExec with the new delegate type and narrows a resize-batches match.
backends-velox/src/main/scala/org/apache/gluten/extension/AppendBatchResizeForShuffleInputAndOutput.scala	Enables shuffle-output resize when CUDF columnar mode is on (with TODO note).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[github-actions]

Run Gluten Clickhouse CI on x86

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

Suppressed comments (3)

gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala:53

• outputPartitioning now comes from delegate.outputPartitioning. When delegate is a ShuffleQueryStageExec, this can differ from the previously used AQEShuffleReadExec(...).outputPartitioning (which is derived from partitionSpecs). That is a semantic change and can affect downstream planning decisions that rely on AQE shuffle-reader partitioning. Consider restoring the previous behavior by deriving outputPartitioning from aqeReader.outputPartitioning (and similarly keep outputs consistent with the AQE reader semantics when the delegate is a stage).

  override def output: Seq[Attribute] = delegate.output

  override lazy val outputPartitioning: Partitioning = delegate.outputPartitioning

gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala:85

• The scaladoc explicitly allows delegate to be ShuffleExchange during canonicalization, but metrics eagerly depends on aqeReader, which will throw for any non-AQEShuffleReadExec/ShuffleQueryStageExec delegate. This can reintroduce canonicalization-time failures if Spark touches metrics (e.g., for explain/UI/debug) on a canonicalized plan. To make canonicalization robust, consider making metrics (and any other accesses that can happen during canonicalization) safe for the ShuffleExchange case (e.g., return an empty metric map or delegate metrics when available), or tighten invariants so a non-AQE/stage delegate cannot reach paths where metrics are evaluated.

  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}.")
    }
  }

  @transient override lazy val metrics: Map[String, SQLMetric] = aqeReader.metrics

backends-velox/src/main/scala/org/apache/spark/sql/execution/AdjustStageExecutionMode.scala:101

• This changes behavior from adjusting VeloxResizeBatchesExec unconditionally to only adjusting it when its child is a shuffle stage / AQE shuffle read. If adjustExecutionMode is expected to consistently propagate stageExecutionMode through the plan, this introduces a special case where the resize node’s stage mode may remain stale while its subtree is updated via the default case. Consider either (a) keeping the unconditional handling, or (b) adding an explicit else branch that still updates the VeloxResizeBatchesExec’s stage mode while recursing, to avoid mixed or inconsistent stage-mode annotations.

      case r: VeloxResizeBatchesExec
          if r.child.isInstanceOf[ShuffleQueryStageExec] ||
            r.child.isInstanceOf[AQEShuffleReadExec] =>
        VeloxResizeBatchesExec(
          adjustExecutionMode(r.child, stageExecutionMode),
          Some(stageExecutionMode))

[github-actions]

Run Gluten Clickhouse CI on x86

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

Suppressed comments (5)

backends-velox/src/main/scala/org/apache/spark/sql/execution/AdjustStageExecutionMode.scala:101

• This narrows the previous VeloxResizeBatchesExec rewrite to only two child types. Any VeloxResizeBatchesExec wrapping other plans will now fall through to the default case and won’t get Some(stageExecutionMode) set anymore, which is a functional behavior change. If the intent is still to propagate stage execution mode through VeloxResizeBatchesExec generally, keep the rewrite unconditional (as before) or add a separate case to preserve/set the stage mode for other children.

      case r: VeloxResizeBatchesExec
          if r.child.isInstanceOf[ShuffleQueryStageExec] ||
            r.child.isInstanceOf[AQEShuffleReadExec] =>
        VeloxResizeBatchesExec(
          adjustExecutionMode(r.child, stageExecutionMode),
          Some(stageExecutionMode))

gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala:81

• The exception message is hard to debug in production because it only prints the simple class name. Consider including delegate.nodeName and/or the full class name, and explicitly listing the expected delegate types (e.g., AQEShuffleReadExec or ShuffleQueryStageExec) plus the canonicalization context. That makes failures actionable if this gets triggered outside canonicalization.

        // The child is Exchange during canonicalization.
        throw new IllegalStateException(
          s"Cannot get aqeReader from delegate class ${delegate.getClass.getSimpleName}.")

gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala:34

• The Scaladoc parameter description is grammatically incomplete and also refers to ShuffleExchange while the concrete Spark plan type appears to be ShuffleExchangeExec (per usage/tests). Please rewrite this line to be precise about the expected delegate plan types and match the actual class naming used in Spark.

 *   AQEShuffleReadExec or ShuffleQueryStageExec. Or ShuffleExchange during canonicalization.

backends-velox/src/test/scala/org/apache/gluten/execution/StageExecutionModeSuite.scala:126

• This error message is misleading (the match is on _.delegate, not the child) and drops important context for diagnosing test failures. Consider changing it to mention delegate and include the actual class (and/or nodeName) of the unexpected plan.

        case _ =>
          throw new IllegalArgumentException("Unexpected child of ColumnarAQEShuffleReadExec")

backends-velox/src/main/scala/org/apache/gluten/extension/AppendBatchResizeForShuffleInputAndOutput.scala:36

• resizeBatchesShuffleOutputEnabled now represents 'explicit shuffle output resize enabled OR CUDF enabled', which is broader than the name implies. To keep intent clear, consider renaming the variable to reflect the combined condition or split into two booleans (one for the config, one for CUDF) and combine them explicitly at the call site.

    // TODO: Move cudf resize batches into shuffle reader.
    val resizeBatchesShuffleOutputEnabled =
      VeloxConfig.get.veloxResizeBatchesShuffleOutput || VeloxConfig.get.enableColumnarCudf

[github-actions]

Run Gluten Clickhouse CI on x86

[github-actions]

Run Gluten Clickhouse CI on x86

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

Suppressed comments (2)

backends-velox/src/main/scala/org/apache/spark/sql/execution/AdjustStageExecutionMode.scala:103

• This guard ignores the new wrapper (ColumnarAQEShuffleReadExec). In plans where VeloxResizeBatchesExec sits on top of ColumnarAQEShuffleReadExec (as asserted in the new test), the case won’t match and VeloxResizeBatchesExec won’t get Some(stageExecutionMode), potentially leaving the resize node in an inconsistent mode configuration. Include r.child.isInstanceOf[ColumnarAQEShuffleReadExec] in the guard (or remove the guard and handle non-shuffle cases explicitly) so shuffle-resize nodes are consistently annotated.

      case r: VeloxResizeBatchesExec
          // TODO: This should be removed after merging resize into native shuffle read.
          // Only change the execution mode for shuffle reader.
          if r.child.isInstanceOf[ShuffleQueryStageExec] ||
            r.child.isInstanceOf[AQEShuffleReadExec] =>
        VeloxResizeBatchesExec(
          adjustExecutionMode(r.child, stageExecutionMode),
          Some(stageExecutionMode))

backends-velox/src/test/scala/org/apache/gluten/execution/StageExecutionModeSuite.scala:114

• This assertion is brittle because canonicalization behavior and exact exchange class (ShuffleExchangeExec vs a different exchange implementation) can vary across Spark versions and planner paths. If the intent is to ensure canonicalization removes query-stages/readers and leaves an exchange-like node, consider asserting on a more stable abstraction (e.g., Spark’s exchange base type) or checking nodeName/isInstanceOf[Exchange] to reduce version-specific failures.

          val canonicalized = reader.canonicalized
          // canonicalized plan before applying query stage optimizer rules.
          assert(canonicalized.children.forall(_.isInstanceOf[ShuffleExchangeExec]))

[github-actions]

Run Gluten Clickhouse CI on x86