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

[tdcmeehan]

What changes were proposed in this pull request?

A DSv2 scan that implements SupportsReportPartitioning reports N partition expressions, and each
of its input partitions reports a key row through HasPartitionKey.partitionKey(). Nothing checks
that a key row actually holds N fields. This adds that check where raw connector key rows first
become a KeyedPartitioning.

• KeyedPartitioning.checkPartitionKeyArity (new) rejects any key row whose numFields differs
  from the number of reported partition expressions, throwing a SparkException that names the
  contract -- matching the sibling connector-contract checks already in replanWithRuntimeFilters.
• KeyedPartitioning.apply(Seq[Expression], Seq[InternalRow]) calls it, so a partitioning built
  from raw rows carries the guarantee for everything downstream.
• The two callers that read raw key rows before a partitioning exists call it themselves:
  • DataSourceV2ScanExecBase.reportedKeyedPartitioning, ahead of the sort it runs on the keys --
    the generated ordering indexes a key at the declared positions, so it has to be checked before
    the sort, not by apply afterwards;
  • PushDownUtils.replanWithRuntimeFilters, on the partitions a source returns from filter().
    Those are re-reported rows that never pass through apply, so a malformed replacement
    partition would otherwise bypass the check entirely.

Why are the changes needed?

HasPartitionKey leaves a key's arity implicit, and nothing downstream reads a key defensively --
Spark indexes it at the arity the scan reported. So an inconsistent key does not fail where it is
produced. It fails later, in interpreted or generated code that names neither the data source nor
the contract it broke, and it fails in four different ways depending on the shape:

reported key	behaviour on master
too wide, >= 2 partitions	ArrayIndexOutOfBoundsException from Murmur3HashFunction (hash.scala:817), via the distinct in KeyedPartitioning.apply
too narrow, ties with another key on its leading fields	ArrayIndexOutOfBoundsException from the generated ordering, in the sort in DataSourceV2ScanExecBase
too narrow, no tie	accepted; fails at whichever positional read comes first, or groups partitions wrong
too wide, exactly 1 partition	accepted (distinct short-circuits); later AssertionError: assertion failed, or the extra field is silently ignored

Reproduction -- one reported partition expression, a second key row with two fields:

KeyedPartitioning(Seq($"a".int), Seq(InternalRow(1), InternalRow(2, 99)))

On master, via the distinct in KeyedPartitioning.apply:

java.lang.ArrayIndexOutOfBoundsException: Index 1 out of bounds for length 1
  at org.apache.spark.sql.catalyst.expressions.InterpretedHashFunction.hash(hash.scala:817)
  at org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper.hashCode(InternalRowComparableWrapper.scala:74)
  at scala.collection.StrictOptimizedSeqOps.distinctBy(StrictOptimizedSeqOps.scala:143)
  ...

With this patch:

org.apache.spark.SparkException: Data source reported a partition key with 2 field(s) from
HasPartitionKey.partitionKey(), but reported 1 partition expression(s). Every reported partition
key must have one field per reported partition expression.

Relation to SPARK-59123 (#58421). That commit replaced a per-key InternalRow.toSeq(dataTypes)
in KeyedPartitioning.reduceKeys with an indexed loop plus a single assert on the head key, and it
can look like it dropped a per-key arity check, since InternalRow.toSeq(fieldTypes) asserts
numFields == fieldTypes.length. It did not. GenericInternalRow overrides toSeq(fieldTypes)
(rows.scala:178) and ignores the types entirely, and partition key rows are GenericInternalRows,
so that assert never ran for them. The two reduceKeys bodies agree on every reachable input, and
the new head assert is strictly stricter in the one single-key case. The missing validation is older
than #58421 and belongs at ingestion rather than in the reduction loop, which is why this centralizes
it at KeyedPartitioning construction instead of adding per-key checks back to the loop.

Does this PR introduce any user-facing change?

Yes, for a DSv2 connector that reports partition keys inconsistent with the partitioning it
reported. In most shapes the query already fails and only the diagnostic improves: an opaque
ArrayIndexOutOfBoundsException or a bare AssertionError becomes an error that names the
contract, the arity it got, and the expressions it was measured against.

Two shapes are newly rejected rather than newly diagnosed, so a connector relying on them would see
a query that completes today start to fail:

• a scan reporting exactly one partition with an over-wide key -- distinct short-circuits on a
  single element, so nothing hashes the row and the extra field is simply never read;
• an under-wide key that never ties with another key on its leading fields and is never read past
  its own width.

Both are contract violations, and the second can mis-group partitions rather than raise anything,
so failing fast at the boundary is the point of the change. No behaviour change for a connector
whose keys match the partitioning it reported.

How was this patch tested?

New tests, both at boundaries a connector actually reaches:

• ShuffleSpecSuite, "a reported partition key of a different arity is rejected at construction" --
  the construction boundary, in the suite that already holds the other KeyedPartitioning
  construction-rejection tests. Covers a too-wide key (an AIOOBE from the hash on master), a
  too-narrow key (accepted on master, the case nothing caught), and that well-formed keys still
  build the partitioning.
• DataSourceV2CatalystRuntimeFilterSuite -- extended the existing "data source that breaks the
  partitioning it reported -> rejected" case with a partition whose key row carries a field the scan
  reported no expression for. This is the runtime-filter re-ingestion path, whose rows never pass
  through KeyedPartitioning.apply; on master it fails with an AIOOBE out of groupBy in
  replanWithRuntimeFilters.

There is no end-to-end SQL reproducer because the in-tree test connectors cannot express the bug:
InMemoryBaseTable and its relatives derive each partition key from the table's own partition
schema, so a key's arity is structurally tied to the reported partitioning. Reaching the case
end-to-end would mean adding a deliberately malformed connector for no other purpose. The
runtime-filter suite above is the connector-shaped test -- real InputPartition / HasPartitionKey
instances through the production replan path -- and it already exists to host exactly this class of
contract violation.

Suites run locally with sbt on JDK 17, all passing:

• catalyst: ShuffleSpecSuite, DistributionSuite
• sql: DataSourceV2CatalystRuntimeFilterSuite, KeyGroupedPartitioningSuite,
  EnsureRequirementsSuite, ValidateRequirementsSuite, GroupPartitionsExecSuite,
  ProjectedOrderingAndPartitioningSuite, SparkThrowableSuite

dev/scalastyle reports 0 errors for every changed module. dev/lint-scala's scalafmt check is
scoped to sql/api and sql/connect, so it does not cover any file here.

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

Generated-by: Isaac

This pull request and its description were written by Isaac.