From 9ceb613a746c1f96aa5fc03d759864f2b178a479 Mon Sep 17 00:00:00 2001 From: Peter Toth Date: Sun, 30 Aug 2026 21:49:25 +0200 Subject: [PATCH 1/2] [SPARK-59123][SQL] Avoid per-key intermediate collections in KeyedPartitioning.projectKeys and reduceKeys `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. 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. No. 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. Generated-by: Claude Code --- .../plans/physical/partitioning.scala | 37 +++++++++++++------ 1 file changed, 25 insertions(+), 12 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala index d8ca230bae17e..450a7a03421f7 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala @@ -854,11 +854,17 @@ object KeyedPartitioning { val projectedDataTypes = positions.map(dataTypes) val comparableKeyWrapperFactory = InternalRowComparableWrapper.getInternalRowComparableWrapperFactory(projectedDataTypes) - val positionsWithTypes = positions.zip(projectedDataTypes) + // A key list is as long as the number of splits the scan reported, so whatever is allocated per + // key is allocated tens of thousands of times. + val positionArray = positions.toArray + val typeArray = projectedDataTypes.toArray val projectedKeys = keys.map { key => - val projectedKey = positionsWithTypes.map { - case (position, dataType) => key.row.get(position, dataType) - }.toArray[Any] + val projectedKey = new Array[Any](positionArray.length) + var i = 0 + while (i < positionArray.length) { + projectedKey(i) = key.row.get(positionArray(i), typeArray(i)) + i += 1 + } comparableKeyWrapperFactory(new GenericInternalRow(projectedKey)) } @@ -872,18 +878,25 @@ object KeyedPartitioning { keys: Seq[InternalRowComparableWrapper], dataTypes: Seq[DataType], reducers: Seq[Option[KeyReducer]]): (Seq[DataType], Seq[InternalRowComparableWrapper]) = { - val reducedDataTypes = dataTypes.zip(reducers).map { - case (_, Some(KeyReducer(reducer: Reducer[Any, Any], _))) => reducer.resultType() - case (t, _) => t + // The `Reducer[Any, Any]` match is erased, so it only ever checks for `Some`. Settling it per + // position keeps it out of the key loop below, and gives the result types with it. + val reducerArray = + reducers.map(_.map(_.reducer.asInstanceOf[Reducer[Any, Any]]).orNull).toArray + val reducedDataTypes = dataTypes.zip(reducerArray).map { + case (t, reducer) => if (reducer == null) t else reducer.resultType() } val comparableKeyWrapperFactory = InternalRowComparableWrapper.getInternalRowComparableWrapperFactory(reducedDataTypes) + val typeArray = dataTypes.toArray val reducedKeys = keys.map { key => - 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 + val reducedKey = new Array[Any](typeArray.length) + var i = 0 + while (i < typeArray.length) { + val value = key.row.get(i, typeArray(i)) + val reducer = reducerArray(i) + reducedKey(i) = if (reducer == null) value else reducer.reduce(value) + i += 1 + } comparableKeyWrapperFactory(new GenericInternalRow(reducedKey)) } From f2f417e77e6aa6f9da7886eac9a75b5fd91df0ae Mon Sep 17 00:00:00 2001 From: Peter Toth Date: Tue, 1 Sep 2026 16:29:17 +0200 Subject: [PATCH 2/2] [SPARK-59123][SQL] Address review: restore the arity assert and reword the hoist comment Keeps the arity check that `InternalRow.toSeq(dataTypes)` carried, as one `assert` before the key loop rather than one per key, and asserts the reducer array's length with it, since the loop indexes both arrays by the same bound where the old `zip` would have truncated. Rewords the `projectKeys` hoist comment to say what the hoist buys, instead of stating the motivation for the whole rewrite above two lines that are not per-key allocations. Generated-by: Claude Code --- .../sql/catalyst/plans/physical/partitioning.scala | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala index 450a7a03421f7..3b39677fa2c8f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala @@ -854,8 +854,8 @@ object KeyedPartitioning { val projectedDataTypes = positions.map(dataTypes) val comparableKeyWrapperFactory = InternalRowComparableWrapper.getInternalRowComparableWrapperFactory(projectedDataTypes) - // A key list is as long as the number of splits the scan reported, so whatever is allocated per - // key is allocated tens of thousands of times. + // Indexed arrays rather than `Seq`s, because the loop below runs once per key and a key list is + // as long as the number of splits the scan reported. val positionArray = positions.toArray val typeArray = projectedDataTypes.toArray val projectedKeys = keys.map { key => @@ -888,6 +888,12 @@ object KeyedPartitioning { val comparableKeyWrapperFactory = InternalRowComparableWrapper.getInternalRowComparableWrapperFactory(reducedDataTypes) val typeArray = dataTypes.toArray + // `InternalRow.toSeq(dataTypes)`, which the loop below replaces, asserted the row's arity once + // per key. All the keys of a partitioning share an arity, so asserting on the first one keeps + // the check. The loop also indexes `reducerArray` by the same bound, where the old `zip` would + // have truncated to the shorter of the two, so that length is asserted with it. + assert(reducerArray.length == typeArray.length) + keys.headOption.foreach(k => assert(k.row.numFields == typeArray.length)) val reducedKeys = keys.map { key => val reducedKey = new Array[Any](typeArray.length) var i = 0