This is an automated email from the ASF dual-hosted git repository.
wenchen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 01464b9ec1ca [SPARK-50124][SQL] LIMIT/OFFSET should preserve data
ordering
01464b9ec1ca is described below
commit 01464b9ec1ca7baca96b1bea0fbb238c25d0eafb
Author: Wenchen Fan <[email protected]>
AuthorDate: Wed Oct 30 10:17:50 2024 -0700
[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 #48661 from cloud-fan/limit.
Authored-by: Wenchen Fan <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
---
.../org/apache/spark/sql/internal/SQLConf.scala | 10 ++
.../execution/InsertSortForLimitAndOffset.scala | 71 ++++++++++++
.../spark/sql/execution/QueryExecution.scala | 2 +
.../execution/adaptive/AdaptiveSparkPlanExec.scala | 2 +
.../InsertSortForLimitAndOffsetSuite.scala | 120 +++++++++++++++++++++
5 files changed, 205 insertions(+)
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index 1c9f5e85d1a0..50f7143145a9 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/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.
*
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala
new file mode 100644
index 000000000000..6c7a9206a8e3
--- /dev/null
+++
b/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
+ }
+ }
+}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
index 6ff2c5d4b9d3..aad905256061 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
+++
b/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,
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
index 77efc4793359..a0a099142930 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala
+++
b/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,
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffsetSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffsetSuite.scala
new file mode 100644
index 000000000000..8d640a1840f4
--- /dev/null
+++
b/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))
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
