Copilot commented on code in PR #12918:
URL: https://github.com/apache/gluten/pull/12918#discussion_r3873408515
##########
gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala:
##########
@@ -57,7 +57,7 @@ case class ColumnarAQEShuffleReadExec(
case a: AQEShuffleReadExec => a.stringArgs
case _ => super.stringArgs
}
- }
+ } ++ Iterator(s"[order=$readerOrder]")
Review Comment:
`ColumnarAQEShuffleReadExec` is being made mutable via `private var
readerOrder` + `setReaderOrder`, and `stringArgs` will always show `[order=0]`
unless the util runs. This is both misleading in explain output and makes the
plan node stateful. Consider representing this as `Option[Int]` (unset by
default) and only appending to `stringArgs` / passing down to shuffle when
defined, so non-GPU/non-assigned plans don’t accidentally look like they have a
real order (and you avoid relying on a mutable field on a `SparkPlan`).
##########
gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala:
##########
@@ -105,7 +110,10 @@ case class ColumnarAQEShuffleReadExec(
}
stage.shuffle match {
case columnarShuffle: ColumnarShuffleExchangeExec =>
- columnarShuffle.getShuffleRDD(aqeReader.partitionSpecs.toArray,
executionMode)
+ columnarShuffle.getShuffleRDD(
+ aqeReader.partitionSpecs.toArray,
+ executionMode,
+ readerOrder)
Review Comment:
`ColumnarAQEShuffleReadExec` is being made mutable via `private var
readerOrder` + `setReaderOrder`, and `stringArgs` will always show `[order=0]`
unless the util runs. This is both misleading in explain output and makes the
plan node stateful. Consider representing this as `Option[Int]` (unset by
default) and only appending to `stringArgs` / passing down to shuffle when
defined, so non-GPU/non-assigned plans don’t accidentally look like they have a
real order (and you avoid relying on a mutable field on a `SparkPlan`).
##########
gluten-substrait/src/main/scala/org/apache/spark/sql/execution/adaptive/ColumnarAQEShuffleReadExec.scala:
##########
@@ -81,9 +81,14 @@ case class ColumnarAQEShuffleReadExec(
s"Cannot get aqeReader from delegate node ${delegate.nodeName}.")
}
}
-
@transient override lazy val metrics: Map[String, SQLMetric] =
aqeReader.metrics
+ private var readerOrder: Int = 0
+
+ def setReaderOrder(readerOrder: Int): Unit = {
+ this.readerOrder = readerOrder
+ }
Review Comment:
`ColumnarAQEShuffleReadExec` is being made mutable via `private var
readerOrder` + `setReaderOrder`, and `stringArgs` will always show `[order=0]`
unless the util runs. This is both misleading in explain output and makes the
plan node stateful. Consider representing this as `Option[Int]` (unset by
default) and only appending to `stringArgs` / passing down to shuffle when
defined, so non-GPU/non-assigned plans don’t accidentally look like they have a
real order (and you avoid relying on a mutable field on a `SparkPlan`).
##########
gluten-substrait/src/main/scala/org/apache/gluten/utils/ShuffleReaderOrderUtil.scala:
##########
@@ -0,0 +1,147 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.gluten.utils
+
+import org.apache.gluten.execution.{BroadcastHashJoinExecTransformerBase,
ShuffledHashJoinExecTransformerBase}
+
+import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight}
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.adaptive.ColumnarAQEShuffleReadExec
+
+import scala.collection.mutable
+
+/**
+ * Computes the order in which Velox will first pull data from each shuffle
reader of a whole stage.
+ */
+object ShuffleReaderOrderUtil {
+
+ /**
+ * Assigns each [[ColumnarAQEShuffleReadExec]] in the stage rooted at
`stageRoot` the order in
+ * which Velox will first pull data from it at runtime (0-based, via
`setReaderOrder`).
+ *
+ * Why the order is not simply the plan-tree order:
+ *
+ * When Gluten hands a whole stage to Velox, Velox's LocalPlanner chops the
operator tree into
+ * "pipelines" (linear chains of operators driven by one source). The cut
points are hash-join
+ * build sides: the probe (streamed) side of a join stays in the current
pipeline, while the build
+ * side becomes a NEW pipeline, appended to the pipeline list at the moment
the planner's
+ * depth-first, probe-side-first walk reaches it. Every pipeline therefore
has one source (a
+ * shuffle read, a scan, or a broadcast input) and may depend on the build
pipelines feeding the
+ * hash joins it contains.
+ *
+ * Gluten then executes the Velox task single-threaded via Task::next(). The
single thread
+ * repeatedly walks the pipeline list from index 0 upward — each full
top-to-bottom walk is called
+ * a "sweep" below. On each sweep, a pipeline whose hash-join build inputs
are not all finished is
+ * blocked (HashProbe reports kWaitForJoinBuild before consuming any input,
so even the pipeline's
+ * source is not touched) and gets skipped; an unblocked pipeline runs,
pulling its source — that
+ * is the moment its shuffle reader is first invoked. A pipeline finishing
mid-sweep is visible to
+ * later-indexed pipelines within the same sweep, but earlier-indexed
pipelines only notice on the
+ * next sweep.
+ *
+ * Net effect: all independent build-side shuffle readers fire first (in
pipeline-list order),
+ * then intermediate probe readers as their builds complete, and the
top-level probe reader fires
+ * 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 = {
+ // One entry per Velox pipeline:
+ // - source: the shuffle reader at the bottom of the pipeline's operator
chain, if any
+ // (pipelines fed by scans or broadcast inputs have none and get no
order number);
+ // - builds: indices of the build pipelines that must finish before this
pipeline may run
+ // (one per hash join contained in this pipeline);
+ // - done: completion flag used while replaying the sweeps.
+ class PipelineSim {
+ var source: Option[ColumnarAQEShuffleReadExec] = None
+ val builds = mutable.ArrayBuffer.empty[Int]
+ var done = false
+ }
+
+ // Ordered as Velox's LocalPlanner orders its driver factories; the index
in this buffer is
+ // the pipeline index the executor sweeps over.
+ val pipelines = mutable.ArrayBuffer.empty[PipelineSim]
+
+ // Walk the plan tree and rebuild the pipeline list, mirroring
LocalPlanner: `pipelineIdx`
+ // is the pipeline the current node belongs to; non-join nodes just stay
in it, joins split.
+ def planPipelines(plan: SparkPlan, pipelineIdx: Int): Unit = {
+ // The probe side is walked first and stays in the current pipeline (so
any joins nested
+ // inside it append THEIR build pipelines before this one); then the
build side is placed
+ // in a fresh pipeline and recorded as a prerequisite of the current
pipeline.
+ def splitBuildPipeline(probe: SparkPlan, build: SparkPlan): Unit = {
+ planPipelines(probe, pipelineIdx)
+ val buildIdx = pipelines.size
+ pipelines += new PipelineSim
+ planPipelines(build, buildIdx)
+ pipelines(pipelineIdx).builds += buildIdx
+ }
+ plan match {
+ // Broadcast builds contain no shuffle reader, but still get a
pipeline: a probe blocked
+ // only on a broadcast build must still wait one sweep, which can
delay its shuffle read
+ // past readers that fire on the first sweep.
+ case bhj: BroadcastHashJoinExecTransformerBase =>
+ bhj.joinBuildSide match {
+ case BuildLeft => splitBuildPipeline(bhj.right, bhj.left)
+ case BuildRight => splitBuildPipeline(bhj.left, bhj.right)
+ }
+
+ // For BuildLeft, Gluten swaps the children when lowering to Velox's
HashJoinNode (whose
+ // build side is always the right source), so the left child is the
build pipeline here.
+ case shj: ShuffledHashJoinExecTransformerBase =>
+ shj.joinBuildSide match {
+ case BuildLeft => splitBuildPipeline(shj.right, shj.left)
+ case BuildRight => splitBuildPipeline(shj.left, shj.right)
+ }
+
+ case c: ColumnarAQEShuffleReadExec =>
+ pipelines(pipelineIdx).source = Some(c)
Review Comment:
`PipelineSim` can hold only a single `source`, but `planPipelines` may visit
multiple `ColumnarAQEShuffleReadExec` nodes within the same `pipelineIdx`
(e.g., operators with multiple children that aren’t modeled as pipeline split
points, or unexpected physical operators). In that case, later visits overwrite
earlier ones and some readers never get an order assigned. A tangible fix is to
track `sources: Seq[ColumnarAQEShuffleReadExec]` (or at minimum detect a second
source and throw/log a clear warning) to avoid silently dropping readers.
##########
backends-velox/src/main/scala/org/apache/gluten/vectorized/ColumnarBatchSerializer.scala:
##########
@@ -137,28 +137,30 @@ private class ColumnarBatchSerializerInstanceImpl(
// `deserializeStream` is currently still used by uniffle shuffle reader.
override def deserializeStream(in: InputStream): DeserializationStream = {
- new TaskDeserializationStream(Iterator((null, in)), None, CPUStageMode)
+ new TaskDeserializationStream(Iterator((null, in)), None, CPUStageMode,
None)
}
def deserializeStreams(
streams: Iterator[(BlockId, InputStream)],
onComplete: () => Unit,
- executionMode: StageExecutionMode = CPUStageMode): DeserializationStream
= {
- new TaskDeserializationStream(streams, Some(onComplete), executionMode)
+ executionMode: StageExecutionMode = CPUStageMode,
+ readerOrder: Option[Int] = None): DeserializationStream = {
+ new TaskDeserializationStream(streams, Some(onComplete), executionMode,
readerOrder)
}
private class TaskDeserializationStream(
streams: Iterator[(BlockId, InputStream)],
onComplete: Option[() => Unit],
- executionMode: StageExecutionMode)
+ executionMode: StageExecutionMode,
+ var readerOrder: Option[Int])
extends DeserializationStream
with TaskResource {
private val streamReader = ShuffleStreamReader(streams)
private val wrappedOut: ClosableIterator[ColumnarBatch] = new
ColumnarBatchOutIterator(
runtime,
jniWrapper
- .read(shuffleReaderHandle, streamReader, executionMode.id))
+ .read(shuffleReaderHandle, streamReader, executionMode.id,
readerOrder.getOrElse(0)))
Review Comment:
This introduces a `logWarning` on the shuffle read hot-path (once per
`TaskDeserializationStream`), which can spam warning logs and add overhead at
scale. Since this appears to be debug/tracing, consider removing it or
downgrading to `logDebug` (and ideally avoid mutating `readerOrder` solely to
control logging).
##########
backends-velox/src/main/scala/org/apache/gluten/vectorized/ColumnarBatchSerializer.scala:
##########
@@ -188,6 +190,10 @@ private class ColumnarBatchSerializerInstanceImpl(
@throws(classOf[EOFException])
override def readValue[T: ClassTag](): T = {
+ if (readerOrder.isDefined) {
+ logWarning(s"Start reading reader order: ${readerOrder.get}")
+ readerOrder = None
+ }
Review Comment:
This introduces a `logWarning` on the shuffle read hot-path (once per
`TaskDeserializationStream`), which can spam warning logs and add overhead at
scale. Since this appears to be debug/tracing, consider removing it or
downgrading to `logDebug` (and ideally avoid mutating `readerOrder` solely to
control logging).
##########
gluten-substrait/src/main/scala/org/apache/gluten/utils/ShuffleReaderOrderUtil.scala:
##########
@@ -0,0 +1,147 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.gluten.utils
+
+import org.apache.gluten.execution.{BroadcastHashJoinExecTransformerBase,
ShuffledHashJoinExecTransformerBase}
+
+import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight}
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.adaptive.ColumnarAQEShuffleReadExec
+
+import scala.collection.mutable
+
+/**
+ * Computes the order in which Velox will first pull data from each shuffle
reader of a whole stage.
+ */
+object ShuffleReaderOrderUtil {
+
+ /**
+ * Assigns each [[ColumnarAQEShuffleReadExec]] in the stage rooted at
`stageRoot` the order in
+ * which Velox will first pull data from it at runtime (0-based, via
`setReaderOrder`).
+ *
+ * Why the order is not simply the plan-tree order:
+ *
+ * When Gluten hands a whole stage to Velox, Velox's LocalPlanner chops the
operator tree into
+ * "pipelines" (linear chains of operators driven by one source). The cut
points are hash-join
+ * build sides: the probe (streamed) side of a join stays in the current
pipeline, while the build
+ * side becomes a NEW pipeline, appended to the pipeline list at the moment
the planner's
+ * depth-first, probe-side-first walk reaches it. Every pipeline therefore
has one source (a
+ * shuffle read, a scan, or a broadcast input) and may depend on the build
pipelines feeding the
+ * hash joins it contains.
+ *
+ * Gluten then executes the Velox task single-threaded via Task::next(). The
single thread
+ * repeatedly walks the pipeline list from index 0 upward — each full
top-to-bottom walk is called
+ * a "sweep" below. On each sweep, a pipeline whose hash-join build inputs
are not all finished is
+ * blocked (HashProbe reports kWaitForJoinBuild before consuming any input,
so even the pipeline's
+ * source is not touched) and gets skipped; an unblocked pipeline runs,
pulling its source — that
+ * is the moment its shuffle reader is first invoked. A pipeline finishing
mid-sweep is visible to
+ * later-indexed pipelines within the same sweep, but earlier-indexed
pipelines only notice on the
+ * next sweep.
+ *
+ * Net effect: all independent build-side shuffle readers fire first (in
pipeline-list order),
+ * then intermediate probe readers as their builds complete, and the
top-level probe reader fires
+ * 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 = {
+ // One entry per Velox pipeline:
+ // - source: the shuffle reader at the bottom of the pipeline's operator
chain, if any
+ // (pipelines fed by scans or broadcast inputs have none and get no
order number);
+ // - builds: indices of the build pipelines that must finish before this
pipeline may run
+ // (one per hash join contained in this pipeline);
+ // - done: completion flag used while replaying the sweeps.
+ class PipelineSim {
+ var source: Option[ColumnarAQEShuffleReadExec] = None
+ val builds = mutable.ArrayBuffer.empty[Int]
+ var done = false
+ }
+
+ // Ordered as Velox's LocalPlanner orders its driver factories; the index
in this buffer is
+ // the pipeline index the executor sweeps over.
+ val pipelines = mutable.ArrayBuffer.empty[PipelineSim]
+
+ // Walk the plan tree and rebuild the pipeline list, mirroring
LocalPlanner: `pipelineIdx`
+ // is the pipeline the current node belongs to; non-join nodes just stay
in it, joins split.
+ def planPipelines(plan: SparkPlan, pipelineIdx: Int): Unit = {
+ // The probe side is walked first and stays in the current pipeline (so
any joins nested
+ // inside it append THEIR build pipelines before this one); then the
build side is placed
+ // in a fresh pipeline and recorded as a prerequisite of the current
pipeline.
+ def splitBuildPipeline(probe: SparkPlan, build: SparkPlan): Unit = {
+ planPipelines(probe, pipelineIdx)
+ val buildIdx = pipelines.size
+ pipelines += new PipelineSim
+ planPipelines(build, buildIdx)
+ pipelines(pipelineIdx).builds += buildIdx
+ }
+ plan match {
+ // Broadcast builds contain no shuffle reader, but still get a
pipeline: a probe blocked
+ // only on a broadcast build must still wait one sweep, which can
delay its shuffle read
+ // past readers that fire on the first sweep.
+ case bhj: BroadcastHashJoinExecTransformerBase =>
+ bhj.joinBuildSide match {
+ case BuildLeft => splitBuildPipeline(bhj.right, bhj.left)
+ case BuildRight => splitBuildPipeline(bhj.left, bhj.right)
+ }
+
+ // For BuildLeft, Gluten swaps the children when lowering to Velox's
HashJoinNode (whose
+ // build side is always the right source), so the left child is the
build pipeline here.
+ case shj: ShuffledHashJoinExecTransformerBase =>
+ shj.joinBuildSide match {
+ case BuildLeft => splitBuildPipeline(shj.right, shj.left)
+ case BuildRight => splitBuildPipeline(shj.left, shj.right)
+ }
+
+ case c: ColumnarAQEShuffleReadExec =>
+ pipelines(pipelineIdx).source = Some(c)
+
+ // Any other operator (project, filter, aggregate, input iterator,
...) is a single-input
+ // link in the current pipeline's chain. NOTE: operators that Velox
would also split into
+ // extra pipelines (union/local exchange, nested-loop join, ...) are
not modeled; if one
+ // appears in a stage, the computed order may not match the runtime
order.
+ case other =>
+ other.children.foreach(planPipelines(_, pipelineIdx))
+ }
Review Comment:
`PipelineSim` can hold only a single `source`, but `planPipelines` may visit
multiple `ColumnarAQEShuffleReadExec` nodes within the same `pipelineIdx`
(e.g., operators with multiple children that aren’t modeled as pipeline split
points, or unexpected physical operators). In that case, later visits overwrite
earlier ones and some readers never get an order assigned. A tangible fix is to
track `sources: Seq[ColumnarAQEShuffleReadExec]` (or at minimum detect a second
source and throw/log a clear warning) to avoid silently dropping readers.
##########
gluten-substrait/src/main/scala/org/apache/gluten/utils/ShuffleReaderOrderUtil.scala:
##########
@@ -0,0 +1,147 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.gluten.utils
+
+import org.apache.gluten.execution.{BroadcastHashJoinExecTransformerBase,
ShuffledHashJoinExecTransformerBase}
+
+import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight}
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.adaptive.ColumnarAQEShuffleReadExec
+
+import scala.collection.mutable
+
+/**
+ * Computes the order in which Velox will first pull data from each shuffle
reader of a whole stage.
+ */
+object ShuffleReaderOrderUtil {
+
+ /**
+ * Assigns each [[ColumnarAQEShuffleReadExec]] in the stage rooted at
`stageRoot` the order in
+ * which Velox will first pull data from it at runtime (0-based, via
`setReaderOrder`).
+ *
+ * Why the order is not simply the plan-tree order:
+ *
+ * When Gluten hands a whole stage to Velox, Velox's LocalPlanner chops the
operator tree into
+ * "pipelines" (linear chains of operators driven by one source). The cut
points are hash-join
+ * build sides: the probe (streamed) side of a join stays in the current
pipeline, while the build
+ * side becomes a NEW pipeline, appended to the pipeline list at the moment
the planner's
+ * depth-first, probe-side-first walk reaches it. Every pipeline therefore
has one source (a
+ * shuffle read, a scan, or a broadcast input) and may depend on the build
pipelines feeding the
+ * hash joins it contains.
+ *
+ * Gluten then executes the Velox task single-threaded via Task::next(). The
single thread
+ * repeatedly walks the pipeline list from index 0 upward — each full
top-to-bottom walk is called
+ * a "sweep" below. On each sweep, a pipeline whose hash-join build inputs
are not all finished is
+ * blocked (HashProbe reports kWaitForJoinBuild before consuming any input,
so even the pipeline's
+ * source is not touched) and gets skipped; an unblocked pipeline runs,
pulling its source — that
+ * is the moment its shuffle reader is first invoked. A pipeline finishing
mid-sweep is visible to
+ * later-indexed pipelines within the same sweep, but earlier-indexed
pipelines only notice on the
+ * next sweep.
+ *
+ * Net effect: all independent build-side shuffle readers fire first (in
pipeline-list order),
+ * then intermediate probe readers as their builds complete, and the
top-level probe reader fires
+ * 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 = {
+ // One entry per Velox pipeline:
+ // - source: the shuffle reader at the bottom of the pipeline's operator
chain, if any
+ // (pipelines fed by scans or broadcast inputs have none and get no
order number);
+ // - builds: indices of the build pipelines that must finish before this
pipeline may run
+ // (one per hash join contained in this pipeline);
+ // - done: completion flag used while replaying the sweeps.
+ class PipelineSim {
+ var source: Option[ColumnarAQEShuffleReadExec] = None
+ val builds = mutable.ArrayBuffer.empty[Int]
+ var done = false
+ }
Review Comment:
`PipelineSim` can hold only a single `source`, but `planPipelines` may visit
multiple `ColumnarAQEShuffleReadExec` nodes within the same `pipelineIdx`
(e.g., operators with multiple children that aren’t modeled as pipeline split
points, or unexpected physical operators). In that case, later visits overwrite
earlier ones and some readers never get an order assigned. A tangible fix is to
track `sources: Seq[ColumnarAQEShuffleReadExec]` (or at minimum detect a second
source and throw/log a clear warning) to avoid silently dropping readers.
##########
gluten-substrait/src/main/scala/org/apache/gluten/utils/ShuffleReaderOrderUtil.scala:
##########
@@ -0,0 +1,147 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.gluten.utils
+
+import org.apache.gluten.execution.{BroadcastHashJoinExecTransformerBase,
ShuffledHashJoinExecTransformerBase}
+
+import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight}
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.adaptive.ColumnarAQEShuffleReadExec
+
+import scala.collection.mutable
+
+/**
+ * Computes the order in which Velox will first pull data from each shuffle
reader of a whole stage.
+ */
+object ShuffleReaderOrderUtil {
+
+ /**
+ * Assigns each [[ColumnarAQEShuffleReadExec]] in the stage rooted at
`stageRoot` the order in
+ * which Velox will first pull data from it at runtime (0-based, via
`setReaderOrder`).
+ *
+ * Why the order is not simply the plan-tree order:
+ *
+ * When Gluten hands a whole stage to Velox, Velox's LocalPlanner chops the
operator tree into
+ * "pipelines" (linear chains of operators driven by one source). The cut
points are hash-join
+ * build sides: the probe (streamed) side of a join stays in the current
pipeline, while the build
+ * side becomes a NEW pipeline, appended to the pipeline list at the moment
the planner's
+ * depth-first, probe-side-first walk reaches it. Every pipeline therefore
has one source (a
+ * shuffle read, a scan, or a broadcast input) and may depend on the build
pipelines feeding the
+ * hash joins it contains.
+ *
+ * Gluten then executes the Velox task single-threaded via Task::next(). The
single thread
+ * repeatedly walks the pipeline list from index 0 upward — each full
top-to-bottom walk is called
+ * a "sweep" below. On each sweep, a pipeline whose hash-join build inputs
are not all finished is
+ * blocked (HashProbe reports kWaitForJoinBuild before consuming any input,
so even the pipeline's
+ * source is not touched) and gets skipped; an unblocked pipeline runs,
pulling its source — that
+ * is the moment its shuffle reader is first invoked. A pipeline finishing
mid-sweep is visible to
+ * later-indexed pipelines within the same sweep, but earlier-indexed
pipelines only notice on the
+ * next sweep.
+ *
+ * Net effect: all independent build-side shuffle readers fire first (in
pipeline-list order),
+ * then intermediate probe readers as their builds complete, and the
top-level probe reader fires
+ * 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 = {
+ // One entry per Velox pipeline:
+ // - source: the shuffle reader at the bottom of the pipeline's operator
chain, if any
+ // (pipelines fed by scans or broadcast inputs have none and get no
order number);
+ // - builds: indices of the build pipelines that must finish before this
pipeline may run
+ // (one per hash join contained in this pipeline);
+ // - done: completion flag used while replaying the sweeps.
+ class PipelineSim {
+ var source: Option[ColumnarAQEShuffleReadExec] = None
+ val builds = mutable.ArrayBuffer.empty[Int]
+ var done = false
+ }
+
+ // Ordered as Velox's LocalPlanner orders its driver factories; the index
in this buffer is
+ // the pipeline index the executor sweeps over.
+ val pipelines = mutable.ArrayBuffer.empty[PipelineSim]
+
+ // Walk the plan tree and rebuild the pipeline list, mirroring
LocalPlanner: `pipelineIdx`
+ // is the pipeline the current node belongs to; non-join nodes just stay
in it, joins split.
+ def planPipelines(plan: SparkPlan, pipelineIdx: Int): Unit = {
+ // The probe side is walked first and stays in the current pipeline (so
any joins nested
+ // inside it append THEIR build pipelines before this one); then the
build side is placed
+ // in a fresh pipeline and recorded as a prerequisite of the current
pipeline.
+ def splitBuildPipeline(probe: SparkPlan, build: SparkPlan): Unit = {
+ planPipelines(probe, pipelineIdx)
+ val buildIdx = pipelines.size
+ pipelines += new PipelineSim
+ planPipelines(build, buildIdx)
+ pipelines(pipelineIdx).builds += buildIdx
+ }
+ plan match {
+ // Broadcast builds contain no shuffle reader, but still get a
pipeline: a probe blocked
+ // only on a broadcast build must still wait one sweep, which can
delay its shuffle read
+ // past readers that fire on the first sweep.
+ case bhj: BroadcastHashJoinExecTransformerBase =>
+ bhj.joinBuildSide match {
+ case BuildLeft => splitBuildPipeline(bhj.right, bhj.left)
+ case BuildRight => splitBuildPipeline(bhj.left, bhj.right)
+ }
+
+ // For BuildLeft, Gluten swaps the children when lowering to Velox's
HashJoinNode (whose
+ // build side is always the right source), so the left child is the
build pipeline here.
+ case shj: ShuffledHashJoinExecTransformerBase =>
+ shj.joinBuildSide match {
+ case BuildLeft => splitBuildPipeline(shj.right, shj.left)
+ case BuildRight => splitBuildPipeline(shj.left, shj.right)
+ }
+
+ case c: ColumnarAQEShuffleReadExec =>
+ pipelines(pipelineIdx).source = Some(c)
+
+ // Any other operator (project, filter, aggregate, input iterator,
...) is a single-input
+ // link in the current pipeline's chain. NOTE: operators that Velox
would also split into
+ // extra pipelines (union/local exchange, nested-loop join, ...) are
not modeled; if one
+ // appears in a stage, the computed order may not match the runtime
order.
+ case other =>
+ other.children.foreach(planPipelines(_, pipelineIdx))
+ }
+ }
+
+ // Pipeline 0 is the output pipeline: the chain from the stage root down
its probe sides.
+ pipelines += new PipelineSim
+ planPipelines(stageRoot, 0)
+
+ // Replay the serial executor: each `while` iteration is one sweep over
the pipeline list.
+ // A pipeline runs once all its build prerequisites are done; running it
assigns the next
+ // order number to its shuffle reader (its source is pulled to exhaustion
at that point).
+ // Marking `done` mid-sweep lets later-indexed pipelines run in the same
sweep, matching the
+ // executor's forward-only scan. `progressed` guards against a malformed
dependency graph.
+ 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
+ }
+ }
+ }
Review Comment:
If `progressed` becomes false while some pipelines are still not `done`, the
method exits silently and leaves some readers at the default order (currently
`0`). This makes mis-ordering hard to diagnose in production. Consider adding
an explicit post-condition check (e.g., if any pipelines remain not done, log a
warning with enough context or fail fast in debug mode) so ordering issues
don’t silently degrade behavior.
##########
gluten-substrait/src/main/scala/org/apache/gluten/utils/ShuffleReaderOrderUtil.scala:
##########
@@ -0,0 +1,147 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.gluten.utils
+
+import org.apache.gluten.execution.{BroadcastHashJoinExecTransformerBase,
ShuffledHashJoinExecTransformerBase}
+
+import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight}
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.adaptive.ColumnarAQEShuffleReadExec
+
+import scala.collection.mutable
+
+/**
+ * Computes the order in which Velox will first pull data from each shuffle
reader of a whole stage.
+ */
+object ShuffleReaderOrderUtil {
+
+ /**
+ * Assigns each [[ColumnarAQEShuffleReadExec]] in the stage rooted at
`stageRoot` the order in
+ * which Velox will first pull data from it at runtime (0-based, via
`setReaderOrder`).
+ *
+ * Why the order is not simply the plan-tree order:
+ *
+ * When Gluten hands a whole stage to Velox, Velox's LocalPlanner chops the
operator tree into
+ * "pipelines" (linear chains of operators driven by one source). The cut
points are hash-join
+ * build sides: the probe (streamed) side of a join stays in the current
pipeline, while the build
+ * side becomes a NEW pipeline, appended to the pipeline list at the moment
the planner's
+ * depth-first, probe-side-first walk reaches it. Every pipeline therefore
has one source (a
+ * shuffle read, a scan, or a broadcast input) and may depend on the build
pipelines feeding the
+ * hash joins it contains.
+ *
+ * Gluten then executes the Velox task single-threaded via Task::next(). The
single thread
+ * repeatedly walks the pipeline list from index 0 upward — each full
top-to-bottom walk is called
+ * a "sweep" below. On each sweep, a pipeline whose hash-join build inputs
are not all finished is
+ * blocked (HashProbe reports kWaitForJoinBuild before consuming any input,
so even the pipeline's
+ * source is not touched) and gets skipped; an unblocked pipeline runs,
pulling its source — that
+ * is the moment its shuffle reader is first invoked. A pipeline finishing
mid-sweep is visible to
+ * later-indexed pipelines within the same sweep, but earlier-indexed
pipelines only notice on the
+ * next sweep.
+ *
+ * Net effect: all independent build-side shuffle readers fire first (in
pipeline-list order),
+ * then intermediate probe readers as their builds complete, and the
top-level probe reader fires
+ * 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 = {
Review Comment:
The new plan simulation logic for reader ordering is non-trivial (hash join
build/probe handling, broadcast vs shuffled join, sweep replay) and currently
isn’t covered by unit tests. Adding Scala tests that build small synthetic
`SparkPlan` trees (BHJ/SHJ with both `BuildLeft` and `BuildRight`, nested
joins, and a case with an unsupported multi-child operator) would help prevent
regressions and validate assumptions about ordering.
--
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]
