peter-toth opened a new pull request, #58420:
URL: https://github.com/apache/spark/pull/58420
### What changes were proposed in this pull request?
`KeyedPartitioning` gains `keyDataTypes`, the schema each `partitionKeys`
row was actually written under, and everything that reads those rows takes its
types from there instead of from the partition expressions.
The two can disagree. With `v2BucketingAllowCompatibleTransforms` a
storage-partitioned join reduces one or both sides' keys onto a common key
space, and the partitioning keeps reporting the expressions it was built from.
Joining an `identity(ts)`-partitioned table to a `years(ts)`-partitioned one
leaves `IntegerType` year values under a `TimestampType`-declared expression. A
row is only readable at its own schema, so the wrappers are the authority:
@transient lazy val keyDataTypes: Seq[DataType] =
partitionKeys.headOption.map(_.dataTypes).getOrElse(expressionDataTypes)
Switched to it: `keyRowOrdering` (and with it `toGrouped` and
`createShuffleSpec`'s subset projection), the instance-level `projectKeys` and
`reduceKeys`, `GroupPartitionsExec`'s projection base types, and the reduce
path in `EnsureRequirements`. `PushDownUtils` changes basis twice, at its
wrapper factory and through `keyRowOrdering` at its sort, and both are no-ops
today: it handles rows a scan has just reported, and a scan's own partitioning
is never reduced. They are switched because those wrappers are compared against
the stored keys, which the switch keeps them comparable with. One drive-by in
the same reduce path: it built its ordering with
`RowOrdering.createNaturalAscendingOrdering` directly, and now calls
`KeyedPartitioning.groupedKeyRowOrdering`, which is where that order is defined.
The expressions' own types stay right in one place. `ShuffleExchangeExec`
builds a `KeyGroupedPartitioner` whose lookup keys come from evaluating the
expressions per row, so both sides of that comparison are declared at
`expressionDataTypes`.
That is also what the second change is about. A reduced partitioning must
not be the layout another child is shuffled onto, whatever types it declares,
so `KeyedShuffleSpec.canCreatePartitioning` refuses the ones it can detect. The
question gets a name next to the two values it relates:
@transient lazy val expressionsDescribeKeyShape: Boolean =
keyDataTypes.corresponds(expressionDataTypes)(
DataType.equalsStructurally(_, _, ignoreNullability = true))
Shapes rather than plain equality, which is the test `HashJoin` already
applies to its join key types. `createPartitioning` puts the other child's
expressions over this side's keys, so a struct key whose two sides name the
field differently reaches the gate with no reducer involved, and refusing it
would cost a shuffle for nothing.
The shape test is only a proxy for the real question, and the comment says
so. A reduction that keeps the type, `bucket(12)` and `bucket(8)` both reducing
onto `bucket(4)`, passes it and still misroutes rows. SPARK-59045 and
SPARK-59121 add the real test.
Two of the switched sites are inert today, for the same reason as the
`PushDownUtils` no-op: `reduceKeys` and `GroupPartitionsExec`'s projection only
see a reduced partitioning when a second reduce lands on a first, which still
fails at planning either way, in `KeyedShuffleSpec.reducers`. That reducer is
bound to the partition expression and then reads a stored key, so it is a
fourth reader of these rows that types cannot fix. It is named in the
`keyDataTypes` scaladoc and belongs to SPARK-59121.
### Why are the changes needed?
Both halves fail a query on the `v2BucketingAllowCompatibleTransforms` path.
Reading a reduced key at the expression's type unboxes an Integer as a Long,
which throws `ClassCastException` at planning. AQE reaches it on its own:
`ValidateRequirements` re-runs `createShuffleSpec` on the already reduced
children, and with the subset opt-in on that projects and sorts the keys. The
first test below is such a query.
The second half is independent of the first, and it fails at execution
rather than at planning. A reduced partitioning is grouped and does satisfy the
distribution, so `EnsureRequirements` offers it as the layout to shuffle
another child onto. `ShuffleExchangeExec` then builds the shuffle's key map by
re-wrapping the stored keys at the expressions' types, and the same unboxing
throws, on the driver as the shuffle is prepared. The gate makes
`EnsureRequirements` fall back to shuffling both sides, which returns the right
rows.
The `EnsureRequirements` reduce path is the third behaviour change. Its
`fold` default now reports the keys' own types, so two sides whose reduced key
spaces genuinely differ raise
`STORAGE_PARTITION_JOIN_INCOMPATIBLE_REDUCED_TYPES` where `master` cast one to
the other and threw `ClassCastException` in the merge.
What the gate does not fix is a reduction that keeps the type. There
`KeyGroupedPartitioner` sends a key it cannot find to `hashCode %
numPartitions`, and rows are lost silently, both before and after this change.
`bucket(12)` joined to `bucket(8)`, then joined to an unpartitioned table,
returns 8 of 12 rows on `master` and still does here.
The reduce and the reporting of reduced keys under the original expressions
arrived in SPARK-47094 (4.0.0). The crashes start in 4.2.0, with the
`KeyedPartitioning` / `GroupPartitionsExec` refactor that derives the key
ordering and types from the reported expressions.
### Does this PR introduce _any_ user-facing change?
Yes, on the opt-in `v2BucketingAllowCompatibleTransforms` path. Queries that
failed with `ClassCastException` at planning, as the shuffle was prepared, or
in the shuffle write, now plan and return rows. One error changes: two sides
with genuinely incompatible reduced types now say so instead of throwing
`ClassCastException`. A connector whose transforms reduce onto a
type-compatible key space could previously lose rows to the partitioner's hash
fallback and now gets a shuffle instead, so its results change from wrong to
right.
### How was this patch tested?
Five new tests, all five failing on `master`:
- `KeyGroupedPartitioningSuite`: `reduced partition keys are read at the
types they were built with`, which throws at planning on `master`, in
`toGrouped` under `createShuffleSpec` under `ValidateRequirements`. `another
child is not shuffled onto reducer-rewritten keys`, which throws on `master` as
the shuffle is prepared. And `incompatible reduced key types are reported
instead of cast`, which throws `ClassCastException` on `master` and the named
error here. The first two also assert the plan the fix is about: no shuffle for
the co-partitioned pair, two shuffles for the child that must not be laid out
on reduced keys.
- `ShuffleSpecSuite`: `createShuffleSpec sorts the projected keys at their
built-with types`, the unit-level form of the first one, same exception on
`master`. And `canCreatePartitioning: KeyedShuffleSpec requires the declared
key shape`, which asserts both directions, the reduced keys refused and the
differently named struct fields allowed.
The two changes are pinned separately by ablation. Removing the
`canCreatePartitioning` clause leaves the first end-to-end test passing and
makes the second one throw again, so neither test stands in for the other.
Tightening the shape test to plain equality fails the allowed direction alone.
291 tests green across `KeyGroupedPartitioningSuite`,
`KeyGroupedPartitioningCatalystRuntimeFilterSuite`, `EnsureRequirementsSuite`,
`ValidateRequirementsSuite`, `ProjectedOrderingAndPartitioningSuite`,
`PlannerSuite`, `ShuffleSpecSuite` and `DistributionSuite`. `dev/lint-scala`
clean.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code
--
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]
