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..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,11 +854,17 @@ object KeyedPartitioning {
     val projectedDataTypes = positions.map(dataTypes)
     val comparableKeyWrapperFactory =
       InternalRowComparableWrapper.getInternalRowComparableWrapperFactory(projectedDataTypes)
-    val positionsWithTypes = positions.zip(projectedDataTypes)
+    // 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 =>
-      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,31 @@ 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
+    // `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 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))
     }