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 06:37:18 UTC
(spark) branch branch-4.0 updated: [SPARK-59050][4.0][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.0
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.0 by this push:
new 01cc00f6bc67 [SPARK-59050][4.0][SQL] SPJ one-side shuffle with out-of-set keys produces wrong results in multi-joins
01cc00f6bc67 is described below
commit 01cc00f6bc67d0bd0c9e46d10e005a18b9b2d5b3
Author: Xiduo You <ul...@gmail.com>
AuthorDate: Tue Sep 8 14:37:02 2026 +0800
[SPARK-59050][4.0][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.0`, reduced to what this branch's SPJ machinery has; it is the `branch-4.1` backport (#58571) applied as is, which holds because the two branches share every piece of SPJ machinery this fix touches. The core fix is unchanged: `KeyGroupedPartitioning` gains `mayContainUnknownPartitionKeys`, set by `KeyGroupedShuffleSpec.createPartitioning` (the one-side shuffle), and `areKeysCompatible` only co-partitions such a side with a subset-keyed (or identically-d [...]
Two e2e tests are adapted around a pre-existing `branch-4.0` limitation, reproduced on the unmodified branch: with `spark.sql.sources.v2.bucketing.pushPartValues.enabled` (default on), `checkKeyGroupCompatible` merges the two sides' partition values and pushes the merged set down without requiring the keyed side to be a plain scan subtree, so a keyed side sitting above a one-side-shuffled join whose partition count differs from the partner's plans a shuffle-free join over unequal part [...]
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.0`: `KeyGroupedPartitioningSuite` (12 `SPARK-59050` tests), `ShuffleSpecSuite` (3 new tests), `DistributionSuite`, `EnsureRequirementsSuite`, `ValidateRequirementsSuite`.
### Was this patch authored or co-authored using generative AI tooling?
Yes. Generated-by: Claude Code.
Closes #58572 from ulysses-you/spj-part-mismatch-4.0.
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 | 23 +-
.../spark/sql/catalyst/ShuffleSpecSuite.scala | 123 +++++
.../execution/exchange/EnsureRequirements.scala | 6 +-
.../execution/exchange/ShuffleExchangeExec.scala | 4 +-
.../DistributionAndOrderingSuiteBase.scala | 2 +-
.../connector/KeyGroupedPartitioningSuite.scala | 552 ++++++++++++++++++++-
.../exchange/EnsureRequirementsSuite.scala | 2 +-
8 files changed, 800 insertions(+), 20 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 9777c31426c0..1c6238e06a59 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
@@ -374,13 +374,34 @@ case class CoalescedHashPartitioning(from: HashPartitioning, partitions: Seq[Coa
* @param originalPartitionValues the original input partition values before any grouping has been
* applied, must be in ascending order, and may contain duplicated
* values
+ * @param mayContainUnknownPartitionKeys Whether the data may contain rows whose partition key is
+ * not among the declared `partitionValues`. `KeyGroupedPartitioner`
+ * routes such rows by a deterministic hash when a side is
+ * re-shuffled onto this partitioning (see
+ * `KeyGroupedShuffleSpec.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 `KeyGroupedShuffleSpec.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 KeyGroupedPartitioning(
expressions: Seq[Expression],
numPartitions: Int,
partitionValues: Seq[InternalRow] = Seq.empty,
originalPartitionValues: Seq[InternalRow] = Seq.empty,
- isPartiallyClustered: Boolean = false) extends HashPartitioningLike {
+ isPartiallyClustered: Boolean = false,
+ mayContainUnknownPartitionKeys: Boolean = false) extends HashPartitioningLike {
// See SPARK-55848. We must check ClusteredDistribution BEFORE delegating to
// super.satisfies0(), because HashPartitioningLike.satisfies0() also matches
@@ -399,7 +420,11 @@ case class KeyGroupedPartitioning(
// 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 (KeyGroupedPartitioning attributes)
requiredClustering.exists(x => attributes.exists(_.semanticEquals(x))) &&
@@ -410,7 +435,10 @@ case class KeyGroupedPartitioning(
}
case o @ OrderedDistribution(_) if SQLConf.get.v2BucketingAllowSorting =>
- 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`).
+ o.areAllClusterKeysMatched(expressions) &&
+ (!mayContainUnknownPartitionKeys || numPartitions == 1)
case _ =>
super.satisfies0(required)
@@ -420,12 +448,22 @@ case class KeyGroupedPartitioning(
override def createShuffleSpec(distribution: ClusteredDistribution): ShuffleSpec = {
val result = KeyGroupedShuffleSpec(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
// `KeyGroupedPartitioning` 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 projectedPartitioning = KeyGroupedPartitioning(expressions, joinKeyPositions,
- partitionValues, originalPartitionValues, isPartiallyClustered)
+ partitionValues, originalPartitionValues, isPartiallyClustered,
+ mayContainUnknownPartitionKeys = mayContainUnknownPartitionKeys)
result.copy(partitioning = projectedPartitioning, joinKeyPositions = Some(joinKeyPositions))
} else {
result
@@ -449,7 +487,8 @@ object KeyGroupedPartitioning {
projectionPositions: Seq[Int],
partitionValues: Seq[InternalRow],
originalPartitionValues: Seq[InternalRow],
- isPartiallyClustered: Boolean): KeyGroupedPartitioning = {
+ isPartiallyClustered: Boolean,
+ mayContainUnknownPartitionKeys: Boolean): KeyGroupedPartitioning = {
val projectedExpressions = projectionPositions.map(expressions(_))
val projectedPartitionValues = partitionValues.map(project(expressions, projectionPositions, _))
val projectedOriginalPartitionValues =
@@ -461,7 +500,8 @@ object KeyGroupedPartitioning {
.map(_.row)
KeyGroupedPartitioning(projectedExpressions, finalPartitionValues.length,
- finalPartitionValues, projectedOriginalPartitionValues, isPartiallyClustered)
+ finalPartitionValues, projectedOriginalPartitionValues, isPartiallyClustered,
+ mayContainUnknownPartitionKeys)
}
def project(
@@ -857,6 +897,46 @@ case class KeyGroupedShuffleSpec(
}
} && 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 `KeyGroupedPartitioning.mayContainUnknownPartitionKeys`).
+ //
+ // The key comparison below must also happen in a single domain: `isExpressionCompatible`
+ // admits two different-but-compatible transforms when `v2BucketingAllowCompatibleTransforms`
+ // is on, and in those cases the two sides' `partitionValues` hold the two transforms'
+ // outputs, 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, 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
+ } && {
+ def valuesOf(kgp: KeyGroupedPartitioning): Seq[InternalRowComparableWrapper] =
+ kgp.partitionValues.map(InternalRowComparableWrapper(_, kgp.expressions))
+ val values = valuesOf(partitioning)
+ val otherValues = valuesOf(other.partitioning)
+ if (partitioning.mayContainUnknownPartitionKeys &&
+ other.partitioning.mayContainUnknownPartitionKeys) {
+ values == otherValues
+ } else if (partitioning.mayContainUnknownPartitionKeys) {
+ val declared = values.toSet
+ otherValues.forall(declared.contains)
+ } else {
+ val declared = otherValues.toSet
+ values.forall(declared.contains)
+ }
+ }
+ } else {
+ true
+ }
}
}
@@ -903,6 +983,10 @@ case class KeyGroupedShuffleSpec(
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]
}
@@ -916,9 +1000,15 @@ case class KeyGroupedShuffleSpec(
te.copy(children = te.children.map(_ => clustering(positionSet.head)))
case (_, positionSet) => clustering(positionSet.head)
}
+ // 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.
KeyGroupedPartitioning(newExpressions,
partitioning.numPartitions,
- partitioning.partitionValues)
+ partitioning.partitionValues,
+ 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 4f3efca4ad0f..87ae810bef27 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.{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,
@@ -391,4 +393,21 @@ class DistributionSuite extends SparkFunSuite {
StatefulOpClusteredDistribution(Seq($"a", $"b", $"c"), 10),
false)
}
+
+ test("SPARK-59050: a marked one-partition layout keeps the global ordering claim") {
+ // The one-partition exemption inside the ordered branch of `satisfies0`: 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 = KeyGroupedPartitioning(Seq(a), 1, Seq(InternalRow(1)))
+ .copy(mayContainUnknownPartitionKeys = true)
+ val markedTwo = KeyGroupedPartitioning(Seq(a), 2, 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 fc5d39fd9c2b..5e3c36e10734 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
@@ -19,9 +19,12 @@ package org.apache.spark.sql.catalyst
import org.apache.spark.{SparkFunSuite, SparkUnsupportedOperationException}
import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.expressions.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.{Reducer, ReducibleFunction, ScalarFunction}
import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{DataType, LongType}
class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
protected def checkCompatible(
@@ -479,4 +482,124 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
"methodName" -> "createPartitioning$",
"className" -> "org.apache.spark.sql.catalyst.plans.physical.ShuffleSpec"))
}
+
+ 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): KeyGroupedShuffleSpec = KeyGroupedShuffleSpec(
+ KeyGroupedPartitioning(Seq(a), keys.length, 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
+ // `KeyGroupedShuffleSpec.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("createShuffleSpec: a marked narrowing projection yields an unusable spec") {
+ val a = $"a".int
+ val b = $"b".int
+ val marked = KeyGroupedPartitioning(Seq(a, b), 2, 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[KeyGroupedShuffleSpec]
+ // 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
+ }
+ } else {
+ null
+ }
+ }
+ val fn = new FakeBucket
+ val a = $"a".long
+ def bucketSpec(numBuckets: Int, hasUnknown: Boolean): KeyGroupedShuffleSpec =
+ KeyGroupedShuffleSpec(
+ KeyGroupedPartitioning(
+ Seq(TransformExpression(fn, Seq(a), Some(numBuckets))), 2,
+ 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")
+ }
+ }
}
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 a6c685e26cc8..02827d39ec7e 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
@@ -296,7 +296,7 @@ case class EnsureRequirements(
private def ensureOrdering(plan: SparkPlan, distribution: Distribution) = {
(plan.outputPartitioning, distribution) match {
- case (p @ KeyGroupedPartitioning(expressions, _, partitionValues, _, _),
+ case (p @ KeyGroupedPartitioning(expressions, _, partitionValues, _, _, _),
d @ OrderedDistribution(ordering)) if p.satisfies(d) =>
val attrs = expressions.flatMap(_.collectLeaves()).map(_.asInstanceOf[Attribute])
val partitionOrdering: Ordering[InternalRow] = {
@@ -329,12 +329,12 @@ case class EnsureRequirements(
reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq, rightExpressions, rightKeys)
.orElse(reorderJoinKeysRecursively(
leftKeys, rightKeys, leftPartitioning, None))
- case (Some(KeyGroupedPartitioning(clustering, _, _, _, _)), _) =>
+ case (Some(KeyGroupedPartitioning(clustering, _, _, _, _, _)), _) =>
val leafExprs = clustering.flatMap(_.collectLeaves())
reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq, leafExprs, leftKeys)
.orElse(reorderJoinKeysRecursively(
leftKeys, rightKeys, None, rightPartitioning))
- case (_, Some(KeyGroupedPartitioning(clustering, _, _, _, _))) =>
+ case (_, Some(KeyGroupedPartitioning(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/exchange/ShuffleExchangeExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala
index 83697b715dfb..7008422da72e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala
@@ -366,7 +366,7 @@ object ShuffleExchangeExec {
ascending = true,
samplePointsPerPartitionHint = SQLConf.get.rangeExchangeSampleSizePerPartition)
case SinglePartition => new ConstantPartitioner
- case k @ KeyGroupedPartitioning(expressions, n, _, _, _) =>
+ case k @ KeyGroupedPartitioning(expressions, n, _, _, _, _) =>
val valueMap = k.uniquePartitionValues.zipWithIndex.map {
case (partition, index) => (partition.toSeq(expressions.map(_.dataType)), index)
}.toMap
@@ -397,7 +397,7 @@ object ShuffleExchangeExec {
val projection = UnsafeProjection.create(sortingExpressions.map(_.child), outputAttributes)
row => projection(row)
case SinglePartition => identity
- case KeyGroupedPartitioning(expressions, _, _, _, _) =>
+ case KeyGroupedPartitioning(expressions, _, _, _, _, _) =>
row => bindReferences(expressions, outputAttributes).map(_.eval(row))
case _ => throw SparkException.internalError(s"Exchange not implemented for $newPartitioning")
}
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 80f360ed216f..a5922e331337 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
@@ -52,7 +52,7 @@ abstract class DistributionAndOrderingSuiteBase
case HashPartitioning(exprs, numPartitions) =>
HashPartitioning(exprs.map(resolveAttrs(_, plan)), numPartitions)
case KeyGroupedPartitioning(clustering, numPartitions, partValues,
- originalPartValues, isPartiallyClustered) =>
+ originalPartValues, isPartiallyClustered, _) =>
KeyGroupedPartitioning(clustering.map(resolveAttrs(_, plan)), numPartitions, partValues,
originalPartValues, isPartiallyClustered)
case PartitioningCollection(partitionings) =>
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 9921eb683724..5e035065b2ad 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,6 +24,7 @@ import org.apache.spark.sql.{DataFrame, Row}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Literal, TransformExpression}
import org.apache.spark.sql.catalyst.plans.physical
+import org.apache.spark.sql.catalyst.plans.physical.KeyGroupedPartitioning
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
@@ -93,13 +94,15 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase {
checkQueryPlan(df, catalystDistribution,
physical.KeyGroupedPartitioning(catalystDistribution.clustering, projectedPositions,
- partitionValues, partitionValues, isPartiallyClustered = false))
+ partitionValues, partitionValues, isPartiallyClustered = false,
+ mayContainUnknownPartitionKeys = false))
// multiple group keys should work too as long as partition keys are subset of them
df = sql(s"SELECT count(*) FROM testcat.ns.$table GROUP BY id, ts")
checkQueryPlan(df, catalystDistribution,
physical.KeyGroupedPartitioning(catalystDistribution.clustering, projectedPositions,
- partitionValues, partitionValues, isPartiallyClustered = false))
+ partitionValues, partitionValues, isPartiallyClustered = false,
+ mayContainUnknownPartitionKeys = false))
}
test("non-clustered distribution: no partition") {
@@ -3249,4 +3252,549 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase {
}
}
}
+
+ 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: KeyGroupedPartitioning =>
+ assert(k.mayContainUnknownPartitionKeys === hasUnknown,
+ s"expected shuffle output mayContainUnknownPartitionKeys=$hasUnknown, got " +
+ s"${k.mayContainUnknownPartitionKeys}:\n$plan")
+ case p =>
+ fail(s"expected a KeyGroupedPartitioning 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
+ // storage-partitions.
+ assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
+ }
+ }
+ }
+
+ 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))
+ }
+
+ // 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.
+ }
+
+ // 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))
+ }
+ }
+ }
+
+ 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). u declares a's key set rather than a superset:
+ // this branch's value-merging path cannot reconcile differing partition counts over a join
+ // subtree (a pre-existing limitation independent of the marker), and the equal-count shape
+ // carries the property this test is about.
+ 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)")
+
+ 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))
+ }
+ }
+ }
+
+ 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))
+ }
+ }
+ }
+
+ 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))
+ }
+ }
+ }
+
+ 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))
+ }
+ }
+ }
+
+
+ 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 costs no extra shuffle") {
+ // The spurious marker on the inner join's collection must cost nothing: consumers read a
+ // joined collection with `exists`, so the accurate member still answers, and the second
+ // storage-partitioned join runs without a re-shuffle. u declares a's bucket values rather
+ // than a finer bucket count: this branch's value-merging path cannot reconcile differing
+ // partition counts over a join subtree (a pre-existing limitation independent of the
+ // marker), and the equal-count shape carries the property this test is about.
+ 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')")
+
+ 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"))
+
+ 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 spurious marker adds no exchange: the first join's one-side shuffle is the only
+ // shuffle and the second join storage-partitions on the accurate member.
+ assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan, Seq(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.
+ // `satisfies0` 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[KeyGroupedPartitioning]),
+ 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 `satisfies0`: 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]. `satisfies0` 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")
+ }
+ }
+ }
+
}
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 259564fcd3a8..a4ad8456c026 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
@@ -1127,7 +1127,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
EnsureRequirements.apply(smjExec) match {
case ShuffledHashJoinExec(_, _, _, _, _,
DummySparkPlan(_, _, left: KeyGroupedPartitioning, _, _),
- ShuffleExchangeExec(KeyGroupedPartitioning(attrs, 4, pv, _, _),
+ ShuffleExchangeExec(KeyGroupedPartitioning(attrs, 4, pv, _, _, _),
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