Skip to content

[SPARK-59279][SQL] Don't re-read the sorted-merge config after GroupPartitionsExec is planned #58543

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
peter-toth wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[SPARK-59279][SQL] Don't re-read the sorted-merge config after GroupPartitionsExec is planned #58543

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 33 additions & 19 deletions ...re/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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](
Expand All @@ -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) {
Expand All @@ -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
Expand All @@ -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 =>
Expand Down Expand Up @@ -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
}
}

Expand Down
124 changes: 96 additions & 28 deletions sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))"
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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",
Expand All @@ -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")
Expand All @@ -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 =
Expand All @@ -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
Expand Down
Loading