This is an automated email from the ASF dual-hosted git repository.
peter-toth 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 c0633e31d625 [SPARK-53322][4.0][SQL] Select a KeyGroupedShuffleSpec
only when join key positions can be fully pushed down
c0633e31d625 is described below
commit c0633e31d62599efc70732355f40777cbf96bc3e
Author: Peter Toth <[email protected]>
AuthorDate: Mon Oct 5 17:12:44 2026 +0200
[SPARK-53322][4.0][SQL] Select a KeyGroupedShuffleSpec only when join key
positions can be fully pushed down
### Differences from #53098
This is the branch-4.0 backport of #53098. It reached master and branch-4.1
(#53132), but not branch-4.0. The description below is the original one. On
this branch:
- `KeyGroupedPartitionedScan` does not exist here.
`canPushDownSPJParamsToScan` asks for a `BatchScanExec` instead, the node
`populateCommonPartitionInfo` pushes the SPJ parameters into. On branch-4.1,
`BatchScanExec` is the only implementer of that trait.
- The tests are the branch-4.1 ones. They conflict only at the end of
`KeyGroupedPartitioningSuite` and `EnsureRequirementsSuite`.
`DummySparkPlanWithBatchScanChild` is already on this branch, copied in by the
SPARK-59080 backport, so its comment that pointed at this ticket is dropped.
- Three SPARK-59050 tests now match branch-4.1:
- Without this change, this branch storage-partitions a keyed side above
a one-side-shuffled join of a parquet table. When the partition counts differ,
the query fails with `Can't zip RDDs with unequal numbers of partitions` or
with the `PartitioningCollection` error below. So the SPARK-59050 backport
(#58572) gave two of these tests equal key sets, and kept the third at 1
shuffle.
- With this change, the first join's output reads a parquet table, so the
second join shuffles one side instead, as on branch-4.1. The rows are the same.
- Measured without this change: two of them fail with those two errors,
and the third gets 1 shuffle instead of 2.
- As on branch-4.1, every leaf under a join side must now be a
`BatchScanExec`. So some joins that storage-partitioned here before now shuffle:
- a keyed side that also reads another kind of source, for example
through a broadcast join with a parquet build side. Measured with every conf at
its default, such a three-table join goes from 0 to 2 shuffles;
- a cached or checkpointed keyed side, even when its partition values
match the other side's.
- Master and branch-4.2 dropped this check later, with the SPARK-55535
refactor.
- The bug reproduces on branch-4.0 without this change, with
`spark.sql.sources.v2.bucketing.enabled` on:
- both sides checkpointed: the join returns 1 of 2 rows;
- one side checkpointed: planning fails with `PartitioningCollection
requires all of its partitionings have the same numPartitions`.
- New tests: three of the four fail on branch-4.0 without this change.
"checkpointed scans avoid shuffles for aggregates" passes, since it needs no
storage-partitioned join. "checkpointed scans can't shuffle other children on
SPJ" fails on the shuffle count only, 2 instead of 3.
- Each half was removed on its own. Without the `candidateSpecs` filter,
"checkpointed scans can't shuffle other children on SPJ" fails. Without the
check in `checkKeyGroupCompatible`, "checkpointed scans aren't used for SPJ"
and "checkpointed scans can be shuffled by children on SPJ" fail.
- Ran the `KeyGroupedPartitioning*` suites, `EnsureRequirementsSuite`,
`PlannerSuite`, `WriteDistributionAndOrderingSuite`,
`BucketedReadWithoutHiveSupportSuite`, the plan stability suites, the
`*JoinSuite` suites, `AdaptiveQueryExecSuite`, `DataSourceV2Suite` and
`CheckpointSuite` on branch-4.0, 1313 tests in all, plus `dev/lint-scala`.
- The backport was prepared with Claude Code. Generated-by: Claude Code
(Claude Opus 5.5)
### What changes were proposed in this pull request?
When a KeyGroupedShuffleSpec is used to shuffle another child of a JOIN, we
must be able to push down JOIN keys or partition values to be able to ensure
that both children have matching partitioning. If one child reports a
KeyGroupedPartitioning but we can't push down these values (for example, if the
child was a key-grouped scan that was checkpointed), then this information
cannot be pushed down to the child scan and we should avoid using this shuffle
spec to shuffle other children.
### Why are the changes needed?
Prevents potential correctness issue when key-grouped partitioning is used
on a checkpointed RDD.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
See test changes.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #59234 from peter-toth/SPARK-53322-checkpoint-pushdown-4.0.
Lead-authored-by: Peter Toth <[email protected]>
Co-authored-by: Chirag Singh <[email protected]>
Signed-off-by: Peter Toth <[email protected]>
---
.../execution/exchange/EnsureRequirements.scala | 31 ++++
.../connector/KeyGroupedPartitioningSuite.scala | 193 ++++++++++++++++++---
.../exchange/EnsureRequirementsSuite.scala | 68 ++++----
3 files changed, 232 insertions(+), 60 deletions(-)
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 e581b78c066d..289f6438d374 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
@@ -140,6 +140,13 @@ case class EnsureRequirements(
// Choose all the specs that can be used to shuffle other children
val candidateSpecs = specs
.filter(_._2.canCreatePartitioning)
+ .filter {
+ // To choose a KeyGroupedShuffleSpec, we must be able to push down
SPJ parameters into
+ // the scan (for join key positions). If these parameters can't be
pushed down, this
+ // spec can't be used to shuffle other children.
+ case (idx, _: KeyGroupedShuffleSpec) =>
canPushDownSPJParamsToScan(children(idx))
+ case _ => true
+ }
.filter(p => !shouldConsiderMinParallelism ||
children(p._1).outputPartitioning.numPartitions >=
conf.defaultNumShufflePartitions)
val bestSpecOpt = if (candidateSpecs.isEmpty) {
@@ -425,6 +432,24 @@ case class EnsureRequirements(
}
}
+ /**
+ * Whether SPJ params can be pushed down to the leaf nodes of a physical
plan. For a plan to be
+ * eligible for SPJ parameter pushdown, all leaf nodes must be a
KeyGroupedPartitioning-aware
+ * scan, i.e. a [[BatchScanExec]], the node `populateCommonPartitionInfo`
pushes them into.
+ *
+ * Notably, if the leaf of `plan` is an [[RDDScanExec]] created by
checkpointing a DSv2 scan, the
+ * reported partitioning will be a [[KeyGroupedPartitioning]], but this plan
will _not_ be
+ * eligible for SPJ parameter pushdown (as the partitioning is static and
can't be easily
+ * re-grouped or padded with empty partitions according to the partition
values on the other side
+ * of the join).
+ */
+ private def canPushDownSPJParamsToScan(plan: SparkPlan): Boolean = {
+ plan.collectLeaves().forall {
+ case _: BatchScanExec => true
+ case _ => false
+ }
+ }
+
/**
* Checks whether two children, `left` and `right`, of a join operator have
compatible
* `KeyGroupedPartitioning`, and can benefit from storage-partitioned join.
@@ -436,6 +461,12 @@ case class EnsureRequirements(
left: SparkPlan,
right: SparkPlan,
requiredChildDistribution: Seq[Distribution]): Option[Seq[SparkPlan]] = {
+ // If SPJ params can't be pushed down to either the left or right side,
it's unsafe to do an
+ // SPJ.
+ if (!canPushDownSPJParamsToScan(left) ||
!canPushDownSPJParamsToScan(right)) {
+ return None
+ }
+
parent match {
case smj: SortMergeJoinExec =>
checkKeyGroupCompatible(left, right, smj.joinType,
requiredChildDistribution)
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 dafb04923058..c76056fb8cc9 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
@@ -35,6 +35,7 @@ import
org.apache.spark.sql.execution.datasources.v2.BatchScanExec
import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation
import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
import org.apache.spark.sql.execution.joins.SortMergeJoinExec
+import org.apache.spark.sql.functions.{col, max}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.internal.SQLConf._
import org.apache.spark.sql.types._
@@ -3625,16 +3626,15 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase {
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.
+ // 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. The join
+ // output also reads the parquet table t, so SPJ parameters cannot be
pushed into it
+ // (SPARK-53322), and the second join shuffles one side instead. Both
exchanges carry the
+ // marker and the answer is complete.
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)")
+ 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")
@@ -3652,7 +3652,7 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase {
val df = sql(query)
checkAnswer(df, Seq(Row(1, "u1"), Row(2, "u2")))
assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
- Seq(true))
+ Seq(true, true))
}
}
}
@@ -3698,9 +3698,10 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase {
}
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.
+ // r (from a RIGHT OUTER JOIN t) has unknown partition keys {1, 2}, and
the downstream u is
+ // keyed on a subset {1}, so the marker's subset rule would accept the
pair. But r also reads
+ // the parquet table t, so SPARK-53322 refuses to push SPJ parameters into
it. The second join
+ // re-shuffles r onto u's layout instead, which is sound for any keys r
holds.
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)")
@@ -3721,9 +3722,8 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase {
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))
+ Seq(true, true))
}
}
}
@@ -3828,16 +3828,15 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase {
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.
+ // joined collection with `exists`, so the accurate member still answers.
The join output
+ // also reads the parquet table t, so SPJ parameters cannot be pushed into
it (SPARK-53322),
+ // and the second join shuffles one side instead, with and without the
marker: two exchanges
+ // either way.
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')")
+ 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")
@@ -3849,7 +3848,7 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase {
|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"))
+ val expected = Seq(Row(0L, "u0"), Row(1L, "u1"), Row(2L, "u2"), Row(3L,
"u3"))
withSQLConf(
SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
@@ -3857,9 +3856,10 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase {
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))
+ // The first join's one-side shuffle plus the second join's
re-shuffle; the spurious
+ // marker adds no third exchange.
+
assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+ Seq(true, true))
}
}
}
@@ -4036,4 +4036,147 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase {
}
}
+ test("SPARK-53322: checkpointed scans avoid shuffles for aggregates") {
+ withTempDir { dir =>
+ spark.sparkContext.setCheckpointDir(dir.getPath)
+ val itemsPartitions = Array(identity("id"))
+ createTable(items, itemsColumns, itemsPartitions)
+ sql(s"INSERT INTO testcat.ns.$items VALUES " +
+ s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " +
+ s"(1, 'aa', 41.0, cast('2020-01-02' as timestamp)), " +
+ s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " +
+ s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp))")
+
+ val scanDF = spark.read.table(s"testcat.ns.$items").checkpoint()
+ val df = scanDF.groupBy("id").agg(max("price").as("res")).select("res")
+ checkAnswer(df.sort("res"), Seq(Row(10.0), Row(15.5), Row(41.0)))
+
+ val shuffles = collectAllShuffles(df.queryExecution.executedPlan)
+ assert(shuffles.isEmpty,
+ "should not contain shuffle when not grouping by partition values")
+ }
+ }
+
+ test("SPARK-53322: checkpointed scans aren't used for SPJ") {
+ withTempDir { dir =>
+ spark.sparkContext.setCheckpointDir(dir.getPath)
+ val itemsPartitions = Array(identity("id"))
+ createTable(items, itemsColumns, itemsPartitions)
+ sql(s"INSERT INTO testcat.ns.$items VALUES " +
+ s"(1, 'aa', 41.0, cast('2020-01-01' as timestamp)), " +
+ s"(2, 'bb', 10.0, cast('2020-01-02' as timestamp)), " +
+ s"(3, 'cc', 15.5, cast('2020-01-03' as timestamp))")
+
+ val purchase_partitions = Array(identity("item_id"))
+ createTable(purchases, purchasesColumns, purchase_partitions)
+ sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
+ s"(1, 40.0, cast('2020-01-01' as timestamp)), " +
+ s"(3, 25.5, cast('2020-01-03' as timestamp)), " +
+ s"(4, 20.0, cast('2020-01-04' as timestamp))")
+
+ for {
+ pushdownValues <- Seq(true, false)
+ checkpointBothScans <- Seq(true, false)
+ } {
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key ->
pushdownValues.toString) {
+ val scanDF1 =
spark.read.table(s"testcat.ns.$items").checkpoint().as("i")
+ val scanDF2 = if (checkpointBothScans) {
+ spark.read.table(s"testcat.ns.$purchases").checkpoint().as("p")
+ } else {
+ spark.read.table(s"testcat.ns.$purchases").as("p")
+ }
+
+ val df = scanDF1
+ .join(scanDF2, col("id") === col("item_id"))
+ .selectExpr("id", "name", "i.price AS purchase_price", "p.price AS
sale_price")
+ .orderBy("id", "purchase_price", "sale_price")
+ checkAnswer(
+ df,
+ Seq(Row(1, "aa", 41.0, 40.0), Row(3, "cc", 15.5, 25.5))
+ )
+ // 1 shuffle for SORT and 2 shuffles for JOIN are expected.
+ assert(collectAllShuffles(df.queryExecution.executedPlan).length ===
3)
+ }
+ }
+ }
+ }
+
+ test("SPARK-53322: checkpointed scans can't shuffle other children on SPJ") {
+ withTempDir { dir =>
+ spark.sparkContext.setCheckpointDir(dir.getPath)
+ val itemsPartitions = Array(identity("id"))
+ createTable(items, itemsColumns, itemsPartitions)
+ sql(s"INSERT INTO testcat.ns.$items VALUES " +
+ s"(1, 'aa', 41.0, cast('2020-01-01' as timestamp)), " +
+ s"(2, 'bb', 10.0, cast('2020-01-02' as timestamp)), " +
+ s"(3, 'cc', 15.5, cast('2020-01-03' as timestamp))")
+
+ createTable(purchases, purchasesColumns, Array.empty)
+ sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
+ s"(1, 40.0, cast('2020-01-01' as timestamp)), " +
+ s"(3, 25.5, cast('2020-01-03' as timestamp)), " +
+ s"(4, 20.0, cast('2020-01-04' as timestamp))")
+
+ Seq(true, false).foreach { pushdownValues =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key ->
pushdownValues.toString) {
+ val scanDF1 =
spark.read.table(s"testcat.ns.$items").checkpoint().as("i")
+ val scanDF2 = spark.read.table(s"testcat.ns.$purchases").as("p")
+
+ val df = scanDF1
+ .join(scanDF2, col("id") === col("item_id"))
+ .selectExpr("id", "name", "i.price AS purchase_price", "p.price AS
sale_price")
+ .orderBy("id", "purchase_price", "sale_price")
+ checkAnswer(
+ df,
+ Seq(Row(1, "aa", 41.0, 40.0), Row(3, "cc", 15.5, 25.5))
+ )
+ // 1 shuffle for SORT and 2 shuffles for JOIN are expected.
+ assert(collectAllShuffles(df.queryExecution.executedPlan).length ===
3)
+ }
+ }
+ }
+ }
+
+ test("SPARK-53322: checkpointed scans can be shuffled by children on SPJ") {
+ withTempDir { dir =>
+ spark.sparkContext.setCheckpointDir(dir.getPath)
+ val itemsPartitions = Array(identity("id"))
+ createTable(items, itemsColumns, itemsPartitions)
+ sql(s"INSERT INTO testcat.ns.$items VALUES " +
+ s"(1, 'aa', 41.0, cast('2020-01-01' as timestamp)), " +
+ s"(2, 'bb', 10.0, cast('2020-01-02' as timestamp)), " +
+ s"(3, 'cc', 15.5, cast('2020-01-03' as timestamp))")
+
+ createTable(purchases, purchasesColumns, Array(identity("item_id")))
+ sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
+ s"(1, 40.0, cast('2020-01-01' as timestamp)), " +
+ s"(3, 25.5, cast('2020-01-03' as timestamp)), " +
+ s"(4, 20.0, cast('2020-01-04' as timestamp))")
+
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true") {
+ val scanDF1 =
spark.read.table(s"testcat.ns.$items").checkpoint().as("i")
+ val scanDF2 = spark.read.table(s"testcat.ns.$purchases").as("p")
+
+ val df = scanDF1
+ .join(scanDF2, col("id") === col("item_id"))
+ .selectExpr("id", "name", "i.price AS purchase_price", "p.price AS
sale_price")
+ .orderBy("id", "purchase_price", "sale_price")
+ checkAnswer(
+ df,
+ Seq(Row(1, "aa", 41.0, 40.0), Row(3, "cc", 15.5, 25.5))
+ )
+
+ // One shuffle for the sort and one shuffle for one side of the JOIN
are expected.
+ assert(collectAllShuffles(df.queryExecution.executedPlan).length === 2)
+ }
+ }
+ }
}
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 4bd5c08b5dee..9ecc94368b9e 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
@@ -92,11 +92,11 @@ class EnsureRequirementsSuite extends SharedSparkSession {
test("reorder should handle KeyGroupedPartitioning") {
// partitioning on the left
- val plan1 = DummySparkPlan(
+ val plan1 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(Seq(
years(exprA), bucket(4, exprB), days(exprC)), 4)
)
- val plan2 = DummySparkPlan(
+ val plan2 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(Seq(
years(exprB), bucket(4, exprA), days(exprD)), 4)
)
@@ -114,7 +114,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
}
// partitioning on the right
- val plan3 = DummySparkPlan(
+ val plan3 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(Seq(
bucket(4, exprD), days(exprA), years(exprC)), 4)
)
@@ -778,9 +778,9 @@ class EnsureRequirementsSuite extends SharedSparkSession {
test("Check with KeyGroupedPartitioning") {
// simplest case: identity transforms
- var plan1 = DummySparkPlan(
- outputPartitioning = KeyGroupedPartitioning(exprA :: exprB :: Nil, 5))
- var plan2 = DummySparkPlan(
+ var plan1 = new DummySparkPlanWithBatchScanChild(
+ KeyGroupedPartitioning(exprA :: exprB :: Nil, 5))
+ var plan2 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(exprA :: exprC :: Nil, 5))
var smjExec = SortMergeJoinExec(
exprA :: exprB :: Nil, exprA :: exprC :: Nil, Inner, None, plan1, plan2)
@@ -794,11 +794,11 @@ class EnsureRequirementsSuite extends SharedSparkSession {
}
// matching bucket transforms from both sides
- plan1 = DummySparkPlan(
+ plan1 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(
bucket(4, exprA) :: bucket(16, exprB) :: Nil, 4)
)
- plan2 = DummySparkPlan(
+ plan2 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(
bucket(4, exprA) :: bucket(16, exprC) :: Nil, 4)
)
@@ -814,11 +814,11 @@ class EnsureRequirementsSuite extends SharedSparkSession {
}
// partition collections
- plan1 = DummySparkPlan(
+ plan1 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(
bucket(4, exprA) :: bucket(16, exprB) :: Nil, 4)
)
- plan2 = DummySparkPlan(
+ plan2 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = PartitioningCollection(Seq(
KeyGroupedPartitioning(bucket(4, exprA) :: bucket(16, exprC) :: Nil,
4),
KeyGroupedPartitioning(bucket(4, exprA) :: bucket(16, exprC) :: Nil,
4))
@@ -844,10 +844,10 @@ class EnsureRequirementsSuite extends SharedSparkSession {
}
// bucket + years transforms from both sides
- plan1 = DummySparkPlan(
+ plan1 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(bucket(4, exprA) ::
years(exprB) :: Nil, 4)
)
- plan2 = DummySparkPlan(
+ plan2 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(bucket(4, exprA) ::
years(exprC) :: Nil, 4)
)
smjExec = SortMergeJoinExec(
@@ -863,11 +863,11 @@ class EnsureRequirementsSuite extends SharedSparkSession {
// by default spark.sql.requireAllClusterKeysForCoPartition is true, so
when there isn't
// exact match on all partition keys, Spark will fallback to shuffle.
- plan1 = DummySparkPlan(
+ plan1 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(
bucket(4, exprA) :: bucket(4, exprB) :: Nil, 4)
)
- plan2 = DummySparkPlan(
+ plan2 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(
bucket(4, exprA) :: bucket(4, exprC) :: Nil, 4)
)
@@ -884,11 +884,11 @@ class EnsureRequirementsSuite extends SharedSparkSession {
}
test(s"KeyGroupedPartitioning with
${REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key} = false") {
- var plan1 = DummySparkPlan(
+ var plan1 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(
bucket(4, exprB) :: years(exprC) :: Nil, 4)
)
- var plan2 = DummySparkPlan(
+ var plan2 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(
bucket(4, exprC) :: years(exprB) :: Nil, 4)
)
@@ -906,11 +906,11 @@ class EnsureRequirementsSuite extends SharedSparkSession {
}
// should also work with distributions with duplicated keys
- plan1 = DummySparkPlan(
+ plan1 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(
bucket(4, exprA) :: years(exprB) :: Nil, 4)
)
- plan2 = DummySparkPlan(
+ plan2 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(
bucket(4, exprA) :: years(exprC) :: Nil, 4)
)
@@ -926,10 +926,10 @@ class EnsureRequirementsSuite extends SharedSparkSession {
}
// both partitioning and distribution have duplicated keys
- plan1 = DummySparkPlan(
+ plan1 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(
years(exprA) :: bucket(4, exprB) :: days(exprA) :: Nil, 5))
- plan2 = DummySparkPlan(
+ plan2 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(
years(exprA) :: bucket(4, exprC) :: days(exprA) :: Nil, 5))
smjExec = SortMergeJoinExec(
@@ -944,11 +944,11 @@ class EnsureRequirementsSuite extends SharedSparkSession {
}
// invalid case: partitioning key positions don't match
- plan1 = DummySparkPlan(
+ plan1 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(
bucket(4, exprA) :: bucket(4, exprB) :: Nil, 4)
)
- plan2 = DummySparkPlan(
+ plan2 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(
bucket(4, exprB) :: bucket(4, exprC) :: Nil, 4)
)
@@ -965,11 +965,11 @@ class EnsureRequirementsSuite extends SharedSparkSession {
}
// invalid case: different number of buckets (we don't support
coalescing/repartitioning yet)
- plan1 = DummySparkPlan(
+ plan1 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(
bucket(4, exprA) :: bucket(4, exprB) :: Nil, 4)
)
- plan2 = DummySparkPlan(
+ plan2 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(
bucket(4, exprA) :: bucket(8, exprC) :: Nil, 4)
)
@@ -985,10 +985,10 @@ class EnsureRequirementsSuite extends SharedSparkSession {
}
// invalid case: partition key positions match but with different
transforms
- plan1 = DummySparkPlan(
+ plan1 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(years(exprA) :: bucket(4,
exprB) :: Nil, 4)
)
- plan2 = DummySparkPlan(
+ plan2 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(days(exprA) :: bucket(4,
exprC) :: Nil, 4)
)
smjExec = SortMergeJoinExec(
@@ -1004,11 +1004,11 @@ class EnsureRequirementsSuite extends
SharedSparkSession {
// invalid case: multiple references in transform
- plan1 = DummySparkPlan(
+ plan1 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(
years(exprA) :: buckets(4, Seq(exprB, exprC)) :: Nil, 4)
)
- plan2 = DummySparkPlan(
+ plan2 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(
years(exprA) :: buckets(4, Seq(exprB, exprC)) :: Nil, 4)
)
@@ -1030,11 +1030,11 @@ class EnsureRequirementsSuite extends
SharedSparkSession {
val rightPartValues = Seq(Array[Any](1, 1), Array[Any](2, 2),
Array[Any](3, 3))
.map(new GenericInternalRow(_))
- var plan1 = DummySparkPlan(
+ var plan1 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(bucket(4, exprB) ::
bucket(8, exprC) :: Nil,
leftPartValues.length, leftPartValues)
)
- var plan2 = DummySparkPlan(
+ var plan2 = new DummySparkPlanWithBatchScanChild(
outputPartitioning = KeyGroupedPartitioning(bucket(4, exprC) ::
bucket(8, exprB) :: Nil,
rightPartValues.length, rightPartValues)
)
@@ -1052,7 +1052,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
}
// With partition collections
- plan1 = DummySparkPlan(outputPartitioning =
+ plan1 = new DummySparkPlanWithBatchScanChild(outputPartitioning =
PartitioningCollection(
Seq(KeyGroupedPartitioning(bucket(4, exprB) :: bucket(8, exprC) ::
Nil,
leftPartValues.length, leftPartValues),
@@ -1076,7 +1076,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
}
// Nested partition collections
- plan2 = DummySparkPlan(outputPartitioning =
+ plan2 = new DummySparkPlanWithBatchScanChild(outputPartitioning =
PartitioningCollection(
Seq(
PartitioningCollection(
@@ -1119,7 +1119,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
val a1 = AttributeReference("a1", IntegerType)()
val partitionValue = Seq(50, 51, 52).map(v =>
InternalRow.fromSeq(Seq(v)))
- val plan1 = DummySparkPlan(outputPartitioning = KeyGroupedPartitioning(
+ val plan1 = new DummySparkPlanWithBatchScanChild(outputPartitioning =
KeyGroupedPartitioning(
identity(a1) :: Nil, 4, partitionValue))
val plan2 = DummySparkPlan(outputPartitioning = SinglePartition)
@@ -1205,8 +1205,6 @@ class EnsureRequirementsSuite extends SharedSparkSession {
TransformExpression(DaysFunction, Seq(expr))
}
- // Backported from SPARK-53322 together with the test below, which needs a
keyed plan whose
- // subtree holds a `BatchScanExec` for the join key positions to be pushed
into.
private class DummySparkPlanWithBatchScanChild(outputPartitioning:
Partitioning)
extends DummySparkPlan(
children = Seq(BatchScanExec(Seq.empty, null, Seq.empty, table = null)),
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
