marin-ma commented on code in PR #12087:
URL: https://github.com/apache/gluten/pull/12087#discussion_r3316329979
##########
backends-velox/src/main/scala/org/apache/gluten/extension/AppendBatchResizeForShuffleInputAndOutput.scala:
##########
@@ -41,65 +42,59 @@ case class AppendBatchResizeForShuffleInputAndOutput()
extends Rule[SparkPlan] {
val range = VeloxConfig.get.veloxResizeBatchesShuffleInputOutputRange
val preferredBatchBytes = VeloxConfig.get.veloxPreferredBatchBytes
+
+ val newPlan = if (resizeBatchesShuffleInputEnabled) {
+ addResizeBatchesForShuffleInput(plan, range.min, range.max,
preferredBatchBytes)
+ } else {
+ plan
+ }
+
+ val resultPlan = if (resizeBatchesShuffleOutputEnabled) {
+ addResizeBatchesForShuffleOutput(newPlan, range.min, range.max,
preferredBatchBytes)
+ } else {
+ newPlan
+ }
+
+ resultPlan
+ }
+
+ private def addResizeBatchesForShuffleInput(
+ plan: SparkPlan,
+ min: Int,
+ max: Int,
+ preferredBatchBytes: Long): SparkPlan = {
plan.transformUp {
case shuffle: ColumnarShuffleExchangeExec
- if resizeBatchesShuffleInputEnabled &&
- shuffle.shuffleWriterType.requiresResizingShuffleInput =>
+ if shuffle.shuffleWriterType.requiresResizingShuffleInput =>
val appendBatches =
- VeloxResizeBatchesExec(shuffle.child, range.min, range.max,
preferredBatchBytes)
+ VeloxResizeBatchesExec(shuffle.child, min, max, preferredBatchBytes)
shuffle.withNewChildren(Seq(appendBatches))
- case a @ AQEShuffleReadExec(
- ShuffleQueryStageExec(_, shuffle: ColumnarShuffleExchangeExec, _),
- _)
- if resizeBatchesShuffleOutputEnabled &&
- shuffle.shuffleWriterType.requiresResizingShuffleOutput =>
- VeloxResizeBatchesExec(a, range.min, range.max, preferredBatchBytes)
- case a @ AQEShuffleReadExec(
- ShuffleQueryStageExec(
- _,
- ReusedExchangeExec(_, shuffle: ColumnarShuffleExchangeExec),
- _),
- _)
- if resizeBatchesShuffleOutputEnabled &&
- shuffle.shuffleWriterType.requiresResizingShuffleOutput =>
- VeloxResizeBatchesExec(a, range.min, range.max, preferredBatchBytes)
- // Since it's transformed in a bottom to up order, so we may first
encounter
- // ShuffeQueryStageExec, which is transformed to
VeloxResizeBatchesExec(ShuffeQueryStageExec),
- // then we see AQEShuffleReadExec
- case a @ AQEShuffleReadExec(
- VeloxResizeBatchesExec(
- s @ ShuffleQueryStageExec(_, shuffle:
ColumnarShuffleExchangeExec, _),
- _,
- _,
- _),
- _)
- if resizeBatchesShuffleOutputEnabled &&
- shuffle.shuffleWriterType.requiresResizingShuffleOutput =>
- VeloxResizeBatchesExec(a.copy(child = s), range.min, range.max,
preferredBatchBytes)
- case a @ AQEShuffleReadExec(
- VeloxResizeBatchesExec(
- s @ ShuffleQueryStageExec(
- _,
- ReusedExchangeExec(_, shuffle: ColumnarShuffleExchangeExec),
- _),
- _,
- _,
- _),
- _)
- if resizeBatchesShuffleOutputEnabled &&
- shuffle.shuffleWriterType.requiresResizingShuffleOutput =>
- VeloxResizeBatchesExec(a.copy(child = s), range.min, range.max,
preferredBatchBytes)
- case s @ ShuffleQueryStageExec(_, shuffle: ColumnarShuffleExchangeExec,
_)
- if resizeBatchesShuffleOutputEnabled &&
- shuffle.shuffleWriterType.requiresResizingShuffleOutput =>
- VeloxResizeBatchesExec(s, range.min, range.max, preferredBatchBytes)
- case s @ ShuffleQueryStageExec(
- _,
- ReusedExchangeExec(_, shuffle: ColumnarShuffleExchangeExec),
- _)
- if resizeBatchesShuffleOutputEnabled &&
- shuffle.shuffleWriterType.requiresResizingShuffleOutput =>
- VeloxResizeBatchesExec(s, range.min, range.max, preferredBatchBytes)
+ }
+ }
+
+ private def addResizeBatchesForShuffleOutput(
+ plan: SparkPlan,
+ min: Int,
+ max: Int,
+ preferredBatchBytes: Long): SparkPlan = {
+ plan match {
+ case s: ShuffleQueryStageExec if requiresResizingShuffleOutput(s) =>
+ VeloxResizeBatchesExec(s, min, max, preferredBatchBytes)
+ case a @ AQEShuffleReadExec(s @ ShuffleQueryStageExec(_, _, _), _)
+ if requiresResizingShuffleOutput(s) =>
+ VeloxResizeBatchesExec(a, min, max, preferredBatchBytes)
+ case other =>
Review Comment:
@wForget The normal entry point for the root node is `case other`. It
traverses the whole plan until a AQE query stage leaf node (plain
`ShuffleQueryStageExec`, or `AQEShuffleReadExec` wrapped `AQEShuffleReadExec`)
is found.
--
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]
