[SPARK-59642][SQL] Validate the arity of connector-reported partition keys

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala

@@ -962,6 +962,7 @@ object KeyedPartitioning {
   def apply(
       expressions: Seq[Expression],
       partitionKeys: Seq[InternalRow]): KeyedPartitioning = {
+    checkPartitionKeyArity(expressions, partitionKeys)
     val factory = InternalRowComparableWrapper
       .getInternalRowComparableWrapperFactory(expressions.map(_.dataType))
     val comparablePartitionKeys = partitionKeys.map(factory)
@@ -972,6 +973,22 @@ object KeyedPartitioning {
       KeyLayout(comparablePartitionKeys, factory.dataTypes, isGrouped, isCollapsed = false))
   }
 
+  // The key's arity is implicit in `HasPartitionKey` and read positionally downstream: a short key
+  // runs off the end as an opaque `ArrayIndexOutOfBoundsException`, while a long key's trailing
+  // fields are silently dropped by readers built over `expressions.length`, grouping too loosely.
+  def checkPartitionKeyArity(
+      expressions: Seq[Expression],
+      partitionKeys: Seq[InternalRow]): Unit = {
+    partitionKeys.foreach { key =>
+      if (key.numFields != expressions.length) {
+        throw new SparkException("Data source reported a partition key with " +
+          s"${key.numFields} field(s) from HasPartitionKey.partitionKey(), but reported " +
+          s"${expressions.length} partition expression(s). Every reported partition key " +
+          "must have one field per reported partition expression.")
+      }
+    }
+  }
+
   /**
    * Concatenates partitionings that agree on their expressions, which is what a `UnionExec` does to
    * its children's partitions. The result reports one key per output partition, so its keys are the

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala

@@ -17,7 +17,7 @@
 
 package org.apache.spark.sql.catalyst
 
-import org.apache.spark.{SparkFunSuite, SparkUnsupportedOperationException}
+import org.apache.spark.{SparkException, SparkFunSuite, SparkUnsupportedOperationException}
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, DirectShufflePartitionID, Expression, TransformExpression}
 import org.apache.spark.sql.catalyst.plans.SQLHelper
@@ -980,6 +980,26 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
     assert(e.getMessage.contains("expected specs to be non-empty"))
   }
 
+  test("SPARK-59642: a reported partition key of a different arity is rejected at construction") {
+    val expressions = Seq($"a".int, $"b".int)
+    val wellFormed = InternalRow(1, 1)
+
+    val tooWide = intercept[SparkException] {
+      KeyedPartitioning(expressions, Seq(wellFormed, InternalRow(2, 2, 99)))
+    }
+    assert(tooWide.getMessage.contains("partition key with 3 field(s)"))
+    assert(tooWide.getMessage.contains("2 partition expression(s)"))
+
+    val tooNarrow = intercept[SparkException] {
+      KeyedPartitioning(expressions, Seq(wellFormed, InternalRow(2)))
+    }
+    assert(tooNarrow.getMessage.contains("partition key with 1 field(s)"))
+    assert(tooNarrow.getMessage.contains("2 partition expression(s)"))
+
+    val partitioning = KeyedPartitioning(expressions, Seq(wellFormed, InternalRow(2, 2)))
+    assert(partitioning.numPartitions === 2)
+  }
+
   test("SPARK-59256: flattening reaches the members of a nested collection") {
     val distribution = ClusteredDistribution(Seq($"a", $"b"))
     val buried = HashShuffleSpec(HashPartitioning(Seq($"a"), 10), distribution)

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala

@@ -109,8 +109,10 @@ trait DataSourceV2ScanExecBase
         // `PartitioningCollection.fromPartitionings` refuses them. See `KeyedPartitioning.apply`
         // for why this is the ordering to sort with.
         val keys = inputPartitions.map(_.asInstanceOf[HasPartitionKey].partitionKey())
-          .sorted(KeyedPartitioning.groupedKeyRowOrdering(exprs.map(_.dataType)))
-        Some(KeyedPartitioning(exprs, keys))
+        // Ahead of the sort, which reads every key at the declared positions.
+        KeyedPartitioning.checkPartitionKeyArity(exprs, keys)
+        val ordering = KeyedPartitioning.groupedKeyRowOrdering(exprs.map(_.dataType))
+        Some(KeyedPartitioning(exprs, keys.sorted(ordering)))
       case _ => None
     }
   }

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala

@@ -306,6 +306,11 @@ object PushDownUtils extends Logging {
                 "filtering")
           }
 
+          // Re-reported rows from `filter()` that never pass through `KeyedPartitioning.apply`.
+          KeyedPartitioning.checkPartitionKeyArity(
+            k.expressions,
+            newPartitions.map(_.asInstanceOf[HasPartitionKey].partitionKey()).toSeq)
+
           val inputMap = k.partitionKeys.groupBy(identity).view.mapValues(_.size)
           val comparableKeyWrapperFactory = InternalRowComparableWrapper
             .getInternalRowComparableWrapperFactory(k.keyDataTypes)

sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala

@@ -655,7 +655,8 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession {
   /**
    * While SPJ is active the scan's partitioning has to survive runtime filtering, so the
    * post-filter partitions still line up with the other side of the join: splits may be pruned,
-   * but the source may not drop a partition key, invent one, or grow a key's split count.
+   * but the source may not drop a partition key, invent one, grow a key's split count, or report a
+   * key of a different shape than the partitioning was built from.
    */
   test("data source that breaks the partitioning it reported -> rejected") {
     val partAttr = AttributeReference("part", IntegerType)()
@@ -680,6 +681,12 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession {
       replanAfterFiltering(Seq(KeyedInputPartition(1), KeyedInputPartition(1)))
     }
     assert(splitsGrown.getMessage.contains("must not report new partitions for a given key"))
+
+    val wrongArity = intercept[SparkException] {
+      replanAfterFiltering(Seq(WideKeyedInputPartition(1, 99)))
+    }
+    assert(wrongArity.getMessage.contains("partition key with 2 field(s)"))
+    assert(wrongArity.getMessage.contains("1 partition expression(s)"))
   }
 
   // ---------------------------------------------------------------------------
@@ -853,6 +860,11 @@ private case class KeyedInputPartition(key: Int) extends InputPartition with Has
   override def partitionKey(): InternalRow = InternalRow(key)
 }
 
+private case class WideKeyedInputPartition(key: Int, extra: Int)
+  extends InputPartition with HasPartitionKey {
+  override def partitionKey(): InternalRow = InternalRow(key, extra)
+}
+
 /**
  * A scan reporting one set of partitions before filtering and another after, so it can break the
  * requirement to preserve the partitioning it originally reported.

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBaseSuite.scala

@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources.v2
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.QueryTest
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.AttributeReference
+import org.apache.spark.sql.connector.read.{Batch, HasPartitionKey, InputPartition, PartitionReaderFactory, Scan}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{IntegerType, StructType}
+
+class DataSourceV2ScanExecBaseSuite extends QueryTest with SharedSparkSession {
+
+  private val exprA = AttributeReference("a", IntegerType)()
+  private val exprB = AttributeReference("b", IntegerType)()
+
+  test("SPARK-59642: a reported key of the wrong arity is rejected before the keys are sorted") {
+    // The keys tie on `a`, so the sort that follows has to read `b` off the one-field key: were the
+    // check to run after it, this would be an ArrayIndexOutOfBoundsException from the ordering.
+    val exec = BatchScanExec(
+      output = Seq(exprA, exprB),
+      scan = new KeyedPartitionsScan(Seq(InternalRow(1, 5), InternalRow(1))),
+      runtimeFilters = Seq.empty,
+      table = null,
+      keyGroupedPartitioning = Some(Seq(exprA, exprB)))
+
+    withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "true") {
+      val e = intercept[SparkException](exec.outputPartitioning)
+      assert(e.getMessage.contains("partition key with 1 field(s)"))
+      assert(e.getMessage.contains("2 partition expression(s)"))
+    }
+  }
+}
+
+private case class KeyedPartition(key: InternalRow) extends InputPartition with HasPartitionKey {
+  override def partitionKey(): InternalRow = key
+}
+
+private class KeyedPartitionsScan(keys: Seq[InternalRow]) extends Scan with Batch {
+  override def readSchema(): StructType =
+    new StructType().add("a", IntegerType).add("b", IntegerType)
+
+  override def toBatch: Batch = this
+
+  override def planInputPartitions(): Array[InputPartition] = keys.map(KeyedPartition(_)).toArray
+
+  override def createReaderFactory(): PartitionReaderFactory =
+    throw new UnsupportedOperationException()
+}