ulysses-you opened a new pull request, #7551:
URL: https://github.com/apache/kyuubi/pull/7551
<!--
Thanks for sending a pull request!
Here are some tips for you:
1. If this is your first time, please read our contributor guidelines:
https://kyuubi.readthedocs.io/en/latest/contributing/code/index.html
2. If the PR is related to an issue in
https://github.com/apache/kyuubi/issues,
add '[KYUUBI #XXXX]' in your PR title, e.g., '[KYUUBI #XXXX] Your PR
title ...'.
3. If the PR is unfinished, add '[WIP]' in your PR title, e.g.,
'[WIP][KYUUBI #XXXX] Your PR title ...'.
-->
### Why are the changes needed?
The `InferRebalanceAndSortOrders` optimization infers rebalance partitioning
and sort columns from the original query (e.g. join keys) to improve the
compression ratio before writing. In practice this can regress performance in
two cases:
1. **Expensive inferred columns.** When the inferred rebalance/sort keys are
non-trivial expressions, evaluating them during the extra shuffle and
local
sort adds CPU cost that can outweigh the compression benefit.
2. **Unwanted local sort.** Some workloads only want the file-layout benefit
from rebalance and do not want the additional local `Sort` before writing.
This PR makes the inference tunable so users can opt out of expensive-column
inference and skip the sort while keeping the rebalance.
Two new configs are added (default behavior is unchanged except that
inference
now restricts to cheap columns by default):
| Config | Default | Description |
|---|---|---|
| `spark.sql.optimizer.inferRebalanceAndSortOrdersWithCheapColumns.enabled`
| `true` | Only infer rebalance/sort columns when all inferred columns are
cheap expressions (attributes, foldable values, or field extractions over cheap
expressions). |
| `spark.sql.optimizer.skipInferRebalanceAndSortOrders.enabled` | `false` |
Only infer the rebalance partition columns and skip inferring the sort orders. |
Both configs only take effect when
`spark.sql.optimizer.inferRebalanceAndSortOrders.enabled` is `true`.
Implementation notes:
- `InferRebalanceAndSortOrders.infer` gains an `onlyInferWithCheapColumns`
parameter. A new `isCheap` helper classifies an expression as cheap when
it is
an `Attribute` / `OuterReference` / `BoundReference`, foldable, or an
`Alias` /
`ExtractValue` whose children are all cheap. When the flag is set and any
inferred column is not cheap, inference returns `None` and the rule falls
back
to a plain rebalance.
- `RebalanceBeforeWritingBase.buildRebalance` reads the two new configs and
skips the local `Sort` when `skipInferRebalanceAndSortOrders.enabled` is
on.
- The change is applied consistently across the Spark `4.1`, `4.0`, and `3.5`
extension modules, and the rule docs table is updated.
### How was this patch tested?
- Added `Skip inferring sort orders` — end-to-end check that enabling
`skipInferRebalanceAndSortOrders.enabled` drops the inferred local `Sort`
while keeping the `RebalancePartitions`.
- Added `Infer rebalance and sort orders only with cheap columns` —
unit-level
check on `InferRebalanceAndSortOrders.infer`: expensive join keys
(`col1 + 1`) are not inferred when the cheap-column restriction is on and
are
inferred when it is off; cheap attribute keys are inferred regardless.
- Tests added to the `RebalanceBeforeWritingSuite` of all three modules.
Ran locally (all green):
```
build/mvn -Pspark-4.1 -Pscala-2.13 test -pl
extensions/spark/kyuubi-extension-spark-4-1 -am
-Dtest=none
-DwildcardSuites=org.apache.spark.sql.RebalanceBeforeWritingSuite # 12/12
build/mvn -Pspark-4.0 -Pscala-2.13 test -pl
extensions/spark/kyuubi-extension-spark-4-0 -am
-Dtest=none
-DwildcardSuites=org.apache.spark.sql.RebalanceBeforeWritingSuite # 12/12
build/mvn -Pspark-3.5 test -pl extensions/spark/kyuubi-extension-spark-3-5
-am
-Dtest=none
-DwildcardSuites=org.apache.spark.sql.RebalanceBeforeWritingSuite # 14/14
```
`dev/reformat` run; no style changes required.
### Was this patch authored or co-authored using generative AI tooling?
Yes.
Assisted-by: 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]
