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))