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/09/08 05:44:24 UTC
(spark) branch branch-4.2 updated: [SPARK-59050][4.2][SQL] SPJ one-side shuffle with out-of-set keys produces wrong results in multi-joins
This is an automated email from the ASF dual-hosted git repository.
ulysses-you pushed a commit to branch branch-4.2
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.2 by this push:
new fc65d768ed8b [SPARK-59050][4.2][SQL] SPJ one-side shuffle with out-of-set keys produces wrong results in multi-joins
fc65d768ed8b is described below
commit fc65d768ed8b3908f316431e7283533cabef9b18
Author: Xiduo You <ul...@gmail.com>
AuthorDate: Tue Sep 8 13:44:10 2026 +0800
[SPARK-59050][4.2][SQL] SPJ one-side shuffle with out-of-set keys produces wrong results in multi-joins
### What changes were proposed in this pull request?
Backport of #58339 to `branch-4.2`, adapted to this branch. The core fix is unchanged: `KeyedPartitioning` gains `mayContainUnknownPartitionKeys`, set by `KeyedShuffleSpec.createPartitioning` (the one-side shuffle), and `areKeysCompatible` only co-partitions such a side with a subset-keyed (or identically-declared and identically-marked) partner; `groupedSatisfies` gives it the full-key rules only; `createShuffleSpec` refuses a narrowing projection of a marked partitioning; `GroupPart [...]
What differs from `master` on this branch:
- The config is spelled `V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS`.
- `isCollapsed` (SPARK-46367) and the `PartitioningCollection` invariant with `fromPartitionings`/`keyedMarkerOf` are not here; consumers read collections with `exists`, so `ShuffledJoin`'s clearing is implemented inline and `GroupPartitionsExec` reads any marked member.
- `UnionExec` and `AliasAwareOutputExpression` are untouched: this branch's union does not merge `KeyedPartitioning`s, and its generic projection drops a `KeyedPartitioning` whose key column leaves the output, so neither wrong-results shape can arise here.
- Consequently the two e2e tests whose `master` expectations rely on the aliasing projection keeping a narrowed claim expect the second join's one-side re-shuffle here, and the union-merge and regrouping repros are not carried.
The commit message carries the full change description with the complete backport notes.
### Why are the changes needed?
A storage-partitioned join whose preserved side is one-side-shuffled (`spark.sql.sources.v2.bucketing.shuffle.enabled`, shipping since 4.0.0) can silently lose matches in a following storage-partitioned join, with no error. Reproduction in #58339.
### Does this PR introduce _any_ user-facing change?
Yes: it fixes a wrong-results bug in `spark.sql.sources.v2.bucketing.shuffle.enabled`. Storage-partitioned joins whose preserved side is one-side-shuffled against a smaller keyed side, followed by another storage-partitioned join on a larger key set, now fall back to shuffles and return correct results.
### How was this patch tested?
The ported suites pass on `branch-4.2`: `KeyGroupedPartitioningSuite` (19 `SPARK-59050` tests), `ShuffleSpecSuite`, `DistributionSuite`, `GroupPartitionsExecSuite`, `EnsureRequirementsSuite`, `ValidateRequirementsSuite`.
### Was this patch authored or co-authored using generative AI tooling?
Yes. Generated-by: Claude Code.
Closes #58570 from ulysses-you/spj-part-mismatch-4.2.
Authored-by: Xiduo You <ul...@gmail.com>
Signed-off-by: Xiduo You <ul...@apache.org>
---
.../sql/catalyst/plans/physical/partitioning.scala | 108 ++-
.../spark/sql/catalyst/DistributionSuite.scala | 22 +-
.../spark/sql/catalyst/ShuffleSpecSuite.scala | 173 +++-
.../datasources/v2/GroupPartitionsExec.scala | 49 +-
.../execution/exchange/EnsureRequirements.scala | 4 +-
.../spark/sql/execution/joins/ShuffledJoin.scala | 43 +-
.../DistributionAndOrderingSuiteBase.scala | 2 +-
.../connector/KeyGroupedPartitioningSuite.scala | 903 +++++++++++++++++++++
.../datasources/v2/GroupPartitionsExecSuite.scala | 156 +++-
.../exchange/EnsureRequirementsSuite.scala | 2 +-
10 files changed, 1439 insertions(+), 23 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 608ad0195f7f..939e5d1deb96 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
@@ -368,6 +368,10 @@ case class CoalescedHashPartitioning(from: HashPartitioning, partitions: Seq[Coa
* - `partitionKeys`: The partition keys, one per partition. May contain duplicates initially
* (ungrouped state), but becomes unique after `GroupPartitionsExec` applies grouping.
*
+ * One exception to "partition `i` holds key `partitionKeys(i)`": when
+ * `mayContainUnknownPartitionKeys` is set, a partition may also hold rows whose key is not
+ * declared at all (see the `@param` below).
+ *
* == Grouping State ==
* A KeyedPartitioning can be in two states:
*
@@ -424,11 +428,33 @@ case class CoalescedHashPartitioning(from: HashPartitioning, partitions: Seq[Coa
* after projection. May contain duplicates when ungrouped.
* @param isGrouped Whether partition keys are unique (no duplicates). Computed on first
* creation, then preserved through copy operations to avoid recomputation.
+ * @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. A node that changes
+ * the declared key set must drop the keyed partitioning, whether
+ * it coarsens it (a key-dropping projection, 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).
*/
case class KeyedPartitioning(
expressions: Seq[Expression],
@transient partitionKeys: Seq[InternalRowComparableWrapper],
- isGrouped: Boolean) extends Expression with Partitioning with Unevaluable {
+ isGrouped: Boolean,
+ mayContainUnknownPartitionKeys: Boolean = false)
+ extends Expression with Partitioning with Unevaluable {
override val numPartitions = partitionKeys.length
override def children: Seq[Expression] = expressions
@@ -551,7 +577,11 @@ case class KeyedPartitioning(
// We'll need to find leaf attributes from the partition expressions first.
val attributes = expressions.flatMap(_.collectLeaves())
- if (SQLConf.get.v2BucketingAllowJoinKeysSubsetOfPartitionKeys) {
+ if (mayContainUnknownPartitionKeys) {
+ // Whole keys co-locate, subsets do not (see the `@param`): a window or aggregate
+ // keyed on a strict subset of the partition columns must still shuffle.
+ attributes.forall(x => requiredClustering.exists(_.semanticEquals(x)))
+ } else if (SQLConf.get.v2BucketingAllowJoinKeysSubsetOfPartitionKeys) {
// check that join keys (required clustering keys)
// overlap with partition keys (KeyedPartitioning attributes)
requiredClustering.exists(x => attributes.exists(_.semanticEquals(x))) &&
@@ -567,7 +597,10 @@ case class KeyedPartitioning(
// keys. This is the local gate. A reduced position always carries a transform and an
// `ORDER BY` cannot name one, so the match below already refuses every case a query can
// reach. Do not read the first clause as redundant.
- expressionsDescribeKeys && o.areAllClusterKeysMatched(expressions)
+ // An out-of-set key can break the ascending sequence of the declared keys, so a marked
+ // layout keeps a global ordering claim only for a single partition (see the `@param`).
+ expressionsDescribeKeys && o.areAllClusterKeysMatched(expressions) &&
+ (!mayContainUnknownPartitionKeys || numPartitions == 1)
case _ =>
false
@@ -577,15 +610,25 @@ case class KeyedPartitioning(
override def createShuffleSpec(distribution: ClusteredDistribution): ShuffleSpec = {
val result = KeyedShuffleSpec(this, distribution)
if (SQLConf.get.v2BucketingAllowJoinKeysSubsetOfPartitionKeys) {
+ val joinKeyPositions = result.keyPositions.map(_.nonEmpty).zipWithIndex.filter(_._1).map(_._2)
+ // The projection coarsens the declared set, which a marked claim cannot survive (see the
+ // `@param`). Return the unprojected spec: its extra expressions map to no clustering key,
+ // so `areKeysCompatible` refuses it as a partner and `canCreatePartitioning` refuses it as
+ // a shuffle template, and the child falls back to the ordinary shuffle.
+ if (mayContainUnknownPartitionKeys && joinKeyPositions.length < expressions.length) {
+ return result
+ }
// If allowing join keys to be subset of clustering keys, we should create a new
// `KeyedPartitioning` here that is grouped on the join keys instead, and use that as
- // the returned shuffle spec.
- val joinKeyPositions = result.keyPositions.map(_.nonEmpty).zipWithIndex.filter(_._1).map(_._2)
+ // the returned shuffle spec. The construction carries the unknown-keys marker across:
+ // only an identity projection reaches here when it is set, the refusal above turns away
+ // the narrowing one.
val projectedExpressions = joinKeyPositions.map(expressions)
val projectedKeys = projectKeys(joinKeyPositions)._2
val distinctProjectedKeys = projectedKeys.distinct
- val projectedPartitioning =
- KeyedPartitioning(projectedExpressions, distinctProjectedKeys, isGrouped = true)
+ val projectedPartitioning = KeyedPartitioning(
+ projectedExpressions, distinctProjectedKeys,
+ isGrouped = true, mayContainUnknownPartitionKeys = mayContainUnknownPartitionKeys)
result.copy(partitioning = projectedPartitioning, joinKeyPositions = Some(joinKeyPositions))
} else {
result
@@ -1145,6 +1188,44 @@ case class KeyedShuffleSpec(
}
} && expressions.zip(otherExpressions).forall {
case (l, r) => isExpressionCompatible(l, r)
+ } && {
+ // An unknown-keyed side co-locates only its declared keys, and the out-of-set routing is
+ // a deterministic hash, so it can pair only with a side whose keys are a subset of those
+ // declared keys (see `KeyedPartitioning.mayContainUnknownPartitionKeys`).
+ //
+ // The key comparison below must also happen in a single domain: `isExpressionCompatible`
+ // admits an `AttributeReference` against a `TransformExpression` (and two different-but-
+ // compatible transforms) when `v2BucketingAllowCompatibleTransforms` is on, and in those
+ // cases the two sides' `partitionKeys` hold raw values on one side and transform outputs on
+ // the other, so the subset test would compare unrelated values. Require the partition
+ // expressions to be the same function per position before comparing keys.
+ //
+ // Two unknown-keyed sides are compatible only when they agree on the declared keys *and*
+ // 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.
+ if (partitioning.mayContainUnknownPartitionKeys ||
+ other.partitioning.mayContainUnknownPartitionKeys) {
+ expressions.zip(otherExpressions).forall {
+ case (_: AttributeReference, _: AttributeReference) => true
+ case (l: TransformExpression, r: TransformExpression) => l.isSameFunction(r)
+ case _ => false
+ } && {
+ if (partitioning.mayContainUnknownPartitionKeys &&
+ other.partitioning.mayContainUnknownPartitionKeys) {
+ partitioning.partitionKeys == other.partitioning.partitionKeys
+ } else if (partitioning.mayContainUnknownPartitionKeys) {
+ val declared = partitioning.partitionKeys.toSet
+ other.partitioning.partitionKeys.forall(declared.contains)
+ } else {
+ val declared = other.partitioning.partitionKeys.toSet
+ partitioning.partitionKeys.forall(declared.contains)
+ }
+ }
+ } else {
+ true
+ }
}
}
@@ -1273,6 +1354,10 @@ case class KeyedShuffleSpec(
override def canCreatePartitioning: Boolean =
SQLConf.get.v2BucketingShuffleEnabled &&
!SQLConf.get.v2BucketingPartiallyClusteredDistributionEnabled &&
+ // 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.
+ keyPositions.forall(_.nonEmpty) &&
partitioning.expressions.forall { e =>
e.isInstanceOf[AttributeReference] || e.isInstanceOf[TransformExpression]
} &&
@@ -1289,7 +1374,14 @@ case class KeyedShuffleSpec(
te.copy(children = te.children.map(_ => clustering(positionSet.head)))
case (_, positionSet) => clustering(positionSet.head)
}
- KeyedPartitioning(newExpressions, partitioning.partitionKeys, partitioning.isGrouped)
+ // The child re-shuffled onto this layout may hold keys outside the declared set, so every
+ // partitioning produced here carries the marker (see the `@param`). Only the shuffle loop
+ // reaches this call, and it carries no same-domain subset proof, so marking is sound; it is
+ // conservative where the child's keys are in fact a known subset (identity key [1] inside
+ // declared [1, 2]), a precision this path does not attempt.
+ KeyedPartitioning(
+ newExpressions, partitioning.partitionKeys, partitioning.isGrouped,
+ mayContainUnknownPartitionKeys = true)
}
}
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 88faf7010318..474718edf26b 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
@@ -20,11 +20,13 @@ package org.apache.spark.sql.catalyst
import org.apache.spark.SparkFunSuite
/* Implicit conversions */
import org.apache.spark.sql.catalyst.dsl.expressions._
-import org.apache.spark.sql.catalyst.expressions.{AttributeReference, CollationAwareMurmur3Hash, Expression, Literal, Pmod}
+import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, CollationAwareMurmur3Hash, Expression, Literal, Pmod, SortOrder}
+import org.apache.spark.sql.catalyst.plans.SQLHelper
import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.IntegerType
-class DistributionSuite extends SparkFunSuite {
+class DistributionSuite extends SparkFunSuite with SQLHelper {
protected def checkSatisfied(
inputPartitioning: Partitioning,
@@ -410,4 +412,20 @@ class DistributionSuite extends SparkFunSuite {
assert(groupedKP.isGrouped)
checkSatisfied(groupedKP, ClusteredDistribution(Seq(x)), true)
}
+ test("SPARK-59050: a marked one-partition layout keeps the global ordering claim") {
+ // The one-partition exemption inside `groupedSatisfies`'s ordered branch: a single partition
+ // holds every row, so an out-of-set key cannot break the cross-partition sequence, while
+ // two partitions can (the e2e `ORDER BY` repro measures that). Positive control: an
+ // always-false gate would shuffle these plans for nothing.
+ val a = AttributeReference("a", IntegerType)()
+ val ordered = OrderedDistribution(Seq(SortOrder(a, Ascending)))
+ val markedOne = KeyedPartitioning(Seq(a), Seq(InternalRow(1)))
+ .copy(mayContainUnknownPartitionKeys = true)
+ val markedTwo = KeyedPartitioning(Seq(a), Seq(InternalRow(1), InternalRow(2)))
+ .copy(mayContainUnknownPartitionKeys = true)
+ withSQLConf(SQLConf.V2_BUCKETING_SORTING_ENABLED.key -> "true") {
+ checkSatisfied(markedOne, ordered, true)
+ checkSatisfied(markedTwo, ordered, false)
+ }
+ }
}
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 69a41b0d24dd..4f01c884d3ec 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
@@ -22,9 +22,9 @@ import org.apache.spark.sql.catalyst.dsl.expressions._
import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, DirectShufflePartitionID, Expression, TransformExpression}
import org.apache.spark.sql.catalyst.plans.SQLHelper
import org.apache.spark.sql.catalyst.plans.physical._
-import org.apache.spark.sql.connector.catalog.functions.ScalarFunction
+import org.apache.spark.sql.connector.catalog.functions.{Reducer, ReducibleFunction, ScalarFunction}
import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.{DataType, IntegerType, StructType}
+import org.apache.spark.sql.types.{DataType, IntegerType, LongType, StructType}
class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
private val passThrough_a_10 = ShufflePartitionIdPassThrough(DirectShufflePartitionID($"a"), 10)
@@ -473,6 +473,175 @@ 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 keyedSpec(
+ keys: Seq[Int],
+ hasUnknown: Boolean = false): KeyedShuffleSpec = KeyedShuffleSpec(
+ KeyedPartitioning(Seq(a), keys.map(k => InternalRow(k)))
+ .copy(mayContainUnknownPartitionKeys = 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
+ // it can only be co-partitioned with a side whose keys are a subset of the declared keys.
+ val unknown12 = keyedSpec(Seq(1, 2), hasUnknown = true)
+ // hasUnknown=true is still compatible with a subset (or equal) partner: every such key is
+ // co-located on both sides.
+ assert(unknown12.areKeysCompatible(keyedSpec(Seq(1))), "subset keys must be compatible")
+ assert(unknown12.areKeysCompatible(keyedSpec(Seq(2))), "another subset key must be compatible")
+ assert(unknown12.areKeysCompatible(keyedSpec(Seq(1, 2))), "equal keys must be compatible")
+ assert(keyedSpec(Seq(1)).areKeysCompatible(unknown12),
+ "compatibility must be symmetric for a subset partner")
+
+ // A larger partner's keys are not all covered by the declared keys.
+ assert(!unknown12.areKeysCompatible(keyedSpec(Seq(1, 2, 3))),
+ "a larger key set must not be compatible with an unknown-keyed partitioning")
+ assert(!keyedSpec(Seq(1, 2, 3)).areKeysCompatible(unknown12),
+ "an unknown-keyed partitioning cannot cover a larger partner's keys")
+
+ // Both sides unknown with the same declared keys: `KeyGroupedPartitioner`'s out-of-set-key
+ // fallback is a deterministic hash of the key, so both sides route those keys to the same
+ // partition and stay compatible, but only when the declared key order also agrees, since a
+ // GroupPartitionsExec regrouping re-labels partitions by each side's declared order. Different
+ // declared keys or a different order are rejected.
+ assert(unknown12.areKeysCompatible(keyedSpec(Seq(1, 2), hasUnknown = true)),
+ "two unknown-keyed partitionings with the same declared keys must be compatible")
+ assert(!unknown12.areKeysCompatible(keyedSpec(Seq(2, 1), hasUnknown = true)),
+ "two unknown-keyed partitionings must agree on the declared key order")
+ assert(!unknown12.areKeysCompatible(keyedSpec(Seq(1, 2, 3), hasUnknown = true)),
+ "two unknown-keyed partitionings with different declared keys must not be compatible")
+
+ // Without the marker, key sets are not compared, only the partition expressions are.
+ assert(keyedSpec(Seq(1, 2)).areKeysCompatible(keyedSpec(Seq(1, 2, 3))),
+ "without unknown keys, different key sets remain expression-compatible")
+ }
+
+ test("areKeysCompatible: unknown partition keys require the same partition functions") {
+ // A stand-in for a connector partition function such as `bucket`: only the canonical name
+ // matters here, since `TransformExpression.isSameFunction` compares names and bucket counts.
+ val bucketFn = new ScalarFunction[Int] {
+ override def inputTypes(): Array[DataType] = Array(LongType)
+ override def resultType(): DataType = IntegerType
+ override def name(): String = "test.bucket"
+ override def canonicalName(): String = "test.bucket"
+ }
+ def bucket(numBuckets: Int, expr: Expression): TransformExpression =
+ TransformExpression(bucketFn, Seq(expr), Some(numBuckets))
+
+ val a = $"a".long
+ val distribution = ClusteredDistribution(Seq(a))
+ def bucketSpec(
+ keys: Seq[Long],
+ hasUnknown: Boolean = false): KeyedShuffleSpec = KeyedShuffleSpec(
+ KeyedPartitioning(Seq(bucket(4, a)), keys.map(k => InternalRow(k)))
+ .copy(mayContainUnknownPartitionKeys = hasUnknown), distribution)
+ def keyedSpec(
+ keys: Seq[Long],
+ hasUnknown: Boolean = false): KeyedShuffleSpec = KeyedShuffleSpec(
+ KeyedPartitioning(Seq(a), keys.map(k => InternalRow(k)))
+ .copy(mayContainUnknownPartitionKeys = 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
+ // on the identity side and bucket ids on the transform side. The unknown-keyed subset test
+ // would compare unrelated numbers, so the marker path must require the same partition
+ // function per position instead.
+ withSQLConf(
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ assert(!keyedSpec(Seq(1), hasUnknown = true).areKeysCompatible(bucketSpec(Seq(0, 1))),
+ "an unknown-keyed identity partitioning must not pair with a transform partitioning")
+ assert(!bucketSpec(Seq(0, 1)).areKeysCompatible(keyedSpec(Seq(1), hasUnknown = true)),
+ "the incompatibility must be symmetric")
+ assert(!keyedSpec(Seq(1), hasUnknown = true).areKeysCompatible(
+ bucketSpec(Seq(0, 1), hasUnknown = true)),
+ "the identity-vs-transform pair must not pair even when both sides are unknown-keyed")
+ // Two unmarked sides keep the pre-existing behavior: the pair stays admissible, and
+ // `EnsureRequirements` computes reducers to reconcile the two key domains.
+ assert(keyedSpec(Seq(1)).areKeysCompatible(bucketSpec(Seq(0, 1))),
+ "unmarked identity-vs-transform pairs remain admissible")
+ }
+ }
+
+ 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)))
+ .copy(mayContainUnknownPartitionKeys = true)
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+ val spec = marked.createShuffleSpec(ClusteredDistribution(Seq(a)))
+ .asInstanceOf[KeyedShuffleSpec]
+ // The refusal returns the unprojected spec whose `b` expression maps to no clustering
+ // key; that position must make the spec unusable as a shuffle template, or
+ // `createPartitioning` would index an empty position set.
+ assert(spec.joinKeyPositions.isEmpty)
+ assert(!spec.canCreatePartitioning)
+ }
+ }
+
+ test("areKeysCompatible: unknown keys require the same function, not just a compatible one") {
+ // A bucket-like reducible function: like the built-in `bucket`, a pair of the same function
+ // with coarser/finer bucket counts is compatible (a reducer exists) but not the same
+ // function. Local to the suite because catalyst has no BucketFunction.
+ class FakeBucket extends ScalarFunction[java.lang.Long]
+ with ReducibleFunction[java.lang.Long, java.lang.Long] {
+ override def inputTypes(): Array[DataType] = Array(LongType)
+ override def resultType(): DataType = LongType
+ override def name(): String = "test.fakeBucket"
+ override def canonicalName(): String = name()
+ override def produceResult(input: InternalRow): java.lang.Long = input.getLong(0)
+ override def reducer(
+ thisNumBuckets: Int,
+ other: ReducibleFunction[_, _],
+ otherNumBuckets: Int): Reducer[java.lang.Long, java.lang.Long] =
+ if (other.isInstanceOf[FakeBucket] && thisNumBuckets != otherNumBuckets &&
+ thisNumBuckets % otherNumBuckets == 0) {
+ new Reducer[java.lang.Long, java.lang.Long] {
+ override def reduce(v: java.lang.Long): java.lang.Long = v % otherNumBuckets
+ override def resultType(): DataType = LongType
+ }
+ } else {
+ null
+ }
+ }
+ val fn = new FakeBucket
+ val a = $"a".long
+ def bucketSpec(numBuckets: Int, hasUnknown: Boolean): KeyedShuffleSpec =
+ KeyedShuffleSpec(
+ KeyedPartitioning(
+ Seq(TransformExpression(fn, Seq(a), Some(numBuckets))),
+ Seq(InternalRow(0L), InternalRow(1L)))
+ .copy(mayContainUnknownPartitionKeys = hasUnknown),
+ ClusteredDistribution(Seq(a)))
+
+ // `allowCompatibleTransforms` lets a differing-bucket-count pair through
+ // `isExpressionCompatible`, and both specs declare the key set {0, 1}, so any
+ // refusal below can only come from the same-function gate on the marker path, not the keys.
+ withSQLConf(
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ assert(!bucketSpec(4, true).areKeysCompatible(bucketSpec(8, true)),
+ "two marked sides must agree on the exact function, not just a compatible one")
+ assert(!bucketSpec(4, true).areKeysCompatible(bucketSpec(8, false)),
+ "a marked side must not pair across bucket counts")
+ assert(!bucketSpec(8, false).areKeysCompatible(bucketSpec(4, true)),
+ "the refusal must be symmetric")
+ // Unmarked, the compatible-transform relaxation still pairs them (pre-existing behavior).
+ assert(bucketSpec(4, false).areKeysCompatible(bucketSpec(8, false)),
+ "unmarked compatible transforms remain admissible")
+ // Positive control: the same marked function with the same keys must stay compatible, so
+ // the refusals above are the function difference's doing, not the marker path always false.
+ assert(bucketSpec(4, true).areKeysCompatible(bucketSpec(4, true)),
+ "identical marked functions remain compatible")
+ }
+ }
+
+
test("createPartitioning: HashShuffleSpec") {
checkCreatePartitioning(
HashShuffleSpec(HashPartitioning(Seq($"a"), 10), ClusteredDistribution(Seq($"a", $"b"))),
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 84f3be562820..ef2ecac5f10b 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
@@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.expressions.codegen.GenerateOrdering
import org.apache.spark.sql.catalyst.plans.QueryPlan
-import org.apache.spark.sql.catalyst.plans.physical.{IdentityReducer, KeyedPartitioning, KeyReducer, Partitioning}
+import org.apache.spark.sql.catalyst.plans.physical.{IdentityReducer, KeyedPartitioning, KeyReducer, Partitioning, UnknownPartitioning}
import org.apache.spark.sql.catalyst.util.{truncatedString, InternalRowComparableWrapper}
import org.apache.spark.sql.execution.{SafeForKWayMerge, SparkPlan, UnaryExecNode}
import org.apache.spark.sql.internal.SQLConf
@@ -86,6 +86,20 @@ case class GroupPartitionsExec(
}
}
+ // A marked claim pins undeclared rows to hash(key) % numPartitions (see
+ // `KeyedPartitioning.mayContainUnknownPartitionKeys`). Only an identity grouping keeps
+ // that relationship: a reorder, a coalesce, or a resize moves those rows, and a
+ // projection or reduction rewrites the keys the claim speaks for (see
+ // `identityGrouping`). Clearing only the marker would misreport the undeclared rows
+ // that remain, so give up the keyed partitioning at the physical output count (one per
+ // group, padding included). The give-up deliberately under-reports: the node still
+ // physically groups the partitions but no longer claims a keyed layout, so a join
+ // planned over it does not see its required distribution satisfied. `identityGrouping`
+ // is a lazy val, so repeated `outputPartitioning` calls scan it at most once.
+ if (keyedPartitionings.exists(_.mayContainUnknownPartitionKeys) && !identityGrouping) {
+ return UnknownPartitioning(groupedPartitions.size)
+ }
+
p.transform {
case k: KeyedPartitioning =>
val projectedExpressions = joinKeyPositions.fold(k.expressions)(_.map(k.expressions))
@@ -103,7 +117,8 @@ case class GroupPartitionsExec(
case None => projectedExpressions
}
KeyedPartitioning(effectiveExpressions, groupedPartitions.map(_._1),
- isGrouped = isGrouped)
+ isGrouped = isGrouped,
+ mayContainUnknownPartitionKeys = k.mayContainUnknownPartitionKeys)
}.asInstanceOf[Partitioning]
case o => o
}
@@ -209,6 +224,36 @@ case class GroupPartitionsExec(
@transient lazy val isGrouped: Boolean = groupedPartitionsTuple._2
+ /**
+ * Whether this node's grouping leaves the declared keys and every partition where they were:
+ * no projection or reduction rewrote the keys, output partition i holds exactly input
+ * partition i, and there is one output per input. That is the only grouping that keeps a
+ * marked layout's undeclared rows at hash(key) % numPartitions. A projection or reduction
+ * re-labels the groups into a different key space, so even a grouping whose indices line up
+ * would pin the claim to keys it no longer declares; `keysRewritten` rejects it up front --
+ * it covers a narrowing projection, a reordering one, and any reducer slot. A reducer slot
+ * is treated as key-changing: a conforming self-reducer cannot rewrite a reachable key
+ * value, so the give-up there loses at most an optimization. A grouping that drops trailing
+ * declared keys still reads identity for every group it keeps, but the partition count
+ * shrinks and the hash modulus with it. The `forall` stops at the first moved partition, so
+ * a reorder or coalesce is rejected without a full scan.
+ */
+ @transient private lazy val identityGrouping: Boolean = {
+ val childKp = child.outputPartitioning
+ .asInstanceOf[Partitioning with Expression]
+ .collectFirst { case k: KeyedPartitioning => k }
+ .getOrElse(
+ throw new SparkException("GroupPartitionsExec requires a child with KeyedPartitioning"))
+ val keysRewritten =
+ joinKeyPositions.exists(_ != childKp.expressions.indices) || reducers.isDefined
+ !keysRewritten &&
+ groupedPartitions.size == child.outputPartitioning.numPartitions &&
+ groupedPartitions.iterator.zipWithIndex.forall {
+ case ((_, Seq(single)), outputIndex) => single == outputIndex
+ case _ => false
+ }
+ }
+
@transient private lazy val hasCoalescing: Boolean = groupedPartitions.exists(_._2.size > 1)
// Whether the child subtree is safe to use with SortedMergeCoalescedRDD (k-way merge).
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 b00e5d92c7b9..3a55862718c6 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
@@ -439,12 +439,12 @@ case class EnsureRequirements(
reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq, rightExpressions, rightKeys)
.orElse(reorderJoinKeysRecursively(
leftKeys, rightKeys, leftPartitioning, None))
- case (Some(KeyedPartitioning(clustering, _, _)), _) =>
+ case (Some(KeyedPartitioning(clustering, _, _, _)), _) =>
val leafExprs = clustering.flatMap(_.collectLeaves())
reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq, leafExprs, leftKeys)
.orElse(reorderJoinKeysRecursively(
leftKeys, rightKeys, None, rightPartitioning))
- case (_, Some(KeyedPartitioning(clustering, _, _))) =>
+ case (_, Some(KeyedPartitioning(clustering, _, _, _))) =>
val leafExprs = clustering.flatMap(_.collectLeaves())
reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq, leafExprs, rightKeys)
.orElse(reorderJoinKeysRecursively(
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 f363156c81e5..6718dc753224 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
@@ -17,9 +17,9 @@
package org.apache.spark.sql.execution.joins
-import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression}
import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, FullOuter, InnerLike, LeftExistence, LeftOuter, LeftSingle, RightOuter}
-import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, Distribution, Partitioning, PartitioningCollection, UnknownPartitioning, UnspecifiedDistribution}
+import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, Distribution, KeyedPartitioning, Partitioning, PartitioningCollection, UnknownPartitioning, UnspecifiedDistribution}
/**
* Holds common logic for join operators by shuffling two child relations
@@ -46,7 +46,14 @@ trait ShuffledJoin extends JoinCodegenSupport {
override def outputPartitioning: Partitioning = joinType match {
case _: InnerLike =>
- PartitioningCollection(Seq(left.outputPartitioning, right.outputPartitioning))
+ // Every `KeyedPartitioning` in the joined collection speaks for the join output's layout,
+ // and `groupedSatisfies` 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 take members'
+ // markers at face value (see `KeyedPartitioning.mayContainUnknownPartitionKeys`).
+ PartitioningCollection(
+ clearUnknownPartitionKeys(Seq(left.outputPartitioning, right.outputPartitioning)))
case LeftOuter | LeftSingle => left.outputPartitioning
case RightOuter => right.outputPartitioning
case FullOuter => UnknownPartitioning(left.outputPartitioning.numPartitions)
@@ -56,6 +63,36 @@ trait ShuffledJoin extends JoinCodegenSupport {
s"ShuffledJoin should not take $x as the JoinType")
}
+ /**
+ * Clears the `mayContainUnknownPartitionKeys` marker of every `KeyedPartitioning` in
+ * `partitionings` when marked and unmarked keyed inputs meet. Only `ShuffledJoin`'s
+ * `InnerLike` arm can mix inputs this way; see the call site for the argument.
+ */
+ private def clearUnknownPartitionKeys(
+ partitionings: Seq[Partitioning]): Seq[Partitioning] = {
+ // Keyless inputs drop out of the `flatMap` rather than reading as unmarked, and a nested
+ // collection answers with one marker per keyed member it holds. The all-unmarked path must
+ // reach no `copy`: `transform`'s `fastEquals` would compare every partition key.
+ def markersOf(p: Partitioning): Seq[Boolean] = p match {
+ case k: KeyedPartitioning => k.mayContainUnknownPartitionKeys :: Nil
+ case c: PartitioningCollection => c.partitionings.flatMap(markersOf)
+ case _ => Nil
+ }
+ val markers = partitionings.flatMap(markersOf)
+ if (markers.isEmpty || markers.forall(_ == markers.head)) {
+ partitionings
+ } else {
+ partitionings.map {
+ case partitioning: Partitioning with Expression if markersOf(partitioning).contains(true) =>
+ partitioning.transform {
+ case k: KeyedPartitioning if k.mayContainUnknownPartitionKeys =>
+ k.copy(mayContainUnknownPartitionKeys = false)
+ }.asInstanceOf[Partitioning]
+ case p => p
+ }
+ }
+ }
+
override def output: Seq[Attribute] = {
joinType match {
case _: InnerLike =>
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala
index d5825432ebbf..0273a5d6dd49 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala
@@ -50,7 +50,7 @@ abstract class DistributionAndOrderingSuiteBase
plan: QueryPlan[T]): Partitioning = partitioning match {
case HashPartitioning(exprs, numPartitions) =>
HashPartitioning(exprs.map(resolveAttrs(_, plan)), numPartitions)
- case KeyedPartitioning(expressions, partitionKeys, isGrouped) =>
+ case KeyedPartitioning(expressions, partitionKeys, isGrouped, _) =>
KeyedPartitioning(expressions.map(resolveAttrs(_, plan)), partitionKeys, isGrouped)
case PartitioningCollection(partitionings) =>
PartitioningCollection(partitionings.map(resolvePartitioning(_, plan)))
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 293a0de0a5d9..dfa7ab0220ad 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
@@ -24,7 +24,9 @@ import org.apache.spark.rdd.SortedMergeCoalescedRDD
import org.apache.spark.sql.{DataFrame, ExplainSuiteHelper, Row}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, ExprId, Literal, TransformExpression}
+import org.apache.spark.sql.catalyst.plans.Inner
import org.apache.spark.sql.catalyst.plans.physical
+import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning
import org.apache.spark.sql.connector.catalog.{Column, Identifier, InMemoryTableCatalog}
import org.apache.spark.sql.connector.catalog.functions._
import org.apache.spark.sql.connector.distributions.Distributions
@@ -5153,4 +5155,905 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
assert(collectAllShuffles(plan).isEmpty, "no shuffle either")
}
}
+
+ /**
+ * Asserts that the plan's shuffles (in tree order) are all `KeyedPartitioning`s carrying the
+ * given `mayContainUnknownPartitionKeys` flags (a `KeyedPartitioning` produced by
+ * `KeyedShuffleSpec.createPartitioning` always carries the marker).
+ */
+ /** Every `KeyedPartitioning` this partitioning reports, flattening collections. */
+ private def flattenKeyedPartitionings(p: physical.Partitioning): Seq[KeyedPartitioning] =
+ p match {
+ case k: KeyedPartitioning => k :: Nil
+ case c: physical.PartitioningCollection => c.partitionings.flatMap(flattenKeyedPartitionings)
+ case _ => Nil
+ }
+
+ private def assertShuffleMayContainUnknownPartitionKeys(
+ plan: SparkPlan,
+ expected: Seq[Boolean]): Unit = {
+ val shuffles = collectAllShuffles(plan)
+ assert(shuffles.size === expected.size,
+ s"expected ${expected.size} shuffles, got ${shuffles.size}:\n$plan")
+ shuffles.zip(expected).foreach { case (shuffle, hasUnknown) =>
+ shuffle.outputPartitioning match {
+ case k: KeyedPartitioning =>
+ assert(k.mayContainUnknownPartitionKeys === hasUnknown,
+ s"expected shuffle output mayContainUnknownPartitionKeys=$hasUnknown, got " +
+ s"${k.mayContainUnknownPartitionKeys}:\n$plan")
+ case p =>
+ fail(s"expected a KeyedPartitioning shuffle, got $p:\n$plan")
+ }
+ }
+ }
+
+
+ test("SPARK-59050: SPJ: one-side shuffle with out-of-set keys loses matches in a following " +
+ "SPJ join") {
+ // a: keyed on id, keys {1, 2}. t: v1 parquet, keys {1, 2, 3}. u: keyed on id, keys {1, 2, 3}.
+ // With shuffle.enabled, a RIGHT OUTER JOIN t shuffles t onto a's declared keys {1, 2}; t's
+ // id=3 row is out-of-set, so the join output's partitioning has unknown keys. A following
+ // storage-partitioned join against u must not trust it and falls back to a shuffle.
+ createTable("a", columns, Array(identity("id")))
+ createTable("u", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES (1, 'a1', NULL), (2, 'a2', NULL)")
+ sql("INSERT INTO testcat.ns.u VALUES (1, 'u1', NULL), (2, 'u2', NULL), (3, 'u3', NULL)")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 't1'), (2, 't2'), (3, 't3')")
+
+ val query =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT t.id AS id FROM testcat.ns.a a RIGHT OUTER JOIN t ON a.id = t.id) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ val expected = Seq(Row(1, "u1"), Row(2, "u2"), Row(3, "u3"))
+
+ // Baseline: no SPJ -> all three rows.
+ withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "false") {
+ checkAnswer(sql(query), expected)
+ }
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ // Two one-side shuffles: t onto a's keys, then the first join's output (unknown-keyed)
+ // onto u's keys. Both are keyed with unknown partition keys; neither join GPEs.
+ assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ assert(collectGroupPartitions(df.queryExecution.executedPlan).isEmpty,
+ s"second join must not storage-partition on an unknown-keyed layout, got: " +
+ df.queryExecution.executedPlan)
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: preserved non-keyed side of outer join falls back to shuffle " +
+ "downstream") {
+ // Same hazard for every outer join type whose preserved side is the non-keyed table: the
+ // one-side shuffle marks the preserved side's partitioning as having unknown keys, so a
+ // downstream storage-partitioned join against a larger key set must fall back to a shuffle.
+ createTable("a", columns, Array(identity("id")))
+ createTable("u", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES (1, 'a1', NULL), (2, 'a2', NULL)")
+ sql("INSERT INTO testcat.ns.u VALUES (1, 'u1', NULL), (2, 'u2', NULL), (3, 'u3', NULL)")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 't1'), (2, 't2'), (3, 't3')")
+
+ val expected = Seq(Row(1, "u1"), Row(2, "u2"), Row(3, "u3"))
+
+ // RIGHT OUTER preserves the non-keyed t on the right.
+ val rightQuery =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT t.id AS id FROM testcat.ns.a a RIGHT OUTER JOIN t ON a.id = t.id) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(rightQuery)
+ checkAnswer(df, expected)
+ assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ assert(collectGroupPartitions(df.queryExecution.executedPlan).isEmpty,
+ s"downstream join must not storage-partition on an unknown-keyed layout, got: " +
+ df.queryExecution.executedPlan)
+ }
+
+ // FULL OUTER exposes UnknownPartitioning, so it is already safe regardless of the shuffle
+ // direction; correctness is the guard.
+ val fullQuery =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT t.id AS id FROM testcat.ns.a a FULL OUTER JOIN t ON a.id = t.id) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(fullQuery)
+ checkAnswer(df, expected)
+ assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ // The downstream join must not storage-partition on the first join's unknown-keyed
+ // layout; FULL OUTER keeps it safe only because the join output exposes
+ // UnknownPartitioning.
+ assert(collectGroupPartitions(df.queryExecution.executedPlan).isEmpty,
+ s"downstream join must not storage-partition on an unknown-keyed layout, got: " +
+ df.queryExecution.executedPlan)
+ }
+
+ // t LEFT OUTER JOIN a preserves the non-keyed t on the left.
+ val leftQuery =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT t.id AS id FROM t LEFT OUTER JOIN testcat.ns.a a ON t.id = a.id) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(leftQuery)
+ checkAnswer(df, expected)
+ assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ assert(collectGroupPartitions(df.queryExecution.executedPlan).isEmpty,
+ s"downstream join must not storage-partition on an unknown-keyed layout, got: " +
+ df.queryExecution.executedPlan)
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: keyed preserved side of outer join still uses the one-side shuffle") {
+ // a (keyed) preserved on the left, t (non-keyed) nullable on the right: t is shuffled onto
+ // a's keys (its partitioning is marked as having unknown keys), but the LEFT OUTER join exposes
+ // only a's accurate partitioning, so the one-side shuffle stays sound and the downstream SPJ
+ // still runs (no shuffle for the second join).
+ createTable("a", columns, Array(identity("id")))
+ createTable("u", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES (1, 'a1', NULL), (2, 'a2', NULL)")
+ sql("INSERT INTO testcat.ns.u VALUES (1, 'u1', NULL), (2, 'u2', NULL), (3, 'u3', NULL)")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 't1'), (2, 't2'), (3, 't3')")
+
+ val query =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT a.id AS id FROM testcat.ns.a a LEFT OUTER JOIN t ON a.id = t.id) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, Seq(Row(1, "u1"), Row(2, "u2")))
+ assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true))
+ assert(collectGroupPartitions(df.queryExecution.executedPlan).nonEmpty,
+ s"downstream join should storage-partition on the accurate keyed layout, got: " +
+ df.queryExecution.executedPlan)
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: one-side shuffle with out-of-set keys loses matches in a following " +
+ "SPJ join (bucket)") {
+ // Same hazard as the identity variant, but the keyed sides are partitioned by bucket(4, id):
+ // a covers buckets {0, 1, 2} (ids 0, 1, 2), while t holds id 3 (bucket 3), which a does
+ // not, so the one-side shuffle misplaces t's bucket-3 row while still declaring a's layout.
+ // `id` is LONG because `BucketFunction` binds its value argument to LongType.
+ val cols = Array(Column.create("id", LongType), Column.create("data", StringType))
+ createTable("a", cols, Array(bucket(4, "id")))
+ createTable("u", cols, Array(bucket(4, "id")))
+ sql("INSERT INTO testcat.ns.a VALUES (0, 'a0'), (1, 'a1'), (2, 'a2')")
+ sql("INSERT INTO testcat.ns.u VALUES (0, 'u0'), (1, 'u1'), (2, 'u2'), (3, 'u3')")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id BIGINT, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (0, 't0'), (1, 't1'), (2, 't2'), (3, 't3')")
+
+ val query =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT t.id AS id FROM testcat.ns.a a RIGHT OUTER JOIN t ON a.id = t.id) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ val expected = Seq(Row(0L, "u0"), Row(1L, "u1"), Row(2L, "u2"), Row(3L, "u3"))
+
+ // Baseline: no SPJ -> all four rows.
+ withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "false") {
+ checkAnswer(sql(query), expected)
+ }
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ assert(collectGroupPartitions(df.queryExecution.executedPlan).isEmpty,
+ s"second join must not storage-partition on an unknown-keyed layout, got: " +
+ df.queryExecution.executedPlan)
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: unknown-keyed partitioning still joins a subset-keyed partner") {
+ // r (from a RIGHT OUTER JOIN t) has unknown partition keys {1, 2}, but the downstream u is
+ // keyed on a subset {1}, so the storage-partitioned join stays compatible and works: every
+ // key u can have is co-located on r's declared layout.
+ createTable("a", columns, Array(identity("id")))
+ createTable("u", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES (1, 'a1', NULL), (2, 'a2', NULL)")
+ sql("INSERT INTO testcat.ns.u VALUES (1, 'u1', NULL)")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 't1'), (2, 't2'), (3, 't3')")
+
+ val query =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT t.id AS id FROM testcat.ns.a a RIGHT OUTER JOIN t ON a.id = t.id) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, Seq(Row(1, "u1")))
+ // Only the first join's one-side shuffle remains; the second join storage-partitions.
+ assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true))
+ assert(collectGroupPartitions(df.queryExecution.executedPlan).nonEmpty,
+ s"subset-keyed partner should still storage-partition join, got: " +
+ df.queryExecution.executedPlan)
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: project dropping a key position drops the unknown-keyed claim") {
+ // The first join's output is keyed on (id, k) and may contain unknown keys (t's rows are all
+ // out-of-set: a holds k=x, t holds k=z). The Project below the second join drops the k
+ // position, so the declared key set coarsens from {(1, x) ... (4, x)} to {1, 2, 3, 4}, and
+ // an out-of-set (id, k) can then land inside the projected declared set. The keyed claim must
+ // be dropped entirely, otherwise the second join trusts the coarsened layout and silently
+ // loses the misplaced rows' matches.
+ val cols = Array(
+ Column.create("id", IntegerType),
+ Column.create("k", StringType),
+ Column.create("data", StringType))
+ createTable("a", cols, Array(identity("id"), identity("k")))
+ createTable("u", cols, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES " +
+ "(1, 'x', 'a1'), (2, 'x', 'a2'), (3, 'x', 'a3'), (4, 'x', 'a4')")
+ sql("INSERT INTO testcat.ns.u VALUES " +
+ "(1, NULL, 'u1'), (2, NULL, 'u2'), (3, NULL, 'u3'), (4, NULL, 'u4')")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, k STRING, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 'z', 't1'), (2, 'z', 't2'), (3, 'z', 't3'), (4, 'z', 't4')")
+
+ val query =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT t.id AS id FROM testcat.ns.a a RIGHT OUTER JOIN t
+ | ON a.id = t.id AND a.k = t.k) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ val expected = Seq(Row(1, "u1"), Row(2, "u2"), Row(3, "u3"), Row(4, "u4"))
+
+ // Baseline: no SPJ -> all four rows.
+ withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "false") {
+ checkAnswer(sql(query), expected)
+ }
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ // The projection drops the unknown-keyed claim, so the second join shuffles: two one-side
+ // shuffles, both keyed with unknown partition keys, and no GroupPartitionsExec.
+ assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ assert(collectGroupPartitions(df.queryExecution.executedPlan).isEmpty,
+ s"second join must not storage-partition on the coarsened layout, got: " +
+ df.queryExecution.executedPlan)
+ }
+ }
+ }
+
+
+ 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
+ // (`id`) instead. A key that was out-of-set in the full key space lands inside the projected
+ // declared set, so the unknown-keyed spec must be refused and the second join must shuffle.
+ val cols = Array(
+ Column.create("id", IntegerType),
+ Column.create("k", StringType),
+ Column.create("data", StringType))
+ createTable("a", cols, Array(identity("id"), identity("k")))
+ createTable("u", cols, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES " +
+ "(1, 'x', 'a1'), (2, 'x', 'a2'), (3, 'x', 'a3'), (4, 'x', 'a4')")
+ sql("INSERT INTO testcat.ns.u VALUES " +
+ "(1, NULL, 'u1'), (2, NULL, 'u2'), (3, NULL, 'u3'), (4, NULL, 'u4')")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, k STRING, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 'z', 't1'), (2, 'z', 't2'), (3, 'z', 't3'), (4, 'z', 't4')")
+
+ val query =
+ """
+ |SELECT r.id, r.k, u.data
+ |FROM (SELECT t.id AS id, t.k AS k FROM testcat.ns.a a RIGHT OUTER JOIN t
+ | ON a.id = t.id AND a.k = t.k) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ val expected = Seq(Row(1, "z", "u1"), Row(2, "z", "u2"), Row(3, "z", "u3"), Row(4, "z", "u4"))
+
+ // Baseline: no SPJ -> all four rows.
+ withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "false") {
+ checkAnswer(sql(query), expected)
+ }
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ // The second join refuses the projected unknown-keyed spec and shuffles instead: the
+ // first join's one-side shuffle plus the re-shuffle of the first join's output.
+ assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: spurious marker of an inner join keeps the reduced SPJ") {
+ // The spurious marker on the inner join's collection is cleared at construction (see
+ // `ShuffledJoin`), so a following reduced storage-partitioned join (bucket(4) onto
+ // bucket(2)) must still work. Reading the marker with `exists` used to make the
+ // GroupPartitionsExec give up with a zero-partition UnknownPartitioning, which threw at
+ // planning when a parent asked for the partitioning.
+ val cols = Array(Column.create("id", LongType), Column.create("data", StringType))
+ createTable("a", cols, Array(bucket(4, "id")))
+ createTable("u", cols, Array(bucket(2, "id")))
+ sql("INSERT INTO testcat.ns.a VALUES (0, 'a0'), (1, 'a1'), (2, 'a2'), (3, 'a3')")
+ sql("INSERT INTO testcat.ns.u VALUES (0, 'u0'), (1, 'u1'), (2, 'u2'), (3, 'u3')")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id BIGINT, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (0, 't0'), (1, 't1'), (2, 't2'), (3, 't3')")
+
+ val query =
+ """
+ |SELECT a.id, u.data
+ |FROM testcat.ns.a a JOIN t ON a.id = t.id
+ |JOIN testcat.ns.u u ON a.id = u.id
+ |""".stripMargin
+ val expected = Seq(Row(0L, "u0"), Row(1L, "u1"), Row(2L, "u2"), Row(3L, "u3"))
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ // The reduced SPJ still runs: the first join's one-side shuffle is the only shuffle.
+ assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan, Seq(true))
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: inner join with in-set rows keeps the sound SPJ downstream") {
+ // Same shape as the one-side-shuffle repro, but the inner join's second side holds only
+ // in-set rows: the marker is spurious and cleared at the join (see `ShuffledJoin`), so it
+ // costs the plan nothing. On this branch the Project between the joins drops the keyed
+ // claim when the `k` column leaves the output, so the second join one-side-shuffles r onto
+ // u's layout: both exchanges carry the marker, and u's side keeps its scan partitioning.
+ val cols = Array(
+ Column.create("id", IntegerType),
+ Column.create("k", StringType),
+ Column.create("data", StringType))
+ createTable("a", cols, Array(identity("id"), identity("k")))
+ createTable("u", cols, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES " +
+ "(1, 'x', 'a1'), (2, 'x', 'a2'), (3, 'x', 'a3'), (4, 'x', 'a4')")
+ sql("INSERT INTO testcat.ns.u VALUES " +
+ "(1, NULL, 'u1'), (2, NULL, 'u2'), (3, NULL, 'u3'), (4, NULL, 'u4')")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, k STRING, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 'x', 't1'), (2, 'x', 't2'), (3, 'x', 't3'), (4, 'x', 't4')")
+
+ val query =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT t.id AS id FROM testcat.ns.a a JOIN t ON a.id = t.id AND a.k = t.k) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ val expected = Seq(Row(1, "u1"), Row(2, "u2"), Row(3, "u3"), Row(4, "u4"))
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: inner join clears the spurious marker on both member orders") {
+ // `t`'s id=3 can match nothing on `a` (keys {1, 2}), so the inner join's marker is spurious
+ // and cleared at construction (see `ShuffledJoin`). If it survived, the subset gate in
+ // `areKeysCompatible` would refuse `u`'s wider key set and cost a shuffle no sibling order
+ // can rescue. Both join orders must plan the expected shape: the single first-join shuffle
+ // plus a `GroupPartitionsExec` on each side of the storage-partitioned second join.
+ createTable("a", columns, Array(identity("id")))
+ createTable("u", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES (1, 'a1', NULL), (2, 'a2', NULL)")
+ sql("INSERT INTO testcat.ns.u VALUES (1, 'u1', NULL), (2, 'u2', NULL), (3, 'u3', NULL)")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 't1'), (2, 't2'), (3, 't3')")
+
+ // `t` first puts the marked member ahead of its unmarked sibling in the collection;
+ // `a` first is the mirror order.
+ for (side <- Seq("t", "a")) {
+ val query = if (side == "t") {
+ """
+ |SELECT t.id, u.data
+ |FROM t JOIN testcat.ns.a a ON a.id = t.id
+ |JOIN testcat.ns.u u ON t.id = u.id
+ |""".stripMargin
+ } else {
+ """
+ |SELECT a.id, u.data
+ |FROM testcat.ns.a a JOIN t ON a.id = t.id
+ |JOIN testcat.ns.u u ON a.id = u.id
+ |""".stripMargin
+ }
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ "spark.sql.autoBroadcastJoinThreshold" -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, Seq(Row(1, "u1"), Row(2, "u2")))
+ assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan, Seq(true))
+ assert(collectGroupPartitions(df.queryExecution.executedPlan).nonEmpty,
+ s"side=$side: the second join should storage-partition, got: " +
+ df.queryExecution.executedPlan)
+ }
+ }
+ }
+ }
+
+ 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.
+ 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).copy(mayContainUnknownPartitionKeys = true),
+ new LocalTableScanExec(Seq(e), Nil, None))
+ def keylessExchange(e: AttributeReference): ShuffleExchangeExec =
+ ShuffleExchangeExec(physical.HashPartitioning(Seq(e), keys.length),
+ new LocalTableScanExec(Seq(e), Nil, None))
+ val joins = Seq(
+ SortMergeJoinExec(attrA :: Nil, attrB :: Nil, Inner, None,
+ markedExchange(attrA), keylessExchange(attrB)),
+ SortMergeJoinExec(attrA :: Nil, attrB :: Nil, Inner, None,
+ keylessExchange(attrB), markedExchange(attrA)))
+ joins.zipWithIndex.foreach { case (join, idx) =>
+ val leaves = flattenKeyedPartitionings(join.outputPartitioning)
+ assert(leaves.size == 1, s"order $idx: $leaves")
+ assert(leaves.forall(_.mayContainUnknownPartitionKeys),
+ s"order $idx: a keyless sibling must not clear the marker: $leaves")
+ }
+ }
+
+ test("SPARK-59050: SPJ: inner join marker clearing reaches nested collections") {
+ // `ShuffledJoin`'s `InnerLike` arm passes each child's partitioning into the joined
+ // collection as reported, so the next inner join can find marked members nested inside a
+ // collection inherited from an all-marked inner join below. SQL can produce that shape:
+ // two one-side-shuffled RIGHT OUTER joins with the same declared keys each expose a marked
+ // output, joining those two keeps an all-marked collection, and a following join against an
+ // accurate keyed table supplies the unmarked sibling (the end-to-end counterpart below
+ // pins that path). The nodes are hand-built here to isolate the shape from planner
+ // choices, through the same exchange-leaf idiom used above.
+ val attrA = AttributeReference("a", IntegerType)()
+ val attrB = AttributeReference("b", IntegerType)()
+ val keys = Seq(InternalRow(1), InternalRow(2))
+ def markedKP(e: AttributeReference): KeyedPartitioning =
+ KeyedPartitioning(Seq(e), keys).copy(mayContainUnknownPartitionKeys = true)
+ def markedExchange(e: AttributeReference): ShuffleExchangeExec =
+ ShuffleExchangeExec(markedKP(e), new LocalTableScanExec(Seq(e), Nil, None))
+ def plainExchange(e: AttributeReference): ShuffleExchangeExec =
+ ShuffleExchangeExec(KeyedPartitioning(Seq(e), keys),
+ new LocalTableScanExec(Seq(e), Nil, None))
+ // The first join: two marked children, no unmarked sibling -> its collection stays marked.
+ val inner1 = SortMergeJoinExec(attrA :: Nil, attrA :: Nil, Inner, None,
+ markedExchange(attrA), markedExchange(attrA))
+ // The second join nests that collection next to an unmarked sibling -> the clearing must
+ // reach inside it.
+ val inner2 = SortMergeJoinExec(attrA :: Nil, attrB :: Nil, Inner, None,
+ inner1, plainExchange(attrB))
+ val cleared = flattenKeyedPartitionings(inner2.outputPartitioning)
+ assert(cleared.size == 3, cleared.toString)
+ assert(cleared.forall(!_.mayContainUnknownPartitionKeys),
+ s"the unmarked sibling must clear every marked member: $cleared")
+ // And a marked sibling of an all-marked nested collection excuses nothing: all stay marked.
+ val inner3 = SortMergeJoinExec(attrA :: Nil, attrB :: Nil, Inner, None,
+ markedExchange(attrA), inner1)
+ val kept = flattenKeyedPartitionings(inner3.outputPartitioning)
+ assert(kept.size == 3, kept.toString)
+ assert(kept.forall(_.mayContainUnknownPartitionKeys),
+ s"an all-marked chain must keep its markers: $kept")
+ }
+
+ test("SPARK-59050: SPJ: marker clearing reaches a nested all-marked collection built by SQL") {
+ // End-to-end counterpart to the hand-built test above. The first two RIGHT OUTER joins
+ // one-side-shuffle `t` and `u` onto `a`'s and `b`'s declared keys {1, 2}; their id=3 rows
+ // are out-of-set, so both join outputs are marked. The inner join of those two outputs
+ // pairs marked against marked on equal key sequences, so its collection stays all-marked.
+ // The final join against the accurate keyed table `w` then supplies the unmarked sibling:
+ // the clearing must reach inside the nested collection, and id=3's rows, dropped by `w`,
+ // must not cost the plan its storage-partitioned final join.
+ createTable("a", columns, Array(identity("id")))
+ createTable("b", columns, Array(identity("id")))
+ createTable("w", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES (1, 'a1', NULL), (2, 'a2', NULL)")
+ sql("INSERT INTO testcat.ns.b VALUES (1, 'b1', NULL), (2, 'b2', NULL)")
+ sql("INSERT INTO testcat.ns.w VALUES (1, 'w1', NULL), (2, 'w2', NULL)")
+
+ withTable("t", "u") {
+ sql("CREATE TABLE t (id INT, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 't1'), (2, 't2'), (3, 't3')")
+ sql("CREATE TABLE u (id INT, data STRING) USING parquet")
+ sql("INSERT INTO u VALUES (1, 'u1'), (2, 'u2'), (3, 'u3')")
+
+ val query =
+ """
+ |SELECT r3.id, w.data
+ |FROM (
+ | SELECT r1.id FROM (
+ | SELECT tt.id AS id FROM testcat.ns.a RIGHT OUTER JOIN t tt ON a.id = tt.id
+ | ) r1
+ | JOIN (
+ | SELECT uu.id AS id FROM testcat.ns.b RIGHT OUTER JOIN u uu ON b.id = uu.id
+ | ) r2 ON r1.id = r2.id
+ |) r3
+ |JOIN testcat.ns.w ON r3.id = w.id
+ |""".stripMargin
+ val expected = Seq(Row(1, "w1"), Row(2, "w2"))
+
+ // Baseline without SPJ.
+ withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "false") {
+ checkAnswer(sql(query), expected)
+ }
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ "spark.sql.autoBroadcastJoinThreshold" -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ // Only the two one-side shuffles carry the marker: the final join storage-partitions
+ // without a re-shuffle (a third shuffle here would mean the marked-vs-marked pairing
+ // was refused).
+ assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ // The clearing reached the nested collection: every keyed leaf of the final output is
+ // unmarked.
+ val leaves = flattenKeyedPartitionings(df.queryExecution.executedPlan.outputPartitioning)
+ assert(leaves.nonEmpty, s"expected keyed leaves in the final output: " +
+ df.queryExecution.executedPlan)
+ assert(leaves.forall(!_.mayContainUnknownPartitionKeys),
+ s"the unmarked sibling must clear the nested collection: $leaves")
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: inner join drops out-of-set rows before the cleared marker is trusted") {
+ // Data-level companion to the spurious-marker tests: `t` here genuinely holds an out-of-set
+ // row (5, 'z') that a's declared keys cannot match. The inner join drops it, so every row
+ // reaching the cleared-marker layout carries a declared key; the downstream storage-
+ // partitioned join must then return exactly the matched rows: ids 1..4, and NOT id=5 even
+ // though u holds it. If the marker were cleared for a shape that keeps out-of-set rows,
+ // this query would silently lose or misplace matches.
+ val cols = Array(
+ Column.create("id", IntegerType),
+ Column.create("k", StringType),
+ Column.create("data", StringType))
+ createTable("a", cols, Array(identity("id"), identity("k")))
+ createTable("u", cols, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.a VALUES " +
+ "(1, 'x', 'a1'), (2, 'x', 'a2'), (3, 'x', 'a3'), (4, 'x', 'a4')")
+ sql("INSERT INTO testcat.ns.u VALUES (1, NULL, 'u1'), (2, NULL, 'u2'), (3, NULL, 'u3'), " +
+ "(4, NULL, 'u4'), (5, NULL, 'u5')")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, k STRING, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 'x', 't1'), (2, 'x', 't2'), (3, 'x', 't3'), " +
+ "(4, 'x', 't4'), (5, 'z', 't5')")
+
+ val query =
+ """
+ |SELECT r.id, u.data
+ |FROM (SELECT t.id AS id FROM t JOIN testcat.ns.a a ON a.id = t.id AND a.k = t.k) r
+ |JOIN testcat.ns.u u ON r.id = u.id
+ |""".stripMargin
+ // id=5 dropped by the inner join (a has no (5, *)); u's id=5 row has no partner.
+ val expected = Seq(Row(1, "u1"), Row(2, "u2"), Row(3, "u3"), Row(4, "u4"))
+
+ // Baseline without SPJ.
+ withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "false") {
+ checkAnswer(sql(query), expected)
+ }
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ // On this branch the Project between the joins drops the keyed claim, so the second
+ // join one-side-shuffles the join output onto u's layout; every row reaching it
+ // carries a declared key, so the re-shuffle is sound.
+ assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: genuine unknown keys survive an inner join on a key subset") {
+ // Adversarial shape against the spurious-marker clearing: the RIGHT OUTER output `r` is a
+ // single unknown-keyed KP carrying a GENUINE unknown row (1, 'zz'); it inner-joins `x`
+ // (a keyed table declaring the same key set) on a strict subset of the key columns; the
+ // join to `u` below matches on the full (id, k). The (1,'zz','xa','uz') match must survive:
+ // clearing markers only applies to mixed collections where every row already carries a
+ // declared key,
+ // and here the planner re-shuffles `r`'s side onto u's layout (the (id)-only claim cannot
+ // pair with the (id,k) spec), routing (1,'zz') by its full tuple into its own partition.
+ val cols = Array(
+ Column.create("id", IntegerType),
+ Column.create("k", StringType),
+ Column.create("data", StringType))
+ createTable("a", cols, Array(identity("id"), identity("k")))
+ createTable("x", cols, Array(identity("id"), identity("k")))
+ createTable("u", cols, Array(identity("id"), identity("k")))
+ sql("INSERT INTO testcat.ns.a VALUES (1, 'x', 'a1'), (2, 'x', 'a2')")
+ sql("INSERT INTO testcat.ns.x VALUES (1, 'x', 'xa'), (2, 'x', 'xb')")
+ sql("INSERT INTO testcat.ns.u VALUES (1, 'x', 'ux'), (2, 'x', 'ub'), (1, 'zz', 'uz')")
+
+ withTable("t") {
+ sql("CREATE TABLE t (id INT, k STRING, data STRING) USING parquet")
+ sql("INSERT INTO t VALUES (1, 'x', 't1'), (2, 'x', 't2'), (1, 'zz', 't3')")
+
+ val query =
+ """
+ |SELECT r.id, r.k, x.data, u.data
+ |FROM (SELECT t.id AS id, t.k AS k FROM testcat.ns.a a RIGHT OUTER JOIN t
+ | ON a.id = t.id AND a.k = t.k) r
+ |JOIN testcat.ns.x x ON r.id = x.id
+ |JOIN testcat.ns.u u ON r.id = u.id AND r.k = u.k
+ |""".stripMargin
+ val expected = Seq(Row(1, "x", "xa", "ux"), Row(2, "x", "xb", "ub"),
+ Row(1, "zz", "xa", "uz"))
+
+ withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "false") {
+ checkAnswer(sql(query), expected)
+ }
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ // Pin the total exchange count so later silent shuffles surface: the marked (id, k)
+ // side refuses to pair on the subset key, cascading one-side shuffles around it.
+ val plan = df.queryExecution.executedPlan
+ val shuffles = collectAllShuffles(plan)
+ assert(shuffles.size == 4, s"expected 4 exchanges, got ${shuffles.size}:\n$plan")
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: a window keyed on a subset of an unknown-keyed layout must shuffle") {
+ // The right-outer join preserves pt's out-of-set row (1, 9): the surviving layout declares
+ // (id, k) pairs (1, 0)..(4, 0) only. A window keyed on `id` alone must not run
+ // partition-locally on it: the declared (1, 0) row and the hash-placed (1, 9) row can
+ // sit in different partitions, splitting the count for id=1 into two groups of 1.
+ // `groupedSatisfies` refuses the subset relaxation for an unknown-keyed layout.
+ val cols = Array(
+ Column.create("id", IntegerType),
+ Column.create("k", IntegerType),
+ Column.create("data", StringType))
+ createTable("pa", cols, Array(identity("id"), identity("k")))
+ sql("INSERT INTO testcat.ns.pa VALUES " +
+ "(1, 0, 'a1'), (2, 0, 'a2'), (3, 0, 'a3'), (4, 0, 'a4')")
+ withTable("pt") {
+ sql("CREATE TABLE pt (id INT, k INT, data STRING) USING parquet")
+ sql("INSERT INTO pt VALUES " +
+ "(1, 0, 't1'), (2, 0, 't2'), (3, 0, 't3'), (4, 0, 't4'), (1, 9, 'u1')")
+ val query =
+ """
+ |SELECT id, k, COUNT(*) OVER (PARTITION BY id) FROM (
+ | SELECT /*+ MERGE */ pt.id AS id, pt.k AS k FROM testcat.ns.pa pa
+ | RIGHT OUTER JOIN pt ON pa.id = pt.id AND pa.k = pt.k) r
+ |""".stripMargin
+ // AQE off pins the plan shape; AQE on replans the same query over stage outputs and must
+ // still return correct counts (the marker survives `ShuffleQueryStageExec`).
+ for (adaptive <- Seq(false, true)) {
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true",
+ "spark.sql.autoBroadcastJoinThreshold" -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) {
+ val df = sql(query)
+ checkAnswer(df, Seq(Row(1, 0, 2), Row(2, 0, 1), Row(3, 0, 1), Row(4, 0, 1),
+ Row(1, 9, 2)))
+ if (!adaptive) {
+ // Two shuffles: the first join's one-side shuffle plus the window's exchange.
+ val plan = df.queryExecution.executedPlan
+ val shuffles = collectAllShuffles(plan)
+ assert(shuffles.size == 2, s"the window must pay an exchange, got:\n$plan")
+ assert(shuffles.exists(!_.outputPartitioning.isInstanceOf[KeyedPartitioning]),
+ s"expected a non-keyed (window) exchange, got:\n$plan")
+ }
+ }
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: a window keyed on the full key of an unknown-keyed layout does not " +
+ "shuffle") {
+ // Positive control for the marked branch of `groupedSatisfies`: whole declared keys
+ // co-locate, so
+ // a window keyed on the COMPLETE partition key of a marked layout still runs
+ // partition-locally. A gate that over-rejects (e.g. also refusing full-key clustering) keeps
+ // every wrong-results test green and only silently adds shuffles; this test would not.
+ val cols = Array(
+ Column.create("id", IntegerType),
+ Column.create("k", IntegerType),
+ Column.create("data", StringType))
+ createTable("qa", cols, Array(identity("id"), identity("k")))
+ sql("INSERT INTO testcat.ns.qa VALUES (1, 0, 'a1'), (2, 0, 'a2'), (3, 0, 'a3')")
+ withTable("qt") {
+ sql("CREATE TABLE qt (id INT, k INT, data STRING) USING parquet")
+ sql("INSERT INTO qt VALUES (1, 0, 't1'), (2, 0, 't2'), (3, 0, 't3'), (1, 9, 'u1')")
+ val query =
+ """
+ |SELECT id, k, COUNT(*) OVER (PARTITION BY id, k) FROM (
+ | SELECT /*+ MERGE */ qt.id AS id, qt.k AS k FROM testcat.ns.qa qa
+ | RIGHT OUTER JOIN qt ON qa.id = qt.id AND qa.k = qt.k) r
+ |""".stripMargin
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ "spark.sql.autoBroadcastJoinThreshold" -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, Seq(Row(1, 0, 1), Row(2, 0, 1), Row(3, 0, 1), Row(1, 9, 1)))
+ val plan = df.queryExecution.executedPlan
+ val shuffles = collectAllShuffles(plan)
+ // Only the first join's one-side shuffle; the window rides the marked layout.
+ assert(shuffles.size == 1, s"the full-key window must not shuffle, got:\n$plan")
+ }
+ }
+ }
+
+ test("SPARK-59050: SPJ: a global ORDER BY over an unknown-keyed layout must " +
+ "range-partition") {
+ // The right-outer join preserves obt's out-of-set id=4, which the KeyGroupedPartitioner
+ // hash-placed into the partition declaring id=1; without a range exchange the global sort
+ // degrades to per-partition sorting and emits [1, 4, 2]. `keysSatisfy` rejects the ordering
+ // claim of a marked layout with more than one partition.
+ createTable("oba", columns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.oba VALUES (1, 'x', NULL), (2, 'x', NULL)")
+ withTable("obt") {
+ sql("CREATE TABLE obt (id INT, data STRING) USING parquet")
+ sql("INSERT INTO obt VALUES (1, 'p1'), (2, 'p2'), (4, 'p4')")
+ val query =
+ """
+ |SELECT /*+ MERGE */ obt.id FROM testcat.ns.oba ba RIGHT OUTER JOIN obt
+ |ON ba.id = obt.id
+ |ORDER BY id
+ |""".stripMargin
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_SORTING_ENABLED.key -> "true",
+ "spark.sql.autoBroadcastJoinThreshold" -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ // compare `collect` directly: checkAnswer would sort both sides and hide the order.
+ val ordered = df.collect().map(_.getInt(0)).toSeq
+ assert(ordered == Seq(1, 2, 4), s"global order broken: $ordered")
+ val plan = df.queryExecution.executedPlan
+ val shuffles = collectAllShuffles(plan)
+ // Two exchanges: the first join's one-side shuffle plus the range partitioning.
+ assert(shuffles.size == 2, s"expected 2 exchanges, got ${shuffles.size}:\n$plan")
+ assert(shuffles.exists(_.outputPartitioning.isInstanceOf[physical.RangePartitioning]),
+ s"a range exchange must precede the global sort, got:\n$plan")
+ }
+ }
+ }
+
+
+ test("SPARK-59050: SPJ: a reducer-free identity regrouping keeps the unknown-keyed claim") {
+ // Positive counterpart to the regrouping test above. The first join one-side-shuffles pt
+ // onto pa's sorted keys [1, 2, 3] and marks the layout (pt's id=4 is out-of-set). The second
+ // join aligns that marked [1, 2, 3] superset with pb's [1, 2] by pushing the merged keys
+ // [1, 2, 3] to both sides. The marked side already holds exactly the merged keys, so its
+ // GroupPartitionsExec is a reducer-free identity grouping and must keep the claim; only the
+ // subset side pads. This pins that the identity branch is reachable, not dead code.
+ val cols = Array(Column.create("id", IntegerType), Column.create("data", StringType))
+ createTable("pa", cols, Array(identity("id")))
+ createTable("pb", cols, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.pa VALUES (1, 'a1'), (2, 'a2'), (3, 'a3')")
+ sql("INSERT INTO testcat.ns.pb VALUES (1, 'b1'), (2, 'b2')")
+ withTable("pt") {
+ sql("CREATE TABLE pt (id INT, data STRING) USING parquet")
+ sql("INSERT INTO pt VALUES (1, 't1'), (2, 't2'), (3, 't3'), (4, 't4')")
+ val query =
+ """
+ |SELECT r.id, pb.data
+ |FROM (SELECT pt.id AS id FROM testcat.ns.pa RIGHT OUTER JOIN pt ON pa.id = pt.id) r
+ |JOIN testcat.ns.pb ON r.id = pb.id
+ |""".stripMargin
+ val expected = Seq(Row(1, "b1"), Row(2, "b2"))
+
+ // Baseline: no SPJ.
+ withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "false") {
+ checkAnswer(sql(query), expected)
+ }
+ withSQLConf(
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ "spark.sql.autoBroadcastJoinThreshold" -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ val plan = df.queryExecution.executedPlan
+ // A GroupPartitionsExec with no reducers whose output partition i holds exactly input
+ // partition i is an identity grouping; the fix keeps the unknown-keyed claim for it.
+ val identityGpe = collectAllGroupPartitions(plan).find { g =>
+ g.reducers.isEmpty && g.groupedPartitions.zipWithIndex.forall {
+ case ((_, inputIndices), outputIndex) =>
+ inputIndices.lengthCompare(1) == 0 && inputIndices.head == outputIndex
+ }
+ }
+ assert(identityGpe.isDefined,
+ s"expected a reducer-free identity GroupPartitionsExec, got:\n$plan")
+ identityGpe.get.outputPartitioning match {
+ case k: KeyedPartitioning =>
+ assert(k.mayContainUnknownPartitionKeys,
+ "the identity grouping must keep the unknown-keyed claim")
+ case other =>
+ fail(s"expected a KeyedPartitioning output, got $other")
+ }
+ }
+ }
+ }
+
}
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 607eddc83bf1..b79cf35ff600 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
@@ -21,11 +21,12 @@ import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeReference, SortOrder, TransformExpression}
import org.apache.spark.sql.catalyst.plans.physical.{KeyedPartitioning, KeyReducer, Partitioning, PartitioningCollection, UnknownPartitioning}
-import org.apache.spark.sql.connector.catalog.functions.{BucketFunction, BucketReducer}
+import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper
+import org.apache.spark.sql.connector.catalog.functions.{BucketFunction, BucketReducer, Reducer}
import org.apache.spark.sql.execution.{DummySparkPlan, LeafExecNode, SafeForKWayMerge}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
-import org.apache.spark.sql.types.IntegerType
+import org.apache.spark.sql.types.{DataType, IntegerType}
class GroupPartitionsExecSuite extends SharedSparkSession {
@@ -190,6 +191,157 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
assert(gpe.outputOrdering === Nil)
}
+ test("SPARK-59050: identity grouping without reducers keeps the child's marker") {
+ // An identity grouping (output partition i holds exactly input partition i) keeps the
+ // unknown-key hash-routing relationship intact, so the marker rides the transform.
+ val child = DummySparkPlan(
+ outputPartitioning = KeyedPartitioning(Seq(exprA), Seq(row(1), row(2)))
+ .copy(mayContainUnknownPartitionKeys = true))
+ val out = GroupPartitionsExec(child)
+ .outputPartitioning.asInstanceOf[KeyedPartitioning]
+ assert(out.mayContainUnknownPartitionKeys, "the marker must survive identity grouping")
+ }
+
+ 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.
+ // 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)))
+ .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")
+ }
+ // 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)))
+ .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")
+ }
+ }
+
+ test("SPARK-59050: an identity grouping that shrinks the partition count drops the claim") {
+ // `alignToExpectedKeys` emits only the expected keys, so a declared key the merged set
+ // does not carry is never emitted. When the dropped key is trailing, every kept group
+ // still reads as the identity, but the partition count shrank and the hash modulus with
+ // it: the child's undeclared rows sit at hash(key) % 3 while the retained claim would
+ // 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)))
+ .copy(mayContainUnknownPartitionKeys = true))
+ def keyOf(a: Int): InternalRowComparableWrapper =
+ InternalRowComparableWrapper(row(a), Seq(exprA))
+ val gpe = GroupPartitionsExec(child,
+ expectedPartitionKeys = Some(Seq(keyOf(1) -> 1, keyOf(2) -> 1)))
+ assert(gpe.groupedPartitions.size === 2, "the trailing key 3 is dropped")
+ gpe.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 a shrink, got $other")
+ }
+ }
+
+ test("SPARK-59050: unknown-keyed child with reducers gives up the keyed claim at the real " +
+ "count") {
+ // The reducer instance of the give-up: a reduction regroups by the reduced keys, a
+ // non-identity grouping, so the unknown-keyed claim is dropped at the physical grouped
+ // count. The node must report `UnknownPartitioning` with its physical grouped count;
+ // reporting zero partitions is what threw once a parent join built a
+ // `PartitioningCollection` over both sides. (`areKeysCompatible` pairs a marked layout
+ // only with same-function partners that need no reducer, so no planner path applies a
+ // reducer to a marked layout; this pins the contract directly.)
+ // The reducer must be real: it is the collapse of [1, 2, 3] to two groups that exercises
+ // the regrouping the give-up answers. A reducer slot that rewrites no key is pinned by the
+ // next test, where it is the slot itself, not a merge, that drops the claim.
+ // Keys [1, 2, 3] reduced mod 2 give [1, 0, 1]: the node emits 2 grouped partitions where
+ // the child had 3, and the give-up count has to be that physical 2.
+ val mod2 = new Reducer[Int, Int] {
+ override def reduce(v: Int): Int = v % 2
+ override def resultType(): DataType = IntegerType
+ override def displayName(): String = "mod2"
+ }
+ val reducer = KeyReducer(mod2, TransformExpression(BucketFunction, Seq(exprA), Some(2)))
+ val partitionKeys = Seq(row(1), row(2), row(3))
+ val child = DummySparkPlan(
+ outputPartitioning = KeyedPartitioning(Seq(exprA), partitionKeys)
+ .copy(mayContainUnknownPartitionKeys = true))
+ val gpe = GroupPartitionsExec(child, reducers = Some(Seq(Some(reducer))))
+
+ assert(gpe.groupedPartitions.size === 2, "mod 2 collapses keys 1 and 3")
+ gpe.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 reduction, got $other")
+ }
+ // A parent join merges the two sides' partitionings into a collection: a count of 0 fails
+ // the uniform-numPartitions requirement, the planning throw reproduced; this call
+ // throwing fails the test.
+ val partner = KeyedPartitioning(Seq(exprB), Seq(row(1), row(2)))
+ PartitioningCollection(Seq(gpe.outputPartitioning, partner))
+ }
+
+ test("SPARK-59050: a grouping that rewrites the declared keys drops the claim") {
+ // `identityGrouping` also asks whether the grouping rewrote the keys: the claim the node
+ // goes on to declare lives in the projected or reduced key space, while the child's
+ // undeclared rows still sit at hash(originalKey) % numPartitions. A reducer slot, a
+ // narrowing projection, or a reordering projection therefore gives up the claim even when
+ // every group keeps its index and the count is unchanged. A conforming self-reducer cannot
+ // rewrite a reachable key (its contract is r(f(x)) = f(x)), so the give-up there loses at
+ // most an optimization; no planner path applies a reducer or a non-identity projection to
+ // 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)))
+ .copy(mayContainUnknownPartitionKeys = true))
+
+ // Reducer slots: `Some(Seq(None, None))` leaves every key and every index unchanged.
+ val reduced = GroupPartitionsExec(child, reducers = Some(Seq(None, None)))
+ assert(reduced.groupedPartitions.size === 2)
+ reduced.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 with reducer slots, got $other")
+ }
+
+ // Narrowing projection: position 0 of keys [(1, 10), (2, 20)] keeps both groups in order.
+ val projected = GroupPartitionsExec(child, joinKeyPositions = Some(Seq(0)))
+ assert(projected.groupedPartitions.size === 2)
+ projected.outputPartitioning match {
+ case u: UnknownPartitioning =>
+ assert(u.numPartitions === 2, "the give-up count must match the physical partitions")
+ case other =>
+ fail("expected the unknown-keyed claim to be dropped on a narrowing projection, " +
+ s"got $other")
+ }
+
+ // Reordering projection: positions Seq(1, 0) re-label every group into the swapped key
+ // space while each group keeps its index and the count, so only the rewrite clause can
+ // catch it. No producer of joinKeyPositions emits anything but ascending positions today;
+ // this pins the predicate directly.
+ val reordered = GroupPartitionsExec(child, joinKeyPositions = Some(Seq(1, 0)))
+ assert(reordered.groupedPartitions.size === 2)
+ assert(reordered.groupedPartitions.zipWithIndex.forall {
+ case ((_, inputIndices), outputIndex) => inputIndices == Seq(outputIndex)
+ }, "the reordering keeps every group at its index")
+ reordered.outputPartitioning match {
+ case u: UnknownPartitioning =>
+ assert(u.numPartitions === 2, "the give-up count must match the physical partitions")
+ case other =>
+ fail("expected the unknown-keyed claim to be dropped on a reordering projection, " +
+ s"got $other")
+ }
+ }
+
test("SPARK-55715: sorted merge config enabled but child not SafeForKWayMerge falls back " +
"to key-expression ordering") {
// DummySparkPlan does not extend SafeForKWayMerge, so childIsSafeForKWayMerge = false and
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 55b8f29f711a..10d8589ec0f3 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
@@ -1137,7 +1137,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
EnsureRequirements.apply(smjExec) match {
case ShuffledHashJoinExec(_, _, _, _, _,
DummySparkPlan(_, _, left: KeyedPartitioning, _, _),
- ShuffleExchangeExec(KeyedPartitioning(attrs, pks, _),
+ ShuffleExchangeExec(KeyedPartitioning(attrs, pks, _, _),
DummySparkPlan(_, _, SinglePartition, _, _), _, _), _) =>
assert(left.expressions == a1 :: Nil)
assert(attrs == a1 :: Nil)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org