peter-toth opened a new pull request, #58262:
URL: https://github.com/apache/spark/pull/58262
### What changes were proposed in this pull request?
This is an alternative to #58245, which fixes the same JIRA by adding the
projection to one of
the two branches below.
`EnsureRequirements` split a child's `KeyedPartitioning`s by `isGrouped` and
then had two
branches that each had to insert a `GroupPartitionsExec`. This PR changes
the classification to
be about what still has to happen to the data:
- `splitKeyedPartitionings` now takes the required distribution and returns
`(other, satisfying, needsGrouping)`. `needsGrouping` pairs each
partitioning with the
partition expression positions the node has to project to.
- A new `projectionPositions` helper computes those positions, from the
`joinKeyPositions` that
`KeyedPartitioning.createShuffleSpec` already derives.
- The four-way match on the required distribution collapses to three cases,
with a single place
that inserts `GroupPartitionsExec`.
`projectionPositions` reports "no projection needed" in two cases. When
projecting would not
merge any partition that coalescing duplicate partition keys alone leaves
apart, every operation
key already lives on a single partition, so the partitioning satisfies the
distribution as it is;
keeping it is also better than projecting, because `KeyedPartitioning([id,
name])` and
`KeyedPartitioning([id])` then describe the same number of partitions and
only the first lets a
downstream operator co-partition on `name` too. And when no partition
expression's *reference* is
a cluster key at all there is no position to project to -
`requireAllClusterKeys` matches at the
expression level, so a transform or nested-field partition expression can
satisfy the distribution
while `KeyedShuffleSpec.keyPositions` yields nothing, and projecting to an
empty position list
would collapse every partition into one.
The classification asks `satisfies` for an already grouped partitioning, not
`groupedSatisfies`,
because only `satisfies` also enforces `Distribution.requiredNumPartitions`
- and `satisfies` is how
this rule used to reach a grouped partitioning. A non-grouped one is left to
`groupedSatisfies`
alone, also as before. Why that requirement is enforced on this path and not
on the others is worth
a separate look.
For an operator that co-partitions more than one child no projection is done
here: the multi-child
block below owns it, through `checkKeyGroupCompatible` for a
storage-partitioned join and through
`withJoinKeyPositions` otherwise. Projecting inline as well would apply the
projection twice, the
second time with positions that index into the unprojected partition
expressions. Note that
`withJoinKeyPositions` is handed the *best* spec's positions for every
child, so that delegation is
only correct when both sides' positions agree - a pre-existing limitation,
unchanged here.
`KeyedPartitioning.satisfies` is not touched, so nothing outside
`EnsureRequirements` changes
behaviour. It does still answer `true` for a partitioning that needs a
projection first, which
means `ValidateRequirements` cannot catch a missing `GroupPartitionsExec`.
Giving that check the
strict test directly, without changing what `satisfies` answers, is left as
a follow-up.
The `KeyedPartitioning` scaladoc is updated as well: it described
`nonGroupedSatisfies` /
`groupedSatisfies` as methods "called on non-grouped KPs", and taught
`isGrouped` as the axis this
PR replaces.
### Why are the changes needed?
With `v2BucketingAllowKeysSubsetOfPartitionKeys` enabled,
`KeyedPartitioning.groupedSatisfies`
only requires that some operation key overlaps the partition attributes. A
partitioning grouped
on `(id, name)` therefore reports that it satisfies
`ClusteredDistribution([id])` while rows
sharing an `id` still sit on separate partitions. A storage-partitioned join
projects the keys
down to the operation keys in `checkKeyGroupCompatible`; for a non-join
operator nothing did.
`isGrouped` is the wrong thing to classify on, because it only says the
*full* partition keys are
unique - it says nothing about whether the *projected* keys are. A
`GroupPartitionsExec` is needed
either to coalesce duplicate partition keys, or to project down to the
operation keys, or both,
and splitting by provenance put those two reasons in different branches.
Both branches returned
wrong results:
1. A grouped partitioning over a superset of the operation keys got no node
at all, so a top-k
window over `PARTITION BY id` on an `(id, name)`-partitioned table
surfaced `id=1` twice, once
per `(1,'aa')`/`(1,'bb')` partition.
2. A non-grouped partitioning over a superset got a node without a
projection, so it coalesced by
the full partition keys and left the operation key split. Any source
reporting more than one
split per partition value hits this: `SUM(price) OVER (PARTITION BY id)`
over the same table
with two splits for `(1,'aa')` returned 25.0 and 20.0 instead of 45.0.
Collapsing the two branches is what makes the second case impossible to
forget again, but it also
opens two ways for a single insertion point to be wrong, which is what the
two "no projection
needed" cases above guard: inserting a node that merges nothing, and
projecting to no position at
all. Neither can happen on `master`, where a grouped partitioning got no
node.
### Does this PR introduce _any_ user-facing change?
Yes, it fixes a data correctness issue. With
`v2BucketingAllowKeysSubsetOfPartitionKeys` enabled, a
single-child operator whose keys are a subset of the partition keys now
produces correct results -
a window `PARTITION BY`, and the single-pass aggregate shapes
(`FlatMapGroupsInBatchExec`,
`ArrowAggregatePythonExec`, `MapGroupsExec`). A two-phase SQL aggregate was
already correct: its
partial `HashAggregate` is a `PartitioningPreservingUnaryExecNode`, so it
narrows `KP([id, name])`
to `KP([id])` before the final aggregate sees it.
A cogroup is unaffected either way, because the projection there is left to
the multi-child block,
exactly as before.
The config is disabled by default, so nothing changes unless it is turned on.
### How was this patch tested?
Added regression tests in `KeyGroupedPartitioningSuite`:
- window top-k over `PARTITION BY` a subset of the partition keys
- window top-k over a duplicated `PARTITION BY` key
- window top-k over union output partitioning
- a plain window over a subset of the partition keys on a non-grouped
`KeyedPartitioning`,
asserting the inserted node projects to the operation key rather than only
coalescing
- no `GroupPartitionsExec` and no shuffle when projecting to the operation
keys merges nothing
and in `EnsureRequirementsSuite`:
- a `FlatMapCoGroupsInPandasExec` over `(n, i)`-partitioned children grouped
on `i`, asserting
both sides are grouped on `i` and not on `n`
- a grouped `KeyedPartitioning` under a `ClusteredDistribution` carrying
`requiredNumPartitions`,
asserting the count is still honoured with a shuffle
- a `years(k)`-partitioned `KeyedPartitioning` under a
`requireAllClusterKeys` distribution,
asserting no node is inserted where there is no position to project to
The three window top-k tests and the non-grouped plain-window test fail
without the
`EnsureRequirements` change. The other four pass on `master` already and
guard this PR's single
insertion point: each of them fails when its own condition is removed and
passes under the other
ablations.
`KeyGroupedPartitioningSuite`, `EnsureRequirementsSuite`, `PlannerSuite`,
`ProjectedOrderingAndPartitioningSuite`, `AdaptiveQueryExecSuite`, the
`execution.exchange` and
`RemoveRedundant*` suites and the `connector` distribution suites pass.
The three window tests come from #58245.
### 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]
