cloud-fan commented on code in PR #59535:
URL: https://github.com/apache/spark/pull/59535#discussion_r4324836901
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:
##########
@@ -1040,6 +1036,36 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
}
}
+ /**
+ * Whether the legs' markers let their keys be merged. A marked leg may hold
rows of a key it
+ * does not declare, and `KeyGroupedPartitioner` routes those into the
partition of some declared
+ * key. Two conditions keep every key's rows in one group, which is what a
clustering reads:
+ *
+ * - a marked leg declares every key the other legs declare, so its rows
of such a key do not
+ * land in a partition the merge names for another key;
+ * - marked legs declare the same keys in the same order, and when more
than one is marked,
+ * each still routes by the hash. A routed leg is grouped, so its keys
are unique and the
+ * list length is the modulus the hash is taken over: the position in
that list is the name a
+ * key no leg declares reaches, and equal lists agree on both. A leg
whose routing is lost
+ * names such a key by an earlier layout's order, which its list does
not show. One marked
+ * leg needs no agreement, since its undeclared rows reach one group
either way.
+ *
+ * The merged layout cannot claim that a partition holds only the key it
names, which
+ * [[KeyedPartitioning.concat]] records.
+ */
+ private def canMergeMarkedKeys(kps: Seq[KeyedPartitioning]): Boolean = {
+ val marked = kps.filter(_.mayContainUnknownPartitionKeys)
+ if (marked.isEmpty) {
+ true
+ } else {
+ val declared = marked.head.partitionKeys
+ marked.tail.forall(_.partitionKeys == declared) &&
+ (marked.size == 1 || marked.forall(_.unknownKeysRouted)) &&
+ kps.filterNot(_.mayContainUnknownPartitionKeys)
+ .forall(_.partitionKeys.forall(declared.contains))
Review Comment:
**Non-blocking (P2):** `declared.contains` scans the marked leg's
`partitionKeys` linearly, once for every key of every unmarked leg. That is U x
D `InternalRowComparableWrapper.equals` calls, where U is the unmarked keys
(one per split when a leg is ungrouped) and D the declared keys.
`outputPartitioning` re-runs `rawPartitioning` on every call for a non-plain
union. It is deliberately not cached, because AQE changes child partition
counts, so EnsureRequirements, ValidateRequirements, stamping and AQE
re-planning each pay the cost again. With D = 1e3 and U = 1e4 that is about 1e7
comparisons per call. Before this PR, a marked leg failed the gate before any
key comparison. `KeyedShuffleSpec.areKeysCompatible` answers the same subset
question with `partitionKeys.toSet`. Building `declared.toSet` once, only when
an unmarked leg exists, keeps this linear and gives the same answer. The
ordered `==` check across marked legs should stay as it is. Only driver
planning time on the opt-in
one-side-shuffle path is affected, not results.
**Verification:**
- **Behavior:** The union merge gate keeps its existing outcomes: a covered
unmarked leg merges, an uncovered key, differing order or multiplicity, two
lost legs, and a routed/lost pair are refused, and a single lost marked leg and
a grouped concatenation still merge.
- **Inspection:** The containment check probes a hashed collection of the
marked leg's declared keys and no longer scans the declared key list once per
unmarked key.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1182,25 +1216,49 @@ object KeyedPartitioning {
KeyLayout(comparablePartitionKeys, factory.dataTypes, isGrouped,
isCollapsed = false))
}
+ /**
+ * The unknown-keys state a layout built from several members carries: the
marked members have to
+ * agree, since one layout cannot carry two descriptions of where an
undeclared row sits, and a
+ * disagreement leaves `ROUTING_LOST`. An unmarked member takes no part.
Asked of the members'
+ * states, so a new origin joins without this changing.
+ */
+ def combine(states: Seq[Option[UnknownKeysOrigin]]):
Option[UnknownKeysOrigin] =
+ states.flatten.distinct match {
+ case Seq() => None
+ case Seq(one) => Some(one)
+ case _ => Some(ROUTING_LOST)
+ }
+
/**
* 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
* children's keys in child order.
*
* Keys repeating across children is not a collapse. Only a child's own
collapse carries over,
* since such a key still stands for several finer-grained ones in the
concatenation.
+ *
+ * One marked child marks the result, and the concatenation reports the
weakest routing of its
+ * legs': ungrouped, it routes an undeclared row per contributing leg rather
than by a hash over
+ * its own count.
*/
def concat(kps: Seq[KeyedPartitioning]): KeyedPartitioning = {
val concatenatedKeys = kps.flatMap(_.partitionKeys)
val ungroupingOrigin =
PartitioningCollection.concatUngrouping(kps.map(_.ungroupingOrigin))
+ // A child that has duplicates of its own puts them in the concatenation
too, which answers
+ // this without walking the keys.
+ val grouped = kps.forall(_.isGrouped) &&
+ concatenatedKeys.distinct.length == concatenatedKeys.length
+ // Only a grouped concatenation, where one marked leg holds every
partition, keeps the routing.
+ val routing = combine(kps.map(_.layout.unknownKeysOrigin)) match {
+ case Some(HASH_ROUTING) if !grouped => Some(ROUTING_LOST)
Review Comment:
**Nit (P3):** This keeps `HASH_ROUTING` whenever the concatenation is
grouped. The comment above states the condition that actually preserves the
modulus: one marked leg holds every partition. The two coincide only because
`UnionExec.canMergeMarkedKeys` runs first and forces every unmarked key into
the marked declared list, so a grouped result there means every other leg is
empty. Without that gate, `concat` over a routed [1, 2, 3] leg and an unmarked
[4, 5] leg is grouped. It reports `HASH_ROUTING` over 5 partitions while the
undeclared rows sit at `hash % 3`. `canCreatePartitioning` would then accept it
as a template and lay a sibling's undeclared rows out by `hash % 5`, losing
matches. No current caller reaches this. Deriving the state from the stated
condition (at most one leg contributes partitions, and that leg is routed)
costs one count and keeps every outcome `UnionExec` sees today.
**Verification:**
- **Regression:** A direct concatenation of a hash-routed marked layout over
keys [1, 2, 3] with an unmarked layout over keys [4, 5] yields a grouped,
marked layout whose routing is lost, so it cannot serve as a shuffle template.
- **Behavior:** A grouped concatenation in which the routed marked leg holds
every partition still reports hash routing, and the existing union merge tests
keep their outcomes.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala:
##########
@@ -2809,7 +2809,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
marked <- Seq(false, true)
} yield {
val kp = KeyedPartitioning(exprs, keys)
- val p = if (marked)
kp.withLayout(_.copy(mayContainUnknownPartitionKeys = true)) else kp
+ val p = if (marked) kp.withLayout(_.copy(unknownKeysOrigin =
Some(HASH_ROUTING))) else kp
Review Comment:
**Nit (P3):** Before this PR, every reachable marked layout was grouped and
hash-routed, so this generator covered the whole marked state space. The PR
adds two more kinds of join input: `ROUTING_LOST` layouts, and ungrouped marked
layouts such as a merged union reports. The generator was only migrated
mechanically, so the sweep never feeds EnsureRequirements either kind. The
SPARK-60277 join test pairs identical {1, 2, 3} key sets, and the
GroupPartitionsExec and ShuffleSpec unit tests check each refusal predicate on
its own. Suppose a later change let a lost-routing side act as a template, or
pair index by index under a non-identity expected-key alignment. No test would
fail. Two additions would close the gap under the existing
`ValidateRequirements` oracle: `Some(ROUTING_LOST)` and an ungrouped marked key
set in the shapes here, and one end-to-end join, with AQE off and on, over a
merged marked union against a side with a different key set.
**Verification:**
- **Regression:** A planner change that pairs index by index with, or lays a
sibling out on, a marked side whose routing is lost fails the generated pairing
sweep.
- **Behavior:** A storage-partitioned join over a merged marked union whose
other side declares a different key set falls back to a valid plan and returns
the expected rows under both AQE settings.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -937,9 +937,10 @@ case class EnsureRequirements(
// The pairing is only worth committing to if both children still declare
the same aligned key
// sequence once the grouping has been pushed into them. A
`GroupPartitionsExec` gives up its
- // keyed claim when it turns out to regroup a layout that pins undeclared
rows to
- // `hash(key) % numPartitions` (see
`KeyLayout.mayContainUnknownPartitionKeys`), and only the
- // node knows the permutation it performs, so that answer arrives after
the pairing was chosen.
+ // keyed claim when the grouping re-labels the keys or replicates them
into expected slots,
Review Comment:
**Nit (P3):** This rewritten sentence matches the new rule. The same block's
last paragraph (955-958), however, still says the `groupIfNeeded` node 'gives
up the claim for the same reason the pushed one does'. That node has no key
rewrite and no expected keys, so it now keeps a marked claim as `ROUTING_LOST`,
and the SPARK-60277 join test depends on that. Other unchanged comments still
state the old rule:
- `createShuffleSpec` (partitioning.scala 1150-1152) justifies its marked
early return by a give-up that no longer happens for the sort-only half.
- `areKeysCompatible` (2316-2317) says a reducing GPE gives the claim up
because it 'is not an identity regrouping'. The cause is now `keysRewritten`.
- The EnsureRequirementsSuite pairing sweep says the failing marked shapes
are those 'where a `GroupPartitionsExec` gives up the keyed claim' while
regrouping (2781-2782), and that 'regrouping one moves them and the node stops
claiming the layout' (2807-2808).
- KeyGroupedPartitioningSuite 9840-9843 still names `keyedMarkerOf`.
A maintainer who trusts these comments could relax the `createShuffleSpec`
early return or the routed checks. That would re-admit a sorted marked layout
whose undeclared rows no longer sit at `hash(key) % numPartitions`.
See **Shared repair plan 1** in the review body.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -447,6 +447,21 @@ case object REPLICATED_FOR_JOIN extends UngroupingOrigin
// A spread an operator's ordering requirement is read through, which pairs
with nothing.
case object SPLIT_FOR_ORDERING extends UngroupingOrigin
+/**
+ * Why a layout may hold rows whose key it does not declare, and what became
of the hash routing
+ * that placed them. `KeyGroupedPartitioner` routes such a row by
+ * `nonNegativeMod(key.hashCode, numPartitions)`, so a marked layout
co-locates whole keys, never a
+ * subset of their columns. Anything that regrouped, sorted or projected the
keys moves the rows off
+ * that routing, and only the co-location survives.
+ */
+sealed trait UnknownKeysOrigin
+
+// An undeclared row sits at `nonNegativeMod(key.hashCode, numPartitions)`;
the relation holds.
+case object HASH_ROUTING extends UnknownKeysOrigin
+
+// A regrouping, a sort or a projection moved the rows off that relation; only
co-location remains.
Review Comment:
**Nit (P3):** No projection produces `ROUTING_LOST`.
`KeyedPartitioning.project` asserts the layout is unmarked (847-848),
`AliasAwareOutputExpression` turns away a narrowing projection of a marked
layout, and a GroupPartitionsExec with `keysRewritten` reports
`UnknownPartitioning`. After a projection, no co-location claim survives at
all. The carriers that do set `ROUTING_LOST` are a non-identity regroup,
`toGrouped`, an ungrouped `concat`, and `combine` over disagreeing members, and
the docs do not mention the last two. Two nearby passages have the same problem.
- Carry rule (1) (506-512) says a union 'expands' the declared set over a
marked leg, which the same rule says must drop the claim. In fact
`canMergeMarkedKeys` never adds a key beyond the marked leg's declaration.
- The comment on the `project` assert (843-846) blames distinct keys
collapsing, but a regroup survives that. The real break is an undeclared row
under another declared key's label projecting to a different value. For
example, project keys [(1, a), (2, b)] to [1, 2] with an undeclared (1, z)
routed into the (2, b) partition. Projected key 1 is then split across two
labels without any collapse.
A maintainer reads these docs before adding a carrier, so the producer list
should match the code.
See **Shared repair plan 1** in the review body.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -817,21 +827,25 @@ private[sql] object GroupPartitionsExec {
case _ => false
}
// A `copy` of the child's layout rather than a fresh one, so the marker
and anything the
- // layout grows later are carried without this having to name them. The
origin is the one
+ // layout grows later are carried without this having to name them, except
for the routing:
+ // only an identity regrouping leaves an undeclared row where it was. The
origin is the one
Review Comment:
**Nit (P3):** The marker and the routing are one field, `unknownKeysOrigin`,
and lines 837-838 and 846 name it explicitly, so it is not carried 'without
this having to name' it. 'The origin is the one field not carried' was true
while `ungroupingOrigin` was the only origin. Now this code decides two origin
fields, so the sentence is false and 'the origin' could mean either one. The
comment would match the code if it said that both are set explicitly:
`unknownKeysOrigin`, which only an identity regrouping keeps, and
`ungroupingOrigin`.
See **Shared repair plan 1** in the review body.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -727,7 +758,7 @@ case class KeyedPartitioning(
*/
override protected def stringArgs: Iterator[Any] =
Iterator(expressions, partitionKeys, isGrouped, isCollapsed,
- mayContainUnknownPartitionKeys) ++ ungroupingOrigin
+ unknownKeysOrigin) ++ ungroupingOrigin
Review Comment:
**Nit (P3):** `flatArguments` (Expression.scala 363-366) flattens only
`Iterable`, and in Scala 2.13 `Option` is only `IterableOnce`, so this slot
renders as `Some(HASH_ROUTING)`, `Some(ROUTING_LOST)` or `None`. Every keyed
exchange from `createPartitioning` now ends in `Some(HASH_ROUTING)`, and every
unmarked SPJ partitioning ends in `None`. The `++ ungroupingOrigin` beside it
flattens to a bare name or nothing, so the two optional layout fields print
differently on the same line. The two case-object families have distinct names,
so flattening this one the same way keeps the trailing slots unambiguous:
```suggestion
Iterator(expressions, partitionKeys, isGrouped, isCollapsed) ++
unknownKeysOrigin ++
ungroupingOrigin
```
**Verification:**
- **Behavior:** Rendering a hash-routed, a routing-lost, and an unmarked
KeyedPartitioning to a string shows the bare state name for the first two and
no routing token for the third.
- **Compatibility:** No plan-stability or explain golden output that renders
KeyedPartitioning changes unexpectedly.
--
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]
