From 06b12e01b5aa3b9c6963f3d08b635c7561a137c6 Mon Sep 17 00:00:00 2001 From: David Mollitor Date: Wed, 16 Sep 2026 16:04:26 +0000 Subject: [PATCH 1/4] [SPARK-59581][SQL] Avoid copying every input row in TakeOrderedAndProject top-N ### What changes were proposed in this pull request? `TakeOrderedAndProjectExec` implements top-N (`ORDER BY ... LIMIT n`). For unsorted input it selects the top-K with `Utils.takeOrdered(iter.map(_.copy()), limit)(ord)`. The `iter.map(_.copy())` copies every input row, even though only `limit` rows are ever retained. The copy is forced by the selector: `org.apache.spark.util.collection.Utils.takeOrdered` wraps Guava's `Ordering.leastOf` / `TopKSelector`, which buffers element references while selecting, and the child's whole-stage-codegen iterator yields a single reused `UnsafeRow` -- so every row must be copied to a detached instance before being offered. For N input rows we allocate N copies to keep `limit`. This replaces the copy-everything-then-select approach with a bounded top-K that copies a row only when it is actually retained. A helper on the `TakeOrderedAndProjectExec` companion builds a bounded max-heap (`java.util.PriorityQueue` with `ord.reverse`, so the head is the current largest of the retained set = the eviction threshold), compares each live row against the head, and calls `copy()` only on insertion: ```scala if (heap.size < num) heap.add(row.copy()) else if (ord.compare(row, heap.peek()) < 0) { heap.poll(); heap.add(row.copy()) } ``` Rows are drained into ascending order to match the sorted output of the previous `Utils.takeOrdered`. This is applied to the three copy-all sites (the `executeCollect` unsorted branch, the `doExecute` per-partition local top-K, and the `doExecute` post-shuffle merge). The already-efficient `orderingSatisfies` fast paths (`_.map(_.copy()).take(limit)`) and the post-limit OFFSET/projection are unchanged. Copies drop from O(N) to O(limit). A per-element heap decision against the exact current threshold enables copy-on-retain, which Guava's reference-buffering design cannot. ### Why are the changes needed? `ORDER BY ... LIMIT n` over large unsorted input is common (top-N, "latest 100", dashboards). JFR profiling of `spark.range(20000000).selectExpr("id","id % 1000 as k").orderBy("k").limit(100)` showed `UnsafeRow.copy()` at ~93.7% of sampled allocation and ~15.5% of CPU -- the operator's cost is almost entirely copying rows that are immediately discarded. Cutting copies from O(N) to O(limit) removes that. Before/after on that query (local): `UnsafeRow.copy()` allocation samples dropped 2532 -> 4, and wall time roughly halved (noop 671 -> 343 ms, collect 650 -> 376 ms). ### Does this PR introduce _any_ user-facing change? No. Both old and new return the `limit` smallest rows by the ordering, in ascending order; which rows win a tie at the exact boundary value is unspecified in both (inherent LIMIT non-determinism). `limit`/OFFSET semantics are identical (the same `limit` is passed and the post-limit `.drop(offset)` is untouched). Worst-case memory is unchanged: `limit` is bounded by `spark.sql.execution.topKSortFallbackThreshold`, and the heap holds at most `limit` rows. ### How was this patch tested? `TakeOrderedAndProjectSuite` (with/without projection, 0/1/10 partitions, 0 and 10k rows, the already-sorted path via both `executeCollect` and `doExecute`), `PlannerSuite` (including the `topKSortFallbackThreshold` and `limit + offset` threshold planner tests), and `InsertSortForLimitAndOffsetSuite` all pass unchanged. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Isaac Co-authored-by: Isaac --- .../apache/spark/sql/execution/limit.scala | 48 +++++++++++++++++-- 1 file changed, 44 insertions(+), 4 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/limit.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/limit.scala index c0fb1c37b2102..855bd416bdb4e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/limit.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/limit.scala @@ -28,7 +28,6 @@ import org.apache.spark.sql.catalyst.util.truncatedString import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.execution.metric.{SQLShuffleReadMetricsReporter, SQLShuffleWriteMetricsReporter} import org.apache.spark.sql.execution.python.HybridRowQueue -import org.apache.spark.util.collection.Utils /** * The operator takes limited number of elements from its child operator. @@ -324,7 +323,9 @@ case class TakeOrderedAndProjectExec( val limited = if (orderingSatisfies) { child.execute().mapPartitionsInternal(_.map(_.copy()).take(limit)).takeOrdered(limit)(ord) } else { - child.execute().mapPartitionsInternal(_.map(_.copy())).takeOrdered(limit)(ord) + child.execute().mapPartitionsInternal { iter => + TakeOrderedAndProjectExec.takeOrderedByCopyOnRetain(iter, limit, ord) + }.takeOrdered(limit)(ord) } val data = if (offset > 0) limited.drop(offset) else limited if (projectList != child.output) { @@ -358,7 +359,7 @@ case class TakeOrderedAndProjectExec( childRDD.mapPartitionsInternal(_.map(_.copy()).take(limit)) } else { childRDD.mapPartitionsInternal { iter => - Utils.takeOrdered(iter.map(_.copy()), limit)(ord) + TakeOrderedAndProjectExec.takeOrderedByCopyOnRetain(iter, limit, ord) } } @@ -372,7 +373,7 @@ case class TakeOrderedAndProjectExec( readMetrics) } singlePartitionRDD.mapPartitionsWithIndexInternal { (idx, iter) => - val limited = Utils.takeOrdered(iter.map(_.copy()), limit)(ord) + val limited = TakeOrderedAndProjectExec.takeOrderedByCopyOnRetain(iter, limit, ord) val topK = if (offset > 0) limited.drop(offset) else limited if (projectList != child.output) { val proj = UnsafeProjection.create(projectList, child.output) @@ -409,3 +410,42 @@ case class TakeOrderedAndProjectExec( override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = copy(child = newChild) } + +object TakeOrderedAndProjectExec { + /** + * Returns the `num` smallest rows of `input` by `ord`, in ascending order. Unlike + * `Utils.takeOrdered(input.map(_.copy()), num)`, this copies a row only when it is actually + * retained in the bounded top-K, rather than copying every input row up front. This matters + * because `input` typically yields a single reused `UnsafeRow`, so only the retained rows need + * a detached copy. + */ + private[execution] def takeOrderedByCopyOnRetain( + input: Iterator[InternalRow], + num: Int, + ord: Ordering[InternalRow]): Iterator[InternalRow] = { + if (num <= 0) { + return Iterator.empty + } + // Max-heap by `ord` (via ord.reverse): the head is the largest of the retained rows, i.e. the + // eviction threshold. A row is copied only when it enters the heap. + val heap = + new java.util.PriorityQueue[InternalRow](math.max(1, math.min(num, 1024)), ord.reverse) + while (input.hasNext) { + val row = input.next() + if (heap.size < num) { + heap.add(row.copy()) + } else if (ord.compare(row, heap.peek()) < 0) { + heap.poll() + heap.add(row.copy()) + } + } + // Drain into ascending order: poll yields the largest remaining row first under ord.reverse. + val result = new Array[InternalRow](heap.size) + var i = result.length + while (i > 0) { + i -= 1 + result(i) = heap.poll() + } + result.iterator + } +} From e0884eeb95d95b3cfc655e0a5cb80c3908adfa5a Mon Sep 17 00:00:00 2001 From: david-mollitor-db Date: Wed, 16 Sep 2026 16:32:19 +0000 Subject: [PATCH 2/4] Benchmark results for org.apache.spark.sql.execution.benchmark.TakeOrderedAndProjectBenchmark (JDK 17, Scala 2.13, split 1 of 1) --- .../benchmarks/TakeOrderedAndProjectBenchmark-results.txt | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/sql/core/benchmarks/TakeOrderedAndProjectBenchmark-results.txt b/sql/core/benchmarks/TakeOrderedAndProjectBenchmark-results.txt index e1db8b2a8c618..ac060ddbc5253 100644 --- a/sql/core/benchmarks/TakeOrderedAndProjectBenchmark-results.txt +++ b/sql/core/benchmarks/TakeOrderedAndProjectBenchmark-results.txt @@ -2,11 +2,11 @@ TakeOrderedAndProject ================================================================================================ -OpenJDK 64-Bit Server VM 17.0.20+8-LTS on Linux 6.17.0-1020-azure -AMD EPYC 9V74 80-Core Processor +OpenJDK 64-Bit Server VM 17.0.20+8-LTS on Linux 6.17.0-1022-azure +INTEL(R) XEON(R) PLATINUM 8573C TakeOrderedAndProject with SMJ: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative --------------------------------------------------------------------------------------------------------------------------------- -TakeOrderedAndProject with SMJ for doExecute 339 352 13 0.0 33898.4 1.0X -TakeOrderedAndProject with SMJ for executeCollect 150 192 37 0.1 15014.3 2.3X +TakeOrderedAndProject with SMJ for doExecute 161 175 13 0.1 16079.4 1.0X +TakeOrderedAndProject with SMJ for executeCollect 91 94 4 0.1 9069.2 1.8X From f55abc9b86ab5912c0f78bf08ee475bd6e111a4e Mon Sep 17 00:00:00 2001 From: david-mollitor-db Date: Wed, 16 Sep 2026 17:03:33 +0000 Subject: [PATCH 3/4] Benchmark results for org.apache.spark.sql.execution.benchmark.TakeOrderedAndProjectBenchmark (JDK 21, Scala 2.13, split 1 of 1) --- .../TakeOrderedAndProjectBenchmark-jdk21-results.txt | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/sql/core/benchmarks/TakeOrderedAndProjectBenchmark-jdk21-results.txt b/sql/core/benchmarks/TakeOrderedAndProjectBenchmark-jdk21-results.txt index 4c0d7446eccfb..7bc203300686d 100644 --- a/sql/core/benchmarks/TakeOrderedAndProjectBenchmark-jdk21-results.txt +++ b/sql/core/benchmarks/TakeOrderedAndProjectBenchmark-jdk21-results.txt @@ -2,11 +2,11 @@ TakeOrderedAndProject ================================================================================================ -OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1020-azure -AMD EPYC 9V74 80-Core Processor +OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1022-azure +INTEL(R) XEON(R) PLATINUM 8573C TakeOrderedAndProject with SMJ: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative --------------------------------------------------------------------------------------------------------------------------------- -TakeOrderedAndProject with SMJ for doExecute 220 232 14 0.0 22030.1 1.0X -TakeOrderedAndProject with SMJ for executeCollect 120 125 6 0.1 12031.5 1.8X +TakeOrderedAndProject with SMJ for doExecute 208 221 22 0.0 20789.0 1.0X +TakeOrderedAndProject with SMJ for executeCollect 82 99 18 0.1 8173.7 2.5X From 66f4a4c1b6af7ebe970411045919a648eaebc2a6 Mon Sep 17 00:00:00 2001 From: david-mollitor-db Date: Wed, 16 Sep 2026 17:17:22 +0000 Subject: [PATCH 4/4] Benchmark results for org.apache.spark.sql.execution.benchmark.TakeOrderedAndProjectBenchmark (JDK 25, Scala 2.13, split 1 of 1) --- .../TakeOrderedAndProjectBenchmark-jdk25-results.txt | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/sql/core/benchmarks/TakeOrderedAndProjectBenchmark-jdk25-results.txt b/sql/core/benchmarks/TakeOrderedAndProjectBenchmark-jdk25-results.txt index 191b1abe619df..b03ab6451b198 100644 --- a/sql/core/benchmarks/TakeOrderedAndProjectBenchmark-jdk25-results.txt +++ b/sql/core/benchmarks/TakeOrderedAndProjectBenchmark-jdk25-results.txt @@ -2,11 +2,11 @@ TakeOrderedAndProject ================================================================================================ -OpenJDK 64-Bit Server VM 25.0.4+7-LTS on Linux 6.17.0-1020-azure -Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz +OpenJDK 64-Bit Server VM 25.0.4+7-LTS on Linux 6.17.0-1022-azure +AMD EPYC 9V74 80-Core Processor TakeOrderedAndProject with SMJ: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative --------------------------------------------------------------------------------------------------------------------------------- -TakeOrderedAndProject with SMJ for doExecute 227 246 21 0.0 22688.5 1.0X -TakeOrderedAndProject with SMJ for executeCollect 102 111 10 0.1 10207.1 2.2X +TakeOrderedAndProject with SMJ for doExecute 236 263 24 0.0 23596.7 1.0X +TakeOrderedAndProject with SMJ for executeCollect 130 151 23 0.1 12990.8 1.8X