[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.

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