ulysses-you opened a new pull request, #7539:
URL: https://github.com/apache/kyuubi/pull/7539
## Why are the changes needed?
`RebalanceBeforeWriting` injects a `RebalancePartitions` operator before
writes (physically a
`ShuffleExchangeExec` with `shuffleOrigin = REBALANCE_PARTITIONS_BY_NONE`)
to merge small files and
avoid the small-file problem. However, in some cases this rebalance shuffle
is pure overhead: once
the upstream data is materialized and its real size is known, the shuffle
brings no benefit and only
adds an extra round of shuffle read/write.
This PR adds an AQE runtime optimizer rule `RemoveRebalanceShuffle` that, at
AQE time (when the
upstream query stage is materialized and real statistics are available),
decides whether such a
not-worthwhile rebalance shuffle can be removed, saving one shuffle.
## What changes were proposed in this pull request?
A new rule `RemoveRebalanceShuffle`, injected into the `AQEOptimizer` via
`injectRuntimeOptimizerRule`, runs on the **logical plan**: it replaces
`RebalancePartitions` with
its child, and the physical planner then re-plans without the shuffle.
**Scope** — only the rebalance feeding the final write is considered:
starting from `WriteFiles`,
the rule walks down through at most `Project` and local `Sort`; any other
operator in between leaves
the rebalance untouched.
**Removal conditions** (all must hold):
1. `spark.sql.adaptive.removeRebalanceShuffle.enabled` is on;
2. the `RebalancePartitions` has no partition expressions (i.e.
`REBALANCE_PARTITIONS_BY_NONE`);
3. its `optAdvisoryPartitionSize` is present and **larger than**
`spark.sql.adaptive.advisoryPartitionSizeInBytes` (a smaller advisory
size sugge
intentionally wants finer partitions, so we respect that).
**Size decision** — the rebalance input is split into groups of sub-plans
whose leaves are all
materialized `LogicalQueryStage` exchange stages (a `Union` yields one group
per br
measured by its runtime `sizeInBytes`. The shuffle is removed only when
every group is removable,
where a group is removable under either:
- **Large data**: the input has no data-reducing operator and the group's
represent
`max(maxStageSize, groupSum / 2)` is larger than
`spark.sql.shuffle.partitions *
removeRebalanceShuffle.smallPartitionSize`; or
- **Small data**: the input has no data-expanding operator and the
representative size is smaller
than `advisoryPartitionSizeInBytes *
removeRebalanceShuffle.tolerableSmallFileNum
Operator classification (logical operators):
- reducing: `Aggregate` / `Filter` / `Sample` / `Offset` / `GlobalLimit` /
`LocalLi
`WindowGroupLimit` / `Inner` join / semi-anti join;
- expanding: `Generate` / `Expand` / non-semi-anti `Join`.
**New configs** (`KyuubiSQLConf`, version 1.12.0):
| Config | Default |
|---|---|
| `spark.sql.adaptive.removeRebalanceShuffle.enabled` | `false` |
| `spark.sql.adaptive.removeRebalanceShuffle.smallPartitionSize` | `96MB` |
| `spark.sql.adaptive.removeRebalanceShuffle.tolerableSmallFileNum` | `5` |
Modules changed (identical across 3.5 / 4.0 / 4.1):
- Add `RemoveRebalanceShuffle.scala` and `RemoveRebalanceShuffleSuite.scala`;
- Add 3 configs in `KyuubiSQLConf.scala`;
- Register the rule in `KyuubiSparkSQLExtension.scala`;
- Document the feature and configs in
`docs/extensions/engines/spark/rules.md`.
## How was this patch tested?
Added `RemoveRebalanceShuffleSuite` (3.5 / 4.0 / 4.1), which asserts on the
presence of
`REBALANCE_PARTITIONS_BY_NONE` shuffle in the executed plan, covering:
- semi join (reducing), small data → removed;
- outer join (expanding, non-reducing), large data with tiny / huge
`smallPartitionSize` → removed /
kept respectively;
- rebalance advisory size ≤ session advisory size → kept;
- inner join (both reducing and expanding) → kept;
- semi join with `tolerableSmallFileNum=0` → kept;
- switch off → kept.
All pass on Spark 3.5:
```
build/mvn -Pspark-3.5 -pl extensions/spark/kyuubi-extension-spark-3-5 -am
test \
-Dtest=none
-DwildcardSuites=org.apache.spark.sql.RemoveRebalanceShuffleSuite
```
No regression in `RebalanceBeforeWritingSuite`. Formatted with
`spotless:apply` across all Spark
profiles.
### Was this patch authored or co-authored using generative AI tooling?
<!--
If generative AI tooling has assisted in authoring or reviewing this patch,
please include
an 'Assisted-by: ' trailer that identifies the agent and model.
For example: 'Assisted-by: Claude:claude-opus-4-7' or 'Assisted-by: Claude
Opus 4.7'.
If no, write 'No'.
Please refer to the [ASF Generative Tooling
Guidance](https://www.apache.org/legal/generative-tooling.html) for details.
-->
co-authored: Claude Opus 4.8
--
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]
