You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by do...@apache.org on 2026/10/16 11:52:07 UTC
(spark) branch branch-4.3 updated: [SPARK-60044][4.3][SQL] Freeze the k-way merge ordering of GroupPartitionsExec when the merge is planned
This is an automated email from the ASF dual-hosted git repository.
dongjoon-hyun pushed a commit to branch branch-4.3
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.3 by this push:
new 268afa2477f9 [SPARK-60044][4.3][SQL] Freeze the k-way merge ordering of GroupPartitionsExec when the merge is planned
268afa2477f9 is described below
commit 268afa2477f9dc537aef5541e4add7c3e94487f3
Author: Peter Toth <pe...@gmail.com>
AuthorDate: Fri Oct 16 04:51:55 2026 -0700
[SPARK-60044][4.3][SQL] Freeze the k-way merge ordering of GroupPartitionsExec when the merge is planned
### Differences from #59306
This is the branch-4.3 backport of #59306. The description below is the original one.
- The bug is on branch-4.3. Without the fix, the new config-change test numbers the rows of `id` 1 `[1, aa, 2]` and `[1, bb, 1]`, as in the example below. The new `debug()` test fails too, since the merge concatenates under a `DebugExec`.
- For that run the three main files were restored from branch-4.3. `GroupPartitionsExecSuite.scala` and `EnsureRequirementsSuite.scala` were restored too, and one line of `KeyGroupedPartitioningSuite.scala` went back to `enableSortedMerge = true`. Those changes need the new field.
- `GroupPartitionsExec` on 4.3 is the older node. It has `expectedPartitionKeys`, `distributePartitions` and a default for every field but `child`, and no factory. `sortedMergeOrdering` replaces `enableSortedMerge` there, with `None` as its default.
- `withKeyPositionsFor`, `checkChildStillMatches`, `sendDriverMetrics` and the factory are not on 4.3. The hunks in `withKeyPositionsFor` and the factory are dropped. So `checkPlannedMergeRuns` is the first check in `doExecute`, and its Scaladoc no longer points to `checkChildStillMatches`.
- The comment on the EXPLAIN string keeps the 4.3 reason for not reading `usesSortedMerge`. On 4.3 that forces `grouping`, which can throw.
- `EnsureRequirements.scala` has no change. On 4.3 `withJoinKeyPositions` builds the node with `copy`, which already keeps the ordering.
- The `withKeyPositionsFor` test in `GroupPartitionsExecSuite` is dropped, since the method is not on 4.3. The rest of `GroupPartitionsExecSuite.scala` conflicted only on imports and on a master-only test helper next to the new ones.
- `KeyGroupedPartitioningSuite.scala` only conflicted on imports. In `EnsureRequirementsSuite.scala` only the new test is taken. The conflict there was master-only tests next to it.
- The new tests already set the configs they need. `partitionKeyOrdering` is off by default on 4.3. `preserveOrderingOnCoalesce` is off by default on every branch. The window is a single-child operator, so it needs no `spark.sql.requireAllClusterKeysForCoPartition`.
- The new `GroupPartitionsExecSuite` and `EnsureRequirementsSuite` tests need the new field, so they cannot run on the unfixed branch. Ablations on the fixed branch:
- With `kWayMergeOrdering` reading the child's ordering, the new `GroupPartitionsExecSuite` test and the config-change test fail, the latter with the wrong row numbers.
- With `withNewChildInternal` dropping the ordering, the new `GroupPartitionsExecSuite` test fails.
- With `outputOrdering` reading the child's ordering, that test and the config-change test fail.
- Without `checkPlannedMergeRuns`, or with `supportsColumnar` reading `usesSortedMerge`, the adjusted SPARK-55715 test fails.
- Without `DebugExec` in the whitelist, the `debug()` test fails with the planned-merge error.
- With `withJoinKeyPositions` dropping the ordering, the new `EnsureRequirementsSuite` test fails.
- Ran `GroupPartitionsExecSuite` (24), `KeyGroupedPartitioningSuite` (190), `EnsureRequirementsSuite` (59), `CombineAdjacentAggregationSuite` (10), `DebuggingSuite` (7) and `DebuggingSuiteAE` (5). All 295 tests pass. `KeyGroupedPartitioningSuite` is the only `KeyGroupedPartitioning*` suite on 4.3.
- `dev/lint-scala` passes.
### What changes were proposed in this pull request?
`GroupPartitionsExec` now freezes the ordering its k-way merge merges by, when the merge is planned:
- `sortedMergeOrdering: Option[Seq[SortOrder]]` replaces the `enableSortedMerge: Boolean` field. `enableSortedMerge` stays, as a method.
- `tryEnableSortedMerge` captures the child's ordering. That is the ordering `EnsureRequirements` checks against the parent's requirement.
- Every rebuild over the same child, or a wrapper of it, passes it on: `copy`, `withNewChildInternal` and `EnsureRequirements.withJoinKeyPositions`. Only `tryEnableSortedMerge` captures it.
- `withKeyPositionsFor` turns a merging node away, since it rebuilds the node over another child. No merging node reaches it today.
- For a merging node, `kWayMergeOrdering` and the merge branch of `outputOrdering` read the frozen ordering instead of `child.outputOrdering`. The frozen ordering is never empty.
- A merging node whose child is no longer safe for the merge fails at execution instead of concatenating. The error names the node outside the `SafeForKWayMerge` whitelist. That needs a rule after `EnsureRequirements` to insert such a node. No built-in rule does.
- The columnar path never merges, so a planned merge never reports itself columnar, whatever its child turns into later.
- `DebugExec` joins the whitelist. Without AQE, the `debug()` helper of `org.apache.spark.sql.execution.debug` wraps every node in it, including the ones under a merging node. Its per-partition state lives in its evaluator, and its accumulators only add.
- `tryEnableSortedMerge` leaves a node that already merges as it is, so the merge is decided once.
- EXPLAIN shows the frozen ordering instead of `SortedMerge: true`, e.g. `SortedMerge: [id#1 ASC NULLS FIRST, name#2 ASC NULLS FIRST]`.
### Why are the changes needed?
SPARK-59279 froze the merge decision, but the node still read the ordering it merges by from `child.outputOrdering`. When the source reports no ordering, the scan derives one from the partition keys under `spark.sql.sources.v2.bucketing.partitionKeyOrdering.enabled`, which it reads live. So turning that conf off after planning emptied the ordering of a planned merge. What the node did then depended on whether it had already read the ordering:
- If it had, it still merged, by an empty ordering, so it treated every row as equal. The new test takes this path without AQE, since it checks the merge's ordering before turning the conf off.
- Otherwise it found the merge infeasible and concatenated the splits. With AQE that always happens, since the result stage's node is made after the change.
Either way, a parent that relied on the merge read its rows out of order.
For example, take a table partitioned by `(identity(id), days(arrive_time), identity(name))`, with `allowKeysSubsetOfPartitionKeys` and `preserveOrderingOnCoalesce` on. The scan derives the ordering `[id, name]`, while Spark sorts the splits by the full key. `ROW_NUMBER() OVER (PARTITION BY id ORDER BY name)` plans a k-way merge and no sort. With the conf turned off before the query runs, the rows of `id` 1 were numbered in split order: `[1, aa, 2]` and `[1, bb, 1]` instead of `[1, aa [...]
### Does this PR introduce _any_ user-facing change?
Yes. A query planned with the k-way merge keeps the merge's ordering when `spark.sql.sources.v2.bucketing.partitionKeyOrdering.enabled` changes after planning. It used to give wrong results then, such as wrong `ROW_NUMBER` values. The bug exists since 4.2.0, which has both the k-way merge (SPARK-55715) and the key-derived ordering (SPARK-56241). EXPLAIN of a planned k-way merge now shows its ordering instead of `SortedMerge: true`.
### How was this patch tested?
- New test in `KeyGroupedPartitioningSuite`, "SPARK-60044: a planned k-way merge keeps its ordering after a later config change". It plans the example above, turns the conf off, then checks the rows and that the merge still reports `[id, name]`, in both AQE modes. It fails without the fix with the wrong row numbers above.
- New test in `EnsureRequirementsSuite`: pushing join key positions into a merging node keeps its frozen ordering.
- New tests in `GroupPartitionsExecSuite`:
- `withKeyPositionsFor` turns a merging node away. It fails without the check.
- `tryEnableSortedMerge` captures the child's ordering. `copy`, `withNewChildren` and `outputOrdering` keep it after the child's ordering becomes empty. It also checks the EXPLAIN string, that the merge is decided once, that deciding again after the change finds no ordering, and that an empty ordering is refused. It fails if `withNewChildInternal` drops the ordering or `outputOrdering` reads the child's.
- A merge over a child that is not safe for it fails at execution on both execute paths, checked in an existing `GroupPartitionsExecSuite` test.
- New test in `KeyGroupedPartitioningSuite`: the `debug()` helper runs a planned k-way merge without AQE. It checks that the merge runs under the `DebugExec` wrapping, and fails if `DebugExec` is not whitelisted.
- The tests that build `GroupPartitionsExec` with the merge on now pass the ordering.
- Ran `GroupPartitionsExecSuite` (35), the `KeyGroupedPartitioning*` suites (234), `EnsureRequirementsSuite` (80), `CombineAdjacentAggregationSuite` (10) and `DebuggingSuite`, plus `dev/lint-scala`.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5.5)
Closes #59508 from peter-toth/SPARK-60044-freeze-kway-merge-ordering-4.3.
Authored-by: Peter Toth <pe...@gmail.com>
Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
.../spark/sql/execution/SafeForKWayMerge.scala | 3 +
.../datasources/v2/GroupPartitionsExec.scala | 128 +++++++++++++-------
.../apache/spark/sql/execution/debug/package.scala | 3 +-
.../connector/KeyGroupedPartitioningSuite.scala | 95 ++++++++++++++-
.../datasources/v2/GroupPartitionsExecSuite.scala | 132 +++++++++++++++------
.../exchange/EnsureRequirementsSuite.scala | 16 +++
6 files changed, 298 insertions(+), 79 deletions(-)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SafeForKWayMerge.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SafeForKWayMerge.scala
index 08be66e3f447..5024a4177e02 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SafeForKWayMerge.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SafeForKWayMerge.scala
@@ -23,5 +23,8 @@ package org.apache.spark.sql.execution
* mutable state in shared plan-node instance fields; all per-partition state must be
* captured inside the partition's iterator closure (e.g. via the
* PartitionEvaluatorFactory pattern).
+ *
+ * A `GroupPartitionsExec` planned with a k-way merge fails at execution if a node below it does
+ * not mix this in, e.g. one that a rule put there after `EnsureRequirements`.
*/
trait SafeForKWayMerge
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 2a9ff7971431..c8d3fb115e5a 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
@@ -51,14 +51,18 @@ import org.apache.spark.sql.vectorized.ColumnarBatch
* @param distributePartitions When true, splits for a key are distributed across the expected
* partitions (padding with empty partitions). When false, all splits
* are replicated to every expected partition for that key.
- * @param enableSortedMerge When true, uses [[SortedMergeCoalescedRDD]] to perform a k-way merge
- * 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).
+ * @param sortedMergeOrdering When set, uses [[SortedMergeCoalescedRDD]] to perform a k-way merge
+ * of the coalesced partitions by this ordering, preserving it
+ * end-to-end. Set by [[EnsureRequirements]] when a parent operator
+ * requires the ordering that this node can satisfy via sorted merge.
+ * This 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). The ordering is captured from the child when the merge
+ * is decided. That is the ordering the parent's requirement was checked
+ * against. The child's ordering can follow a config too, so a later
+ * read could merge by a shorter or empty one (SPARK-60044). The
+ * captured ordering is never empty.
*/
case class GroupPartitionsExec(
child: SparkPlan,
@@ -66,7 +70,7 @@ case class GroupPartitionsExec(
@transient expectedPartitionKeys: Option[Seq[(InternalRowComparableWrapper, Int)]] = None,
@transient reducers: Option[Seq[Option[KeyReducer]]] = None,
@transient distributePartitions: Boolean = false,
- @transient enableSortedMerge: Boolean = false
+ @transient sortedMergeOrdering: Option[Seq[SortOrder]] = None
) extends UnaryExecNode {
override def outputPartitioning: Partitioning = {
@@ -318,8 +322,8 @@ case class GroupPartitionsExec(
// not the one belonging to the just-exhausted partition. If that last sorter is still being
// actively read by another partition in the k-way merge, freeing it causes a use-after-free.
//
- // To be conservative, we use a whitelist: unknown node types fall through to unsafe, causing
- // a fallback to simple sequential coalescing. Only node types explicitly confirmed to store no
+ // To be conservative, we use a whitelist. An unknown node type counts as unsafe, so the merge
+ // is not planned and the node concatenates. Only node types explicitly confirmed to store no
// per-partition state in shared (plan node) instance fields are permitted.
//
// --- Interaction with late plan-tree mutations ---
@@ -330,33 +334,47 @@ case class GroupPartitionsExec(
// 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.
+ // whitelist so the evaluation stays consistent. So is the `DebugExec` that the `debug()` helper
+ // (`org.apache.spark.sql.execution.debug`) wraps every node in without AQE. A rule that breaks
+ // this under a planned merge fails the query at execution, see `checkPlannedMergeRuns`.
@transient private lazy val childIsSafeForKWayMerge: Boolean =
!child.exists {
case _: SafeForKWayMerge => false
case _ => true
}
- /** 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
+ // A merge by an empty ordering would run but sort nothing.
+ require(sortedMergeOrdering.forall(_.nonEmpty), "A k-way merge needs an ordering to merge by")
- /** Whether this node performs the k-way merge. No config term, see `enableSortedMerge`. */
- @transient private lazy val usesSortedMerge: Boolean =
- enableSortedMerge && hasCoalescing && kWayMergeIsFeasible
+ /**
+ * Whether the k-way merge was planned for this node. See `sortedMergeOrdering`. Driver-only, like
+ * the other accessors of `@transient` fields: a deserialized copy holds no ordering.
+ */
+ def enableSortedMerge: Boolean = sortedMergeOrdering.isDefined
+
+ /** Whether this node was planned to k-way merge partitions it coalesces. */
+ @transient private lazy val mergeIsPlanned: Boolean = enableSortedMerge && hasCoalescing
+
+ /** Whether this node performs the k-way merge. No config term, see `sortedMergeOrdering`. */
+ @transient private lazy val usesSortedMerge: Boolean = mergeIsPlanned && childIsSafeForKWayMerge
/**
* 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`.
+ * coalesces partitions, the child reports an ordering and its subtree is safe for the merge. The
+ * only read of `preserveOrderingOnCoalesce`. The copy keeps the ordering checked here, see
+ * `sortedMergeOrdering`. A node that already merges is left as it is, so the merge is decided
+ * once.
*/
def tryEnableSortedMerge(): Option[GroupPartitionsExec] = {
- Option.when(conf.v2BucketingPreserveOrderingOnCoalesceEnabled && hasCoalescing &&
- kWayMergeIsFeasible) {
- val newGroupPartitions = copy(enableSortedMerge = true)
- newGroupPartitions.copyTagsFrom(this)
- newGroupPartitions
- }
+ // The ordering is read once, so the ordering frozen is the one checked.
+ Option.when(!enableSortedMerge && conf.v2BucketingPreserveOrderingOnCoalesceEnabled &&
+ hasCoalescing && childIsSafeForKWayMerge)(child.outputOrdering)
+ .filter(_.nonEmpty)
+ .map { ordering =>
+ val newGroupPartitions = copy(sortedMergeOrdering = Some(ordering))
+ newGroupPartitions.copyTagsFrom(this)
+ newGroupPartitions
+ }
}
/**
@@ -366,9 +384,26 @@ case class GroupPartitionsExec(
* serialized with the RDD in every task -- is dropped.
*/
private[v2] def kWayMergeOrdering: Seq[SortOrder] =
- child.outputOrdering.map(_.copy(sameOrderExpressions = Seq.empty))
+ sortedMergeOrdering.get.map(_.copy(sameOrderExpressions = Seq.empty))
+
+ /**
+ * Fails a planned merge that would not run. The parent's required ordering was checked against
+ * the merge, so no sort was added above it. A rule after `EnsureRequirements` can still put a
+ * node outside the `SafeForKWayMerge` whitelist below this one. This node would then concatenate
+ * and hand the parent unsorted rows. No built-in rule does that, see `childIsSafeForKWayMerge`.
+ * This makes the failure explicit if one starts.
+ */
+ private def checkPlannedMergeRuns(): Unit = {
+ if (mergeIsPlanned && !childIsSafeForKWayMerge) {
+ val blocker = child.find(!_.isInstanceOf[SafeForKWayMerge]).map(_.nodeName)
+ throw SparkException.internalError(
+ "GroupPartitionsExec was planned with a k-way merge, but its child holds " +
+ s"${blocker.getOrElse("a node")}, which does not mix in SafeForKWayMerge")
+ }
+ }
override protected def doExecute(): RDD[InternalRow] = {
+ checkPlannedMergeRuns()
if (groupedPartitions.isEmpty) {
sparkContext.emptyRDD
} else if (usesSortedMerge) {
@@ -385,9 +420,16 @@ case class GroupPartitionsExec(
}
}
- override def supportsColumnar: Boolean = child.supportsColumnar && !usesSortedMerge
+ // The columnar path never merges, so a planned merge stays row-based. That is decided by the plan
+ // rather than by the live child, so the columnar transition lands below the node, and `doExecute`
+ // either merges or fails in `checkPlannedMergeRuns`.
+ override def supportsColumnar: Boolean = child.supportsColumnar && !mergeIsPlanned
override protected def doExecuteColumnar(): RDD[ColumnarBatch] = {
+ if (mergeIsPlanned) {
+ throw SparkException.internalError(
+ "GroupPartitionsExec was planned with a k-way merge, which the columnar path cannot run")
+ }
if (groupedPartitions.isEmpty) {
sparkContext.emptyRDD
} else {
@@ -408,9 +450,9 @@ case class GroupPartitionsExec(
// `DataSourceV2ScanExecBase` already prepended).
child.outputOrdering
} 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
+ // Coalescing with sorted merge: SortedMergeCoalescedRDD performs a k-way merge by the
+ // planned ordering, so that ordering is preserved end-to-end. See `sortedMergeOrdering`.
+ sortedMergeOrdering.get
} else {
// Coalescing: multiple input partitions are merged into one output partition. The child's
// within-partition ordering is lost due to concatenation -- for example, if two input
@@ -419,9 +461,11 @@ case class GroupPartitionsExec(
// 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.
+ // The config below and the child's ordering stay per-call reads. They only change how much
+ // of an ordering that holds either way is reported. So a late read can never claim more
+ // than this node delivers. The merge is different, because its config and its ordering
+ // change what this node produces. It reads the config once and freezes the ordering when it
+ // is planned, see `sortedMergeOrdering`.
outputPartitioning match {
case p: Partitioning with Expression
if reducers.isEmpty && conf.v2BucketingPreserveKeyOrderingOnCoalesceEnabled =>
@@ -445,20 +489,22 @@ case class GroupPartitionsExec(
override protected def stringArgs: Iterator[Any] = planSummaryParts(Int.MaxValue)
- private def planSummaryParts(joinKeyMaxFields: Int): Iterator[String] = {
+ private def planSummaryParts(maxFields: Int): Iterator[String] = {
val joinKeyStr = joinKeyPositions.map { p =>
- s"JoinKeyPositions: ${truncatedString(p, "[", ", ", "]", joinKeyMaxFields)}"
+ s"JoinKeyPositions: ${truncatedString(p, "[", ", ", "]", maxFields)}"
}.iterator
val expectedStr = expectedPartitionKeys.map(ks => s"ExpectedPartitionKeys: ${ks.size}")
val reducersStr = reducers.map { seq =>
val names = seq.map(_.map(_.reducer.displayName()).getOrElse("identity"))
- s"Reducers: ${truncatedString(names, "[", ", ", "]", joinKeyMaxFields)}"
+ s"Reducers: ${truncatedString(names, "[", ", ", "]", maxFields)}"
}
val distributeStr = Iterator(s"DistributePartitions: $distributePartitions")
- // 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")
+ // Rendered from the constructor field `sortedMergeOrdering`, 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. The ordering is shown
+ // since it can differ from the one the child reports now.
+ val sortedMergeStr = Iterator("SortedMerge: " +
+ sortedMergeOrdering.fold("false")(truncatedString(_, "[", ", ", "]", maxFields)))
joinKeyStr ++ expectedStr ++ reducersStr ++ distributeStr ++ sortedMergeStr
}
}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/debug/package.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/debug/package.scala
index b408f2823e16..b068f8fbf5ec 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/debug/package.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/debug/package.scala
@@ -229,7 +229,8 @@ package object debug {
override def value: java.util.Set[T] = _set
}
- case class DebugExec(child: SparkPlan) extends UnaryExecNode with CodegenSupport {
+ case class DebugExec(child: SparkPlan)
+ extends UnaryExecNode with CodegenSupport with SafeForKWayMerge {
def output: Seq[Attribute] = child.output
/**
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 4ea60e3aa3e1..4ff0d3142c51 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
@@ -21,7 +21,7 @@ import java.util.Collections
import org.apache.spark.{SparkConf, SparkException}
import org.apache.spark.rdd.SortedMergeCoalescedRDD
-import org.apache.spark.sql.{DataFrame, ExplainSuiteHelper, Row}
+import org.apache.spark.sql.{DataFrame, ExplainSuiteHelper, QueryTest, Row}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, ExprId, Literal, TransformExpression}
import org.apache.spark.sql.catalyst.plans.Inner
@@ -3144,7 +3144,8 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
// A k-way merge binds the scan's ordering against the scan output. A row-based scan like this
// one always gets a `ProjectExec`, which truncates the ordering, but a columnar scan may get
// nothing between it and the merge that does. So the merge is built right over the scan here.
- val merged = GroupPartitionsExec(scan, enableSortedMerge = true).execute()
+ val merged =
+ GroupPartitionsExec(scan, sortedMergeOrdering = Some(scan.outputOrdering)).execute()
assert(merged.isInstanceOf[SortedMergeCoalescedRDD[_]])
val rows = merged.map(r => (r.getInt(0), r.getString(1))).collect().toSeq
val expected = Seq((1, "aa"), (1, "bb"), (2, "cc"))
@@ -5607,7 +5608,7 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
assert(merging.length == 1, "expected one k-way merge")
assert(merging.head.child.output.exists(_.name == "arrive_time"),
"test setup: arrive_time stays in the scan output")
- assert(merging.head.child.outputOrdering.map(_.child.sql) == expectedOrdering)
+ assert(merging.head.outputOrdering.map(_.child.sql) == expectedOrdering)
assert(merging.head.execute().isInstanceOf[SortedMergeCoalescedRDD[_]])
}
@@ -5748,7 +5749,8 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
// 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)
+ // Without `checkToRDD`, so only the committed plan runs, not a newly planned one.
+ QueryTest.checkAnswer(df, orderedIdJoinRows, checkToRDD = false)
// Re-collected, because under AQE `executedPlan` only reaches the final plan's nodes
// once the query has run.
val executed =
@@ -5764,6 +5766,91 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
}
}
+ /**
+ * The setup of the SPARK-60044 tests. Creates `items`, which reports no ordering, so the scan
+ * derives [id, name] from its partition keys and drops the transform between them. The splits
+ * are sorted by the full key, so within id 1 the 'bb' split comes before the 'aa' one. Returns a
+ * window query that groups by id and orders by name. Under `kWayMergeWindowConfs` a k-way merge
+ * delivers that ordering, so EnsureRequirements adds no SortExec. The query selects
+ * `arrive_time`, so the transform's column stays in the scan output.
+ */
+ private def createKWayMergeWindowQuery(): String = {
+ createTable(items, itemsColumns, Array(identity("id"), days("arrive_time"), identity("name")))
+ sql(s"INSERT INTO testcat.ns.$items VALUES " +
+ "(1, 'bb', 10.0, cast('2021-01-01' as timestamp)), " +
+ "(1, 'aa', 20.0, cast('2022-01-01' as timestamp)), " +
+ "(2, 'cc', 30.0, cast('2021-01-01' as timestamp))")
+ s"""SELECT id, name, arrive_time, ROW_NUMBER() OVER (PARTITION BY id ORDER BY name) AS rn
+ |FROM testcat.ns.$items
+ |""".stripMargin
+ }
+
+ private val kWayMergeWindowConfs = Seq(
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true",
+ SQLConf.V2_BUCKETING_PARTITION_KEY_ORDERING_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key -> "true")
+
+ private val kWayMergeWindowRows = Seq(
+ Row(1, "aa", Timestamp.valueOf("2022-01-01 00:00:00"), 1),
+ Row(1, "bb", Timestamp.valueOf("2021-01-01 00:00:00"), 2),
+ Row(2, "cc", Timestamp.valueOf("2021-01-01 00:00:00"), 1))
+
+ test("SPARK-60044: a planned k-way merge keeps its ordering after a later config change") {
+ // Turning partitionKeyOrdering off after planning empties the scan's ordering. The planned
+ // merge must still merge by [id, name]. Otherwise the window numbers the rows of id 1 in split
+ // order.
+ //
+ // Both AQE modes are covered, see the SPARK-59279 test above.
+ val query = createKWayMergeWindowQuery()
+ Seq(true, false).foreach { aqeEnabled =>
+ val confs = (SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString) +:
+ kWayMergeWindowConfs
+ withSQLConf(confs: _*) {
+ val df = sql(query)
+ // Force the plan without executing it, so the k-way merge is decided under this config.
+ val plan = df.queryExecution.executedPlan
+ val merging = collectAllGroupPartitions(plan).filter(_.enableSortedMerge)
+ assert(merging.length == 1, "expected one k-way merge")
+ assert(merging.head.outputOrdering.map(_.child.sql) == Seq("id", "name"))
+ assert(collect(plan) { case s: SortExec => s }.isEmpty,
+ "the k-way merge satisfies the window's ordering, so no SortExec should be added")
+
+ // Flip only the key ordering config. The plan above is already committed.
+ withSQLConf(SQLConf.V2_BUCKETING_PARTITION_KEY_ORDERING_ENABLED.key -> "false") {
+ assert(merging.head.child.outputOrdering.isEmpty,
+ "test setup: the scan no longer derives an ordering from its keys")
+ // Without `checkToRDD`, so only the committed plan runs, not a newly planned one.
+ QueryTest.checkAnswer(df, kWayMergeWindowRows, checkToRDD = false)
+ // Re-collected, because under AQE `executedPlan` only reaches the final plan's nodes
+ // once the query has run.
+ val executed = collectAllGroupPartitions(df.queryExecution.executedPlan)
+ assert(executed.map(_.outputOrdering.map(_.child.sql)) == Seq(Seq("id", "name")),
+ "the planned k-way merge must keep the ordering it was planned with")
+ }
+ }
+ }
+ }
+
+ test("SPARK-60044: the debug() helper runs a planned k-way merge") {
+ // Without AQE, `debug()` wraps every node of the executed plan in a `DebugExec`, including the
+ // nodes under a merging `GroupPartitionsExec`. That wrapper must not stop the planned merge.
+ import org.apache.spark.sql.execution.debug._
+ val query = createKWayMergeWindowQuery()
+ val confs = (SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") +: kWayMergeWindowConfs
+ withSQLConf(confs: _*) {
+ val df = sql(query)
+ val plan = df.queryExecution.executedPlan
+ assert(collectAllGroupPartitions(plan).exists(_.enableSortedMerge),
+ "test setup: the plan has a k-way merge")
+ df.debug()
+
+ // `debug()` returns nothing, so the merge is also checked under a `DebugExec` child.
+ val planned = collectAllGroupPartitions(plan).filter(_.enableSortedMerge).head
+ val wrapped = planned.withNewChildren(Seq(DebugExec(planned.child)))
+ assert(wrapped.execute().isInstanceOf[SortedMergeCoalescedRDD[_]])
+ }
+ }
+
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 3d14fac3a656..e0c252fabb9a 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
@@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, Keye
import org.apache.spark.sql.catalyst.util.DateTimeTestUtils.date
import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper
import org.apache.spark.sql.connector.catalog.functions.{BucketFunction, BucketReducer, Reducer, YearsFunction}
-import org.apache.spark.sql.execution.{DummySparkPlan, LeafExecNode, SafeForKWayMerge}
+import org.apache.spark.sql.execution.{DummySparkPlan, InputAdapter, LeafExecNode, SafeForKWayMerge}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.sql.types.{DataType, IntegerType, TimestampType}
@@ -116,21 +116,56 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
}
test("SPARK-58324: k-way merge ordering drops sameOrderExpressions") {
- // The child ordering carries sameOrderExpressions (planner metadata). The k-way merge
+ // The planned ordering carries sameOrderExpressions (planner metadata). The k-way merge
// comparator only needs the sort key, so kWayMergeOrdering keeps child/direction/nullOrdering
// but drops sameOrderExpressions, so LazyRowOrdering does not serialize them with the RDD.
- val childOrdering = Seq(SortOrder(exprA, Ascending, Seq(exprB, exprC)))
+ val plannedOrdering = Seq(SortOrder(exprA, Ascending, Seq(exprB, exprC)))
val child = DummySparkPlan(
- outputPartitioning = KeyedPartitioning(Seq(exprA), Seq(row(1), row(2), row(1))),
- outputOrdering = childOrdering)
- val gpe = GroupPartitionsExec(child)
+ outputPartitioning = KeyedPartitioning(Seq(exprA), Seq(row(1), row(2), row(1))))
+ val gpe = GroupPartitionsExec(child, sortedMergeOrdering = Some(plannedOrdering))
- assert(child.outputOrdering.head.sameOrderExpressions.nonEmpty, "test setup")
val merged = gpe.kWayMergeOrdering
assert(merged.map(so => (so.child, so.direction)) === Seq((exprA, Ascending)))
assert(merged.forall(_.sameOrderExpressions.isEmpty))
}
+ test("SPARK-60044: a planned merge keeps the ordering it was planned with") {
+ val partitionKeys = Seq(row(1), row(2), row(1))
+ val childOrdering = Seq(SortOrder(exprA, Ascending), SortOrder(exprC, Ascending))
+ val child = ConfOrderedLeaf(childOrdering, KeyedPartitioning(Seq(exprA), partitionKeys))
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTITION_KEY_ORDERING_ENABLED.key -> "true") {
+ val unplanned = GroupPartitionsExec(child)
+ val merging = unplanned.tryEnableSortedMerge().get
+ assert(merging.sortedMergeOrdering.contains(childOrdering))
+ // EXPLAIN shows the ordering the merge was planned with.
+ val planned = childOrdering.mkString("[", ", ", "]")
+ assert(merging.simpleString(10).contains(s"SortedMerge: $planned"))
+ // The merge is decided once.
+ assert(merging.tryEnableSortedMerge().isEmpty)
+
+ withSQLConf(SQLConf.V2_BUCKETING_PARTITION_KEY_ORDERING_ENABLED.key -> "false") {
+ assert(child.outputOrdering.isEmpty, "test setup: the child's ordering follows the config")
+ // The planned node and two rebuilds of it, all read after the change. `InputAdapter`
+ // stands in for a wrapper that `CollapseCodegenStages` inserts after planning.
+ Seq(merging, merging.copy(), merging.withNewChildren(Seq(InputAdapter(child))))
+ .map(_.asInstanceOf[GroupPartitionsExec])
+ .foreach { node =>
+ assert(node.outputOrdering === childOrdering)
+ assert(node.kWayMergeOrdering === childOrdering)
+ }
+ // Deciding again, even on the same instance, finds no ordering to merge by.
+ assert(unplanned.tryEnableSortedMerge().isEmpty)
+ }
+ }
+ // An empty merge ordering is refused.
+ intercept[IllegalArgumentException] {
+ GroupPartitionsExec(child, sortedMergeOrdering = Some(Seq.empty))
+ }
+ }
+
test("SPARK-59995: the k-way merge ordering falls back to interpreted evaluation") {
// The scan drops every sort order over a transform, so none reaches the merge. Here a transform
// only stands in for a sort key without generated code. `TransformExpression` generates none,
@@ -405,11 +440,10 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
}
}
- 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 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).
+ test("SPARK-55715: a planned merge over a child that is not SafeForKWayMerge reports only " +
+ "the key ordering and fails at execution") {
+ // DummySparkPlan does not extend SafeForKWayMerge, so the node cannot merge even with a
+ // planned merge ordering. outputOrdering therefore keeps only the key-expression orders.
val partitionKeys = Seq(row(1), row(2), row(1))
val childOrdering = Seq(SortOrder(exprA, Ascending), SortOrder(exprC, Ascending))
val child = DummySparkPlan(
@@ -419,19 +453,31 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
assert(!GroupPartitionsExec(child).groupedPartitions.forall(_._2.size <= 1),
"expected coalescing")
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
+ val merging = GroupPartitionsExec(child, sortedMergeOrdering = Some(childOrdering))
+ val ordering = merging.outputOrdering
assert(ordering.length === 1)
assert(ordering.head.child === exprA)
+ // A planned merge that cannot run fails rather than concatenating (SPARK-60044). The error
+ // names the node outside the whitelist. The columnar path never merges, so it always fails.
+ val error = intercept[SparkException](merging.execute())
+ assert(error.getMessage.contains("holds DummySparkPlan"))
+ val columnarError = intercept[SparkException](merging.executeColumnar())
+ assert(columnarError.getMessage.contains("the columnar path cannot run"))
+
+ // The columnar side follows the plan rather than the live child, so a columnar child
+ // outside the whitelist does not make a planned merge columnar.
+ val columnarChild = UnsafeColumnarLeaf(KeyedPartitioning(Seq(exprA), partitionKeys))
+ assert(GroupPartitionsExec(columnarChild).supportsColumnar, "test setup")
+ assert(!GroupPartitionsExec(columnarChild, sortedMergeOrdering = Some(childOrdering))
+ .supportsColumnar)
}
}
- 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 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.
+ test("SPARK-55715: coalescing with a planned merge returns full child ordering") {
+ // Key 1 appears on partitions 0 and 2, causing coalescing. The child, DummyLeafSparkPlan,
+ // mixes in SafeForKWayMerge, so childIsSafeForKWayMerge = true. With a planned merge ordering
+ // 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(
@@ -440,19 +486,20 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
assert(!GroupPartitionsExec(child).groupedPartitions.forall(_._2.size <= 1),
"expected coalescing")
- assert(GroupPartitionsExec(child, enableSortedMerge = true).outputOrdering === childOrdering)
+ assert(GroupPartitionsExec(child, sortedMergeOrdering = Some(childOrdering))
+ .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.
+ // Without a planned merge ordering 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
+ test("SPARK-59279: the planned merge ordering decides the k-way merge, not the config") {
+ // The config is the planner's input, and `sortedMergeOrdering` records what the planner
+ // decided under it. Once it 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))
@@ -464,13 +511,13 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
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")
- val unflagged = GroupPartitionsExec(child)
- assert(unflagged.outputOrdering !== childOrdering,
- s"config=$configEnabled: without the flag there is no k-way merge to report, and the " +
- "config cannot supply one")
+ val merging = GroupPartitionsExec(child, sortedMergeOrdering = Some(childOrdering))
+ assert(merging.outputOrdering === childOrdering,
+ s"config=$configEnabled: the planned merge alone must report the full ordering")
+ val unplanned = GroupPartitionsExec(child)
+ assert(unplanned.outputOrdering !== childOrdering,
+ s"config=$configEnabled: without a planned merge there is no k-way merge to report, " +
+ "and the config cannot supply one")
}
}
}
@@ -566,3 +613,22 @@ private case class DummyLeafSparkPlan(
throw new UnsupportedOperationException
override def output: Seq[Attribute] = Seq.empty
}
+
+/** A columnar leaf that does not mix in `SafeForKWayMerge`. */
+private case class UnsafeColumnarLeaf(override val outputPartitioning: Partitioning)
+ extends LeafExecNode {
+ override def supportsColumnar: Boolean = true
+ override protected def doExecute(): RDD[InternalRow] = throw new UnsupportedOperationException
+ override def output: Seq[Attribute] = Seq.empty
+}
+
+/** A leaf whose ordering follows a config, the way a V2 scan's key-derived ordering does. */
+private case class ConfOrderedLeaf(
+ ordering: Seq[SortOrder],
+ override val outputPartitioning: Partitioning) extends LeafExecNode with SafeForKWayMerge {
+ override def outputOrdering: Seq[SortOrder] =
+ if (conf.v2BucketingPartitionKeyOrderingEnabled) ordering else Nil
+ override protected def doExecute(): RDD[InternalRow] =
+ throw new UnsupportedOperationException
+ override def output: Seq[Attribute] = Seq.empty
+}
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 2536e120e393..319f00785cf1 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
@@ -2344,6 +2344,22 @@ class EnsureRequirementsSuite extends SharedSparkSession {
s"EnsureRequirements must satisfy the required distribution:\n${newChild.treeString}")
}
+ test("SPARK-60044: pushing join key positions into a merging node keeps its ordering") {
+ val keys = Seq(InternalRow(1, 1), InternalRow(1, 2))
+ // Safe for the merge, so the rebuilt node can merge, which `Seq(0)` makes it do.
+ val leaf = DummyLeafSafeForKWayMerge(
+ outputPartitioning = KeyedPartitioning(Seq(exprA, exprB), keys))
+ val frozen = Seq(SortOrder(exprA, Ascending), SortOrder(exprB, Ascending))
+ val merging = GroupPartitionsExec(leaf, sortedMergeOrdering = Some(frozen))
+ assert(leaf.outputOrdering.isEmpty, "test setup: the child reports no ordering now")
+ EnsureRequirements.withJoinKeyPositions(merging, Seq(0)) match {
+ case g: GroupPartitionsExec =>
+ assert(g.sortedMergeOrdering.contains(frozen))
+ assert(g.outputOrdering == frozen, "the rebuilt node merges by the frozen ordering")
+ case other => fail(s"expected a rebuilt GroupPartitionsExec, got $other")
+ }
+ }
+
private def anyGpeEnabled(plan: SparkPlan): Boolean =
plan.collectFirst { case gpe: GroupPartitionsExec if gpe.enableSortedMerge => true }.isDefined
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org