This is an automated email from the ASF dual-hosted git repository.
marin-ma pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new c1f72ca95e [GLUTEN-11524][VL] Fix ColumnarAQEShuffleReadExec (#12691)
c1f72ca95e is described below
commit c1f72ca95e45ccc1785b0df9d3f2fc3954979703
Author: Rong Ma <[email protected]>
AuthorDate: Wed Aug 5 15:11:02 2026 +0100
[GLUTEN-11524][VL] Fix ColumnarAQEShuffleReadExec (#12691)
---
.../org/apache/gluten/config/VeloxConfig.scala | 2 +-
...AppendBatchResizeForShuffleInputAndOutput.scala | 4 +-
.../sql/execution/AdjustStageExecutionMode.scala | 16 ++---
.../gluten/execution/MiscOperatorSuite.scala | 45 ++++++++++++-
.../gluten/execution/StageExecutionModeSuite.scala | 29 +++++----
docs/velox-configuration.md | 2 +-
.../adaptive/ColumnarAQEShuffleReadExec.scala | 76 +++++++++++-----------
7 files changed, 111 insertions(+), 63 deletions(-)
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala
b/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala
index b742e3a5a0..fade4402cf 100644
--- a/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala
+++ b/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala
@@ -822,7 +822,7 @@ object VeloxConfig extends ConfigRegistry {
.createWithDefault(false)
val CUDF_ENABLE_VALIDATION =
-
buildStaticConf("spark.gluten.sql.columnar.backend.velox.cudf.enableValidation")
+ buildConf("spark.gluten.sql.columnar.backend.velox.cudf.enableValidation")
.doc(
"Heuristics you can apply to validate a cuDF/GPU plan and only offload
when " +
"the entire stage can be fully and profitably executed on GPU")
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/extension/AppendBatchResizeForShuffleInputAndOutput.scala
b/backends-velox/src/main/scala/org/apache/gluten/extension/AppendBatchResizeForShuffleInputAndOutput.scala
index 872ab1db25..fbed0f5f5c 100644
---
a/backends-velox/src/main/scala/org/apache/gluten/extension/AppendBatchResizeForShuffleInputAndOutput.scala
+++
b/backends-velox/src/main/scala/org/apache/gluten/extension/AppendBatchResizeForShuffleInputAndOutput.scala
@@ -31,7 +31,9 @@ case class
AppendBatchResizeForShuffleInputAndOutput(isAdaptiveContext: Boolean)
extends Rule[SparkPlan] {
override def apply(plan: SparkPlan): SparkPlan = {
val resizeBatchesShuffleInputEnabled =
VeloxConfig.get.veloxResizeBatchesShuffleInput
- val resizeBatchesShuffleOutputEnabled =
VeloxConfig.get.veloxResizeBatchesShuffleOutput
+ // TODO: Move cudf resize batches into shuffle reader.
+ val resizeBatchesShuffleOutputEnabled =
+ VeloxConfig.get.veloxResizeBatchesShuffleOutput ||
VeloxConfig.get.enableColumnarCudf
if (!resizeBatchesShuffleInputEnabled &&
!resizeBatchesShuffleOutputEnabled) {
return plan
}
diff --git
a/backends-velox/src/main/scala/org/apache/spark/sql/execution/AdjustStageExecutionMode.scala
b/backends-velox/src/main/scala/org/apache/spark/sql/execution/AdjustStageExecutionMode.scala
index adadc67df7..adddc630eb 100644
---
a/backends-velox/src/main/scala/org/apache/spark/sql/execution/AdjustStageExecutionMode.scala
+++
b/backends-velox/src/main/scala/org/apache/spark/sql/execution/AdjustStageExecutionMode.scala
@@ -85,21 +85,21 @@ object AdjustStageExecutionMode extends Logging {
// TODO: support BroadcastQueryStageExec.
case aqeShuffleRead @ AQEShuffleReadExec(s @ ShuffleQueryStageExec(_, _,
_), _)
if s.shuffle.isInstanceOf[ColumnarShuffleExchangeExec] =>
- ColumnarAQEShuffleReadExec(
- Left(aqeShuffleRead),
- stageExecutionMode)
+ ColumnarAQEShuffleReadExec(aqeShuffleRead, stageExecutionMode)
case queryStageExec: ShuffleQueryStageExec
if queryStageExec.shuffle.isInstanceOf[ColumnarShuffleExchangeExec]
=>
- ColumnarAQEShuffleReadExec(
- Right(queryStageExec),
- stageExecutionMode)
+ ColumnarAQEShuffleReadExec(queryStageExec, stageExecutionMode)
case shuffle: ColumnarShuffleExchangeExec =>
shuffle
.copy(mapperStageMode = Some(stageExecutionMode))
.withNewChildren(Seq(adjustExecutionMode(shuffle.child,
stageExecutionMode)))
- case resizeBatches: VeloxResizeBatchesExec =>
+ 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(resizeBatches.child, stageExecutionMode),
+ adjustExecutionMode(r.child, stageExecutionMode),
Some(stageExecutionMode))
case _ =>
plan.withNewChildren(plan.children.map(adjustExecutionMode(_,
stageExecutionMode)))
diff --git
a/backends-velox/src/test/scala/org/apache/gluten/execution/MiscOperatorSuite.scala
b/backends-velox/src/test/scala/org/apache/gluten/execution/MiscOperatorSuite.scala
index 6fad06dc0b..6cac6ccb1c 100644
---
a/backends-velox/src/test/scala/org/apache/gluten/execution/MiscOperatorSuite.scala
+++
b/backends-velox/src/test/scala/org/apache/gluten/execution/MiscOperatorSuite.scala
@@ -23,7 +23,7 @@ import org.apache.spark.SparkConf
import org.apache.spark.shuffle.GlutenShuffleUtils
import org.apache.spark.sql.{DataFrame, Row}
import org.apache.spark.sql.execution._
-import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper,
AQEShuffleReadExec, ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper,
AQEShuffleReadExec, ColumnarAQEShuffleReadExec, ShuffleQueryStageExec}
import org.apache.spark.sql.execution.joins.BaseJoinExec
import org.apache.spark.sql.execution.window.WindowExec
import org.apache.spark.sql.functions._
@@ -2199,6 +2199,49 @@ class MiscOperatorSuite extends
VeloxWholeStageTransformerSuite with AdaptiveSpa
})
}
+ test("Check VeloxResizeBatches is added in ShuffleRead when cuDF is
enabled") {
+ Seq(true, false).foreach(
+ coalesceEnabled => {
+ withSQLConf(
+ GlutenConfig.COLUMNAR_CUDF_ENABLED.key -> "true",
+ VeloxConfig.CUDF_ENABLE_VALIDATION.key -> "false",
+ VeloxConfig.COLUMNAR_VELOX_RESIZE_BATCHES_SHUFFLE_OUTPUT.key ->
"false",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "10",
+ SQLConf.COALESCE_PARTITIONS_ENABLED.key -> coalesceEnabled.toString
+ ) {
+ runQueryAndCompare(
+ "SELECT l_orderkey, count(1) from lineitem group by
l_orderkey".stripMargin) {
+ df =>
+ val executedPlan = getExecutedPlan(df)
+ if (coalesceEnabled) {
+ //
VeloxResizeBatches(AQEShuffleRead(ShuffleQueryStage(ColumnarShuffleExchange)))
+ assert(executedPlan.sliding(4).exists {
+ case Seq(
+ _: ColumnarShuffleExchangeExec,
+ _: ShuffleQueryStageExec,
+ ColumnarAQEShuffleReadExec(AQEShuffleReadExec(_, _),
_),
+ _: VeloxResizeBatchesExec
+ ) =>
+ true
+ case _ => false
+ })
+ } else {
+ //
VeloxResizeBatches(ShuffleQueryStage(ColumnarShuffleExchange))
+ assert(executedPlan.sliding(4).exists {
+ case Seq(
+ _: ColumnarShuffleExchangeExec,
+ _: ShuffleQueryStageExec,
+ ColumnarAQEShuffleReadExec(ShuffleQueryStageExec(_, _,
_), _),
+ _: VeloxResizeBatchesExec) =>
+ true
+ case _ => false
+ })
+ }
+ }
+ }
+ })
+ }
+
test("RowToVeloxColumnar preferredBatchBytes") {
Seq("1", "80", "100000000").foreach(
preferredBatchBytes => {
diff --git
a/backends-velox/src/test/scala/org/apache/gluten/execution/StageExecutionModeSuite.scala
b/backends-velox/src/test/scala/org/apache/gluten/execution/StageExecutionModeSuite.scala
index 50458d43b6..6efaeae753 100644
---
a/backends-velox/src/test/scala/org/apache/gluten/execution/StageExecutionModeSuite.scala
+++
b/backends-velox/src/test/scala/org/apache/gluten/execution/StageExecutionModeSuite.scala
@@ -21,7 +21,8 @@ import org.apache.gluten.config.{GlutenConfig, VeloxConfig}
import org.apache.spark.SparkConf
import org.apache.spark.sql.Row
import org.apache.spark.sql.execution.ColumnarShuffleExchangeExec
-import org.apache.spark.sql.execution.adaptive.{ColumnarAQEShuffleReadExec,
ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec,
ColumnarAQEShuffleReadExec, ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
import org.apache.spark.sql.internal.SQLConf
class StageExecutionModeSuite extends VeloxWholeStageTransformerSuite {
@@ -108,25 +109,27 @@ class StageExecutionModeSuite extends
VeloxWholeStageTransformerSuite {
shuffleReaders.foreach {
reader =>
+ val canonicalized = reader.canonicalized
+ // canonicalized plan before applying query stage optimizer rules.
+
assert(canonicalized.children.forall(_.isInstanceOf[ShuffleExchangeExec]))
assert(
reader.executionMode == MockGPUStageMode,
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 {
- exchange =>
+ shuffleStages.foreach {
+ shuffleStage =>
+
assert(shuffleStage.shuffle.isInstanceOf[ColumnarShuffleExchangeExec])
+ val exchange =
shuffleStage.shuffle.asInstanceOf[ColumnarShuffleExchangeExec]
assert(
!exchange.mapperStageMode.contains(MockGPUStageMode),
s"Expected CPU mapper stage, but got ${exchange.mapperStageMode}")
diff --git a/docs/velox-configuration.md b/docs/velox-configuration.md
index 9ef1b3ba4b..8f80bb2d38 100644
--- a/docs/velox-configuration.md
+++ b/docs/velox-configuration.md
@@ -23,7 +23,7 @@ nav_order: 16
| spark.gluten.sql.columnar.backend.velox.cudf.batchSize
| 🔄 Dynamic | 2147483647 | Cudf input batch size after shuffle
reader
[...]
| spark.gluten.sql.columnar.backend.velox.cudf.concurrentGpuTasks
| âš“ Static | 1 | The number of concurrent GPU tasks to
run.
[...]
| spark.gluten.sql.columnar.backend.velox.cudf.enableTableScan
| âš“ Static | false | Enable cudf table scan
[...]
-| spark.gluten.sql.columnar.backend.velox.cudf.enableValidation
| âš“ Static | true | Heuristics you can apply to validate
a cuDF/GPU plan and only offload when the entire stage can be fully and
profitably executed on GPU
[...]
+| spark.gluten.sql.columnar.backend.velox.cudf.enableValidation
| 🔄 Dynamic | true | Heuristics you can apply to validate a
cuDF/GPU plan and only offload when the entire stage can be fully and
profitably executed on GPU
[...]
| spark.gluten.sql.columnar.backend.velox.cudf.memoryPercent
| âš“ Static | 50 | The initial percent of GPU memory to
allocate for memory resource for one thread.
[...]
| spark.gluten.sql.columnar.backend.velox.cudf.memoryResource
| âš“ Static | async | GPU RMM memory resource.
[...]
| spark.gluten.sql.columnar.backend.velox.cudf.shuffleMaxPrefetchBytes
| 🔄 Dynamic | 1028MB | Maximum bytes to prefetch in CPU
memory during GPU shuffle read while waiting for GPU available.
[...]
diff --git
a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala
b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala
index 28c3ce2dfb..02f7517cc6 100644
---
a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala
+++
b/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
- override def stringArgs: Iterator[Any] = aqeReader.stringArgs
+ override protected def stringArgs: Iterator[Any] = {
+ delegate match {
+ case a: AQEShuffleReadExec => a.stringArgs
+ case _ => super.stringArgs
+ }
+ }
- @transient override lazy val metrics: Map[String, SQLMetric] =
aqeReader.metrics
+ override protected def withNewChildInternal(newChild: SparkPlan):
ColumnarAQEShuffleReadExec = {
+ delegate match {
+ case a: AQEShuffleReadExec => copy(delegate =
a.withNewChildren(Seq(newChild)))
+ case _ => copy(delegate = newChild)
+ }
+ }
- private def isCoalescedSpec(spec: ShufflePartitionSpec) = {
- val method =
classOf[AQEShuffleReadExec].getDeclaredMethod("isCoalescedSpec")
- method.setAccessible(true)
- method.invoke(aqeReader, spec).asInstanceOf[Boolean]
+ 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 node ${delegate.nodeName}.")
+ }
}
+ @transient override lazy val metrics: Map[String, SQLMetric] =
aqeReader.metrics
+
private def shuffleStage = {
val method = classOf[AQEShuffleReadExec].getDeclaredMethod("shuffleStage")
method.setAccessible(true)
@@ -89,7 +99,8 @@ case class ColumnarAQEShuffleReadExec(
private lazy val shuffleRDD: RDD[_] = {
shuffleStage match {
case Some(stage) =>
- if (isAQEShuffleRead) {
+ // Only send driver metrics if it's a wrapper for AQEShuffleRead.
+ if (delegate.isInstanceOf[AQEShuffleReadExec]) {
sendDriverMetrics()
}
stage.shuffle match {
@@ -108,15 +119,4 @@ case class ColumnarAQEShuffleReadExec(
override protected def doExecuteColumnar(): RDD[ColumnarBatch] = {
shuffleRDD.asInstanceOf[RDD[ColumnarBatch]]
}
-
- override protected def withNewChildInternal(newChild: SparkPlan):
ColumnarAQEShuffleReadExec = {
- if (isAQEShuffleRead) {
- copy(delegate =
-
Left(delegate.left.get.withNewChildren(Seq(newChild)).asInstanceOf[AQEShuffleReadExec]))
- } else {
- copy(delegate =
- Right(
-
delegate.right.get.withNewChildren(Seq(newChild)).asInstanceOf[ShuffleQueryStageExec]))
- }
- }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
