[SPARK-59123][SQL] Avoid per-key intermediate collections in KeyedPartitioning.projectKeys and reduceKeys

[peter-toth]

What changes were proposed in this pull request?

KeyedPartitioning.projectKeys and KeyedPartitioning.reduceKeys walk every partition key of a storage-partitioned join, and each one built several throwaway collections to end up with one array. Both now hoist what does not depend on the key and fill a single Array[Any] with an indexed loop.

projectKeys materialised an intermediate Seq per key and copied it into an array, destructuring a Tuple2 per position:

val projectedKey = positionsWithTypes.map {
case (position, dataType) => key.row.get(position, dataType)
}.toArray[Any]

reduceKeys did the same four times over. key.row.toSeq(dataTypes) allocated an array and an ArraySeq wrapper, zip(reducers) a sequence of tuples, map a third sequence, and toArray the array that was wanted in the first place:

val keyValues = key.row.toSeq(dataTypes)
val reducedKey = keyValues.zip(reducers).map {
case (v, Some(KeyReducer(reducer: Reducer[Any, Any], _))) => reducer.reduce(v)
case (v, _) => v
}.toArray

Which positions have a reducer is now settled once, outside the key loop, so the erased Some(KeyReducer(reducer: Reducer[Any, Any], _)) type test runs once per position instead of once per key value, and it gives the reduced data types with it. Unwrapping the KeyReducer moves out of the loop with it.

reduceKeys keeps the arity check that InternalRow.toSeq(dataTypes) carried, as one assert before the loop rather than one per key. It asserts the reducer array's length with it, because the loop indexes both arrays by the same bound where the old zip would have truncated to the shorter. Both call sites keep the two equal today, so this pins that rather than changing anything.

Why are the changes needed?

A key list is as long as the number of splits the scan reported, so tens of thousands is ordinary, and anything allocated per key is allocated that many times. projectKeys runs over all of them on every EnsureRequirements and ValidateRequirements pass whenever spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled is on, not only on the reducer path.

Measured on a KeyedPartitioning with 50k keys of 12 positions, six IntegerType and six StringType, projecting two of each, with a reducer on one position, 20 evaluations after a warm-up:

	before	after
projectKeys	105-111 ms	35-39 ms
reduceKeys	227-252 ms	66-95 ms

One thing I tried and dropped: hoisting the type dispatch out of the loop with InternalRow.getAccessor, the way BoundReference does. It measured slower, 63-79 ms for projectKeys, because these rows are GenericInternalRows, whose get(ordinal, dataType) ignores the requested type and reads the array directly. The accessor only adds a closure call and a null-check wrapper on top of that. PhysicalDataType.apply per value is on the UnsafeRow path, which partition keys are not.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

No new test: the two bodies are rewritten, not changed in behaviour, and both are on the path of the existing storage-partitioned-join tests. 305 tests green across KeyGroupedPartitioningSuite, KeyGroupedPartitioningCatalystRuntimeFilterSuite, EnsureRequirementsSuite, ValidateRequirementsSuite, ProjectedOrderingAndPartitioningSuite, PlannerSuite, ShuffleSpecSuite and DistributionSuite. dev/lint-scala clean.

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

Generated-by: Claude Code

[peter-toth]

cc @dongjoon-hyun

[dongjoon-hyun]

LGTM. Both rewrites look behavior-preserving to me:

• projectKeys: the same reads, just without the intermediate Seq and the per-position tuple.
• reduceKeys: key.row.toSeq(dataTypes) is exactly the same indexed get(i, dataTypes(i)) loop, and the Reducer[Any, Any] test is erased, so settling it per position is sound.
• The null sentinel in reducerArray is safe: TransformExpression.reducers wraps the connector's reducer in Option(res), so a KeyReducer never carries a null reducer.
• The new loop bounds on dataTypes.length while indexing reducerArray, where the old zip truncated to the shorter of the two. Both call sites keep them equal -- EnsureRequirements takes the reducers from reducersBothWays over the specs' (already projected) expressions, after areKeysCompatible checked equal arity, and GroupPartitionsExec projects the data types with the same joinKeyPositions -- so the difference is not reachable today.

Three minor comments:

1. The PR description quotes stale code: it shows case (v, Some(reducer: Reducer[Any, Any])) => reducer.reduce(v), but master has case (v, Some(KeyReducer(reducer: Reducer[Any, Any], _))). Same for the resultType() snippet. Worth refreshing before merge.

2. reduceKeys drops the assert(numFields == fieldTypes.length) that InternalRow.toSeq(fieldTypes) carried. No impact outside -ea, but it is an unmentioned side effect of not going through toSeq; one assert outside the key loop would restore it for free.

3. In projectKeys, the comment "whatever is allocated per key is allocated tens of thousands of times" sits above positionArray / typeArray, which are hoisted out of the key loop and are not per-key allocations. It reads as the motivation for the whole rewrite rather than for those two lines -- moving it above keys.map, or rewording it to say what the hoist buys, would be clearer.

Agreed that no new test is needed here.

[peter-toth]

Thanks @dongjoon-hyun. All three taken, in f2f417e.

1. The description was stale and that one is on me. I refreshed the commit message when I rebased onto [SPARK-59045][SQL] Fix SPJ ClassCastException when reducer changes partition key data type #58335 and forgot the description. Both snippets now read Some(KeyReducer(reducer: Reducer[Any, Any], _)), and so does the sentence about the erased type test.

2. Restored, as one assert before the key loop instead of one per key. I asserted reducerArray.length against the same bound there too, since that is the length mismatch you named in the LGTM list, and the comment records that neither is reachable from today's two call sites.

3. Reworded to say what the hoist buys. The scale argument stays, moved next to the loop that pays it.

On the reachability: agreed, and thanks for tracing both call sites rather than taking the equality on trust. The assert felt like the right weight for it, so a third caller finds out at the boundary rather than through an ArrayIndexOutOfBoundsException or a silently short key.

The description's suite total moved from 286 to 305 with the rebase, since KeyGroupedPartitioningSuite grew in #58335. Re-ran the same eight suites on the new head.

[ulysses-you]

thank you @peter-toth

[peter-toth]

Merge Summary:

• merged into master 9e0bea0
• merged into branch-4.x b52ba30

Posted by merge_spark_pr.py

[peter-toth]

Thank you @dongjoon-hyun and @ulysses-you for the review.