refine logic

Author: peter-toth

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala

@@ -191,22 +191,21 @@ case class GroupPartitionsExec(
       child.outputOrdering
     } else {
       // Coalescing: multiple input partitions are merged into one output partition. The child's
-      // within-partition ordering is lost due to concatenation, so we re-derive ordering purely
-      // from the key expressions. For example, if two input partitions both belong to the same
-      // output group (same partition key value) and hold [1, 3] and [2, 5] respectively (each
-      // sorted ascending), concatenating them yields [1, 3, 2, 5] which is no longer sorted.
-      // A join may embed multiple `KeyedPartitioning`s (one per join side) within a single
-      // expression tree; they share the same partitionKeys but carry different expressions.
-      // Collect them all and expose each position's equivalent expressions via
-      // `sameOrderExpressions` so the planner can use any of them for ordering checks.
+      // within-partition ordering is lost due to concatenation -- for example, if two input
+      // partitions both belong to the same output group (same partition key value) and hold
+      // [1, 3] and [2, 5] respectively (each sorted ascending), concatenating them yields
+      // [1, 3, 2, 5] which is no longer sorted. Only sort orders over partition key expressions
+      // (which are constant across all merged partitions) remain valid.
       outputPartitioning match {
         case p: Partitioning with Expression if reducers.isEmpty =>
           // Without reducers all merged partitions share the same original key value, so the key
-          // expressions remain constant within the output partition.
+          // expressions remain constant within the output partition. The child's outputOrdering
+          // should already be in sync with the partitioning (either reported by the source or
+          // derived from it in DataSourceV2ScanExecBase), so we only need to keep the sort orders
+          // whose expression is a partition key expression -- all others are lost by concatenation.
           val keyedPartitionings = p.collect { case k: KeyedPartitioning => k }
-          keyedPartitionings.map(_.expressions).transpose.map { exprs =>
-            SortOrder(exprs.head, Ascending, sameOrderExpressions = exprs.tail)
-          }
+          val keyExprs = ExpressionSet(keyedPartitionings.flatMap(_.expressions))
+          child.outputOrdering.filter(order => keyExprs.contains(order.child))
         case _ =>
           // With reducers, merged partitions share only the reduced key, not the original key
           // expressions, which can take different values within the output partition.

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala

@@ -28,6 +28,7 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
 
   private val exprA = AttributeReference("a", IntegerType)()
   private val exprB = AttributeReference("b", IntegerType)()
+  private val exprC = AttributeReference("c", IntegerType)()
 
   private def row(a: Int): InternalRow = InternalRow.fromSeq(Seq(a))
   private def row(a: Int, b: Int): InternalRow = InternalRow.fromSeq(Seq(a, b))
@@ -45,10 +46,12 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
     assert(gpe.outputOrdering === childOrdering)
   }
 
-  test("SPARK-56241: coalescing without reducers returns key-derived ordering") {
+  test("SPARK-56241: coalescing without reducers keeps key-expression orders from child") {
     // Key 1 appears on partitions 0 and 2, causing coalescing.
     val partitionKeys = Seq(row(1), row(2), row(1))
-    val child = DummySparkPlan(outputPartitioning = KeyedPartitioning(Seq(exprA), partitionKeys))
+    val child = DummySparkPlan(
+      outputPartitioning = KeyedPartitioning(Seq(exprA), partitionKeys),
+      outputOrdering = Seq(SortOrder(exprA, Ascending)))
     val gpe = GroupPartitionsExec(child)
 
     assert(!gpe.groupedPartitions.forall(_._2.size <= 1), "expected coalescing")
@@ -59,11 +62,12 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
     assert(ordering.head.sameOrderExpressions.isEmpty)
   }
 
-  test("SPARK-56241: coalescing without reducers returns one SortOrder per key expression") {
+  test("SPARK-56241: coalescing without reducers keeps one SortOrder per key expression") {
     // Multi-key partition: key (1,10) appears on partitions 0 and 2, causing coalescing.
     val partitionKeys = Seq(row(1, 10), row(2, 20), row(1, 10))
     val child = DummySparkPlan(
-      outputPartitioning = KeyedPartitioning(Seq(exprA, exprB), partitionKeys))
+      outputPartitioning = KeyedPartitioning(Seq(exprA, exprB), partitionKeys),
+      outputOrdering = Seq(SortOrder(exprA, Ascending), SortOrder(exprB, Ascending)))
     val gpe = GroupPartitionsExec(child)
 
     assert(!gpe.groupedPartitions.forall(_._2.size <= 1), "expected coalescing")
@@ -75,13 +79,16 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
     assert(ordering(1).sameOrderExpressions.isEmpty)
   }
 
-  test("SPARK-56241: coalescing join case exposes sameOrderExpressions across join sides") {
+  test("SPARK-56241: coalescing join case preserves sameOrderExpressions from child") {
     // PartitioningCollection wraps two KeyedPartitionings (one per join side), sharing the same
-    // partition keys. Key 1 coalesces partitions 0 and 2.
+    // partition keys. Key 1 coalesces partitions 0 and 2. The child (e.g. SortMergeJoinExec)
+    // already carries sameOrderExpressions linking both sides' key expressions.
     val partitionKeys = Seq(row(1), row(2), row(1))
     val leftKP = KeyedPartitioning(Seq(exprA), partitionKeys)
     val rightKP = KeyedPartitioning(Seq(exprB), partitionKeys)
-    val child = DummySparkPlan(outputPartitioning = PartitioningCollection(Seq(leftKP, rightKP)))
+    val child = DummySparkPlan(
+      outputPartitioning = PartitioningCollection(Seq(leftKP, rightKP)),
+      outputOrdering = Seq(SortOrder(exprA, Ascending, sameOrderExpressions = Seq(exprB))))
     val gpe = GroupPartitionsExec(child)
 
     assert(!gpe.groupedPartitions.forall(_._2.size <= 1), "expected coalescing")
@@ -91,6 +98,22 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
     assert(ordering.head.sameOrderExpressions === Seq(exprB))
   }
 
+  test("SPARK-56241: coalescing drops non-key sort orders from child") {
+    // exprA is the partition key; exprC is a non-key sort order the child also reports
+    // (e.g. a secondary sort within each partition). After coalescing, exprC ordering is lost
+    // by concatenation, so only the exprA order should survive.
+    val partitionKeys = Seq(row(1), row(2), row(1))
+    val child = DummySparkPlan(
+      outputPartitioning = KeyedPartitioning(Seq(exprA), partitionKeys),
+      outputOrdering = Seq(SortOrder(exprA, Ascending), SortOrder(exprC, Ascending)))
+    val gpe = GroupPartitionsExec(child)
+
+    assert(!gpe.groupedPartitions.forall(_._2.size <= 1), "expected coalescing")
+    val ordering = gpe.outputOrdering
+    assert(ordering.length === 1)
+    assert(ordering.head.child === exprA)
+  }
+
   test("SPARK-56241: coalescing with reducers returns empty ordering") {
     // When reducers are present, the original key expressions are not constant within the merged
     // partition, so outputOrdering falls back to the default (empty).