tdcmeehan opened a new pull request, #58899:
URL: https://github.com/apache/spark/pull/58899
### 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`, under a named error condition.
- `KeyedPartitioning.checkPartitionKeyArity` (new) rejects any key row whose
`numFields` differs
from the number of reported partition expressions, with
`DATA_SOURCE_PARTITION_KEY_ARITY_MISMATCH` (`SQLSTATE: KD000`).
- `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:
```scala
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:
```
[DATA_SOURCE_PARTITION_KEY_ARITY_MISMATCH] The data source reported a
partition key with 2 field(s)
from `HasPartitionKey.partitionKey()`, but reported 1 partition
expression(s): ["a"]. Every reported
partition key must have one field per reported partition expression.
SQLSTATE: KD000
```
**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 `GenericInternalRow`s,
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.
--
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]
