[Spark] Optimize InferRebalanceAndSortOrders to support skip non-cheap keys and local sort

[ulysses-you]

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