From a3a1889b314c7989a15cdde7da594727a2ac4c33 Mon Sep 17 00:00:00 2001 From: Peter Toth Date: Sat, 5 Sep 2026 10:17:14 +0200 Subject: [PATCH] [SPARK-59279][SQL] Don't re-read the sorted-merge config after GroupPartitionsExec is planned ### What changes were proposed in this pull request? `GroupPartitionsExec.canUseSortedMerge` is replaced by two members, and the config read moves into the planner's own method. - `kWayMergeIsFeasible` holds the two live terms, the child having an ordering and the child subtree being `SafeForKWayMerge`. - `usesSortedMerge` is what `doExecute`, `supportsColumnar` and `outputOrdering` ask, and it carries no config term. - `tryEnableSortedMerge` reads the config itself, so it appears once in the file, at the only place that decides anything with it. It reads it through `conf` rather than `SQLConf.get`. The `enableSortedMerge` scaladoc now states the contract, that the flag is the decision rather than a hint. EXPLAIN shows the flag too, since it is now the only thing that says whether a node k-way merges. ### Why are the changes needed? A sort-merge join silently drops rows when `spark.sql.sources.v2.bucketing.preserveOrderingOnCoalesce.enabled` is turned off between planning and execution. Both tables identity-partitioned on the join key, both reporting a two-column ordering, two splits per key so `GroupPartitionsExec` coalesces: val df = sql("SELECT /*+ MERGE(i, p) */ i.id, i.name FROM items i JOIN purchases p " + "ON p.item_id = i.id AND p.time = i.arrive_time") df.queryExecution.executedPlan // planned with the config on, no SortExec below the join spark.conf.set("spark.sql.sources.v2.bucketing.preserveOrderingOnCoalesce.enabled", "false") df.collect() // 3 rows instead of 5 At planning, `tryEnableSortedMerge()` finds the config on and returns `copy(enableSortedMerge = true)`. That copy reports the child's full ordering, so `EnsureRequirements` adds no `SortExec` under the join. The copy is a new instance, so its own `canUseSortedMerge` is still unevaluated. At execution `doExecute` forces it, now under the new config value, and builds a plain `CoalescedRDD` instead. The concatenated partitions are no longer sorted, and the sort-merge join walks them as if they were. `outputOrdering` read the same member, so the node and the plan above it ended up disagreeing about what the node delivers. `supportsColumnar` read it too, which is a second route to the same lost rows. Under AQE, `ApplyColumnarRulesAndInsertTransitions` runs when the result stage is created, so it sees the flipped config. A columnar child would then make `supportsColumnar` true and route to `doExecuteColumnar`, which only ever builds a plain `CoalescedRDD`, while the join above had already been planned against the merged ordering. That route is closed by the same change and is not covered by a test, because the suite has no columnar V2 source. Measured on `branch-4.3` and `branch-4.2` as well, same 3 rows of 5 in both, with the same two rows lost. `branch-4.1` has no `GroupPartitionsExec`, so the path does not exist there. ### Why this shape `enableSortedMerge` is already the record of the planner's decision, and only `tryEnableSortedMerge` sets it, after checking the config. So the execution side does not need the config at all. `childIsSafeForKWayMerge` does have to stay live. `ApplyColumnarRulesAndInsertTransitions` and `CollapseCodegenStages` insert nodes under the child after `EnsureRequirements` ran. Those wrappers are all whitelisted, which is what keeps the answer stable, and the live check is the fallback if something outside the whitelist ever appears. Splitting the terms apart is what lets the config term go while that guard stays. Two alternatives were rejected. Snapshotting the config into a `private val` at construction, the way `SortExec` does with `enableRadixSort`, fixes the same case, but `SparkPlan.conf` is the live conf and a `val` freezes at construction rather than at the decision, so the window reopens for any node copied after `EnsureRequirements` ran. Dropping the `&& canUseSortedMerge` re-check outright would take `childIsSafeForKWayMerge` with it. Reading `conf` instead of `SQLConf.get` is not what fixes this, and that was measured rather than assumed. `SparkPlan.conf` is `session.sessionState.conf`, the session's live mutable `SQLConf`, which is the same object `spark.conf.set` and `withSQLConf` mutate. The read does switch to `conf` here, on separate grounds. It now happens only on the driver during planning, so the node's own session conf is the right one to ask, and `SparkPlan.conf` falls back to `SQLConf.get` when the node has no session. It also makes the file consistent, since `outputOrdering` already read `conf`. `outputOrdering`'s other config, `preserveKeyOrderingOnCoalesce`, is deliberately still a live read, and the file now says why. It gates whether to *report* an ordering that holds either way, so a late read only makes the node claim less than it delivers. The sorted-merge config gates whether the merge *happens*. ### Does this PR introduce _any_ user-facing change? Yes, it fixes wrong results. A query whose plan was built with the config on keeps the k-way merge and returns the right rows when the config is turned off before it runs. One consequence worth stating. The config is documented as a cost knob, and turning it off no longer stops a merge in a plan that is already built. That includes an `InMemoryRelation`'s cached plan, which can outlive many config changes in a session. Re-planning is what picks the new value up. ### How was this patch tested? A new test in `KeyGroupedPartitioningSuite`, "SPARK-59279: a planned k-way merge survives a later config change". It forces the plan with the config on, asserts there is no `SortExec` below the sort-merge join, then turns the config off and executes the same `DataFrame`. It covers both AQE modes, because they freeze the plan at different points. AQE builds it in `AdaptiveSparkPlanExec.initialPlan`, and without AQE `prepareForExecution` does. Without the fix both modes return 3 of the 5 rows, dropping `[1,aa]` and `[2,cc]`. One existing test changed, plus a new unit test beside it. `GroupPartitionsExecSuite`'s "SPARK-55715: coalescing with enableSortedMerge = true returns full child ordering" wrapped its flag assertion in `withSQLConf(preserveOrderingOnCoalesce -> true)`. That wrapper is inert now, because `outputOrdering` no longer reads that config, so it is dropped and the assertion runs at the config's default of `false`. A new "SPARK-59279: enableSortedMerge decides the k-way merge, not the config" states the invariant directly, as a grid over the config with and without the flag. The sibling test's name lost the words "sorted merge config enabled", for the same reason. All three fail without the fix. The new integration test's fixture was byte-identical to the two tests above it, so it is extracted into `createOrderedIdTables` plus `orderedIdJoinRows`, next to the suite's other fixture helpers. "SPARK-56549: k-way merge enabled only when parent requires ordering" now uses it too. The third copy, under "SPARK-55715", is left alone to keep this diff small. "SPARK-55992: GroupPartitions string in simple and extended explain" now expects the `SortedMerge` field. `KeyGroupedPartitioningSuite`, `GroupPartitionsExecSuite`, `EnsureRequirementsSuite`, `SortedMergeCoalescedRDDSuite`, `PlannerSuite` and `ProjectedOrderingAndPartitioningSuite`, 313 tests. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 5) --- .../datasources/v2/GroupPartitionsExec.scala | 52 +++++--- .../KeyGroupedPartitioningSuite.scala | 124 ++++++++++++++---- .../v2/GroupPartitionsExecSuite.scala | 53 +++++--- 3 files changed, 162 insertions(+), 67 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala index eb0142c5a01bd..503fe172e9e7a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala @@ -29,7 +29,6 @@ import org.apache.spark.sql.catalyst.plans.QueryPlan import org.apache.spark.sql.catalyst.plans.physical.{IdentityReducer, KeyedPartitioning, KeyReducer, Partitioning} import org.apache.spark.sql.catalyst.util.{truncatedString, InternalRowComparableWrapper} import org.apache.spark.sql.execution.{SafeForKWayMerge, SparkPlan, UnaryExecNode} -import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.DataType import org.apache.spark.sql.vectorized.ColumnarBatch @@ -57,6 +56,10 @@ import org.apache.spark.sql.vectorized.ColumnarBatch * of the coalesced partitions, preserving the child's output ordering * end-to-end. Set by [[EnsureRequirements]] when a parent operator * requires the ordering that this node can satisfy via sorted merge. + * This flag is the decision, not a hint. Nothing below re-reads the + * config that produced it. A node planned with the merge would otherwise + * concatenate instead, and a sort-merge join above it would lose rows + * (SPARK-59279). */ case class GroupPartitionsExec( child: SparkPlan, @@ -272,28 +275,33 @@ case class GroupPartitionsExec( // // `outputOrdering` is first evaluated during EnsureRequirements (which decides whether to // add SortExec), but the child plan tree changes afterwards when - // ApplyColumnarRulesAndInsertTransitions and CollapseCodegenStages insert wrapper nodes. The - // correctness of this code relies on all such insertable nodes (WholeStageCodegenExec, - // InputAdapter, ColumnarToRowExec) being in the SafeForKWayMerge whitelist so the evaluation - // stays consistent. + // ApplyColumnarRulesAndInsertTransitions and CollapseCodegenStages insert wrapper nodes. Each + // such change mints a fresh node through `withNewChildInternal`, so this is recomputed against + // the new child. The correctness of this code relies on all such insertable nodes + // (WholeStageCodegenExec, InputAdapter, ColumnarToRowExec) being in the SafeForKWayMerge + // whitelist so the evaluation stays consistent. @transient private lazy val childIsSafeForKWayMerge: Boolean = !child.exists { case _: SafeForKWayMerge => false case _ => true } - @transient private lazy val canUseSortedMerge: Boolean = - SQLConf.get.v2BucketingPreserveOrderingOnCoalesceEnabled && - child.outputOrdering.nonEmpty && - childIsSafeForKWayMerge + /** Whether a k-way merge would work at all, leaving aside whether it is switched on. */ + @transient private lazy val kWayMergeIsFeasible: Boolean = + child.outputOrdering.nonEmpty && childIsSafeForKWayMerge + + /** Whether this node performs the k-way merge. No config term, see `enableSortedMerge`. */ + @transient private lazy val usesSortedMerge: Boolean = + enableSortedMerge && hasCoalescing && kWayMergeIsFeasible /** - * Returns a copy of this node with k-way merge enabled if it is feasible: the config is on, - * the child has an ordering, the child subtree is `SafeForKWayMerge`, and this node actually - * coalesces partitions. + * Returns a copy of this node with k-way merge enabled, when the config is on, this node + * coalesces partitions and the merge is feasible. The only read of + * `preserveOrderingOnCoalesce`. */ def tryEnableSortedMerge(): Option[GroupPartitionsExec] = { - Option.when(hasCoalescing && canUseSortedMerge) { + Option.when(conf.v2BucketingPreserveOrderingOnCoalesceEnabled && hasCoalescing && + kWayMergeIsFeasible) { val newGroupPartitions = copy(enableSortedMerge = true) newGroupPartitions.copyTagsFrom(this) newGroupPartitions @@ -312,7 +320,7 @@ case class GroupPartitionsExec( override protected def doExecute(): RDD[InternalRow] = { if (groupedPartitions.isEmpty) { sparkContext.emptyRDD - } else if (hasCoalescing && enableSortedMerge && canUseSortedMerge) { + } else if (usesSortedMerge) { val partitionCoalescer = new GroupedPartitionCoalescer(groupedPartitions.map(_._2)) val rowOrdering = new LazyCodeGenOrdering(kWayMergeOrdering, child.output) new SortedMergeCoalescedRDD[InternalRow]( @@ -326,8 +334,7 @@ case class GroupPartitionsExec( } } - override def supportsColumnar: Boolean = - child.supportsColumnar && !(hasCoalescing && enableSortedMerge && canUseSortedMerge) + override def supportsColumnar: Boolean = child.supportsColumnar && !usesSortedMerge override protected def doExecuteColumnar(): RDD[ColumnarBatch] = { if (groupedPartitions.isEmpty) { @@ -349,7 +356,7 @@ case class GroupPartitionsExec( // within-partition ordering is fully preserved (including any key-derived ordering that // `DataSourceV2ScanExecBase` already prepended). child.outputOrdering - } else if (enableSortedMerge && canUseSortedMerge) { + } else if (usesSortedMerge) { // Coalescing with sorted merge: SortedMergeCoalescedRDD performs a k-way merge using the // child's ordering, so the full within-partition ordering is preserved end-to-end. child.outputOrdering @@ -360,6 +367,10 @@ case class GroupPartitionsExec( // sorted ascending by the data column), concatenating them yields (A,1),(A,3),(A,2),(A,5) // which is no longer sorted by the data column. Only sort orders over partition key // expressions remain valid -- they evaluate to the same value (A) in every merged partition. + // + // The config below stays a per-call read. It gates only whether to report an ordering that + // holds either way, so a late read can never claim more than this node delivers. The merge + // config is different, because what it gates changes what this node produces. outputPartitioning match { case p: Partitioning with Expression if reducers.isEmpty && conf.v2BucketingPreserveKeyOrderingOnCoalesceEnabled => @@ -395,8 +406,11 @@ case class GroupPartitionsExec( s"Reducers: ${truncatedString(names, "[", ", ", "]", joinKeyMaxFields)}" } val distributeStr = Iterator(s"DistributePartitions: $distributePartitions") - joinKeyStr ++ expectedStr ++ reducersStr ++ distributeStr - + // Rendered from the constructor field, as `DistributePartitions` above is. Not from + // `usesSortedMerge`, because that forces `grouping`, which can throw, and this method feeds + // `simpleString`, which `treeString` calls on error paths. + val sortedMergeStr = Iterator(s"SortedMerge: $enableSortedMerge") + joinKeyStr ++ expectedStr ++ reducersStr ++ distributeStr ++ sortedMergeStr } } 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 6a911ce652615..a0a63b1b05777 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 @@ -591,6 +591,40 @@ class KeyGroupedPartitioningSuite "JOIN testcat.ns.bucket8 b8 ON b12.id = b8.id " + s"JOIN testcat.ns.bucket$third b ON b12.id = b.id") + /** + * Creates `items` and `purchases` identity-partitioned on the join key, each reporting a + * two-column ordering. Keys 1 and 2 sit on two splits per side, so a join over them makes + * `GroupPartitionsExec` coalesce. Rows are inserted so that concatenating a key's splits violates + * the reported ordering, which is what a k-way merge has to repair. + */ + private def createOrderedIdTables(): Unit = { + val itemOrdering = Array( + sort(FieldReference("id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST), + sort(FieldReference("arrive_time"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST)) + createTable(items, itemsColumns, Array(identity("id")), itemOrdering) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + "(2, 'cc', 30.0, cast('2023-06-15' as timestamp)), " + + "(1, 'bb', 20.0, cast('2022-03-10' as timestamp)), " + + "(3, 'dd', 40.0, cast('2024-01-01' as timestamp)), " + + "(1, 'aa', 10.0, cast('2021-05-20' as timestamp)), " + + "(2, 'ee', 50.0, cast('2025-09-01' as timestamp))") + + val purchaseOrdering = Array( + sort(FieldReference("item_id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST), + sort(FieldReference("time"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST)) + createTable(purchases, purchasesColumns, Array(identity("item_id")), purchaseOrdering) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + "(2, 50.0, cast('2025-09-01' as timestamp)), " + + "(1, 10.0, cast('2021-05-20' as timestamp)), " + + "(3, 40.0, cast('2024-01-01' as timestamp)), " + + "(2, 30.0, cast('2023-06-15' as timestamp)), " + + "(1, 20.0, cast('2022-03-10' as timestamp))") + } + + /** What a join of the `createOrderedIdTables` tables on both ordering columns returns. */ + private val orderedIdJoinRows = + Seq(Row(1, "aa"), Row(1, "bb"), Row(2, "cc"), Row(2, "ee"), Row(3, "dd")) + /** The `(id, ts)` rows the `withReducedTsJoinLegs` tables are filled from, one per year. */ private val row2020 = "(0, cast('2020-01-01' as timestamp))" private val row2021 = "(1, cast('2021-01-03' as timestamp))" @@ -4124,10 +4158,10 @@ class KeyGroupedPartitioningSuite |""".stripMargin) val simpleAndExtendedKeyword = "GroupPartitions JoinKeyPositions: [0] ExpectedPartitionKeys: 2 " + - "Reducers: [BucketReducer(2)] DistributePartitions: false" + "Reducers: [BucketReducer(2)] DistributePartitions: false SortedMerge: false" val formattedKeyword = "Arguments: JoinKeyPositions: [0], ExpectedPartitionKeys: 2, " + - "Reducers: [BucketReducer(2)], DistributePartitions: false" + "Reducers: [BucketReducer(2)], DistributePartitions: false, SortedMerge: false" checkKeywordsExistsInExplain(df, SimpleMode, simpleAndExtendedKeyword) checkKeywordsExistsInExplain(df, ExtendedMode, simpleAndExtendedKeyword) checkKeywordsExistsInExplain(df, FormattedMode, formattedKeyword) @@ -4601,32 +4635,9 @@ class KeyGroupedPartitioningSuite } test("SPARK-56549: k-way merge enabled only when parent requires ordering") { - // Both tables are partitioned by id/item_id and report a two-column ordering. - // Key 1 appears on two splits on each side, so GroupPartitionsExec must coalesce. - // // Dynamic gate: with the config enabled, k-way merge must be activated only when the parent // actually requires ordering (SMJ), and must stay off when the parent does not (hash join). - val itemOrdering = Array( - sort(FieldReference("id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST), - sort(FieldReference("arrive_time"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST)) - createTable(items, itemsColumns, Array(identity("id")), itemOrdering) - sql(s"INSERT INTO testcat.ns.$items VALUES " + - "(2, 'cc', 30.0, cast('2023-06-15' as timestamp)), " + - "(1, 'bb', 20.0, cast('2022-03-10' as timestamp)), " + - "(3, 'dd', 40.0, cast('2024-01-01' as timestamp)), " + - "(1, 'aa', 10.0, cast('2021-05-20' as timestamp)), " + - "(2, 'ee', 50.0, cast('2025-09-01' as timestamp))") - - val purchaseOrdering = Array( - sort(FieldReference("item_id"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST), - sort(FieldReference("time"), SortDirection.ASCENDING, NullOrdering.NULLS_FIRST)) - createTable(purchases, purchasesColumns, Array(identity("item_id")), purchaseOrdering) - sql(s"INSERT INTO testcat.ns.$purchases VALUES " + - "(2, 50.0, cast('2025-09-01' as timestamp)), " + - "(1, 10.0, cast('2021-05-20' as timestamp)), " + - "(3, 40.0, cast('2024-01-01' as timestamp)), " + - "(2, 30.0, cast('2023-06-15' as timestamp)), " + - "(1, 20.0, cast('2022-03-10' as timestamp))") + createOrderedIdTables() withSQLConf( SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", @@ -4638,7 +4649,7 @@ class KeyGroupedPartitioningSuite |FROM testcat.ns.$items i |JOIN testcat.ns.$purchases p ON p.item_id = i.id AND p.time = i.arrive_time |""".stripMargin) - checkAnswer(hashDf, Seq(Row(1, "aa"), Row(1, "bb"), Row(2, "cc"), Row(2, "ee"), Row(3, "dd"))) + checkAnswer(hashDf, orderedIdJoinRows) val hashPlan = hashDf.queryExecution.executedPlan assert(collect(hashPlan) { case j: ShuffledHashJoinExec => j }.nonEmpty, "expected ShuffledHashJoinExec") @@ -4660,7 +4671,7 @@ class KeyGroupedPartitioningSuite |FROM testcat.ns.$items i |JOIN testcat.ns.$purchases p ON p.item_id = i.id AND p.time = i.arrive_time |""".stripMargin) - checkAnswer(smjDf, Seq(Row(1, "aa"), Row(1, "bb"), Row(2, "cc"), Row(2, "ee"), Row(3, "dd"))) + checkAnswer(smjDf, orderedIdJoinRows) val smjPlan = smjDf.queryExecution.executedPlan assert(collectAllShuffles(smjPlan).isEmpty, "should not shuffle for compatible partitioning") val smjCoalescing = @@ -4675,6 +4686,63 @@ class KeyGroupedPartitioningSuite } } + test("SPARK-59279: a planned k-way merge survives a later config change") { + // The plan is built with the config on, so the k-way merge delivers the child's full ordering + // and EnsureRequirements adds no SortExec below the sort-merge join. Turning the config off + // afterwards must not change what the already planned node does. If it did, the join would read + // concatenated partitions as if they were still sorted and silently drop rows. + // + // Both AQE modes are covered, because a fresh node instance is minted at a different point in + // each. Without AQE, `CollapseCodegenStages` puts a `WholeStageCodegenExec` under this node + // during `prepareForExecution`. With AQE, the wrappers go in when the result stage is created, + // which is after the flip below. + createOrderedIdTables() + + Seq(true, false).foreach { aqeEnabled => + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString, + SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", + SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key -> "true") { + val df = sql( + s""" + |${selectWithMergeJoinHint("i", "p")} + |i.id, i.name + |FROM testcat.ns.$items i + |JOIN testcat.ns.$purchases p ON p.item_id = i.id AND p.time = i.arrive_time + |""".stripMargin) + // Force the plan without executing it, so the k-way merge is decided under this config. + val plan = df.queryExecution.executedPlan + assert(collectAllShuffles(plan).isEmpty, "should not shuffle for compatible partitioning") + val coalescing = + collectAllGroupPartitions(plan).filter(_.groupedPartitions.exists(_._2.size > 1)) + assert(coalescing.nonEmpty, "expected coalescing GroupPartitionsExec") + coalescing.foreach { gp => + assert(gp.enableSortedMerge, + "sort-merge join requires ordering: enableSortedMerge must be true") + } + val smjs = collect(plan) { case j: SortMergeJoinExec => j } + assert(smjs.nonEmpty, "expected SortMergeJoinExec") + assert(smjs.flatMap(_.children).forall(c => collect(c) { case s: SortExec => s }.isEmpty), + "the k-way merge satisfies the ordering, so no SortExec should be added") + + // Flip only the ordering config. The plan above is already committed. + withSQLConf(SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key -> "false") { + checkAnswer(df, orderedIdJoinRows) + // Re-collected, because under AQE `executedPlan` only reaches the final plan's nodes + // once the query has run. + val executed = + collectAllGroupPartitions(df.queryExecution.executedPlan) + .filter(_.groupedPartitions.exists(_._2.size > 1)) + assert(executed.nonEmpty, "expected coalescing GroupPartitionsExec") + executed.foreach { gp => + assert(gp.execute().isInstanceOf[SortedMergeCoalescedRDD[_]], + "the planned k-way merge must not be dropped by a config change") + } + } + } + } + } + test("SPARK-46367: partition key alias in subquery projects KeyedPartitioning") { // A subquery that renames a partition key (id -> pk) creates a ProjectExec between the scan and // the join. This test verifies that KeyedPartitioning expressions are correctly projected diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala index d8ea11c6ea37e..a6ce1cc17a5ad 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala @@ -223,10 +223,10 @@ class GroupPartitionsExecSuite extends SharedSparkSession { assert(gpe.outputOrdering === Nil) } - test("SPARK-55715: sorted merge config enabled but child not SafeForKWayMerge falls back " + + test("SPARK-55715: enableSortedMerge with a child that is not SafeForKWayMerge falls back " + "to key-expression ordering") { - // DummySparkPlan does not extend SafeForKWayMerge, so childIsSafeForKWayMerge = false and - // canUseSortedMerge = false even with enableSortedMerge = true. outputOrdering must + // DummySparkPlan does not extend SafeForKWayMerge, so childIsSafeForKWayMerge = false and the + // k-way merge is not feasible even with enableSortedMerge = true. outputOrdering must // therefore fall back to key-expression filtering (not return the full child ordering). val partitionKeys = Seq(row(1), row(2), row(1)) val childOrdering = Seq(SortOrder(exprA, Ascending), SortOrder(exprC, Ascending)) @@ -236,9 +236,7 @@ class GroupPartitionsExecSuite extends SharedSparkSession { assert(!GroupPartitionsExec(child).groupedPartitions.forall(_._2.size <= 1), "expected coalescing") - withSQLConf( - SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key -> "true", - SQLConf.V2_BUCKETING_PRESERVE_KEY_ORDERING_ON_COALESCE_ENABLED.key -> "true") { + withSQLConf(SQLConf.V2_BUCKETING_PRESERVE_KEY_ORDERING_ON_COALESCE_ENABLED.key -> "true") { // Even though enableSortedMerge = true, the child is not safe for k-way merge, // so only key-expression orders survive (non-key exprC is dropped). val ordering = GroupPartitionsExec(child, enableSortedMerge = true).outputOrdering @@ -249,9 +247,9 @@ class GroupPartitionsExecSuite extends SharedSparkSession { test("SPARK-55715: coalescing with enableSortedMerge = true returns full child ordering") { // Key 1 appears on partitions 0 and 2, causing coalescing. The child is a LeafExecNode so - // childIsSafeForKWayMerge = true. With enableSortedMerge = true and the config enabled, - // canUseSortedMerge = true and the full child ordering (including the non-key exprC) must be - // returned, not just the subset of key-expression orders. + // childIsSafeForKWayMerge = true. With enableSortedMerge = true the node performs the k-way + // merge, so the full child ordering (including the non-key exprC) must be returned, not just + // the subset of key-expression orders. val partitionKeys = Seq(row(1), row(2), row(1)) val childOrdering = Seq(SortOrder(exprA, Ascending), SortOrder(exprC, Ascending)) val child = DummyLeafSparkPlan( @@ -260,22 +258,37 @@ class GroupPartitionsExecSuite extends SharedSparkSession { assert(!GroupPartitionsExec(child).groupedPartitions.forall(_._2.size <= 1), "expected coalescing") - withSQLConf(SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key -> "true") { - assert(GroupPartitionsExec(child).outputOrdering !== childOrdering, - "config alone should not enable k-way merge; enableSortedMerge must be set by planner") - assert(GroupPartitionsExec(child, enableSortedMerge = true).outputOrdering === childOrdering) - } - withSQLConf( - SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key -> "false", - SQLConf.V2_BUCKETING_PRESERVE_KEY_ORDERING_ON_COALESCE_ENABLED.key -> "true") { - // Sorted-merge config disabled, key-ordering config enabled: only key-expression orders - // survive simple concatenation (non-key exprC is dropped). - val ordering = GroupPartitionsExec(child, enableSortedMerge = true).outputOrdering + assert(GroupPartitionsExec(child, enableSortedMerge = true).outputOrdering === childOrdering) + withSQLConf(SQLConf.V2_BUCKETING_PRESERVE_KEY_ORDERING_ON_COALESCE_ENABLED.key -> "true") { + // Without the flag there is no k-way merge, so only key-expression orders survive simple + // concatenation and the non-key exprC is dropped. + val ordering = GroupPartitionsExec(child).outputOrdering assert(ordering.length === 1) assert(ordering.head.child === exprA) } } + test("SPARK-59279: enableSortedMerge decides the k-way merge, not the config") { + // The config is the planner's input, and `enableSortedMerge` records what the planner decided + // under it. Once the flag is set the config no longer matters, because the plan above was + // built on the ordering the merge delivers. + val partitionKeys = Seq(row(1), row(2), row(1)) + val childOrdering = Seq(SortOrder(exprA, Ascending), SortOrder(exprC, Ascending)) + val child = DummyLeafSparkPlan( + outputPartitioning = KeyedPartitioning(Seq(exprA), partitionKeys), + outputOrdering = childOrdering) + + Seq(true, false).foreach { configEnabled => + withSQLConf( + SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key -> + configEnabled.toString) { + val flagged = GroupPartitionsExec(child, enableSortedMerge = true) + assert(flagged.outputOrdering === childOrdering, + s"config=$configEnabled: the flag alone must keep the full ordering") + } + } + } + test("SPARK-56549: tryEnableSortedMerge returns Some when conditions are met") { val partitionKeys = Seq(row(1), row(2), row(1)) val childOrdering = Seq(SortOrder(exprA, Ascending), SortOrder(exprC, Ascending))