peter-toth opened a new pull request, #58421:
URL: https://github.com/apache/spark/pull/58421
### 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(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(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.
### 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
`GenericInternalRow`s, 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. 286 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
--
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]
