This is an automated email from the ASF dual-hosted git repository.
pan3793 pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 46c1826289e0 [SPARK-57881][SQL] UnionExec outputPartitioning supports
KeyedPartitioning
46c1826289e0 is described below
commit 46c1826289e05f87b8d935d76c8d629e8e64e0ae
Author: Cheng Pan <[email protected]>
AuthorDate: Tue Jul 7 11:47:15 2026 +0800
[SPARK-57881][SQL] UnionExec outputPartitioning supports KeyedPartitioning
### What changes were proposed in this pull request?
Extend `spark.sql.unionOutputPartitioning` to `KeyedPartitioning`, so a
`UNION ALL` of V2 storage-partitioned tables can do a shuffle-free SPJ.
Changes in `UnionExec`:
- `comparePartitioning`: add a `KeyedPartitioning` case matching on
partition *expressions* (keys not compared — children carry different key sets
that are merged below).
- `outputPartitioning`: when all children report compatible
`KeyedPartitioning`s, emit a merged one whose `partitionKeys` is the
concatenation of the children's keys (one per physical partition), recomputing
`isGrouped` and propagating `isNarrowed` (sticky: set if any child is narrowed).
- `doExecute`: a `KeyedPartitioning` union uses the plain concatenating
`UnionRDD`, not the index-based `SQLPartitioningAwareUnionRDD` (which is only
correct for `HashPartitioning`/`SinglePartition` and would mix keys).
The merged descriptor is consumed by the existing SPJ logic in
`EnsureRequirements`: duplicate keys (overlapping child sets,
`isGrouped=false`) are coalesced by `GroupPartitionsExec`; join-leg key-set
mismatches are reconciled via the push-part-values path.
### Why are the changes needed?
A `UNION ALL` of V2 bucketed tables reports `UnknownPartitioning` today, so
any join on the partition key shuffles — defeating SPJ for sharded/bucketed
sources.
```sql
SELECT /*+ MERGE(u, t3) */ u.id, u.data, t3.data
FROM (SELECT id, data FROM t1 UNION ALL SELECT id, data FROM t2) u
JOIN t3 ON u.id = t3.id;
```
Before: both sides shuffle. After: the union reports a merged
`KeyedPartitioning` over `id`; SPJ runs shuffle-free (with
`GroupPartitionsExec` coalescing overlapping keys).
### Does this PR introduce _any_ user-facing change?
Results are unchanged. The executed plan may drop a previously inserted
shuffle.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: GLM 5.2
Closes #56961 from pan3793/SPARK-57881.
Authored-by: Cheng Pan <[email protected]>
Signed-off-by: Cheng Pan <[email protected]>
(cherry picked from commit f8ee9ef678794e086970d66d0631e16e4bdd3798)
Signed-off-by: Cheng Pan <[email protected]>
---
.../sql/execution/basicPhysicalOperators.scala | 57 +++-
.../spark/sql/DataFrameSetOperationsSuite.scala | 44 +++-
.../connector/KeyGroupedPartitioningSuite.scala | 288 ++++++++++++++++++++-
3 files changed, 379 insertions(+), 10 deletions(-)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
index 40c27b9e77f5..4e25bf8bef40 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
@@ -923,6 +923,13 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
(left, right) match {
case (SinglePartition, SinglePartition) => true
case (l: HashPartitioningLike, r: HashPartitioningLike) => l == r
+ // For `KeyedPartitioning`, only the partition expressions must match
(the other child's
+ // expressions have already been remapped to the first child's
attributes by
+ // `prepareOutputPartitioning`). The partition keys are intentionally
not compared here:
+ // children typically carry different key sets, and `outputPartitioning`
merges them.
+ case (l: KeyedPartitioning, r: KeyedPartitioning) =>
+ l.expressions.length == r.expressions.length &&
+ l.expressions.zip(r.expressions).forall { case (le, re) =>
le.semanticEquals(re) }
// Note: two `RangePartitioning`s with even same ordering and number of
partitions
// are not equal, because they might have different partition bounds.
case _ => false
@@ -938,6 +945,28 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
// Take the output attributes of this union and map the partitioner to
them.
val attributeMap = children.head.output.zip(output).toMap
partitioner match {
+ case headKp: KeyedPartitioning =>
+ // A `UnionExec` concatenates its children's partitions in order
(one child's
+ // partitions after another's), so the merged `KeyedPartitioning`
carries the
+ // concatenation of the children's partition keys, one key per
physical output
+ // partition. Children usually hold different key sets, so the
merged keys often
+ // contain duplicates and `isGrouped` is false; a downstream
`GroupPartitionsExec`
+ // regroups partitions that share a key. The children's
expressions have already
+ // been remapped to the first child's attributes by
`prepareOutputPartitioning`;
+ // here they are remapped to the union's output attributes.
+ val mergedKeys = partitionings.flatMap {
+ case k: KeyedPartitioning => k.partitionKeys
+ case _ => return super.outputPartitioning
+ }
+ val mergedExpressions = headKp.expressions.map(_.transform {
+ case a: Attribute if attributeMap.contains(a) => attributeMap(a)
+ })
+ val isGrouped = mergedKeys.distinct.size == mergedKeys.size
+ val isNarrowed = partitionings.exists {
+ case k: KeyedPartitioning => k.isNarrowed
+ case _ => false
+ }
+ KeyedPartitioning(mergedExpressions, mergedKeys, isGrouped,
isNarrowed)
case e: Expression =>
e.transform {
case a: Attribute if attributeMap.contains(a) => attributeMap(a)
@@ -954,6 +983,10 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
// True when the codegen path applies: `outputPartitioning` is
`UnknownPartitioning`,
// and `unionedInputRDD` matches the semantics of `sparkContext.union(...)`
in `doExecute`.
+ // A `KeyedPartitioning` union also uses `sparkContext.union(...)` in
`doExecute`, but
+ // codegen is disabled for it (`supportCodegenFailureReason` reports
"partitioning-aware"):
+ // the per-partition key descriptor is consumed by a downstream
`GroupPartitionsExec`, and
+ // keeping these unions out of whole-stage codegen matches the
`HashPartitioning` union case.
private def isPlainUnion: Boolean =
outputPartitioning.isInstanceOf[UnknownPartitioning]
// Per-child projection from the child's output to the union's output. The
wrapped
@@ -1160,14 +1193,22 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
override def usedInputs: AttributeSet = AttributeSet.empty
protected override def doExecute(): RDD[InternalRow] = {
- if (isPlainUnion) {
- sparkContext.union(children.map(_.execute()))
- } else {
- // This union has a known partitioning, i.e., its children have the same
partitioning
- // in semantics so this union can choose not to change the partitioning
by using a
- // custom partitioning aware union RDD.
- val nonEmptyRdds =
children.map(_.execute()).filter(!_.partitions.isEmpty)
- new SQLPartitioningAwareUnionRDD(sparkContext, nonEmptyRdds,
outputPartitioning.numPartitions)
+ outputPartitioning match {
+ case _: UnknownPartitioning | _: KeyedPartitioning =>
+ // An `UnknownPartitioning` union simply concatenates its children. A
+ // `KeyedPartitioning` union does the same: its merged partition keys
describe the
+ // concatenated layout (one key per physical partition), and a
downstream
+ // `GroupPartitionsExec` regroups partitions that share a key. This
differs from an
+ // index-co-locatable partitioning (e.g. `HashPartitioning`), where a
partitioning-aware
+ // union RDD interleaves same-index partitions across children.
+ sparkContext.union(children.map(_.execute()))
+ case _ =>
+ // This union has a known, index-co-locatable partitioning, i.e., its
children have the
+ // same partitioning in semantics so this union can choose not to
change the partitioning
+ // by using a custom partitioning aware union RDD.
+ val nonEmptyRdds =
children.map(_.execute()).filter(!_.partitions.isEmpty)
+ new SQLPartitioningAwareUnionRDD(
+ sparkContext, nonEmptyRdds, outputPartitioning.numPartitions)
}
}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala
index d838ba4c234f..7e9a7ad74579 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala
@@ -22,7 +22,8 @@ import java.util.Locale
import org.apache.spark.sql.catalyst.optimizer.RemoveNoopUnion
import org.apache.spark.sql.catalyst.plans.logical.Union
-import org.apache.spark.sql.catalyst.plans.physical.UnknownPartitioning
+import org.apache.spark.sql.catalyst.plans.physical.{KeyedPartitioning,
UnknownPartitioning}
+import org.apache.spark.sql.connector.catalog.InMemoryCatalog
import org.apache.spark.sql.execution.{SparkPlan, UnionExec}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
@@ -1617,6 +1618,47 @@ class DataFrameSetOperationsSuite extends
SharedSparkSession with AdaptiveSparkP
}
}
+ test("SPARK-57881: union partitioning - keyed partitioning") {
+ withSQLConf("spark.sql.catalog.testcat" ->
classOf[InMemoryCatalog].getName) {
+ sql("CREATE TABLE testcat.ns.t1 (id bigint, data string) PARTITIONED BY
(id)")
+ sql("INSERT INTO testcat.ns.t1 VALUES (1, 'a1'), (2, 'a2')")
+ sql("CREATE TABLE testcat.ns.t2 (id bigint, data string) PARTITIONED BY
(id)")
+ sql("INSERT INTO testcat.ns.t2 VALUES (2, 'b2'), (3, 'b3')")
+
+ def unionDF: DataFrame = sql(
+ """SELECT id, data FROM testcat.ns.t1
+ |UNION ALL
+ |SELECT id, data FROM testcat.ns.t2
+ |""".stripMargin)
+
+ val correctResult = withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key ->
"false") {
+ unionDF.collect()
+ }
+
+ Seq(true, false).foreach { enabled =>
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.UNION_OUTPUT_PARTITIONING.key -> enabled.toString) {
+ val union = unionDF
+ val unionExec = union.queryExecution.executedPlan.collect { case u:
UnionExec => u }
+ assert(unionExec.size == 1)
+
+ val partitioning = unionExec.head.outputPartitioning
+ if (enabled) {
+ // The two children report compatible KeyedPartitionings (both
identity(id)); the
+ // union merges their partition keys into a single
KeyedPartitioning.
+ assert(partitioning.isInstanceOf[KeyedPartitioning],
+ "union of compatible KeyedPartitionings should report a merged
KeyedPartitioning")
+ } else {
+ assert(partitioning.isInstanceOf[UnknownPartitioning])
+ }
+
+ checkAnswer(union, correctResult)
+ }
+ }
+ }
+ }
+
test("SPARK-53550: union partitioning should compare canonicalized
attributes") {
withSQLConf(
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
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 711f6dbdcdb1..e765b8630189 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
@@ -30,7 +30,14 @@ import org.apache.spark.sql.connector.catalog.functions._
import org.apache.spark.sql.connector.distributions.Distributions
import org.apache.spark.sql.connector.expressions._
import org.apache.spark.sql.connector.expressions.Expressions._
-import org.apache.spark.sql.execution.{ExtendedMode, FormattedMode,
RDDScanExec, SimpleMode, SortExec, SparkPlan}
+import org.apache.spark.sql.execution.{
+ ExtendedMode,
+ FormattedMode,
+ RDDScanExec,
+ SimpleMode,
+ SortExec,
+ SparkPlan,
+ UnionExec}
import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec,
DataSourceV2ScanRelation, GroupPartitionsExec}
import org.apache.spark.sql.execution.exchange.{ShuffleExchangeExec,
ShuffleExchangeLike}
import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec,
SortMergeJoinExec}
@@ -4183,4 +4190,283 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase with
}
}
}
+
+ test("SPARK-57881: storage-partitioned join leverages union output
KeyedPartitioning to " +
+ "avoid shuffle") {
+ val cols = Array(
+ Column.create("id", LongType),
+ Column.create("data", StringType))
+ val partitions = Array(identity("id"))
+ withTable("t1", "t2", "t3") {
+ createTable("t1", cols, partitions)
+ sql("INSERT INTO testcat.ns.t1 VALUES (1, 'a1'), (2, 'a2')")
+ createTable("t2", cols, partitions)
+ sql("INSERT INTO testcat.ns.t2 VALUES (2, 'b2'), (3, 'b3')")
+ createTable("t3", cols, partitions)
+ sql("INSERT INTO testcat.ns.t3 VALUES (1, 'c1'), (2, 'c2'), (3, 'c3')")
+
+ // Disable AQE for a deterministic, fully-planned tree to inspect.
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") {
+ val df = sql(
+ """SELECT /*+ MERGE(u, t3) */ u.id, u.data, t3.data AS t3data
+ |FROM (
+ | SELECT id, data FROM testcat.ns.t1
+ | UNION ALL
+ | SELECT id, data FROM testcat.ns.t2
+ |) u
+ |JOIN testcat.ns.t3 ON u.id = t3.id
+ |""".stripMargin)
+ val plan = df.queryExecution.executedPlan
+ // The union reports a KeyedPartitioning over `id`, which the SMJ
leverages for a
+ // storage-partitioned join, so no shuffle is needed.
+ assert(collectShuffles(plan).isEmpty)
+ checkAnswer(df,
+ Seq(Row(1, "a1", "c1"), Row(2, "a2", "c2"), Row(2, "b2", "c2"),
Row(3, "b3", "c3")))
+ }
+ }
+ }
+
+ test("SPARK-57881: storage-partitioned join over union: compatible
expressions, " +
+ "disjoint child partition keys") {
+ // Both children are partitioned by identity(id) (compatible expressions)
but hold
+ // disjoint key sets: t1=[1,2], t2=[3,4]. The union merges the keys into
[1,2,3,4];
+ // because no key repeats across children, the merged KeyedPartitioning is
already
+ // grouped, so no GroupPartitionsExec is needed on the union side. t3
carries the same
+ // keys in the same order, so SPJ matches the two legs directly without a
shuffle.
+ val cols = Array(Column.create("id", LongType), Column.create("data",
StringType))
+ val partitions = Array(identity("id"))
+ withTable("t1", "t2", "t3") {
+ createTable("t1", cols, partitions)
+ sql("INSERT INTO testcat.ns.t1 VALUES (1, 'a1'), (2, 'a2')")
+ createTable("t2", cols, partitions)
+ sql("INSERT INTO testcat.ns.t2 VALUES (3, 'b3'), (4, 'b4')")
+ createTable("t3", cols, partitions)
+ sql("INSERT INTO testcat.ns.t3 VALUES (1, 'c1'), (2, 'c2'), (3, 'c3'),
(4, 'c4')")
+
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") {
+ val df = sql(
+ """SELECT /*+ MERGE(u, t3) */ u.id, u.data, t3.data AS t3data
+ |FROM (
+ | SELECT id, data FROM testcat.ns.t1
+ | UNION ALL
+ | SELECT id, data FROM testcat.ns.t2
+ |) u
+ |JOIN testcat.ns.t3 ON u.id = t3.id
+ |""".stripMargin)
+ val plan = df.queryExecution.executedPlan
+
+ // The merged descriptor carries the concatenation of the children's
keys; disjoint
+ // keys leave it grouped, which the SMJ consumes directly.
+ val union = collect(plan) { case u: UnionExec => u }.head
+ val kp =
union.outputPartitioning.asInstanceOf[physical.KeyedPartitioning]
+ assert(kp.numPartitions == 4, "one key per physical partition of the
concatenation")
+ assert(kp.isGrouped, "disjoint child keys merge without duplicates")
+
+ assert(collectShuffles(plan).isEmpty, "no shuffle: merged grouped keys
match t3")
+ assert(collectGroupPartitions(plan).isEmpty,
+ "no GroupPartitionsExec: merged keys are already grouped and
aligned")
+ checkAnswer(df, Seq(
+ Row(1, "a1", "c1"), Row(2, "a2", "c2"), Row(3, "b3", "c3"), Row(4,
"b4", "c4")))
+ }
+ }
+ }
+
+ test("SPARK-57881: storage-partitioned join over union: compatible
expressions, " +
+ "union keys are a strict subset of the other leg") {
+ // Expressions are compatible (both identity(id)), but the join legs carry
different
+ // partition key sets: the union (t1=[1,2] UNION t2=[2,3]) groups to
[1,2,3] while t3
+ // holds [1,2,3,4,5]. The merged KeyedPartitioning has a duplicate key (2
from both
+ // children), so isGrouped=false and EnsureRequirements inserts a
GroupPartitionsExec.
+ // With pushPartValues enabled, SPJ computes the superset [1,2,3,4,5] and
pads the union
+ // side with empty partitions for the missing keys 4 and 5, avoiding a
shuffle. With
+ // pushPartValues disabled the key mismatch cannot be reconciled, so both
legs shuffle.
+ val cols = Array(Column.create("id", LongType), Column.create("data",
StringType))
+ val partitions = Array(identity("id"))
+ withTable("t1", "t2", "t3") {
+ createTable("t1", cols, partitions)
+ sql("INSERT INTO testcat.ns.t1 VALUES (1, 'a1'), (2, 'a2')")
+ createTable("t2", cols, partitions)
+ sql("INSERT INTO testcat.ns.t2 VALUES (2, 'b2'), (3, 'b3')")
+ createTable("t3", cols, partitions)
+ sql("INSERT INTO testcat.ns.t3 VALUES (1, 'c1'), (2, 'c2'), (3, 'c3'),
(4, 'c4'), (5, 'c5')")
+
+ Seq(true, false).foreach { pushPartValues =>
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key ->
pushPartValues.toString) {
+ val df = sql(
+ """SELECT /*+ MERGE(u, t3) */ u.id, u.data, t3.data AS t3data
+ |FROM (
+ | SELECT id, data FROM testcat.ns.t1
+ | UNION ALL
+ | SELECT id, data FROM testcat.ns.t2
+ |) u
+ |JOIN testcat.ns.t3 ON u.id = t3.id
+ |""".stripMargin)
+ val plan = df.queryExecution.executedPlan
+
+ // The merged descriptor is ungrouped regardless of the
pushPartValues flag, since
+ // key 2 is duplicated across the two children.
+ val union = collect(plan) { case u: UnionExec => u }.head
+ val kp =
union.outputPartitioning.asInstanceOf[physical.KeyedPartitioning]
+ assert(!kp.isGrouped, "overlapping child keys merge with duplicates")
+
+ val shuffles = collectShuffles(plan)
+ val groupPartitions = collectGroupPartitions(plan)
+ if (pushPartValues) {
+ assert(shuffles.isEmpty, "no shuffle: superset of keys pushed to
both legs")
+ assert(groupPartitions.nonEmpty &&
+ groupPartitions.forall(_.outputPartitioning.numPartitions === 5),
+ "both legs aligned to the 5-key superset")
+ } else {
+ assert(shuffles.length == 2,
+ "both legs shuffled when keys mismatch and pushPartValues is
off")
+ assert(groupPartitions.isEmpty,
+ "GroupPartitionsExec is dropped once a shuffle is inserted")
+ }
+ // Inner join: keys 4 and 5 have no match on the union side.
+ checkAnswer(df, Seq(
+ Row(1, "a1", "c1"), Row(2, "a2", "c2"), Row(2, "b2", "c2"), Row(3,
"b3", "c3")))
+ }
+ }
+ }
+ }
+
+ test("SPARK-57881: storage-partitioned join over union: compatible
expressions, " +
+ "the other leg is a strict subset of the union keys") {
+ // Expressions compatible (identity(id)); partition keys mismatch in the
other direction:
+ // the union (t1=[1,2] UNION t2=[2,3,4]) groups to [1,2,3,4] while t3 only
holds [2,3].
+ // SPJ pushes the superset [1,2,3,4] to t3, padding keys 1 and 4 with
empty partitions.
+ // No shuffle.
+ val cols = Array(Column.create("id", LongType), Column.create("data",
StringType))
+ val partitions = Array(identity("id"))
+ withTable("t1", "t2", "t3") {
+ createTable("t1", cols, partitions)
+ sql("INSERT INTO testcat.ns.t1 VALUES (1, 'a1'), (2, 'a2')")
+ createTable("t2", cols, partitions)
+ sql("INSERT INTO testcat.ns.t2 VALUES (2, 'b2'), (3, 'b3'), (4, 'b4')")
+ createTable("t3", cols, partitions)
+ sql("INSERT INTO testcat.ns.t3 VALUES (2, 'c2'), (3, 'c3')")
+
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") {
+ val df = sql(
+ """SELECT /*+ MERGE(u, t3) */ u.id, u.data, t3.data AS t3data
+ |FROM (
+ | SELECT id, data FROM testcat.ns.t1
+ | UNION ALL
+ | SELECT id, data FROM testcat.ns.t2
+ |) u
+ |JOIN testcat.ns.t3 ON u.id = t3.id
+ |""".stripMargin)
+ val plan = df.queryExecution.executedPlan
+
+ assert(collectShuffles(plan).isEmpty, "no shuffle: superset pushed to
the t3 leg")
+ assert(collectGroupPartitions(plan).nonEmpty &&
+
collectGroupPartitions(plan).forall(_.outputPartitioning.numPartitions === 4),
+ "both legs aligned to the 4-key superset")
+ // Inner join: only ids 2 and 3 match.
+ checkAnswer(df, Seq(
+ Row(2, "a2", "c2"), Row(2, "b2", "c2"), Row(3, "b3", "c3")))
+ }
+ }
+ }
+
+ test("SPARK-57881: storage-partitioned join over union: bucket transform
partitioning") {
+ val cols = Array(Column.create("id", LongType), Column.create("data",
StringType))
+ val partitions = Array(bucket(4, "id"))
+ withTable("t1", "t2", "t3") {
+ createTable("t1", cols, partitions)
+ sql("INSERT INTO testcat.ns.t1 VALUES (1, 'a1'), (2, 'a2')")
+ createTable("t2", cols, partitions)
+ sql("INSERT INTO testcat.ns.t2 VALUES (2, 'b2'), (3, 'b3')")
+ createTable("t3", cols, partitions)
+ sql("INSERT INTO testcat.ns.t3 VALUES (1, 'c1'), (2, 'c2'), (3, 'c3')")
+
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true") {
+ val df = sql(
+ """SELECT /*+ MERGE(u, t3) */ u.id, u.data, t3.data AS t3data
+ |FROM (
+ | SELECT id, data FROM testcat.ns.t1
+ | UNION ALL
+ | SELECT id, data FROM testcat.ns.t2
+ |) u
+ |JOIN testcat.ns.t3 ON u.id = t3.id
+ |""".stripMargin)
+ val plan = df.queryExecution.executedPlan
+
+ // The union reports a KeyedPartitioning whose expression is the
`bucket(4, id)` transform.
+ val union = collect(plan) { case u: UnionExec => u }.head
+ val kp =
union.outputPartitioning.asInstanceOf[physical.KeyedPartitioning]
+ assert(kp.expressions.length == 1 &&
kp.expressions.head.isInstanceOf[TransformExpression],
+ "merged KeyedPartitioning carries the bucket transform expression")
+
+ assert(collectShuffles(plan).isEmpty, "no shuffle: SPJ over the bucket
transform")
+ checkAnswer(df,
+ Seq(Row(1, "a1", "c1"), Row(2, "a2", "c2"), Row(2, "b2", "c2"),
Row(3, "b3", "c3")))
+ }
+ }
+ }
+
+ test("SPARK-57881: storage-partitioned join over union: a union leg is
entirely " +
+ "runtime-pruned") {
+ // The merged descriptor is built from each leg's unfiltered
`inputPartitions`, while the union
+ // RDD concatenates each leg's `filteredPartitions` (pruned splits kept as
`None`). Here dynamic
+ // partition filtering prunes the entire t1 leg (only t3 ids [3, 4]
survive), so this guards
+ // that the per-leg partition-count == partitionKeys.length alignment
holds under pruning.
+ val cols = Array(Column.create("id", LongType), Column.create("data",
StringType))
+ val partitions = Array(identity("id"))
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true",
+ SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true",
+ SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "false",
+ SQLConf.DYNAMIC_PARTITION_PRUNING_FALLBACK_FILTER_RATIO.key -> "10") {
+ withTable("t1", "t2", "t3") {
+ createTable("t1", cols, partitions)
+ sql("INSERT INTO testcat.ns.t1 VALUES (1, 'a1'), (2, 'a2')")
+ createTable("t2", cols, partitions)
+ sql("INSERT INTO testcat.ns.t2 VALUES (3, 'b3'), (4, 'b4')")
+ createTable("t3", cols, partitions)
+ sql("INSERT INTO testcat.ns.t3 VALUES (1, 'c1'), (2, 'c2'), (3, 'c3'),
(4, 'c4')")
+
+ val df = sql(
+ """SELECT /*+ MERGE(u, t3) */ u.id, u.data, t3.data AS t3data
+ |FROM (
+ | SELECT id, data FROM testcat.ns.t1
+ | UNION ALL
+ | SELECT id, data FROM testcat.ns.t2
+ |) u
+ |JOIN testcat.ns.t3 ON u.id = t3.id
+ |WHERE t3.data IN ('c3', 'c4')
+ |""".stripMargin)
+ val plan = df.queryExecution.executedPlan
+
+ // The merged descriptor carries the concatenation of both legs'
unfiltered keys
+ // ([1, 2] ++ [3, 4]), independent of runtime pruning.
+ val union = collect(plan) { case u: UnionExec => u }.head
+ val kp =
union.outputPartitioning.asInstanceOf[physical.KeyedPartitioning]
+ assert(kp.numPartitions == 4,
+ "merged descriptor keeps one key per unfiltered physical partition
of both legs")
+
+ assert(collectShuffles(plan).isEmpty, "no shuffle: merged grouped keys
match t3")
+
+ // Force execution, then verify the t1 leg is entirely runtime-pruned
(all `None`) while
+ // its keys still live in the merged descriptor above.
+ checkAnswer(df, Seq(Row(3, "b3", "c3"), Row(4, "b4", "c4")))
+ val unionScans = collectScans(union)
+ assert(unionScans.exists(_.filteredPartitions.forall(_.isEmpty)),
+ "one union leg must be entirely pruned to None while its keys remain
in the descriptor")
+ }
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
