You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ul...@apache.org on 2026/10/21 12:16:26 UTC
(spark) branch master updated: [SPARK-60277][SQL] Preserve keyed partitioning through unions with unknown keys
This is an automated email from the ASF dual-hosted git repository.
ulysses-you pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 2f58f6239635 [SPARK-60277][SQL] Preserve keyed partitioning through unions with unknown keys
2f58f6239635 is described below
commit 2f58f62396354367a2485ac82446048b738c600b
Author: Xiduo You <ul...@gmail.com>
AuthorDate: Wed Oct 21 20:16:11 2026 +0800
[SPARK-60277][SQL] Preserve keyed partitioning through unions with unknown keys
### What changes were proposed in this pull request?
Allow `UnionExec` to concatenate marked keyed layouts when:
- every unmarked leg declares only keys covered by the marked legs;
- marked legs have identical declared key lists, including order and multiplicity;
- if multiple legs are marked, all retain hash routing.
Replace the Boolean field on `KeyLayout` with `unknownKeysOrigin`. `HASH_ROUTING` records that undeclared keys follow `nonNegativeMod(key.hashCode, numPartitions)` for the current layout; `ROUTING_LOST` retains the grouping claim without promising a partition index. The Boolean accessor remains derived from the origin.
`KeyedPartitioning.concat` preserves hash routing only for a grouped result with at most one nonempty leg. Ordinary `GroupPartitionsExec` regrouping retains keyed clustering, while key rewriting and non-identity expected-key alignment still drop the claim. Shuffle templates and pairing of two marked layouts require hash routing; a marked layout may still pair with an accurate layout whose declared keys are a subset.
Exclude marked layouts from the constant-key ordering shortcut when GPE concatenates partitions. Undeclared rows can have different key values within a merged partition, so the shortcut could incorrectly remove a window's required sort. No-coalescing and planned sorted-merge ordering remain unchanged.
The union containment check builds a set once when unmarked legs exist, rather than scanning the marked key list for each unmarked key. Plan strings show bare routing-origin names.
Non-identity expected-key alignment remains conservative, so joins with differing declared key sets may retain the base's shuffle count. `toGrouped` also conservatively loses routing even when its key list is unchanged.
### Why are the changes needed?
A union leg from a one-side-shuffled outer join may hold undeclared keys, even when its declared keys cover every key the other union legs declare. Rejecting all such merges adds an unnecessary downstream shuffle.
For example, let `a` declare identity-partitioned keys `{1, 2, 3}`, `s` declare `{1, 2}`, and `t` be unpartitioned and contain an undeclared key `4`:
```sql
SELECT id, count(*) FROM (
SELECT t.id AS id FROM a RIGHT OUTER JOIN t ON a.id = t.id
UNION ALL
SELECT id FROM s
) GROUP BY id
```
The merge is safe: declared keys reach their named groups, and undeclared keys occur only in the marked leg. The aggregate can group the union's partitions instead of shuffling them, leaving the join's one-side shuffle as the only exchange. Multiple marked legs are admitted only when their routing proves that equal undeclared keys reach the same declared-key group.
### Does this PR introduce _any_ user-facing change?
Yes. Eligible marked keyed unions retain partitioning and can avoid downstream shuffles. Such unions use the existing partitioning-aware execution path rather than whole-stage union codegen fusion, as unmarked keyed unions already do.
Results are unchanged. No configuration or public API is added.
### How was this patch tested?
Tests cover:
- safe and rejected marked union merges, including uncovered keys, differing order or multiplicity, and lost routing;
- downstream aggregate, join and window results, including AQE on/off joins and a different-declared-key-set shuffle fallback;
- two SQL regressions for differently ordered marked declarations and two lost-routing marked unions, with AQE-off exchange counts of 3 and 4;
- a locally sorted marked union feeding a window, with preserve-key-ordering on/off controls and assertions that the required upper sort remains;
- routing propagation, shared layouts, concat's local routing invariant, shuffle-template eligibility, and marked pairing;
- an expanded EnsureRequirements pairing sweep covering unmarked, hash-routed, routing-lost and ungrouped layouts;
- bare origin names in plan strings and ordering preservation without coalescing or with a planned sorted merge.
The SQL regressions detect weakened guards: comparing marked key sets instead of ordered lists produces 76 aggregate groups instead of 40; removing the marked/marked routing guard produces 4 join groups instead of 40. Removing the marked coalescing ordering guard gives both rows for `id=1` window count `1` instead of `2`.
Passed 59 Catalyst tests across `ShuffleSpecSuite` and `DistributionSuite`, and 545 SQL tests across `GroupPartitionsExecSuite`, `EnsureRequirementsSuite`, `KeyGroupedPartitioningSuite`, `ValidateRequirementsSuite`, `UnionCodegenSuite`, `DataFrameSetOperationsSuite`, `ProjectedOrderingAndPartitioningSuite`, and `ExplainSuite`. Catalyst and SQL main/test scalastyle passed.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code; OpenAI GPT-6.1
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Closes #59535 from ulysses-you/union-marked-keyed-leg.
Authored-by: Xiduo You <ul...@gmail.com>
Signed-off-by: Xiduo You <ul...@apache.org>
---
.../sql/catalyst/plans/physical/partitioning.scala | 193 +++++---
.../spark/sql/catalyst/DistributionSuite.scala | 29 +-
.../spark/sql/catalyst/ShuffleSpecSuite.scala | 59 ++-
.../sql/execution/basicPhysicalOperators.scala | 46 +-
.../datasources/v2/GroupPartitionsExec.scala | 60 +--
.../execution/exchange/EnsureRequirements.scala | 20 +-
.../spark/sql/execution/joins/ShuffledJoin.scala | 11 +-
.../connector/KeyGroupedPartitioningSuite.scala | 501 ++++++++++++++++++++-
.../datasources/v2/GroupPartitionsExecSuite.scala | 98 ++--
.../exchange/EnsureRequirementsSuite.scala | 39 +-
10 files changed, 861 insertions(+), 195 deletions(-)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
index ad6d67a19d9d..f4d77b0b51dc 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
@@ -447,6 +447,23 @@ 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. Non-identity regrouping and `toGrouped` lose routing, as does
+ * concatenation with duplicate keys or multiple nonempty legs. Combining different origins also
+ * records `ROUTING_LOST`.
+ * Key-rewriting projections and reductions cannot retain a marked keyed claim.
+ */
+sealed trait UnknownKeysOrigin
+
+// An undeclared row sits at `nonNegativeMod(key.hashCode, numPartitions)`; the relation holds.
+case object HASH_ROUTING extends UnknownKeysOrigin
+
+// The layout no longer guarantees the hash-index relation; full-key grouping remains possible.
+case object ROUTING_LOST extends UnknownKeysOrigin
+
/**
* The physical layout of the partitions a [[KeyedPartitioning]] describes, which is everything
* about them except the expressions naming them.
@@ -473,34 +490,33 @@ case object SPLIT_FOR_ORDERING extends UngroupingOrigin
* partitioning this layout was derived from onto the same key, so one key here
* can stand for several of the original ones. Sticky. See
* [[KeyedPartitioning]]'s "Key Collapse" for what it gates and how it travels.
- * @param mayContainUnknownPartitionKeys Whether the data may contain rows whose partition key is
- * not among the declared `partitionKeys`. `KeyGroupedPartitioner`
- * routes such rows by a deterministic hash when a side is
- * re-shuffled onto this partitioning (see
- * `KeyedShuffleSpec.createPartitioning`), so co-location holds
- * for whole keys only: two marked partitionings declaring the
- * same keys in the same order and using the same partition
- * function per position still pair (equal undeclared keys hash
- * to the same partition), but a row of an undeclared key sits in
- * the partition of some other declared key and need not be
- * co-located with rows sharing only a subset of its columns.
- * `satisfies` and `KeyedShuffleSpec.areKeysCompatible` therefore
- * accept a marked partitioning only for full-key clustering,
- * never for a subset of its partition columns and never for a
- * global ordering across several partitions. Two carry rules:
- * (1) a node that changes the declared key set must drop the
- * keyed partitioning, whether it coarsens it (a key-dropping
- * projection, a key-changing reduction, a join-key projection)
- * or expands it over a marked leg (a union, where another leg
- * may declare exactly the key that leg holds out-of-set);
- * (2) marker agreement comes free with the layout, since the
- * members of a [[PartitioningCollection]] share it by reference
- * and `fromPartitionings` merges one canonical layout, so
- * consumers may read any member. `ShuffledJoin`'s `InnerLike`
- * arm, the only site that meets a marked input with an unmarked
- * one, clears the markers first. That is precision, not the
- * guarantee: without it the merge would spread the spurious
- * marker onto the accurate side.
+ * @param unknownKeysOrigin Whether the data may contain rows whose partition key is not among the
+ * declared `partitionKeys`, and whether the layout still says where they
+ * sit. `None` when every row's key is declared. See
+ * [[UnknownKeysOrigin]]. `KeyGroupedPartitioner` routes such rows by a
+ * deterministic hash when a side is re-shuffled onto this partitioning
+ * (see `KeyedShuffleSpec.createPartitioning`), so co-location holds for
+ * whole keys only: two marked partitionings declaring the same keys in
+ * the same order and still routing by that hash still pair (equal
+ * undeclared keys hash to the same partition), but a row of an undeclared
+ * key sits in the partition of some other declared key and need not be
+ * co-located with rows sharing only a subset of its columns.
+ * `satisfies` and `KeyedShuffleSpec` therefore accept a marked
+ * partitioning only for full-key clustering, never for a subset of its
+ * partition columns and never for a global ordering across several
+ * partitions. Two carry rules:
+ * (1) key-rewriting projections and reductions drop the marked claim.
+ * A union retains it only when the legs prove full-key grouping remains
+ * possible, without adding keys beyond the marked declaration. Hash
+ * routing survives only a grouped concatenation with at most one nonempty
+ * leg; other concatenations and non-identity regrouping lose routing;
+ * (2) state agreement comes free with the layout, since the members of a
+ * [[PartitioningCollection]] share it by reference and
+ * `fromPartitionings` merges one canonical layout, so consumers may read
+ * any member. `ShuffledJoin`'s `InnerLike` arm clears the marker when
+ * marked and unmarked inputs meet. That is precision, not the guarantee:
+ * without it the merge would spread the spurious state onto the accurate
+ * side.
* @param ungroupingOrigin Why the keys repeat on purpose, when a producer built them that way
* rather than leaving them for the grouping node that would settle them.
* A clustering refuses an ungrouped layout whatever this says: on its
@@ -516,7 +532,7 @@ case class KeyLayout(
dataTypes: Seq[DataType],
isGrouped: Boolean,
isCollapsed: Boolean,
- mayContainUnknownPartitionKeys: Boolean = false,
+ unknownKeysOrigin: Option[UnknownKeysOrigin] = None,
ungroupingOrigin: Option[UngroupingOrigin] = None) {
// The rows carry the types they are compared at, so the pair is checked against itself rather
@@ -526,6 +542,19 @@ case class KeyLayout(
s"A KeyLayout's dataTypes ($dataTypes) must be the types its keys are compared at " +
s"(${partitionKeys.head.dataTypes})")
+ def mayContainUnknownPartitionKeys: Boolean = unknownKeysOrigin.isDefined
+
+ /**
+ * Whether an undeclared row still sits at `nonNegativeMod(key.hashCode, numPartitions)`. A
+ * marked layout answers true only while its key list is the one the data was routed over; every
+ * carrier that changes the list records `ROUTING_LOST` instead.
+ */
+ def unknownKeysRouted: Boolean = unknownKeysOrigin.forall(_ == HASH_ROUTING)
+
+ /** The unknown-keys state after a change that moved the keys, which gives the routing up. */
+ def unknownKeysRoutingLost: Option[UnknownKeysOrigin] =
+ unknownKeysOrigin.map(_ => ROUTING_LOST)
+
/**
* Whether `other` describes the same partition keys as this layout, which is the question every
* caller comparing two layouts' keys is asking.
@@ -707,6 +736,8 @@ case class KeyedPartitioning(
def isGrouped: Boolean = layout.isGrouped
def isCollapsed: Boolean = layout.isCollapsed
def mayContainUnknownPartitionKeys: Boolean = layout.mayContainUnknownPartitionKeys
+ def unknownKeysOrigin: Option[UnknownKeysOrigin] = layout.unknownKeysOrigin
+ def unknownKeysRouted: Boolean = layout.unknownKeysRouted
def ungroupingOrigin: Option[UngroupingOrigin] = layout.ungroupingOrigin
/** This partitioning over a changed layout, e.g. `withLayout(_.copy(isGrouped = false))`. */
@@ -726,8 +757,8 @@ case class KeyedPartitioning(
* string is the only place a reader sees why a marked partitioning still shuffles.
*/
override protected def stringArgs: Iterator[Any] =
- Iterator(expressions, partitionKeys, isGrouped, isCollapsed,
- mayContainUnknownPartitionKeys) ++ ungroupingOrigin
+ Iterator(expressions, partitionKeys, isGrouped, isCollapsed) ++ unknownKeysOrigin ++
+ ungroupingOrigin
override protected def withNewChildrenInternal(
newChildren: IndexedSeq[Expression]): KeyedPartitioning =
@@ -783,8 +814,8 @@ case class KeyedPartitioning(
* Projects this partitioning onto the key positions in `positions`, whose order becomes the order
* of the projected expressions and key fields. The projected `isGrouped` and `isCollapsed` are
* computed together, because dropping a position can map keys that were distinct here onto the
- * same key. Keeping every position changes no key, so both answers are inherited as they are,
- * with no pass over the keys. `GroupPartitionsExec` decides the same two answers from the key
+ * same key. Keeping every position in its original order returns this partitioning without
+ * scanning keys. `GroupPartitionsExec` decides the same two answers from the key
* groups it keeps, so it does not use this.
*/
def project(positions: Seq[Int]): KeyedPartitioning = {
@@ -810,9 +841,10 @@ case class KeyedPartitioning(
}
}
// A `copy` of the layout rather than a fresh one, so every field this method does not decide
- // is carried, including the unknown-keys marker and anything the layout grows later. A
- // projection that coarsens the declared key set cannot keep a marked claim, and
- // `createShuffleSpec` refuses that case before it gets here, but this does not depend on it.
+ // is carried, including anything the layout grows later. A marked projection cannot keep
+ // co-location: an undeclared row can project to a different key than its partition's label.
+ assert(!mayContainUnknownPartitionKeys,
+ "A key-dropping projection cannot be applied to a layout that may hold undeclared rows")
copy(
expressions = positions.map(expressions),
layout = layout.copy(
@@ -828,7 +860,8 @@ case class KeyedPartitioning(
val uniqueKeys = if (isGrouped) partitionKeys else partitionKeys.distinct
// The keys are unique afterwards, so a reason to be ungrouped has been spent.
withLayout(_.copy(partitionKeys = uniqueKeys.sorted(keyOrdering), isGrouped = true,
- ungroupingOrigin = None))
+ ungroupingOrigin = None,
+ unknownKeysOrigin = layout.unknownKeysRoutingLost))
}
/**
@@ -1129,10 +1162,9 @@ case class KeyedPartitioning(
val joinKeyPositions = result.keyPositions.map(_.nonEmpty).zipWithIndex.filter(_._1).map(_._2)
// A marked claim is never projected. The claim says the undeclared rows sit at
// `hash(key) % numPartitions` over *these* keys in *this* order, and both halves of the
- // projection below break it: narrowing coarsens the declared set, and `toGrouped` sorts it,
- // which sends a consumer's undeclared rows where this side does not hold them. Producing
- // either would also need a `GroupPartitionsExec` on this side, and that node gives the keyed
- // claim up rather than perform a non-identity regrouping of a marked layout.
+ // projection below can break it: narrowing loses co-location, and `toGrouped` sorts it,
+ // which can change undeclared rows' partition indices. Plain regrouping retains clustering
+ // but loses routing, so its result cannot be a shuffle template.
//
// What the unprojected spec is good for depends on the layout. Where every partition
// expression covers a clustering key, so the projection would only have re-sorted, the spec
@@ -1201,6 +1233,19 @@ 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
@@ -1208,18 +1253,30 @@ object KeyedPartitioning {
*
* 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. Hash routing survives only when the combined origin is
+ * `HASH_ROUTING`, the result is grouped, and at most one leg is nonempty. Other marked
+ * concatenations lose routing: each leg routed undeclared rows over its own partition 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 || kps.count(_.numPartitions > 0) > 1 =>
+ Some(ROUTING_LOST)
+ case other => other
+ }
kps.head.withLayout(_.copy(
ungroupingOrigin = ungroupingOrigin,
partitionKeys = concatenatedKeys,
- // A child that has duplicates of its own puts them in the concatenation too, which answers
- // this without walking the keys.
- isGrouped = kps.forall(_.isGrouped) &&
- concatenatedKeys.distinct.length == concatenatedKeys.length,
- isCollapsed = kps.exists(_.isCollapsed)))
+ isGrouped = grouped,
+ isCollapsed = kps.exists(_.isCollapsed),
+ unknownKeysOrigin = routing))
}
def supportsExpressions(expressions: Seq[Expression]): Boolean = {
@@ -1624,13 +1681,12 @@ object PartitioningCollection {
}
/**
- * The uniform `mayContainUnknownPartitionKeys` marker of `p`'s keyed members, read from its
- * first keyed member, the same one `checkKeyedPartitioningInvariant` compares against.
- * `None` when `p` holds no [[KeyedPartitioning]] at all, distinct from the `Some(false)` of
- * an unmarked keyed one: there is no claim there to answer about.
+ * The uniform `unknownKeysOrigin` state of `p`'s keyed members, read from its first keyed
+ * member, the same one `checkKeyedPartitioningInvariant` compares against. `None` when `p` holds
+ * no [[KeyedPartitioning]], and when its keyed member holds no undeclared rows.
*/
- private[sql] def keyedMarkerOf(p: Partitioning): Option[Boolean] =
- representativeOf(p).map(_.mayContainUnknownPartitionKeys)
+ private[sql] def keyedUnknownKeysOf(p: Partitioning): Option[UnknownKeysOrigin] =
+ representativeOf(p).flatMap(_.layout.unknownKeysOrigin)
/**
* Builds a [[PartitioningCollection]], unifying the [[KeyLayout]] reference across all
@@ -1641,17 +1697,19 @@ object PartitioningCollection {
* Note: this can't be implemented with `TreeNode.transform`.
*/
def fromPartitionings(partitionings: Seq[Partitioning]): PartitioningCollection = {
- // See the class doc for why the flags are normalized by OR rather than required to agree. One
+ // See the class doc for why the flags are normalized rather than required to agree. One
// representative per member is enough, because every collection agrees internally by this same
// construction, and only a member that disagrees is rebuilt.
val representatives = partitionings.flatMap(representativeOf)
val anyCollapsed = representatives.exists(_.isCollapsed)
- val anyUnknownKeys = representatives.exists(_.mayContainUnknownPartitionKeys)
-
- // The canonical layout is the first one that already agrees with the OR'd flags, so the members
- // that agree keep their instance and only the ones that disagree are rebuilt. Anchoring on the
- // first representative instead would push *every* member off the `eq` path whenever that one
- // happened to need a flag fixed, since no existing layout is `eq` a fresh copy.
+ // The weakest state wins: any member that may hold undeclared rows marks the joint layout,
+ // and one that lost the routing leaves it unable to say where they sit.
+ val anyRouting = KeyedPartitioning.combine(representatives.map(_.layout.unknownKeysOrigin))
+
+ // The canonical layout is the first one that already agrees with the normalized flags, so the
+ // members that agree keep their instance and only the ones that disagree are rebuilt. Anchoring
+ // on the first representative instead would push *every* member off the `eq` path whenever that
+ // one happened to need a flag fixed, since no existing layout is `eq` a fresh copy.
//
// A partitioning with no `KeyedPartitioning` in it has nothing to normalize, and one that
// already holds the canonical layout is returned as it is. That is what keeps repeated
@@ -1663,10 +1721,10 @@ object PartitioningCollection {
PartitioningCollection.concatUngrouping(representatives.map(_.layout.ungroupingOrigin))
val canonicalLayout = representatives.map(_.layout)
.find(l => l.isCollapsed == anyCollapsed &&
- l.mayContainUnknownPartitionKeys == anyUnknownKeys &&
+ l.unknownKeysOrigin == anyRouting &&
l.ungroupingOrigin == concatUngrouping)
.orElse(representatives.headOption.map(_.layout.copy(
- isCollapsed = anyCollapsed, mayContainUnknownPartitionKeys = anyUnknownKeys,
+ isCollapsed = anyCollapsed, unknownKeysOrigin = anyRouting,
ungroupingOrigin = concatUngrouping)))
.orNull
@@ -2270,20 +2328,22 @@ case class KeyedShuffleSpec(
// their order: the out-of-set keys hash to the same-index partition on both sides, and a
// `GroupPartitionsExec` regrouping re-labels each partition by that side's declared key, so
// a differing declared order would push the out-of-set keys into different output
- // partitions and lose their matches.
+ // partitions and lose their matches. They also have to still route by that hash: a side a
+ // regrouping moved off it reports `ROUTING_LOST`, and two such sides cannot pair this way.
if (unknownKeys) {
// Comparing the declared keys also asks that neither side's keys were reduced, since
// reduced keys are in a space neither expression describes. A marked layout never carries
// them: `canCreatePartitioning` requires `expressionsDescribeKeys` before
// `KeyedShuffleSpec.createPartitioning` can mark anything, a `GroupPartitionsExec` that
- // reduces is not an identity regrouping and so gives the keyed claim up, and
+ // reduces rewrites keys and so gives the keyed claim up, and
// `ShuffledJoin.clearUnknownPartitionKeys` strips the marker a `PartitioningCollection`
// would otherwise OR onto a merged layout. Asked rather than asserted, so that a producer
// this misses costs a shuffle rather than the query.
partitioning.expressionsDescribeKeys && other.partitioning.expressionsDescribeKeys && {
if (partitioning.mayContainUnknownPartitionKeys &&
other.partitioning.mayContainUnknownPartitionKeys) {
- partitioning.partitionKeys == other.partitioning.partitionKeys
+ partitioning.unknownKeysRouted && other.partitioning.unknownKeysRouted &&
+ partitioning.partitionKeys == other.partitioning.partitionKeys
} else if (partitioning.mayContainUnknownPartitionKeys) {
val declared = partitioning.partitionKeys.toSet
other.partitioning.partitionKeys.forall(declared.contains)
@@ -2494,6 +2554,9 @@ case class KeyedShuffleSpec(
// satisfy the distribution in a `GroupPartitionsExec` before it builds any spec -- so do not
// read this clause as redundant.
partitioning.isGrouped &&
+ // The sibling laid out here routes its undeclared rows by `hash % numPartitions`, so it
+ // lands where this side's own undeclared rows sit only while this layout still does too.
+ partitioning.unknownKeysRouted &&
// Every partition expression must map to a clustering key, otherwise `createPartitioning`
// cannot rewrite it. This also keeps the unprojected spec returned by `createShuffleSpec`
// for a marked narrowing projection from being chosen as the best spec.
@@ -2537,7 +2600,7 @@ case class KeyedShuffleSpec(
// One `copy`, so no intermediate node is built and discarded.
partitioning.copy(
expressions = newExpressions,
- layout = partitioning.layout.copy(mayContainUnknownPartitionKeys = true))
+ layout = partitioning.layout.copy(unknownKeysOrigin = Some(HASH_ROUTING)))
}
}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/DistributionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/DistributionSuite.scala
index 5ca1ead5d22e..da20e40e043b 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/DistributionSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/DistributionSuite.scala
@@ -516,9 +516,9 @@ class DistributionSuite extends SparkFunSuite with SQLHelper {
val a = AttributeReference("a", IntegerType)()
val ordered = OrderedDistribution(Seq(SortOrder(a, Ascending)))
val markedOne = KeyedPartitioning(Seq(a), Seq(InternalRow(1)))
- .withLayout(_.copy(mayContainUnknownPartitionKeys = true))
+ .withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING)))
val markedTwo = KeyedPartitioning(Seq(a), Seq(InternalRow(1), InternalRow(2)))
- .withLayout(_.copy(mayContainUnknownPartitionKeys = true))
+ .withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING)))
withSQLConf(SQLConf.V2_BUCKETING_SORTING_ENABLED.key -> "true") {
checkSatisfied(markedOne, ordered, true)
checkSatisfied(markedTwo, ordered, false)
@@ -529,14 +529,23 @@ class DistributionSuite extends SparkFunSuite with SQLHelper {
val x = AttributeReference("x", IntegerType)()
val y = AttributeReference("y", IntegerType)()
val marked = KeyedPartitioning(Seq(x), Seq(InternalRow(1), InternalRow(2)))
- .withLayout(_.copy(mayContainUnknownPartitionKeys = true))
+ .withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING)))
val plain = KeyedPartitioning(Seq(y), Seq(InternalRow(1), InternalRow(2)))
val combined = PartitioningCollection.fromPartitionings(Seq(marked, plain))
- // The conservative direction: an unmarked member must never excuse marked data, because
- // the OR'd-on marker only costs a shuffle while the AND'd-off one could cost correctness.
+ // The conservative direction: an unmarked member must never excuse marked data, because a
+ // merged marker only costs a shuffle while dropping one could cost correctness.
val members = combined.partitionings.map(_.asInstanceOf[KeyedPartitioning])
assert(members.forall(_.mayContainUnknownPartitionKeys), members.toString)
assert(members.last.layout eq members.head.layout)
+
+ // A member that lost the routing makes the joint layout lose it as well.
+ val lost = KeyedPartitioning(Seq(x), Seq(InternalRow(1), InternalRow(2)))
+ .withLayout(_.copy(unknownKeysOrigin = Some(ROUTING_LOST)))
+ val withLost = PartitioningCollection.fromPartitionings(Seq(lost, plain))
+ val lostMembers = withLost.partitionings.map(_.asInstanceOf[KeyedPartitioning])
+ assert(lostMembers.forall(_.mayContainUnknownPartitionKeys), lostMembers.toString)
+ assert(lostMembers.forall(!_.unknownKeysRouted), lostMembers.toString)
+ assert(lostMembers.last.layout eq lostMembers.head.layout)
}
test("SPARK-59050: PartitioningCollection requires members to agree on the marker") {
@@ -544,11 +553,11 @@ class DistributionSuite extends SparkFunSuite with SQLHelper {
val y = AttributeReference("y", IntegerType)()
val keys = Seq(InternalRow(1), InternalRow(2))
val base = KeyedPartitioning(Seq(x), keys)
- val marked = base.withLayout(_.copy(mayContainUnknownPartitionKeys = true))
+ val marked = base.withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING)))
// The marker lives on the layout, so two members that disagree on it hold two layouts, and the
// one `eq` on the layout refuses them. There is no clause of its own to reach.
val disagree =
- KeyedPartitioning(Seq(y), marked.layout.copy(mayContainUnknownPartitionKeys = false))
+ KeyedPartitioning(Seq(y), marked.layout.copy(unknownKeysOrigin = None))
val err = intercept[IllegalArgumentException] {
PartitioningCollection(Seq(marked, disagree))
}
@@ -559,8 +568,8 @@ class DistributionSuite extends SparkFunSuite with SQLHelper {
val x = AttributeReference("x", IntegerType)()
val y = AttributeReference("y", IntegerType)()
val keys = Seq(InternalRow(1), InternalRow(2))
- val marked =
- KeyedPartitioning(Seq(x), keys).withLayout(_.copy(mayContainUnknownPartitionKeys = true))
+ val marked = KeyedPartitioning(Seq(x), keys)
+ .withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING)))
val plainNested = PartitioningCollection.fromPartitionings(
Seq(KeyedPartitioning(Seq(y), keys)))
val combined = PartitioningCollection.fromPartitionings(Seq(marked, plainNested))
@@ -581,7 +590,7 @@ class DistributionSuite extends SparkFunSuite with SQLHelper {
val keys = Seq(InternalRow(1), InternalRow(2))
val base = KeyedPartitioning(Seq(x), keys)
val markedNested = PartitioningCollection.fromPartitionings(Seq(
- base.withLayout(_.copy(mayContainUnknownPartitionKeys = true))))
+ base.withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING)))))
val plain = base.copy(expressions = Seq(y))
// The nested collection is internally uniform, so its own construction passes. The outer
// constructor must still catch its representative against the unmarked sibling, and the marker
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala
index 7b20a1337d3e..d4d5d50f325c 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala
@@ -172,7 +172,7 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
// A marked layout is never offered through a projection: the projecting half of `keysSatisfy`
// excludes a marked member, so the keys an offered marked one carries are the keys its claim
// is over.
- val marked = source.withLayout(_.copy(mayContainUnknownPartitionKeys = true))
+ val marked = source.withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING)))
assert(PartitioningCollection.specsForPairing(marked, cd).isEmpty,
"a marked member is not offered through a projection")
}
@@ -813,11 +813,13 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
test("areKeysCompatible: unknown partition keys only allow a subset of the declared keys") {
val a = $"a".int
val distribution = ClusteredDistribution(Seq(a))
+ def routing(hasUnknown: Boolean): Option[UnknownKeysOrigin] =
+ Option.when(hasUnknown)(HASH_ROUTING)
def keyedSpec(
keys: Seq[Int],
hasUnknown: Boolean = false): KeyedShuffleSpec = KeyedShuffleSpec(
KeyedPartitioning(Seq(a), keys.map(k => InternalRow(k)))
- .withLayout(_.copy(mayContainUnknownPartitionKeys = hasUnknown)), distribution)
+ .withLayout(_.copy(unknownKeysOrigin = routing(hasUnknown))), distribution)
// A partitioning with unknown partition keys (e.g. a side re-shuffled onto a keyed layout by
// `KeyedShuffleSpec.createPartitioning`) only guarantees co-location for its declared keys, so
@@ -873,16 +875,18 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
val a = $"a".long
val distribution = ClusteredDistribution(Seq(a))
+ def routing(hasUnknown: Boolean): Option[UnknownKeysOrigin] =
+ Option.when(hasUnknown)(HASH_ROUTING)
def bucketSpec(
keys: Seq[Long],
hasUnknown: Boolean = false): KeyedShuffleSpec = KeyedShuffleSpec(
KeyedPartitioning(Seq(bucket(4, a)), keys.map(k => InternalRow(k)))
- .withLayout(_.copy(mayContainUnknownPartitionKeys = hasUnknown)), distribution)
+ .withLayout(_.copy(unknownKeysOrigin = routing(hasUnknown))), distribution)
def keyedSpec(
keys: Seq[Long],
hasUnknown: Boolean = false): KeyedShuffleSpec = KeyedShuffleSpec(
KeyedPartitioning(Seq(a), keys.map(k => InternalRow(k)))
- .withLayout(_.copy(mayContainUnknownPartitionKeys = hasUnknown)), distribution)
+ .withLayout(_.copy(unknownKeysOrigin = routing(hasUnknown))), distribution)
// `isExpressionCompatible` admits an identity-vs-transform pair when compatible transforms
// are allowed, but then the two sides' partition keys live in different domains: raw values
@@ -1061,11 +1065,52 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
}
}
+ test("SPARK-60277: concat preserves routing only for one grouped nonempty leg") {
+ val a = $"a".int
+ def kp(keys: Int*): KeyedPartitioning =
+ KeyedPartitioning(Seq(a), keys.map(InternalRow(_)))
+ val marked = kp(1, 2, 3).withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING)))
+ val accurate = kp(4, 5)
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ Seq(Seq(marked, accurate), Seq(accurate, marked)).foreach { legs =>
+ val out = KeyedPartitioning.concat(legs)
+ assert(out.isGrouped && out.mayContainUnknownPartitionKeys && !out.unknownKeysRouted)
+ assert(!KeyedShuffleSpec(out, ClusteredDistribution(Seq(a))).canCreatePartitioning)
+ }
+ val unchanged = KeyedPartitioning.concat(Seq(kp(), marked, kp()))
+ assert(unchanged.isGrouped && unchanged.unknownKeysOrigin.contains(HASH_ROUTING))
+ assert(KeyedShuffleSpec(unchanged, ClusteredDistribution(Seq(a))).canCreatePartitioning)
+ }
+ val repeated = kp(1, 1, 2).withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING)))
+ assert(!KeyedPartitioning.concat(Seq(kp(), repeated)).unknownKeysRouted)
+ }
+
+ test("SPARK-60277: keyed plan strings show bare unknown-key origins") {
+ val base = KeyedPartitioning(Seq($"a".int), Seq(InternalRow(1)))
+ assert(!base.toString.contains("None"))
+ Seq(HASH_ROUTING, ROUTING_LOST).foreach { origin =>
+ val rendered = base.withLayout(_.copy(unknownKeysOrigin = Some(origin))).toString
+ assert(rendered.contains(origin.toString))
+ assert(!rendered.contains("Some("))
+ }
+ }
+
+ test("SPARK-60277: toGrouped gives the routing up") {
+ val a = $"a".int
+ val marked = KeyedPartitioning(Seq(a), Seq(InternalRow(1), InternalRow(1), InternalRow(2)))
+ .withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING)))
+ assert(!marked.isGrouped, "test setup: the duplicate key keeps the layout ungrouped")
+ val grouped = marked.toGrouped
+ assert(grouped.isGrouped && grouped.mayContainUnknownPartitionKeys &&
+ !grouped.unknownKeysRouted,
+ s"dedup and sort move the keys off the routing: $grouped")
+ }
+
test("createShuffleSpec: a marked narrowing projection yields an unusable spec") {
val a = $"a".int
val b = $"b".int
val marked = KeyedPartitioning(Seq(a, b), Seq(InternalRow(1, 2), InternalRow(3, 4)))
- .withLayout(_.copy(mayContainUnknownPartitionKeys = true))
+ .withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING)))
withSQLConf(
SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
@@ -1086,7 +1131,7 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
// refused would only have re-sorted the keys.
val marked = KeyedPartitioning(Seq(a),
Seq(InternalRow(3), InternalRow(4), InternalRow(1), InternalRow(2)))
- .withLayout(_.copy(mayContainUnknownPartitionKeys = true))
+ .withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING)))
assert(marked.isGrouped, "test setup: the keys are distinct, so the layout is grouped")
withSQLConf(
SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
@@ -1110,7 +1155,7 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
KeyedPartitioning(
Seq(TransformExpression(fn, Seq(a), Some(numBuckets))),
Seq(InternalRow(0L), InternalRow(1L)))
- .withLayout(_.copy(mayContainUnknownPartitionKeys = hasUnknown)),
+ .withLayout(_.copy(unknownKeysOrigin = Option.when(hasUnknown)(HASH_ROUTING))),
ClusteredDistribution(Seq(a)))
// `allowCompatibleTransforms` lets a differing-bucket-count pair through
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
index 070d06451312..d8314b8807a0 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
@@ -1000,13 +1000,9 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup
if (partitionings.forall(_.isInstanceOf[KeyedPartitioning])) {
val kps = partitionings.map(_.asInstanceOf[KeyedPartitioning])
val headKp = kps.head
- // To merge, the `KeyedPartitioning`s must agree on the partition expressions and no leg
- // may carry the marker: the merged set declares the other legs' keys, so an out-of-set row
- // can sit in the wrong partition for the merged claim (rule (1) of the
- // `KeyedPartitioning.mayContainUnknownPartitionKeys` doc). Unlike a join, a union keeps
- // every leg's rows: an unmarked leg does not excuse a marked one.
- val compatible = kps.forall(kp =>
- !kp.mayContainUnknownPartitionKeys && comparePartitioning(kp, headKp))
+ // To merge, the `KeyedPartitioning`s must agree on the partition expressions, and a leg
+ // carrying the marker must declare what the others declare.
+ val compatible = kps.forall(kp => comparePartitioning(kp, headKp)) && canMergeMarkedKeys(kps)
if (compatible) {
return KeyedPartitioning.concat(kps)
}
@@ -1040,6 +1036,42 @@ 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)) && {
+ val unmarked = kps.filterNot(_.mayContainUnknownPartitionKeys)
+ if (unmarked.isEmpty) {
+ true
+ } else {
+ val declaredSet = declared.toSet
+ unmarked.forall(_.partitionKeys.forall(declaredSet.contains))
+ }
+ }
+ }
+ }
+
/**
* True when this union behaves as a plain concatenation, so `unionedInputRDD` matches the
* semantics of `sparkContext.union(...)` in `unionRDDs`. It satisfies the partitioning gate on
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala
index bed5ae47dd0a..809896399a7e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala
@@ -431,8 +431,8 @@ case class GroupPartitionsExec(
// within-partition ordering is lost due to concatenation -- for example, if two input
// partitions both share key=A and hold rows (A,1),(A,3) and (A,2),(A,5) respectively (each
// sorted ascending by the data column), concatenating them yields (A,1),(A,3),(A,2),(A,5)
- // which is no longer sorted by the data column. Only sort orders over partition key
- // expressions remain valid -- they evaluate to the same value (A) in every merged partition.
+ // which is no longer sorted by the data column. For unmarked layouts without reducers, sort
+ // orders over partition keys remain valid: they evaluate to the same value in each partition.
//
// The config below and the child's ordering stay per-call reads. They only change how much
// of an ordering that holds either way is reported. So a late read can never claim more
@@ -441,9 +441,10 @@ case class GroupPartitionsExec(
// is planned, see `sortedMergeOrdering`.
outputPartitioning match {
case p: Partitioning with Expression
- if reducers.isEmpty && conf.v2BucketingPreserveKeyOrderingOnCoalesceEnabled =>
- // Without reducers all merged partitions share the same original key value, so the sort
- // orders on key expressions still hold. The transform keys match nothing here, since
+ if reducers.isEmpty && conf.v2BucketingPreserveKeyOrderingOnCoalesceEnabled &&
+ PartitioningCollection.keyedUnknownKeysOf(p).isEmpty =>
+ // Undeclared rows can have different key values within a merged partition.
+ // Unmarked layouts without reducers keep key orders. Transform keys match nothing here:
// `DataSourceV2ScanExecBase.outputOrdering` drops every sort order over a transform.
val keyedPartitionings = p.collect { case k: KeyedPartitioning => k }
val keyExprs = ExpressionSet(keyedPartitionings.flatMap(_.expressions))
@@ -566,6 +567,9 @@ case class GroupPartitionsExec(
* claim into a different key space or changes the modulus. A reducer slot counts
* as key-changing even where a conforming self-reducer could not rewrite a
* reachable key value, so the give-up there loses at most an optimization.
+ * @param keysRewritten whether the grouping re-labels the groups, by projecting the keys onto
+ * cluster-key positions or by reducing them, so an undeclared row is grouped
+ * under the image of a key the clustering does not read.
* @param numPrunedPartitions the alignment's effect on the reads of the child's splits, 0 outside
* the alignment path. See `alignToExpectedKeys`.
* @param numReplicatedPartitionReads as above.
@@ -574,6 +578,7 @@ private[sql] case class PartitionGrouping(
partitions: Seq[(InternalRowComparableWrapper, Seq[Int])],
layout: KeyLayout,
isIdentity: Boolean,
+ keysRewritten: Boolean,
numPrunedPartitions: Int,
numReplicatedPartitionReads: Int) {
@@ -616,7 +621,9 @@ private[sql] object GroupPartitionsExec {
val grouping = computeGrouping(
childKp, joinKeyPositions, expectedPartitionKeys, reducers, ungroupingOrigin)
GroupPartitionsExec(child, grouping,
- computeOutputPartitioning(childExpr, grouping, joinKeyPositions, reducers),
+ computeOutputPartitioning(
+ childExpr, grouping, joinKeyPositions, reducers,
+ hasExpectedKeys = expectedPartitionKeys.isDefined),
childPartitioning, joinKeyPositions, expectedPartitionKeys.map(_.size), reducers,
ungroupingOrigin, sortedMergeOrdering)
}
@@ -631,7 +638,8 @@ private[sql] object GroupPartitionsExec {
p: Partitioning with Expression,
grouping: PartitionGrouping,
joinKeyPositions: Option[Seq[Int]],
- reducers: Option[Seq[Option[KeyReducer]]]): Partitioning = {
+ reducers: Option[Seq[Option[KeyReducer]]],
+ hasExpectedKeys: Boolean): Partitioning = {
// There can be multiple `KeyedPartitioning`s in an output partitioning of a join, but they
// can only differ in `expressions`, since they share one `KeyLayout` (enforced by
// `PartitioningCollection`). So the grouping is computed once, and the layout it describes
@@ -642,18 +650,20 @@ private[sql] object GroupPartitionsExec {
// single-side-transform reducers; for the both-sides-reduce shape no single transform
// describes the keys, so the reduce marks it (see `KeyedShuffleSpec.reducersBothWays`).
//
- // A marked claim pins undeclared rows to hash(key) % numPartitions (see [[KeyLayout]]'s
- // `mayContainUnknownPartitionKeys`), and only an identity grouping keeps that relationship (see
- // `PartitionGrouping.isIdentity`). Clearing the marker alone would misreport the undeclared
- // rows that remain, so the keyed partitioning goes and only the physical output count is
- // reported, one per group with padding included, which is the count a parent's
- // `PartitioningCollection` needs for uniformity. `EnsureRequirements` asks this before it
- // commits to a pairing and falls back to a shuffle when it happens.
+ // A marked layout keeps its keyed claim through a plain regrouping: coalescing the partitions
+ // that share a key leaves every key's rows, declared or not, in the group named for the
+ // declared key they reached, and full-key clustering is all a `ClusteredDistribution` reads.
+ // `computeGrouping` reports that layout as `ROUTING_LOST`, since the regrouping moved an
+ // undeclared key's group off `hash(key) % numPartitions`.
+ //
+ // A key rewrite or non-identity expected-key alignment re-labels or replicates the groups, so
+ // those shapes give the keyed claim up, and only the physical output count is reported, which
+ // is the count a parent's `PartitioningCollection` needs for uniformity.
//
- // The marker read off `p` is the same bit the `else` branch publishes through
- // `grouping.layout`, since `computeGrouping` copies the child's layout and carries the
- // marker over. A change to that copy has to move this guard with it.
- if (PartitioningCollection.keyedMarkerOf(p).contains(true) && !grouping.isIdentity) {
+ // The state read off `p` is the one `computeGrouping` copies into `grouping.layout`; a change
+ // to that copy has to move this guard with it.
+ if (PartitioningCollection.keyedUnknownKeysOf(p).isDefined &&
+ (grouping.keysRewritten || (hasExpectedKeys && !grouping.isIdentity))) {
UnknownPartitioning(grouping.partitions.size)
} else {
// One instance for every member, so they share it by reference. It already carries the
@@ -827,13 +837,10 @@ private[sql] object GroupPartitionsExec {
case ((_, Seq(single)), outputIndex) => single == outputIndex
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
- // field not carried: the output claims what this node's own producer stamped, never what
- // the child's claimed, and only while the shape holds the claim -- a grouping that settles
- // every key leaves nothing repeating, so it spends the claim the way `toGrouped` does.
- // Built here, where the child is in scope, so `computeOutputPartitioning` has nothing left
- // to re-apply.
+ // Copy the child's layout, deciding both origins explicitly. Only identity regrouping keeps
+ // unknown-key routing. Ungrouping uses this node's stamp, cleared when every key is grouped.
+ val unknownKeysOrigin =
+ if (isIdentity) childKp.layout.unknownKeysOrigin else childKp.layout.unknownKeysRoutingLost
PartitionGrouping(
partitions,
childKp.layout.copy(
@@ -841,8 +848,9 @@ private[sql] object GroupPartitionsExec {
dataTypes = reducedDataTypes,
isGrouped = isGrouped,
isCollapsed = isCollapsed,
+ unknownKeysOrigin = unknownKeysOrigin,
ungroupingOrigin = if (isGrouped) None else ungroupingOrigin),
- isIdentity, numPrunedPartitions, numReplicatedPartitionReads)
+ isIdentity, keysRewritten, numPrunedPartitions, numReplicatedPartitionReads)
}
}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
index 27196074288c..45c4036bb8a0 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
@@ -657,10 +657,9 @@ case class EnsureRequirements(
l.joinKeyPositions.isEmpty && r.joinKeyPositions.isEmpty
// Whether the merged key list below may be narrowed to what the join type allows. Not when
- // either side is marked. Narrowing drops groups that side holds, so its regrouping stops being
- // the identity, and `GroupPartitionsExec` gives up a marked claim it does not regroup
- // identically. That costs the pair its join at the gate at the end of this method, which is a
- // poor trade for pruning a few groups.
+ // either side is marked. Narrowing drops groups that side holds, so the expected-key alignment
+ // is non-identity and `GroupPartitionsExec` gives up the marked claim. That costs the pair its
+ // join at the final gate, a poor trade for pruning a few groups.
//
// Narrowing is the only thing here that can shrink the list below a marked side's own keys: the
// merging arms take the union, and `KeyedShuffleSpec.areKeysCompatible` pairs a marked layout
@@ -938,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,
+ // which a layout holding undeclared rows cannot survive (see `KeyLayout.unknownKeysOrigin`),
+ // and only the node knows the permutation it performs, so that answer arrives after the pairing
+ // was chosen.
// Asking before returning is what keeps the join from skipping both shuffles for a child that
// no longer satisfies its distribution, which is a plan `ValidateRequirements` rejects and
// every AQE rule that needs a valid plan then refuses to touch.
@@ -952,10 +952,8 @@ case class EnsureRequirements(
// `KeyLayout.describesSameKeys` asks exactly that, and it compares the key types as well as
// the rows.
//
- // Both routes to `newLeft`/`newRight` build, so both are asked. On the base this was the push
- // branch's question alone, because the other route left the children exactly as the pairing
- // read them. Here `groupIfNeeded` can wrap a side that is not grouped, and that node gives up
- // the claim for the same reason the pushed one does.
+ // Check both pushed alignment and plain grouping. Plain grouping retains a marked keyed claim
+ // but loses routing; pushed key rewrites or non-identity expected alignment drop the claim.
val (newLeft, newRight) =
pushed.getOrElse((groupIfNeeded(rawLeft), groupIfNeeded(rawRight)))
def declaredLayout(plan: SparkPlan): Option[KeyLayout] =
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledJoin.scala
index 7d9d7adc15e7..4e483746dd51 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledJoin.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledJoin.scala
@@ -73,8 +73,8 @@ trait ShuffledJoin extends JoinCodegenSupport {
// `PartitioningCollection` invariant), and `keysSatisfy` admits a marked side into it only
// for full-key join keys, where the join equality ties every column of its claim: a row of
// an undeclared key matches nothing on the accurate side and is filtered, so a marked
- // member alongside an unmarked one is spurious. Clear it at the one site that can mix
- // them, so consumers can take members' markers at face value (see
+ // member alongside an unmarked one is spurious. Clear it here so consumers can take
+ // members' markers at face value (see
// `KeyedPartitioning.mayContainUnknownPartitionKeys`).
PartitioningCollection.fromPartitionings(
clearUnknownPartitionKeys(Seq(left.outputPartitioning, right.outputPartitioning)))
@@ -89,9 +89,8 @@ trait ShuffledJoin extends JoinCodegenSupport {
/**
* Clears the `mayContainUnknownPartitionKeys` marker of every `KeyedPartitioning` in
- * `partitionings` when marked and unmarked keyed inputs meet; within a collection the
- * constructor makes that unrepresentable. Only `ShuffledJoin`'s `InnerLike` arm can mix
- * inputs this way; see the call site for the argument.
+ * `partitionings` when marked and unmarked keyed inputs meet in an `InnerLike` join.
+ * See the call site for why the marker can be cleared.
*/
private def clearUnknownPartitionKeys(
partitionings: Seq[Partitioning]): Seq[Partitioning] = {
@@ -118,7 +117,7 @@ trait ShuffledJoin extends JoinCodegenSupport {
partitionings.zip(representatives).map {
case (partitioning: Partitioning with Expression, Some(representative))
if representative.mayContainUnknownPartitionKeys =>
- val copied = representative.layout.copy(mayContainUnknownPartitionKeys = false)
+ val copied = representative.layout.copy(unknownKeysOrigin = None)
val cleared = unmarkedLayout
.filter(l => (l.partitionKeys eq copied.partitionKeys) && l == copied)
.getOrElse(copied)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
index 51504692a265..a65048a03357 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
@@ -27,9 +27,9 @@ import org.apache.spark.sql.{AnalysisException, DataFrame, ExplainSuiteHelper, Q
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, Descending, ExprId, Literal, TransformExpression}
import org.apache.spark.sql.catalyst.expressions.aggregate.Complete
-import org.apache.spark.sql.catalyst.plans.{Cross, ExistenceJoin, Inner, JoinType, LeftAnti, LeftSemi, LeftSingle}
+import org.apache.spark.sql.catalyst.plans.{Cross, ExistenceJoin, Inner, JoinType, LeftAnti, LeftSemi, LeftSingle, RightOuter}
import org.apache.spark.sql.catalyst.plans.physical
-import org.apache.spark.sql.catalyst.plans.physical.{KeyedPartitioning, PartitioningCollection, REPLICATED_FOR_JOIN, SPLIT_FOR_JOIN}
+import org.apache.spark.sql.catalyst.plans.physical.{HASH_ROUTING, KeyedPartitioning, PartitioningCollection, REPLICATED_FOR_JOIN, ROUTING_LOST, SPLIT_FOR_JOIN}
import org.apache.spark.sql.connector.catalog.{Column, Identifier, InMemoryBaseTable, InMemoryCatalystRuntimeFilterCatalog, InMemoryTable, InMemoryTableCatalog}
import org.apache.spark.sql.connector.catalog.constraints.Constraint
import org.apache.spark.sql.connector.catalog.functions._
@@ -8914,10 +8914,10 @@ class KeyGroupedPartitioningSuite
// `name` away drops a key position but loses no distinct key, so nothing collapsed. purchases
// reports no partitioning, so with v2BucketingShuffleEnabled its side is shuffled using the
// projected partitioning as the template, still true under SPARK-59050: the shuffled output
- // carries the unknown-keys marker, the join's right-outer arm exposes it, and the union drops
- // the keyed merge, so the final aggregate re-shuffles instead of grouping the union's
- // partitions. What stays exercised for 58316: the template carries no collapse and needs no
- // opt-in either way.
+ // carries the unknown-keys marker and the join's right-outer arm exposes it. That leg declares
+ // every key t3 declares, so the union merges the keys and carries the marker onto them, and
+ // the final aggregate groups the union's partitions instead of re-shuffling them. What stays
+ // exercised for 58316: the template carries no collapse and needs no opt-in either way.
createTable(items, itemsColumns, Array(identity("id"), identity("name")))
sql(s"INSERT INTO testcat.ns.$items VALUES " +
s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " +
@@ -8967,13 +8967,20 @@ class KeyGroupedPartitioningSuite
"the ids were unique, so dropping `name` collapsed nothing for the template")
assert(keyed.head.mayContainUnknownPartitionKeys,
"the re-shuffled output may hold keys outside the declared set (SPARK-59050)")
- // The marked leg reaches the union through the right-outer arm, so the keyed merge is
- // dropped and the final aggregate pays its own exchange.
+ // The marked leg reaches the union through the right-outer arm. It declares {1, 2, 3},
+ // which covers t3's {1, 2}, so the keys merge and the marker travels onto them; the
+ // aggregate then groups those partitions rather than paying an exchange.
val union = collect(plan) { case u: UnionExec => u }.head
- assert(union.outputPartitioning.isInstanceOf[physical.UnknownPartitioning],
- s"the unknown-keyed leg drops the keyed merge, got ${union.outputPartitioning}")
- assert(collectAllShuffles(plan).size == 2,
- s"allowSubset=$allowSubset: template re-shuffle plus the aggregate's exchange")
+ val merged = union.outputPartitioning match {
+ case kp: physical.KeyedPartitioning => kp
+ case other => fail(s"the union should merge the legs' keys, got $other")
+ }
+ assert(merged.mayContainUnknownPartitionKeys,
+ s"the merged layout must carry the marked leg's marker, got $merged")
+ assert(collectAllGroupPartitions(plan).size == 1,
+ s"allowSubset=$allowSubset: the aggregate groups the union's partitions:\n$plan")
+ assert(collectAllShuffles(plan).size == 1,
+ s"allowSubset=$allowSubset: the template re-shuffle is the only exchange:\n$plan")
}
}
}
@@ -9481,6 +9488,463 @@ class KeyGroupedPartitioningSuite
}
}
+ test("SPARK-60277: SPJ: a union merges a marked leg that declares what the other legs " +
+ "declare") {
+ // The marked leg is `a RIGHT OUTER JOIN t`: t reports no partitioning, so its side is
+ // shuffled onto a's declared keys {1, 2, 3}, and the right-outer arm exposes that marked
+ // layout, holding t's id=4 rows out-of-set. The second leg `s` declares {1, 2}, a subset, so
+ // the keys merge and the marker travels onto them; the out-of-set rows ride one partition of
+ // the marked leg, so they stay together in its group, which is what the counts assert.
+ createTable("a", columns, Array(identity("id")))
+ createTable("s", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES (1, 'a1', NULL), (2, 'a2', NULL), (3, 'a3', NULL)")
+ sql("INSERT INTO testcat.ns.s VALUES (1, 's1', NULL), (2, 's2', NULL)")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 't1'), (2, 't2'), (4, 't4'), (4, 't4b')")
+
+ val query =
+ """
+ |SELECT id, count(*) AS cnt
+ |FROM (SELECT t.id AS id FROM testcat.ns.a a RIGHT OUTER JOIN t ON a.id = t.id
+ | UNION ALL
+ | SELECT id FROM testcat.ns.s)
+ |GROUP BY id
+ |""".stripMargin
+ val expected = Seq(Row(1, 2), Row(2, 2), Row(4, 2))
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ val plan = df.queryExecution.executedPlan
+ val union = collect(plan) { case u: UnionExec => u }.head
+ val merged = union.outputPartitioning match {
+ case kp: physical.KeyedPartitioning => kp
+ case other => fail(s"the union should merge the legs' keys, got $other")
+ }
+ assert(merged.mayContainUnknownPartitionKeys,
+ s"the merged layout must carry the marked leg's marker, got $merged")
+ assert(!merged.unknownKeysRouted,
+ s"an ungrouped concatenation gives the routing up, got $merged")
+ // The regrouping leaves every key's rows in one group, so the plan satisfies the
+ // aggregate's distribution; that is what keeps AQE's rules from discarding it.
+ assert(ValidateRequirements.validate(plan),
+ s"the plan must satisfy its requirements:\n$plan")
+ assert(collectAllShuffles(plan).size == 1,
+ s"the join's one-side shuffle should be the only exchange:\n$plan")
+ }
+
+ // Negative control: `s` declares the very key the marked leg holds out-of-set, so the
+ // merged claim would promise co-location for id=4 that the marked leg cannot honor. The
+ // keys must not merge, and the aggregate pays its own exchange.
+ sql("DROP TABLE IF EXISTS testcat.ns.s")
+ createTable("s", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.s VALUES (4, 's4', NULL)")
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, Seq(Row(1, 1), Row(2, 1), Row(4, 3)))
+ val plan = df.queryExecution.executedPlan
+ val union = collect(plan) { case u: UnionExec => u }.head
+ assert(union.outputPartitioning.isInstanceOf[physical.UnknownPartitioning],
+ s"a leg declaring an out-of-set key must stop the merge, got " +
+ union.outputPartitioning)
+ assert(collectAllShuffles(plan).size == 2,
+ s"the one-side shuffle plus the aggregate's own exchange:\n$plan")
+ }
+ }
+ }
+
+ test("SPARK-60277: SPJ: a join above a merged marked union pairs when the keys line up") {
+ // The merged layout is marked and ungrouped, which `resolveChild` admits and wraps in a
+ // `GroupPartitionsExec`. Coalescing the partitions that share a key keeps every key's rows in
+ // the group named for it, so the node reports a keyed layout whose routing is lost, and the
+ // join pairs it with the other side's scan instead of re-shuffling it. The marked leg's
+ // out-of-set id=4 rows sit in a declared key's group, and the scan holds no id=4 row, so they
+ // match nothing.
+ createTable("ja", columns, Array(identity("id")))
+ createTable("js", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.ja VALUES (1, 'a1', NULL), (2, 'a2', NULL), (3, 'a3', NULL)")
+ sql("INSERT INTO testcat.ns.js VALUES (1, 's1', NULL), (2, 's2', NULL)")
+
+ withTable("jt") {
+ sql("CREATE TABLE jt (id INT, data STRING) USING parquet")
+ sql("INSERT INTO jt VALUES (1, 't1'), (2, 't2'), (4, 't4'), (4, 't4b')")
+
+ val query =
+ """
+ |SELECT u.id, count(*) AS cnt FROM (
+ | SELECT jt.id AS id FROM testcat.ns.ja a RIGHT OUTER JOIN jt ON a.id = jt.id
+ | UNION ALL
+ | SELECT id FROM testcat.ns.js) u
+ |JOIN testcat.ns.ja a2 ON u.id = a2.id
+ |GROUP BY u.id
+ |""".stripMargin
+
+ Seq(false, true).foreach { aqe =>
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val df = sql(query)
+ checkAnswer(df, Seq(Row(1, 2), Row(2, 2)))
+ if (!aqe) {
+ // The merge is what puts the union's layout in front of the join's grouping; under
+ // AQE the plan is rewritten around the stages, so the counts above carry that arm.
+ val plan = df.queryExecution.executedPlan
+ val union = collect(plan) { case u: UnionExec => u }.head
+ union.outputPartitioning match {
+ case kp: physical.KeyedPartitioning =>
+ assert(kp.mayContainUnknownPartitionKeys,
+ s"the merged layout must carry the marked leg's state, got $kp")
+ case other =>
+ fail(s"the union should merge the legs' keys, got $other")
+ }
+ val grouped = collectAllGroupPartitions(plan)
+ assert(grouped.size == 1, s"the join pairs through the grouping:\n$plan")
+ grouped.head.outputPartitioning match {
+ case kp: physical.KeyedPartitioning =>
+ assert(kp.isGrouped && kp.mayContainUnknownPartitionKeys &&
+ !kp.unknownKeysRouted,
+ s"the grouping keeps the claim and loses the routing, got $kp")
+ case other =>
+ fail(s"the grouping should keep the keyed claim, got $other")
+ }
+ assert(collectAllShuffles(plan).size == 1,
+ s"the join pairs instead of re-shuffling the union side:\n$plan")
+ }
+ val smaller = sql(query.replace("JOIN testcat.ns.ja a2", "JOIN testcat.ns.js a2"))
+ checkAnswer(smaller, Seq(Row(1, 2), Row(2, 2)))
+ if (!aqe) {
+ val plan = smaller.queryExecution.executedPlan
+ assert(ValidateRequirements.validate(plan))
+ assert(collectAllShuffles(plan).size == 2,
+ s"different declared key sets must fall back to a valid shuffle plan:\n$plan")
+ }
+ }
+ }
+ }
+ }
+
+ test("SPARK-60277: a window above a merged marked union keeps its sort") {
+ // The grouping coalesces the partitions of one declared key, and one of them may also hold an
+ // undeclared key's rows, so the key ordering is not constant across the merged partition. The
+ // window needs the sort its required ordering asks for.
+ createTable("wa", columns, Array(identity("id")))
+ createTable("ws", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.wa VALUES (1, 'a1', NULL), (2, 'a2', NULL), (3, 'a3', NULL)")
+ sql("INSERT INTO testcat.ns.ws VALUES (1, 's1', NULL), (2, 's2', NULL), (3, 's3', NULL)")
+
+ withTable("wt") {
+ sql("CREATE TABLE wt (id INT, data STRING) USING parquet")
+ sql("INSERT INTO wt VALUES (1, 't1'), (2, 't2'), (3, 't3'), (4, 't4')")
+
+ val union =
+ """
+ |SELECT wt.id AS id FROM testcat.ns.wa a RIGHT OUTER JOIN wt ON a.id = wt.id
+ |UNION ALL
+ |SELECT id FROM testcat.ns.ws
+ |""".stripMargin
+ val expected = Seq(
+ Row(1, 2), Row(1, 2), Row(2, 2), Row(2, 2), Row(3, 2), Row(3, 2), Row(4, 1))
+
+ Seq(false, true).foreach { preserveKeyOrdering =>
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_PRESERVE_KEY_ORDERING_ON_COALESCE_ENABLED.key ->
+ preserveKeyOrdering.toString) {
+ val df = sql(union).sortWithinPartitions("id")
+ .selectExpr("id", "count(1) OVER (PARTITION BY id) AS cnt")
+ val mismatch = QueryTest.sameRows(expected, df.collect().toSeq)
+ assert(mismatch.isEmpty, mismatch.getOrElse(""))
+ val plan = df.queryExecution.executedPlan
+ val grouping = collectAllGroupPartitions(plan)
+ assert(grouping.size == 1, s"the union should be grouped once:\n$plan")
+ assert(grouping.head.groupedPartitions.exists(_._2.size > 1))
+ assert(grouping.head.outputPartitioning.asInstanceOf[physical.KeyedPartitioning]
+ .mayContainUnknownPartitionKeys)
+ assert(grouping.head.outputOrdering.isEmpty)
+ assert(collect(plan) {
+ case s: SortExec if unwrapWrappers(s.child) == grouping.head => s
+ }.size == 1, s"the window needs a sort above the grouping:\n$plan")
+ assert(ValidateRequirements.validate(plan))
+ }
+ }
+ }
+ }
+
+ test("SPARK-60277: SPJ: two marked legs route the keys they do not declare alike") {
+ // Both legs are marked and declare {1, 2, 3}, so each is shuffled onto those keys and each
+ // right-outer arm exposes the marked layout. id=4 is out-of-set on both, and both route it by
+ // the same hash over a count the shared declaration fixes, so its rows are counted once; a
+ // different routing would split them, which the count asserts.
+ createTable("a", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES (1, 'a1', NULL), (2, 'a2', NULL), (3, 'a3', NULL)")
+
+ withTable("t1", "t2") {
+ sql("CREATE TABLE t1 (id INT, data STRING) USING parquet")
+ sql("INSERT INTO t1 VALUES (1, 'x1'), (4, 'x4'), (4, 'x4b')")
+ sql("CREATE TABLE t2 (id INT, data STRING) USING parquet")
+ sql("INSERT INTO t2 VALUES (2, 'y2'), (4, 'y4'), (5, 'y5')")
+
+ val query =
+ """
+ |SELECT id, count(*) AS cnt
+ |FROM (SELECT t1.id AS id FROM testcat.ns.a a RIGHT OUTER JOIN t1 ON a.id = t1.id
+ | UNION ALL
+ | SELECT t2.id AS id FROM testcat.ns.a a RIGHT OUTER JOIN t2 ON a.id = t2.id)
+ |GROUP BY id
+ |""".stripMargin
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, Seq(Row(1, 1), Row(2, 1), Row(4, 3), Row(5, 1)))
+ val plan = df.queryExecution.executedPlan
+ val union = collect(plan) { case u: UnionExec => u }.head
+ val merged = union.outputPartitioning match {
+ case kp: physical.KeyedPartitioning => kp
+ case other => fail(s"the union should merge the legs' keys, got $other")
+ }
+ assert(merged.mayContainUnknownPartitionKeys,
+ s"the merged layout must carry the legs' marker, got $merged")
+ assert(collectAllShuffles(plan).size == 2,
+ s"one one-side shuffle per leg should be the whole exchange count:\n$plan")
+ assert(collectAllGroupPartitions(plan).nonEmpty,
+ s"the aggregate should group the union's partitions:\n$plan")
+ }
+ }
+ }
+
+ Seq("different-orders", "lost-pairing", "lost-union").foreach { shape =>
+ val scenario = shape match {
+ case "different-orders" =>
+ "aggregating marked legs with different declared orders preserves undeclared rows"
+ case "lost-pairing" =>
+ "joining marked unions with lost routing preserves undeclared matches"
+ case "lost-union" =>
+ "a union rejects marked legs with identical declarations but different routing"
+ }
+ test(s"SPARK-60277: SPJ: $scenario") {
+ Seq(
+ "a1" -> Seq(3, 4),
+ "a2" -> Seq(1, 2),
+ "a3" -> Seq(1, 2, 3, 4),
+ "s1" -> Seq(1, 2),
+ "s2" -> Seq(1, 2)).foreach { case (name, ids) =>
+ createTable(name, columns, Array(identity("id")))
+ val values = ids.map(id => s"($id, '$name$id', NULL)").mkString(", ")
+ sql(s"INSERT INTO testcat.ns.$name VALUES $values")
+ }
+
+ withTable("t1", "t2") {
+ Seq("t1", "t2").foreach { name =>
+ sql(s"CREATE TABLE $name (id INT) USING parquet")
+ sql(s"INSERT INTO $name SELECT CAST(id AS INT) FROM range(1, 41)")
+ }
+ val producers =
+ """
+ |WITH leg1 AS (
+ | SELECT t1.id FROM (
+ | SELECT id FROM testcat.ns.a1
+ | UNION ALL
+ | SELECT id FROM testcat.ns.a2) a
+ | RIGHT OUTER JOIN t1 ON a.id = t1.id),
+ |leg2 AS (
+ | SELECT t2.id FROM testcat.ns.a3 a
+ | RIGHT OUTER JOIN t2 ON a.id = t2.id)
+ |""".stripMargin
+ val query = producers + (if (shape == "lost-pairing") {
+ """
+ |SELECT l.id, count(*) AS cnt FROM (
+ | SELECT id FROM leg1
+ | UNION ALL
+ | SELECT id FROM testcat.ns.s1) l
+ |JOIN (
+ | SELECT id FROM leg2
+ | UNION ALL
+ | SELECT id FROM testcat.ns.s2) r ON l.id = r.id
+ |GROUP BY l.id
+ |""".stripMargin
+ } else if (shape == "lost-union") {
+ """
+ |SELECT id, count(*) AS cnt FROM (
+ | SELECT id FROM (
+ | SELECT id, count(*) AS c FROM (
+ | SELECT id FROM leg1
+ | UNION ALL
+ | SELECT id FROM testcat.ns.s1) merged
+ | GROUP BY id) grouped
+ | UNION ALL
+ | SELECT id FROM leg2) u
+ |GROUP BY id
+ |""".stripMargin
+ } else {
+ """
+ |SELECT id, count(*) AS cnt FROM (
+ | SELECT id FROM leg1
+ | UNION ALL
+ | SELECT id FROM leg2) u
+ |GROUP BY id
+ |""".stripMargin
+ })
+ val expected = (1 to 40).map { id =>
+ val count = if (shape != "lost-pairing") {
+ 2L
+ } else if (id <= 2) {
+ 4L
+ } else {
+ 1L
+ }
+ Row(id, count)
+ }
+
+ Seq(false, true).foreach { aqe =>
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ if (!aqe) {
+ val plan = df.queryExecution.executedPlan
+ val markedProducers = collect(plan) {
+ case j: ShuffledJoin if j.joinType == RightOuter =>
+ PartitioningCollection.representativeOf(j.outputPartitioning).get
+ }
+ assert(markedProducers.size == 2, s"both producers must remain:\n$plan")
+ assert(markedProducers.forall(_.unknownKeysRouted),
+ s"the producers must retain hash routing: $markedProducers")
+ val orders = markedProducers.map(_.partitionKeys.map(_.row.getInt(0)))
+ assert(orders.toSet == Set(Seq(3, 4, 1, 2), Seq(1, 2, 3, 4)),
+ s"the same declared key set must arrive in different orders: $orders")
+ val unions = collect(plan) { case u: UnionExec => u }
+ if (shape == "lost-union") {
+ val outer = unions.head
+ val legs = outer.children.map { child =>
+ PartitioningCollection.representativeOf(child.outputPartitioning).get
+ }
+ assert(legs.forall(_.partitionKeys.map(_.row.getInt(0)) == Seq(1, 2, 3, 4)),
+ s"the outer legs must have identical declarations: $legs")
+ assert(legs.map(_.unknownKeysOrigin).toSet ==
+ Set(Some(ROUTING_LOST), Some(HASH_ROUTING)),
+ s"only one outer leg must retain routing: $legs")
+ assert(outer.outputPartitioning.isInstanceOf[physical.UnknownPartitioning],
+ s"the outer union must reject lost/routed marked legs:\n$plan")
+ assert(ValidateRequirements.validate(plan))
+ } else if (shape == "lost-pairing") {
+ val markedUnions = unions.flatMap { u =>
+ PartitioningCollection.representativeOf(u.outputPartitioning)
+ .filter(_.mayContainUnknownPartitionKeys)
+ }
+ assert(markedUnions.size == 2,
+ s"both single-marked-leg unions must merge:\n$plan")
+ assert(markedUnions.forall(kp => !kp.isGrouped &&
+ kp.layout.unknownKeysOrigin.contains(ROUTING_LOST)),
+ s"concatenation must lose both unions' routing: $markedUnions")
+ assert(collectAllShuffles(plan).size == 4,
+ s"two producer shuffles plus two join-side shuffles are required:\n$plan")
+ } else {
+ assert(unions.size == 2, s"the producer and outer unions must remain:\n$plan")
+ assert(unions.head.outputPartitioning.isInstanceOf[physical.UnknownPartitioning],
+ s"the outer union must reject differing declared orders:\n$plan")
+ assert(collectAllShuffles(plan).size == 3,
+ s"two producer shuffles plus the aggregate shuffle are required:\n$plan")
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+
+ test("SPARK-60277: marked legs must declare their keys in one order to be merged") {
+ // Both legs send a key K no leg declares to `nonNegativeMod(K.hashCode, 3)`, so legs
+ // declaring [1, 2, 3] and [2, 3, 1] send it to a different named group and K's rows split.
+ // The gate compares the key lists, not the key sets; a differing multiplicity is refused for
+ // the same reason. Hand-built to isolate order and routing from planner choices.
+ val attrA = AttributeReference("a", IntegerType)()
+ val attrB = AttributeReference("b", IntegerType)()
+ def rows(values: Int*): Seq[InternalRow] = values.map(InternalRow(_))
+ def markedLeg(
+ e: AttributeReference,
+ keys: Seq[InternalRow],
+ routed: Boolean = true): ShuffleExchangeExec =
+ ShuffleExchangeExec(
+ KeyedPartitioning(Seq(e), keys)
+ .withLayout(_.copy(
+ unknownKeysOrigin = Some(if (routed) HASH_ROUTING else ROUTING_LOST))),
+ new LocalTableScanExec(Seq(e), Nil, None, false))
+ def accurateLeg(e: AttributeReference, keys: Seq[InternalRow]): ShuffleExchangeExec =
+ ShuffleExchangeExec(
+ KeyedPartitioning(Seq(e), keys),
+ new LocalTableScanExec(Seq(e), Nil, None, false))
+
+ withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") {
+ val sameOrder = UnionExec(Seq(
+ markedLeg(attrA, rows(1, 2, 3)), markedLeg(attrB, rows(1, 2, 3))))
+ sameOrder.outputPartitioning match {
+ case kp: physical.KeyedPartitioning =>
+ assert(kp.mayContainUnknownPartitionKeys, s"the merged layout must stay marked: $kp")
+ assert(kp.numPartitions === 6, s"the merge concatenates the legs' keys: $kp")
+ case other => fail(s"one order should merge, got $other")
+ }
+
+ Seq(
+ "a different order" -> rows(2, 3, 1),
+ "a different multiplicity" -> rows(1, 2, 2, 3)).foreach { case (label, keys) =>
+ val union = UnionExec(Seq(markedLeg(attrA, rows(1, 2, 3)), markedLeg(attrB, keys)))
+ assert(union.outputPartitioning.isInstanceOf[physical.UnknownPartitioning],
+ s"$label must stop the merge, got ${union.outputPartitioning}")
+ }
+
+ // A lost leg names a key no leg declares by an earlier layout's order, which its key list
+ // does not show, so two lost legs cannot be shown to agree on where it goes.
+ val lostUnion = UnionExec(Seq(
+ markedLeg(attrA, rows(1, 2, 3), routed = false),
+ markedLeg(attrB, rows(1, 2, 3), routed = false)))
+ assert(lostUnion.outputPartitioning.isInstanceOf[physical.UnknownPartitioning],
+ s"a lost routing must stop the merge, got ${lostUnion.outputPartitioning}")
+
+ // A routed and a lost leg cannot be shown to agree either.
+ val mixedUnion = UnionExec(Seq(
+ markedLeg(attrA, rows(1, 2, 3)),
+ markedLeg(attrB, rows(1, 2, 3), routed = false)))
+ assert(mixedUnion.outputPartitioning.isInstanceOf[physical.UnknownPartitioning],
+ s"a routed and a lost leg must stop the merge, got ${mixedUnion.outputPartitioning}")
+
+ // One marked leg needs no agreement: its undeclared rows reach one group either way.
+ val oneLostUnion = UnionExec(Seq(
+ markedLeg(attrA, rows(1, 2, 3), routed = false), accurateLeg(attrB, rows(1))))
+ oneLostUnion.outputPartitioning match {
+ case kp: physical.KeyedPartitioning =>
+ assert(kp.mayContainUnknownPartitionKeys && !kp.unknownKeysRouted,
+ s"a single marked leg merges and keeps its state, got $kp")
+ case other => fail(s"a single marked leg should merge, got $other")
+ }
+
+ // A grouped concatenation keeps the routing: the marked leg holds every partition.
+ val groupedUnion = UnionExec(Seq(markedLeg(attrA, rows(1, 2, 3)), accurateLeg(attrB, Nil)))
+ groupedUnion.outputPartitioning match {
+ case kp: physical.KeyedPartitioning =>
+ assert(kp.isGrouped && kp.mayContainUnknownPartitionKeys && kp.unknownKeysRouted,
+ s"a grouped concatenation keeps the hash routing, got $kp")
+ case other => fail(s"a grouped concatenation should merge, got $other")
+ }
+ }
+ }
+
test("SPARK-59050: SPJ: join-key projection of an unknown-keyed layout drops the claim") {
// Like the key-dropping-project repro, but the projection keeps both key positions: the
// coarsening happens when the second join projects the declared keys down to its join key
@@ -9653,16 +10117,16 @@ class KeyGroupedPartitioningSuite
test("SPARK-59050: SPJ: inner join keeps the marker beside a keyless sibling") {
// A keyless sibling makes no keyed claim about its rows, so it cannot prove the marked
// side's out-of-set rows away: the join equality can still match them, and they stay where
- // the claim pins them. `keyedMarkerOf` answers `None` for such an input and the clearing
- // must read that as "no unmarked keyed member", not as "unmarked" -- mapping `None` to
- // `false` would clear the marker here while every sibling that carries a
- // `KeyedPartitioning` stays green.
+ // the claim pins them. `representativeOf` returns `None` for a keyless input, which drops
+ // out of the marker comparison rather than supplying an unmarked keyed member. Treating it
+ // as unmarked would incorrectly clear the marker.
val attrA = AttributeReference("a", IntegerType)()
val attrB = AttributeReference("b", IntegerType)()
val keys = Seq(InternalRow(1), InternalRow(2))
def markedExchange(e: AttributeReference): ShuffleExchangeExec =
ShuffleExchangeExec(
- KeyedPartitioning(Seq(e), keys).withLayout(_.copy(mayContainUnknownPartitionKeys = true)),
+ KeyedPartitioning(Seq(e), keys)
+ .withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING))),
new LocalTableScanExec(Seq(e), Nil, None, false))
def keylessExchange(e: AttributeReference): ShuffleExchangeExec =
ShuffleExchangeExec(physical.HashPartitioning(Seq(e), keys.length),
@@ -9694,7 +10158,8 @@ class KeyGroupedPartitioningSuite
val attrB = AttributeReference("b", IntegerType)()
val keys = Seq(InternalRow(1), InternalRow(2))
def markedKP(e: AttributeReference): KeyedPartitioning =
- KeyedPartitioning(Seq(e), keys).withLayout(_.copy(mayContainUnknownPartitionKeys = true))
+ KeyedPartitioning(Seq(e), keys)
+ .withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING)))
def markedExchange(e: AttributeReference): ShuffleExchangeExec =
ShuffleExchangeExec(markedKP(e), new LocalTableScanExec(Seq(e), Nil, None, false))
def plainExchange(e: AttributeReference): ShuffleExchangeExec =
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala
index 18fe4e93c796..26990a954c4a 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala
@@ -22,7 +22,7 @@ import org.apache.spark.rdd.RDD
import org.apache.spark.serializer.JavaSerializer
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeReference, CodegenObjectFactoryMode, SortOrder, TransformExpression}
-import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, KeyedPartitioning, KeyReducer, OrderedDistribution, Partitioning, PartitioningCollection, REPLICATED_FOR_JOIN, SPLIT_FOR_JOIN, UnknownPartitioning}
+import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, HASH_ROUTING, KeyedPartitioning, KeyedShuffleSpec, KeyReducer, OrderedDistribution, Partitioning, PartitioningCollection, REPLICATED_FOR_JOIN, SPLIT_FOR_JOIN, UnknownPartitioning}
import org.apache.spark.sql.catalyst.util.DateTimeTestUtils.date
import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper
import org.apache.spark.sql.connector.catalog.functions.{BucketFunction, BucketReducer, DaysFunctionWithToYearsReducerWithLongResult, DaysToYearsReducerWithLongResult, Reducer, YearsFunction, YearsFunctionWithToYearsReducerWithLongResult}
@@ -370,35 +370,60 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
// unknown-key hash-routing relationship intact, so the marker rides the transform.
val child = DummySparkPlan(
outputPartitioning = KeyedPartitioning(Seq(exprA), Seq(row(1), row(2)))
- .withLayout(_.copy(mayContainUnknownPartitionKeys = true)))
+ .withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING))))
val out = GroupPartitionsExec(child)
.outputPartitioning.asInstanceOf[KeyedPartitioning]
assert(out.mayContainUnknownPartitionKeys, "the marker must survive identity grouping")
+ assert(out.unknownKeysRouted, "an identity grouping leaves the hash routing in place")
}
- test("SPARK-59050: non-identity grouping without reducers drops the unknown-keyed claim") {
- // A regrouping that reorders or coalesces the partitions moves the rows an unknown-key claim
- // pins to hash(key) % numPartitions, so the claim cannot survive; clearing only the marker
- // would misreport the undeclared rows that remain.
+ test("SPARK-59050: non-identity grouping keeps the clustering claim and loses the routing") {
+ // A regrouping that reorders or coalesces the partitions leaves every key's rows in the group
+ // named for it, which is all a clustering reads, so the keyed claim survives; it only moves an
+ // undeclared row off `hash(key) % numPartitions`, which `ROUTING_LOST` records and the two
+ // routing readers refuse.
+ def markedKp(keys: Seq[InternalRow]): KeyedPartitioning =
+ KeyedPartitioning(Seq(exprA), keys)
+ .withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING)))
+ def markedChild(keys: Seq[InternalRow]): DummySparkPlan =
+ DummySparkPlan(outputPartitioning = markedKp(keys))
+
// Coalesce: keys [1, 2, 1] merge the two key-1 partitions.
- val coalesced = DummySparkPlan(
- outputPartitioning = KeyedPartitioning(Seq(exprA), Seq(row(1), row(2), row(1)))
- .withLayout(_.copy(mayContainUnknownPartitionKeys = true)))
- GroupPartitionsExec(coalesced).outputPartitioning match {
- case u: UnknownPartitioning =>
- assert(u.numPartitions === 2, "the give-up count must match the physical partitions")
- case other =>
- fail(s"expected the unknown-keyed claim to be dropped on coalesce, got $other")
- }
+ val coalesced = GroupPartitionsExec(markedChild(Seq(row(1), row(2), row(1))))
+ val coalescedOut = coalesced.outputPartitioning.asInstanceOf[KeyedPartitioning]
+ assert(coalescedOut.isGrouped && coalescedOut.numPartitions === 2)
+ assert(coalescedOut.mayContainUnknownPartitionKeys, "the undeclared rows remain")
+ assert(!coalescedOut.unknownKeysRouted, "the regrouping moved them off the hash")
+
// Reorder: keys [2, 1] sort to [1, 2], swapping partitions 0 and 1.
- val reordered = DummySparkPlan(
- outputPartitioning = KeyedPartitioning(Seq(exprA), Seq(row(2), row(1)))
- .withLayout(_.copy(mayContainUnknownPartitionKeys = true)))
- GroupPartitionsExec(reordered).outputPartitioning match {
- case u: UnknownPartitioning =>
- assert(u.numPartitions === 2, "the give-up count must match the physical partitions")
- case other =>
- fail(s"expected the unknown-keyed claim to be dropped on reorder, got $other")
+ val reordered = GroupPartitionsExec(markedChild(Seq(row(2), row(1))))
+ val reorderedOut = reordered.outputPartitioning.asInstanceOf[KeyedPartitioning]
+ assert(reorderedOut.isGrouped && reorderedOut.numPartitions === 2)
+ assert(reorderedOut.mayContainUnknownPartitionKeys && !reorderedOut.unknownKeysRouted)
+
+ // The routing readers refuse the reported layout: it cannot be the template a sibling is
+ // shuffled onto, nor pair index by index with another marked side. A routed layout of the same
+ // shape passes both, which is what makes these assertions discriminating.
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ val distribution = ClusteredDistribution(Seq(exprA))
+ val routed = markedKp(Seq(row(1), row(2)))
+ val routedSpec = KeyedShuffleSpec(routed, distribution)
+ assert(routedSpec.canCreatePartitioning, "control: a routed layout can serve as a template")
+ assert(!KeyedShuffleSpec(coalescedOut, distribution).canCreatePartitioning,
+ "a layout whose routing is lost cannot")
+
+ assert(routedSpec.areKeysCompatible(
+ KeyedShuffleSpec(routed.copy(), distribution), allowReduce = false),
+ "control: two routed sides with the same keys pair")
+ assert(!KeyedShuffleSpec(coalescedOut, distribution).areKeysCompatible(
+ routedSpec, allowReduce = false),
+ "an unrouted side cannot pair index by index with another marked side")
+ assert(!routedSpec.areKeysCompatible(
+ KeyedShuffleSpec(coalescedOut, distribution), allowReduce = false),
+ "a routed side must refuse a lost partner too")
+ assert(!KeyedShuffleSpec(coalescedOut, distribution).areKeysCompatible(
+ KeyedShuffleSpec(coalescedOut, distribution), allowReduce = false),
+ "two lost sides cannot pair either")
}
}
@@ -410,7 +435,7 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
// promise hash(key) % 2. The count check catches what the per-group check cannot.
val child = DummySparkPlan(
outputPartitioning = KeyedPartitioning(Seq(exprA), Seq(row(1), row(2), row(3)))
- .withLayout(_.copy(mayContainUnknownPartitionKeys = true)))
+ .withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING))))
def keyOf(a: Int): InternalRowComparableWrapper =
InternalRowComparableWrapper(row(a), Seq(exprA))
val gpe = GroupPartitionsExec(child,
@@ -447,7 +472,7 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
val partitionKeys = Seq(row(1), row(2), row(3))
val child = DummySparkPlan(
outputPartitioning = KeyedPartitioning(Seq(exprA), partitionKeys)
- .withLayout(_.copy(mayContainUnknownPartitionKeys = true)))
+ .withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING))))
val gpe = GroupPartitionsExec(child, reducers = Some(Seq(Some(reducer))))
assert(gpe.groupedPartitions.size === 2, "mod 2 collapses keys 1 and 3")
@@ -475,7 +500,7 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
// a marked layout, so these shapes are pinned here directly.
val child = DummySparkPlan(
outputPartitioning = KeyedPartitioning(Seq(exprA, exprB), Seq(row(1, 10), row(2, 20)))
- .withLayout(_.copy(mayContainUnknownPartitionKeys = true)))
+ .withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING))))
// Reducer slots: `Some(Seq(None, None))` leaves every key and every index unchanged.
val reduced = GroupPartitionsExec(child, reducers = Some(Seq(None, None)))
@@ -573,6 +598,27 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
}
}
+ test("SPARK-60277: marked layouts preserve ordering only without concatenation") {
+ val ordering = Seq(SortOrder(exprA, Ascending), SortOrder(exprC, Ascending))
+ def child(keys: Seq[InternalRow]): DummyLeafSparkPlan = DummyLeafSparkPlan(
+ outputPartitioning = KeyedPartitioning(Seq(exprA), keys)
+ .withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING))),
+ outputOrdering = ordering)
+
+ val unmerged = GroupPartitionsExec(child(Seq(row(1), row(2))))
+ assert(unmerged.groupedPartitions.forall(_._2.size == 1))
+ assert(unmerged.outputOrdering === ordering)
+
+ val repeated = child(Seq(row(1), row(2), row(1)))
+ withSQLConf(SQLConf.V2_BUCKETING_PRESERVE_KEY_ORDERING_ON_COALESCE_ENABLED.key -> "true") {
+ val concatenated = GroupPartitionsExec(repeated)
+ assert(concatenated.groupedPartitions.exists(_._2.size > 1))
+ assert(concatenated.outputOrdering.isEmpty)
+ assert(GroupPartitionsExec(repeated, sortedMergeOrdering = Some(ordering))
+ .outputOrdering === ordering)
+ }
+ }
+
test("SPARK-59279: the planned merge ordering decides the k-way merge, not the config") {
// The config is the planner's input, and `sortedMergeOrdering` records what the planner
// decided under it. Once it is set the config no longer matters, because the plan above was
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala
index 8a1697cce80e..f79ea2efe15c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala
@@ -2802,7 +2802,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
// behind. `KeyedPartitioning.apply` keeps the order it is given.
val marked = KeyedPartitioning(Seq(exprA),
Seq(InternalRow(3), InternalRow(4), InternalRow(1), InternalRow(2)))
- .withLayout(_.copy(mayContainUnknownPartitionKeys = true))
+ .withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING)))
assert(marked.isGrouped, "test setup: the keys are distinct, so the layout is grouped")
val smj = SortMergeJoinExec(Seq(exprA), Seq(exprB), Inner, None,
@@ -2830,17 +2830,15 @@ class EnsureRequirementsSuite extends SharedSparkSession {
test("SPARK-59272: every plan the rule produces for a generated pair of shapes validates") {
// A generated sweep asserting the property rather than a count: every plan this rule produces
- // for a pair of shapes has to satisfy both children's required distribution. The shapes that
- // make it fail are the marked ones, where a `GroupPartitionsExec` gives up the keyed claim
- // while regrouping, which is an answer that arrives after the child's fate was decided, at
- // either of the two sites that insert such a node.
+ // for a pair of shapes has to satisfy both children's required distribution. Plain regrouping
+ // keeps a marked layout keyed, with lost routing if partitions move. Rewritten keys or a
+ // non-identity expected-key alignment drop the keyed claim, so either insertion site must
+ // account for the grouping node's output before deciding the child's fate.
//
// `v2BucketingPartiallyClusteredDistributionEnabled` is deliberately not one of the cells.
// That family leaves both children keyed but ungrouped on purpose, so neither satisfies its
// distribution on its own and the oracle below rejects a plan this rule is right to produce.
- // A marked layout cannot occur with it either: `KeyedShuffleSpec.createPartitioning` is the
- // only producer of the marker, and it is only ever reached through `canCreatePartitioning`,
- // which requires the config to be off.
+ // The ungrouped input shapes below still exercise plain regrouping with this config off.
def shapes(a: Attribute): Seq[(String, KeyedPartitioning)] = {
val keySets = Seq(
"12" -> Seq(InternalRow(1), InternalRow(2)),
@@ -2849,7 +2847,9 @@ class EnsureRequirementsSuite extends SharedSparkSession {
// Not ascending, which is what a one-side shuffle onto a union's key order leaves behind.
// `KeyedPartitioning.apply` keeps the order it is given, and for a marked layout that
// order is part of the claim.
- "312" -> Seq(InternalRow(3), InternalRow(1), InternalRow(2)))
+ "312" -> Seq(InternalRow(3), InternalRow(1), InternalRow(2)),
+ // Duplicate keys separated by another key force regrouping on either side of the join.
+ "121-interleaved" -> Seq(InternalRow(1), InternalRow(2), InternalRow(1)))
val exprSets = Seq(
"id" -> Seq[Expression](a),
"bucket4" -> Seq[Expression](bucket(4, a)),
@@ -2857,13 +2857,14 @@ class EnsureRequirementsSuite extends SharedSparkSession {
for {
(keyName, keys) <- keySets
(exprName, exprs) <- exprSets
- // A marked layout pins rows outside the declared keys to `hash(key) % numPartitions`, so
- // regrouping one moves them and the node stops claiming the layout.
- marked <- Seq(false, true)
+ // Hash routing locates undeclared keys; lost routing retains only full-key clustering.
+ (originName, origin) <- Seq(
+ "unmarked" -> None,
+ "hash-routing" -> Some(HASH_ROUTING),
+ "routing-lost" -> Some(ROUTING_LOST))
} yield {
- val kp = KeyedPartitioning(exprs, keys)
- val p = if (marked) kp.withLayout(_.copy(mayContainUnknownPartitionKeys = true)) else kp
- (s"$exprName/$keyName${if (marked) "/marked" else ""}", p)
+ val kp = KeyedPartitioning(exprs, keys).withLayout(_.copy(unknownKeysOrigin = origin))
+ (s"$exprName/$keyName/$originName", kp)
}
}
@@ -2871,8 +2872,8 @@ class EnsureRequirementsSuite extends SharedSparkSession {
// `Inner` and `LeftOuter`, so they would repeat those cells exactly.
val joinTypes = Seq(Inner, LeftOuter, RightOuter, FullOuter)
val configCells = for {
- // Without this one no marked layout is reachable and no keyed spec can serve as a shuffle
- // template, so the second grouping site is never entered.
+ // Without keyed shuffle no keyed spec can serve as a shuffle template, so the second
+ // grouping site is never entered. Marked inputs can still come from a union.
keyedShuffle <- Seq(false, true)
pushPartValues <- Seq(false, true)
subsetKeys <- Seq(false, true)
@@ -2922,11 +2923,11 @@ class EnsureRequirementsSuite extends SharedSparkSession {
// The two join types put the marker on opposite sides on purpose, because that is the side
// `rank` reads for the one-sided arms.
val marked = KeyedPartitioning(Seq(exprA), Seq(InternalRow(1), InternalRow(2), InternalRow(3)))
- .withLayout(_.copy(mayContainUnknownPartitionKeys = true))
+ .withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING)))
val plain = KeyedPartitioning(Seq(exprB), Seq(InternalRow(1), InternalRow(2)))
val markedRight = KeyedPartitioning(
Seq(exprB), Seq(InternalRow(1), InternalRow(2), InternalRow(3)))
- .withLayout(_.copy(mayContainUnknownPartitionKeys = true))
+ .withLayout(_.copy(unknownKeysOrigin = Some(HASH_ROUTING)))
val plainLeft = KeyedPartitioning(Seq(exprA), Seq(InternalRow(1), InternalRow(2)))
Seq(
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org