[GLUTEN-11524][VL] Set shuffle reader order for gpu async shuffle read - #12918
Draft
marin-ma wants to merge 2 commits into
Draft
[GLUTEN-11524][VL] Set shuffle reader order for gpu async shuffle read#12918marin-ma wants to merge 2 commits into
marin-ma wants to merge 2 commits into
Conversation
Contributor
There was a problem hiding this comment.
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.
This PR introduces a “shuffle reader order” (priority) derived from the Velox pipeline execution order and propagates it through Spark/Gluten shuffle reader plumbing down to the native Velox GPU async shuffle reader, aiming to better align prefetching with actual consumption order.
Changes:
- Add plan-based computation of Velox-first-read order for
ColumnarAQEShuffleReadExecand attach it as a per-reader priority. - Thread
readerOrderthrough Spark shuffle RDD/reader APIs into JNI and nativeShuffleReader::read(...). - Extend Velox shuffle reader interfaces to accept a priority and pass it into the GPU async hash shuffle deserializer.
Reviewed changes
Copilot reviewed 23 out of 23 changed files in this pull request and generated 10 comments.
Show a summary per file
| File | Description |
|---|---|
| gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala | Adds mutable readerOrder, exposes it in explain output, and passes it into shuffle RDD creation. |
| gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ShuffledColumnarBatchRDD.scala | Adds optional readerOrder and forwards it to shuffle manager reader creation. |
| gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarShuffleExchangeExec.scala | Extends getShuffleRDD API to accept a reader order and wrap it as Option. |
| gluten-substrait/src/main/scala/org/apache/spark/shuffle/sort/ColumnarShuffleManager.scala | Threads readerOrder into the columnar shuffle reader generator. |
| gluten-substrait/src/main/scala/org/apache/spark/shuffle/GlutenShuffleUtils.scala | Adds readerOrder to shuffle reader parameter plumbing. |
| gluten-substrait/src/main/scala/org/apache/spark/shuffle/GlutenShuffleReaderWrapper.scala | Extends reader parameters case class to include readerOrder. |
| gluten-substrait/src/main/scala/org/apache/gluten/utils/ShuffleReaderOrderUtil.scala | New utility to simulate Velox pipeline “sweeps” and assign per-shuffle-read order. |
| gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageTransformer.scala | Calls the order assignment before whole-stage transform when offloadCuda is enabled. |
| gluten-arrow/src/main/java/org/apache/gluten/vectorized/ShuffleReaderJniWrapper.java | Extends JNI read signature to include readerOrder. |
| cpp/core/jni/JniWrapper.cc | Forwards readerOrder from JNI into native ShuffleReader::read. |
| cpp/core/shuffle/ShuffleReader.h | Extends shuffle reader interface to accept priority. |
| cpp/velox/shuffle/VeloxShuffleReader.h / .cc | Plumbs priority into deserializer creation and uses it for GPU async reader. |
| cpp/velox/shuffle/VeloxGpuAsyncShuffleReader.h / .cc | Adds priority to GPU async deserializer constructor and stores it. |
| backends-velox/src/main/scala/org/apache/spark/shuffle/ColumnarShuffleReader.scala | Forwards readerOrder into columnar batch deserialization. |
| backends-velox/src/main/scala/org/apache/gluten/vectorized/ColumnarBatchSerializerInstance.scala | Extends API to accept optional readerOrder. |
| backends-velox/src/main/scala/org/apache/gluten/vectorized/ColumnarBatchSerializer.scala | Passes readerOrder into JNI read and adds a log on first use. |
| backends-velox/src/test/scala/org/apache/spark/shuffle/ColumnarShuffleReaderSuite.scala | Updates unit test call sites for the new readerOrder parameter. |
| cpp/velox/tests/VeloxShuffleWriterTest.cc / VeloxGpuShuffleWriterTest.cc / cpp/velox/benchmarks/GenericBenchmark.cc | Updates call sites for new native read(..., priority) signature (plus formatting changes). |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| case _ => super.stringArgs | ||
| } | ||
| } | ||
| } ++ Iterator(s"[order=$readerOrder]") |
Comment on lines
+86
to
+90
| private var readerOrder: Int = 0 | ||
|
|
||
| def setReaderOrder(readerOrder: Int): Unit = { | ||
| this.readerOrder = readerOrder | ||
| } |
Comment on lines
+113
to
+116
| columnarShuffle.getShuffleRDD( | ||
| aqeReader.partitionSpecs.toArray, | ||
| executionMode, | ||
| readerOrder) |
Comment on lines
+67
to
+71
| class PipelineSim { | ||
| var source: Option[ColumnarAQEShuffleReadExec] = None | ||
| val builds = mutable.ArrayBuffer.empty[Int] | ||
| var done = false | ||
| } |
Comment on lines
+108
to
+109
| case c: ColumnarAQEShuffleReadExec => | ||
| pipelines(pipelineIdx).source = Some(c) |
Comment on lines
+115
to
+117
| case other => | ||
| other.children.foreach(planPipelines(_, pipelineIdx)) | ||
| } |
Comment on lines
+129
to
+145
| var readerOrder = 0 | ||
| var progressed = true | ||
| while (progressed && pipelines.exists(!_.done)) { | ||
| progressed = false | ||
| pipelines.foreach { | ||
| p => | ||
| if (!p.done && p.builds.forall(pipelines(_).done)) { | ||
| p.source.foreach { | ||
| reader => | ||
| reader.setReaderOrder(readerOrder) | ||
| readerOrder += 1 | ||
| } | ||
| p.done = true | ||
| progressed = true | ||
| } | ||
| } | ||
| } |
| runtime, | ||
| jniWrapper | ||
| .read(shuffleReaderHandle, streamReader, executionMode.id)) | ||
| .read(shuffleReaderHandle, streamReader, executionMode.id, readerOrder.getOrElse(0))) |
Comment on lines
+193
to
+196
| if (readerOrder.isDefined) { | ||
| logWarning(s"Start reading reader order: ${readerOrder.get}") | ||
| readerOrder = None | ||
| } |
| * last. This method reproduces the pipeline list and replays the sweeps to compute each reader's | ||
| * first-read order without running anything. | ||
| */ | ||
| def assign(stageRoot: SparkPlan): Unit = { |
|
Run Gluten Clickhouse CI on x86 |
marin-ma
marked this pull request as draft
August 28, 2026 10:33
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes are proposed in this pull request?
Currently the async shuffle reader threads fetch and deserialize shuffle data in the order the readers are created in Spark. In a Spark stage this creation order is usually different from the order the Velox pipeline actually reads them: Velox runs the build side of each hash join before the probe side, so build inputs are read first and the main probe input last.
This PR computes the read order from the plan and passes it to the native reader as the prefetch priority.
In theory matching the two orders should help, but no significant perf improvement on TPC-DS q95 so far.
How was this patch tested?
Manually testing on Cuda nodes.
Was this patch authored or co-authored using generative AI tooling?
gluten-substrait/src/main/scala/org/apache/gluten/utils/ShuffleReaderOrderUtil.scalais generated by ClaudeRelated issue: #11524