ulysses-you opened a new pull request, #7587: URL: https://github.com/apache/kyuubi/pull/7587
### Why are the changes needed? Two fixes to `InferRebalanceAndSortOrders` (applied across the Spark 3.5, 4.0, 4.1 and 4.2 extension modules): 1. **Support cast-wrapped join keys.** When join key columns have mismatched types (e.g. via a `UNION ALL` branch that widens `int` to `bigint`), the analyzer inserts a cast on the join key such as `cast(genre_tag_id as bigint)`. The outer projection's alias map is keyed by the unary-peeled attribute, so the full cast expression never matched and the key was silently dropped by the `outputSet` filter — only the un-cast keys survived, weakening the inferred rebalance/sort. This adds `mapThroughAlias`, which retries the alias lookup with the unary-peeled key (symmetric with how `getAliasMap` builds its keys), so a cast-wrapped key maps back to its aliased output attribute. `Cast` is also treated as cheap in `isCheap` for keys that survive without being aliased away. 2. **Gate the cheap-column check by `maxColumns`.** `INFER_REBALANCE_AND_SORT_ORDERS_MAX_COLUMNS` is now threaded into `infer` and the truncation is applied *before* the cheap-column check, so only the first `maxColumns` inferred columns need to be cheap (previously an expensive key beyond the limit could discard the whole inference). ### How was this patch tested? Added unit tests in `RebalanceBeforeWritingSuite` for each affected module: - cast-wrapped join key resolves to its aliased output attribute, - a `Cast` over a cheap column is treated as cheap, - `maxColumns` gates the cheap-column check. Verified locally with: - Spark 3.5: `build/mvn -Pspark-3.5 test -pl extensions/spark/kyuubi-extension-spark-3-5 -am -DwildcardSuites=org.apache.spark.sql.RebalanceBeforeWritingSuite` - Spark 4.0 / 4.1 / 4.2 (JDK 17, `-Pscala-2.13`): same suite per module. All suites pass; `dev/reformat` applied. ### Was this patch authored or co-authored using generative AI tooling? Assisted-by: Claude: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]
