[SPARK-50124][SQL] LIMIT/OFFSET should preserve data ordering

### What changes were proposed in this pull request?

This PR fixes a long-standing correctness bug for LIMIT/OFFSET. Spark has 3 execution paths for LIMIT:
1. If LIMIT is the root node, use `CollectLimitExec` to eliminate the shuffle and get the first N records on the driver side.
2. If there is a global sort under LIMIT, use `TakeOrderedAndProjectExec` to avoid a full sort using top-N algorithm
3. Otherwise, use `GlobalLimitExec` to get the first N records via a single-partition shuffle.

The third execution path has a problem: Spark shuffle reader fetches shuffle blocks in random order and can't preserve the data ordering. It's usually OK when the data order doesn't matter, but the second execution path is not always picked because it has a trigger condition: the LIMIT N must be smaller than TOP_K_SORT_FALLBACK_THRESHOLD .

This bug is more likely to happen for OFFSET because it doesn't have the top-K optimization.

An ideal solution is to have an order preserving mode for the shuffle reader, which is non-trivial work. This PR proposes a workaround: add a local sort after the single-partition shuffle to restore the data ordering.

### Why are the changes needed?

correctness fix

### Does this PR introduce _any_ user-facing change?

Yes, now LIMIT/OFFSET with global sort can preserve the data ordering

### How was this patch tested?

new tests

### Was this patch authored or co-authored using generative AI tooling?

no

Closes apache#48661 from cloud-fan/limit.

Authored-by: Wenchen Fan <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -5207,6 +5207,16 @@ object SQLConf {
     .booleanConf
     .createWithDefault(false)
 
+  val ORDERING_AWARE_LIMIT_OFFSET = buildConf("spark.sql.orderingAwareLimitOffset")
+    .internal()
+    .doc("When set to true, a local sort will be inserted between GlobalLimitExec and " +
+      "single-partition ShuffleExchangeExec, if the underlying plan produces sorted data. " +
+      "This is because shuffle reader in Spark fetches shuffle blocks in a random order and " +
+      "can not preserve the data ordering, while LIMIT/OFFSET must preserve ordering.")
+    .version("4.0.0")
+    .booleanConf
+    .createWithDefault(true)
+
   /**
    * Holds information about keys that have been deprecated.
    *

sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala

@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution
+
+import org.apache.spark.sql.catalyst.expressions.SortOrder
+import org.apache.spark.sql.catalyst.plans.physical.SinglePartition
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.adaptive.{AQEShuffleReadExec, ShuffleQueryStageExec}
+import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * When LIMIT/OFFSET is the root node, Spark plans it as CollectLimitExec which preserves the data
+ * ordering. However, when OFFSET/LIMIT is not the root node, Spark uses GlobalLimitExec which
+ * shuffles all the data into one partition and then gets a slice of it. Unfortunately, the shuffle
+ * reader fetches shuffle blocks in a random order and can not preserve the data ordering, which
+ * violates the requirement of LIMIT/OFFSET.
+ *
+ * This rule inserts an extra local sort before LIMIT/OFFSET to preserve the data ordering.
+ * TODO: add a order preserving mode in the shuffle reader.
+ */
+object InsertSortForLimitAndOffset extends Rule[SparkPlan] {
+  override def apply(plan: SparkPlan): SparkPlan = {
+    if (!conf.getConf(SQLConf.ORDERING_AWARE_LIMIT_OFFSET)) return plan
+
+    plan transform {
+      case l @ GlobalLimitExec(
+          _,
+          SinglePartitionShuffleWithGlobalOrdering(ordering),
+          _) =>
+        val newChild = SortExec(ordering, global = false, child = l.child)
+        l.withNewChildren(Seq(newChild))
+    }
+  }
+
+  object SinglePartitionShuffleWithGlobalOrdering {
+    def unapply(plan: SparkPlan): Option[Seq[SortOrder]] = plan match {
+      case ShuffleExchangeExec(SinglePartition, SparkPlanWithGlobalOrdering(ordering), _, _) =>
+        Some(ordering)
+      case p: AQEShuffleReadExec => unapply(p.child)
+      case p: ShuffleQueryStageExec => unapply(p.plan)
+      case _ => None
+    }
+  }
+
+  // Note: this is not implementing a generalized notion of "global order preservation", but just
+  // tackles the regular ORDER BY semantics with optional LIMIT (top-K).
+  object SparkPlanWithGlobalOrdering {
+    def unapply(plan: SparkPlan): Option[Seq[SortOrder]] = plan match {
+      case p: SortExec if p.global => Some(p.sortOrder)
+      case p: LocalLimitExec => unapply(p.child)
+      case p: WholeStageCodegenExec => unapply(p.child)
+      case _ => None
+    }
+  }
+}

sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala

@@ -518,6 +518,8 @@ object QueryExecution {
       PlanSubqueries(sparkSession),
       RemoveRedundantProjects,
       EnsureRequirements(),
+      // This rule must be run after `EnsureRequirements`.
+      InsertSortForLimitAndOffset,
       // `ReplaceHashWithSortAgg` needs to be added after `EnsureRequirements` to guarantee the
       // sort order of each node is checked to be valid.
       ReplaceHashWithSortAgg,

sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala

@@ -145,6 +145,8 @@ case class AdaptiveSparkPlanExec(
       CoalesceBucketsInJoin,
       RemoveRedundantProjects,
       ensureRequirements,
+      // This rule must be run after `EnsureRequirements`.
+      InsertSortForLimitAndOffset,
       AdjustShuffleExchangePosition,
       ValidateSparkPlan,
       ReplaceHashWithSortAgg,

sql/core/src/test/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffsetSuite.scala

@@ -0,0 +1,120 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution
+
+import org.apache.spark.sql.QueryTest
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+
+class InsertSortForLimitAndOffsetSuite extends QueryTest
+  with SharedSparkSession
+  with AdaptiveSparkPlanHelper {
+  import testImplicits._
+
+  private def assertHasTopKSort(plan: SparkPlan): Unit = {
+    assert(find(plan) {
+      case _: TakeOrderedAndProjectExec => true
+      case _ => false
+    }.isDefined)
+  }
+
+  private def assertHasCollectLimitExec(plan: SparkPlan): Unit = {
+    assert(find(plan) {
+      case _: CollectLimitExec => true
+      case _ => false
+    }.isDefined)
+  }
+
+  private def assertHasGlobalLimitExec(plan: SparkPlan): Unit = {
+    assert(find(plan) {
+      case _: GlobalLimitExec => true
+      case _ => false
+    }.isDefined)
+  }
+
+  private def hasLocalSort(plan: SparkPlan): Boolean = {
+    find(plan) {
+      case GlobalLimitExec(_, s: SortExec, _) => !s.global
+      case _ => false
+    }.isDefined
+  }
+
+  test("root LIMIT preserves data ordering with top-K sort") {
+    val df = spark.range(10).orderBy($"id" % 8).limit(2)
+    df.collect()
+    val physicalPlan = df.queryExecution.executedPlan
+    assertHasTopKSort(physicalPlan)
+    // Extra local sort is not needed for LIMIT with top-K sort optimization.
+    assert(!hasLocalSort(physicalPlan))
+  }
+
+  test("middle LIMIT preserves data ordering with top-K sort") {
+    val df = spark.range(10).orderBy($"id" % 8).limit(2).distinct()
+    df.collect()
+    val physicalPlan = df.queryExecution.executedPlan
+    assertHasTopKSort(physicalPlan)
+    // Extra local sort is not needed for LIMIT with top-K sort optimization.
+    assert(!hasLocalSort(physicalPlan))
+  }
+
+  test("root LIMIT preserves data ordering with CollectLimitExec") {
+    withSQLConf(SQLConf.TOP_K_SORT_FALLBACK_THRESHOLD.key -> "1") {
+      val df = spark.range(10).orderBy($"id" % 8).limit(2)
+      df.collect()
+      val physicalPlan = df.queryExecution.executedPlan
+      assertHasCollectLimitExec(physicalPlan)
+      // Extra local sort is not needed for root LIMIT
+      assert(!hasLocalSort(physicalPlan))
+    }
+  }
+
+  test("middle LIMIT preserves data ordering with the extra sort") {
+    withSQLConf(
+      SQLConf.TOP_K_SORT_FALLBACK_THRESHOLD.key -> "1",
+      // To trigger the bug, we have to disable the coalescing optimization. Otherwise we use only
+      // one partition to read the range-partition shuffle and there is only one shuffle block for
+      // the final single-partition shuffle, random fetch order is no longer an issue.
+      SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false") {
+      val df = spark.range(10).orderBy($"id" % 8).limit(2).distinct()
+      df.collect()
+      val physicalPlan = df.queryExecution.executedPlan
+      assertHasGlobalLimitExec(physicalPlan)
+      // Extra local sort is needed for middle LIMIT
+      assert(hasLocalSort(physicalPlan))
+    }
+  }
+
+  test("root OFFSET preserves data ordering with CollectLimitExec") {
+    val df = spark.range(10).orderBy($"id" % 8).offset(2)
+    df.collect()
+    val physicalPlan = df.queryExecution.executedPlan
+    assertHasCollectLimitExec(physicalPlan)
+    // Extra local sort is not needed for root OFFSET
+    assert(!hasLocalSort(physicalPlan))
+  }
+
+  test("middle OFFSET preserves data ordering with the extra sort") {
+    val df = spark.range(10).orderBy($"id" % 8).offset(2).distinct()
+    df.collect()
+    val physicalPlan = df.queryExecution.executedPlan
+    assertHasGlobalLimitExec(physicalPlan)
+    // Extra local sort is needed for middle OFFSET
+    assert(hasLocalSort(physicalPlan))
+  }
+}