[SPARK-57993][SQL] Support advisory partition size in rebalance hint

[wForget]

What changes were proposed in this pull request?

This PR adds support for a new REBALANCE_BY_SIZE SQL hint.

The new hint works like REBALANCE, but accepts an optional advisory partition size in bytes as its first parameter, followed by optional partition columns. For example:

SELECT /*+ REBALANCE_BY_SIZE(134217728) */ * FROM t;
SELECT /*+ REBALANCE_BY_SIZE(134217728, c) */ * FROM t;

The implementation resolves REBALANCE_BY_SIZE to RebalancePartitions with optAdvisoryPartitionSize, so AQE can use the hint-specific advisory partition size for the rebalance shuffle.

Why are the changes needed?

A common use case of REBALANCE is small file control before writing query results. Users may want to insert a rebalance before the write and set a larger advisory partition size for the final output stage, so each output partition/file has a reasonable target size.

Today this requires changing the session-level spark.sql.adaptive.advisoryPartitionSizeInBytes configuration. That is too coarse-grained because it may also affect preceding AQE stages, including shuffle coalescing and read/computation parallelism before the final write stage.

REBALANCE_BY_SIZE provides a query-level way to apply a target partition size only to the rebalance shuffle, which is useful for output file size control without changing the advisory size used by earlier stages.

Does this PR introduce any user-facing change?

Yes, a new REBALANCE_BY_SIZE hint.

How was this patch tested?

Added unit tests

Was this patch authored or co-authored using generative AI tooling?

Docs and tests generated-by: Codex (GPT-5)

[wForget]

@ulysses-you @pan3793 Could you please take a look? Thanks.

[ulysses-you]

I'm fine with this new hint, cc @cloud-fan if you have other comments

[cloud-fan]

1 blocking, 0 non-blocking, 1 nit.
Clean, well-scoped extension that reuses existing plumbing; the one substantive issue is a missing positive-value guard that every peer path enforces.

Correctness (1)

• CoalesceHintUtils.scala:43: advisory size not validated as positive — REBALANCE_BY_SIZE(0) / negative via the DataFrame API are accepted into a degenerate plan — see inline

Nits (1)

• CoalesceHintUtils.scala:26-29: the CoalesceHintUtils object Scaladoc still lists only "COALESCE, REPARTITION, REPARTITION_BY_RANGE and REBALANCE" and omits REBALANCE_BY_SIZE, which the object now handles. (Line is outside the diff, so noting it here.)

[cloud-fan]

The advisory size isn't checked for positivity here, so REBALANCE_BY_SIZE(0) is accepted (0 parses as an IntegerLiteral), and df.hint("REBALANCE_BY_SIZE", -1) passes a negative Literal straight through. A size of 0 reaches OptimizeSkewInRebalancePartitions, where it's used as both the skew threshold and the split target size, so every non-empty partition is treated as skewed and split down to one-partition-per-map-block, silently.

Every peer that handles this value rejects non-positive: REBALANCE's num path via require(numPartitions > 0), the write producer via if (partitionSize > 0) Some else None (DistributionAndOrderingUtils), and the session config via checkValue(_ > 0). I'd match the write producer and coerce a non-positive value to the session default, and add a test for REBALANCE_BY_SIZE(0).

[wForget]

@cloud-fan Thanks for the review. I’ve added validation to ensure the advisory size is positive.

[cloud-fan]

1 addressed, 1 remaining, 1 new. (1 new = 0 newly introduced, 1 late catch — my own miss from round 1.)

0 blocking, 0 non-blocking, 2 nits.
The round-1 blocking finding (non-positive advisory size) is properly fixed and well-tested; only two stale-Scaladoc nits remain, both outside the diff.

Nits (2)

• CoalesceHintUtils.scala:26-29: the CoalesceHintUtils object Scaladoc still lists only "COALESCE, REPARTITION, REPARTITION_BY_RANGE and REBALANCE" and omits REBALANCE_BY_SIZE, which the object now handles. (Round-1 nit, still open; line is outside the diff.)
• ResolveHints.scala:185: the ResolveCoalesceHints object Scaladoc has the same stale list — "COALESCE, REPARTITION, REPARTITION_BY_RANGE and REBALANCE" — and should also mention REBALANCE_BY_SIZE. (Line is outside the diff.)

Verification

Confirmed the round-1 fix: getAdvisorySizeOfPartitions now rejects a non-positive advisory size (result._1.exists(_ <= 0), CoalesceHintUtils.scala:55-59) with INVALID_REBALANCE_BY_SIZE_HINT_PARAMETER. This runs on both the SQL hint and the DataFrame .hint("REBALANCE_BY_SIZE", -1) path since both flow through ResolveCoalesceHints, so REBALANCE_BY_SIZE(0) and negatives no longer reach the degenerate OptimizeSkewInRebalancePartitions plan. The rewrite preserves the child row multiset (output = child.output), advisory size only influences AQE shuffle partitioning, and the AQE test asserts the hint size (150) overrides the session default (10000) at the rebalance shuffle.