[SPARK] Infer cast-wrapped join keys in InferRebalanceAndSortOrders

[ulysses-you]

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