spark-test commented on code in PR #59832:
URL: https://github.com/apache/spark/pull/59832#discussion_r4392428924
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -643,9 +643,16 @@ case class KeyedPartitioning(
// the narrowing one.
val projectedExpressions = joinKeyPositions.map(expressions)
val projectedKeys = projectKeys(joinKeyPositions)._2
- val distinctProjectedKeys = projectedKeys.distinct
+ // Sort the distinct projected keys the same way `GroupPartitionsExec`
does. When only the
Review Comment:
**Nit (P3):** This now always declares the projected keys distinct and
sorted, but the `KeyedPartitioning` class doc in this file still says the
opposite. Usage Forms item 2 (lines 362-366) says that inside
`KeyedShuffleSpec` the `partitionKeys` "may not be in sorted order" because the
spec can project them by join key positions, and that `EnsureRequirements`
matches "the unordered keys from both sides". The `@param partitionKeys` text
(lines 427-429) says they "may be unsorted after projection".
On this branch that is more than stale wording. `ShuffleExchangeExec` still
builds its `KeyGroupedPartitioner` from `toGrouped` (sorted) keys, so unsorted
declared keys are exactly what produced the wrong rows this PR fixes. Could you
update both passages to say that a projected spec's keys are distinct and
sorted like `GroupPartitionsExec` output? The part about `EnsureRequirements`
pushing a common ordered key set down to `GroupPartitionsExec` is still
accurate. Master's newer wording ("not guaranteed after a projection") does not
carry over, because master's keyed shuffle keeps the declared order instead of
re-sorting.
**Verification:**
- **Inspection:** Read the updated KeyedPartitioning scaladoc and confirm
that both the KeyedShuffleSpec usage-form description and the partitionKeys
parameter description state that projected spec keys are distinct and sorted
like GroupPartitionsExec output, and that no remaining sentence in the class
doc describes them as unsorted or unordered.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -6999,4 +6999,76 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase with
}
}
+ test("SPARK-58988: v2 bucketed table with subset join keys joining v1
table") {
+ // The v2 table is partitioned by an extra identity key `dt` plus
`bucket(16, c1)`, while the
+ // join is only on `c1`. allowJoinKeysSubsetOfPartitionKeys lets the join
key `c1` be a subset
+ // of the partition keys `[dt, bucket(16, c1)]`, so EnsureRequirements
projects the keyed side
+ // to `[bucket(16, c1)]`. v2BucketingShuffleEnabled then re-shuffles only
the v1 side using that
+ // projected KeyedPartitioning. The v2 side's keys are sorted by
GroupPartitionsExec, while the
+ // keys re-used for the v1 side kept their first-occurrence order from
createShuffleSpec.
+ val cols = Array(
+ Column.create("c1", LongType),
+ Column.create("c2", StringType),
+ Column.create("dt", StringType))
+ val partitions = Array(identity("dt"), bucket(16, "c1"))
+
+ createTable("iceberg_t2", cols, partitions)
+ sql("INSERT INTO testcat.ns.iceberg_t2 VALUES (2, 'cc', '2020'), (1, 'aa',
'2021')")
+
+ withTable("t1") {
+ sql("CREATE TABLE t1 (c1 BIGINT, c2 STRING) USING parquet")
+ sql("INSERT INTO t1 VALUES (1, 'aa'), (2, 'cc')")
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = sql("SELECT * FROM testcat.ns.iceberg_t2 t0 JOIN t1 ON t0.c1
= t1.c1")
+ val plan = df.queryExecution.executedPlan
+ // Only the v1 side is re-shuffled; the v2 side is regrouped onto the
join key instead.
+ assert(collectShuffles(plan).length == 1)
+ assert(collectGroupPartitions(plan).length == 1)
+ checkAnswer(df, Seq(
+ Row(1L, "aa", "2021", 1L, "aa"),
+ Row(2L, "cc", "2020", 2L, "cc")))
+ }
+ }
+ }
+
+
+ test("SPARK-58988: one-side shuffle onto subset join keys whose projected
keys are unsorted") {
+ // `sa` is partitioned by (store, dept) and joined on dept only. Projected
onto dept, its keys
+ // come out as 20, 10. `sb` is re-shuffled onto that projected spec, and
`ShuffleExchangeExec`
+ // routes it by the sorted keys. With AQE, the join is planned again over
the materialized
+ // shuffle, and the `GroupPartitionsExec` planned there reads its
partitions in the declared
+ // order. Both joins returned wrong rows with AQE on before the declared
keys were sorted.
+ val cols = Array(
+ Column.create("store", IntegerType),
+ Column.create("dept", IntegerType),
+ Column.create("data", StringType))
+ createTable("sa", cols, Array(identity("store"), identity("dept")))
+ sql("INSERT INTO testcat.ns.sa VALUES (1, 20, 'a1'), (2, 10, 'a2')")
+ createTable("sb", cols, Array.empty)
+ sql("INSERT INTO testcat.ns.sb VALUES (0, 10, 'b10'), (0, 20, 'b20')")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ Seq("JOIN", "LEFT JOIN").foreach { joinType =>
+ // Selecting `sa.store` keeps `sa`'s keyed partitioning from being
pruned away.
+ val df = sql(
+ s"""
+ |SELECT /*+ MERGE(sa, sb) */ sa.dept, sa.data, sb.data, sa.store
+ |FROM testcat.ns.sa $joinType testcat.ns.sb ON sa.dept = sb.dept
+ |""".stripMargin)
+ assert(collectAllShuffles(df.queryExecution.executedPlan).length ==
1,
Review Comment:
**Nit (P3):** With `aqe = true`, this assertion runs before the query
executes, so `df.queryExecution.executedPlan` is still
`AdaptiveSparkPlanExec`'s initial plan (`executedPlan` returns
`currentPhysicalPlan`, which is `initialPlan` until execution). The re-planning
over the materialized `sb` stage that the test comment describes happens later,
inside `checkAnswer`, and nothing checks that final plan.
Today `checkAnswer` still catches the regression. But suppose a later change
made AQE re-plan this join so that it no longer reads the `sb` stage by its
declared keys, for example by also shuffling `sa`. The AQE iteration would then
pass with or without the sort, and this assertion would not notice. Moving it
after `checkAnswer`, as the AQE on/off loop around line 5011 already does,
makes it check the plan that actually ran. `collectAllShuffles` still counts
the `sb` exchange there through its `ShuffleQueryStageExec`.
**Verification:**
- **Regression:** Run the KeyGroupedPartitioningSuite test 'SPARK-58988:
one-side shuffle onto subset join keys whose projected keys are unsorted' with
the fix and confirm that the AQE-on and AQE-off iterations both pass, with the
shuffle-count assertion now reading the executed plan. Then revert only the key
sort in KeyedPartitioning.createShuffleSpec and confirm that the AQE-on
iterations still fail in checkAnswer.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
