You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@auron.apache.org by sl...@apache.org on 2026/08/16 02:46:52 UTC
(auron) branch master updated: [AURON #2457] Fix per-partition job submission in NativeCollectLimit (#2458)
This is an automated email from the ASF dual-hosted git repository.
slfan1989 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/auron.git
The following commit(s) were added to refs/heads/master by this push:
new 3ac177910 [AURON #2457] Fix per-partition job submission in NativeCollectLimit (#2458)
3ac177910 is described below
commit 3ac1779100f38d8d2e145dac2324b2bf48481d61
Author: linfeng <33...@users.noreply.github.com>
AuthorDate: Sun Aug 16 10:46:46 2026 +0800
[AURON #2457] Fix per-partition job submission in NativeCollectLimit (#2458)
# Which issue does this PR close?
Closes #2457
# Rationale for this change
`NativeCollectLimitBase.executeCollect()` currently uses
`toLocalIterator`, which submits a separate Spark job for every scanned
partition.
For empty or highly selective inputs, this can produce many unnecessary
jobs and make driver-side scheduling dominate query execution.
# What changes are included in this PR?
- Replace the `toLocalIterator` collection loop with `executeTake`.
- Add a regression test to verify that partitions are scanned in
batches.
# Are there any user-facing changes?
Queries over empty or sparse partitions may submit fewer Spark jobs.
# How was this patch tested?
Added a regression test to `AuronExecSuite`.
# Was this patch authored or co-authored using generative AI tooling?
- [x] Yes
- [ ] No
If yes, include: `Generated-by: GPT-5`
ASF guidance: https://www.apache.org/legal/generative-tooling.html
Co-authored-by: Shilun Fan <sl...@apache.org>
Signed-off-by: Shilun Fan <sl...@apache.org>
---
.../org/apache/auron/exec/AuronExecSuite.scala | 30 ++++++++++++++++++++++
.../org/apache/spark/sql/AuronQueryTest.scala | 4 +++
.../auron/plan/NativeCollectLimitBase.scala | 11 +-------
3 files changed, 35 insertions(+), 10 deletions(-)
diff --git a/spark-extension-shims-spark/src/test/scala/org/apache/auron/exec/AuronExecSuite.scala b/spark-extension-shims-spark/src/test/scala/org/apache/auron/exec/AuronExecSuite.scala
index eac82b493..0a6f36255 100644
--- a/spark-extension-shims-spark/src/test/scala/org/apache/auron/exec/AuronExecSuite.scala
+++ b/spark-extension-shims-spark/src/test/scala/org/apache/auron/exec/AuronExecSuite.scala
@@ -24,6 +24,7 @@ import org.apache.spark.sql.catalyst.expressions.WindowExpression
import org.apache.spark.sql.catalyst.expressions.WindowSpecDefinition
import org.apache.spark.sql.execution.auron.plan.{NativeCollectLimitExec, NativeGlobalLimitExec, NativeLocalLimitExec, NativeTakeOrderedExec}
import org.apache.spark.sql.execution.auron.plan.NativeWindowExec
+import org.apache.spark.sql.internal.SQLConf
import org.apache.auron.BaseAuronSQLSuite
import org.apache.auron.util.AuronTestUtils
@@ -43,6 +44,35 @@ class AuronExecSuite extends AuronQueryTest with BaseAuronSQLSuite {
}
}
+ test("CollectLimit batches partition scans") {
+ withTempPath { path =>
+ spark.range(0, 80, 1, 8).write.parquet(path.getCanonicalPath)
+
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ "spark.sql.files.maxPartitionBytes" -> "4096",
+ "spark.sql.limit.initialNumPartitions" -> "1") {
+ val df = spark.read.parquet(path.getCanonicalPath).where("id < 0").limit(1)
+ val collectLimit = collectFirst(df.queryExecution.executedPlan) {
+ case exec: NativeCollectLimitExec => exec
+ }.get
+ val numPartitions = collectLimit.child.execute().getNumPartitions
+ assert(numPartitions > 1)
+
+ val jobGroup = s"collect-limit-${System.nanoTime()}"
+ spark.sparkContext.setJobGroup(jobGroup, "test CollectLimit job count")
+ try {
+ assert(collectLimit.executeCollect().isEmpty)
+ waitUntilListenerBusEmpty()
+ val jobCount = spark.sparkContext.statusTracker.getJobIdsForGroup(jobGroup).length
+ assert(jobCount > 0 && jobCount < numPartitions)
+ } finally {
+ spark.sparkContext.clearJobGroup()
+ }
+ }
+ }
+ }
+
test("CollectLimit with offset") {
if (AuronTestUtils.isSparkV34OrGreater) {
withTempView("t1") {
diff --git a/spark-extension-shims-spark/src/test/scala/org/apache/spark/sql/AuronQueryTest.scala b/spark-extension-shims-spark/src/test/scala/org/apache/spark/sql/AuronQueryTest.scala
index 263f06231..5540b38e3 100644
--- a/spark-extension-shims-spark/src/test/scala/org/apache/spark/sql/AuronQueryTest.scala
+++ b/spark-extension-shims-spark/src/test/scala/org/apache/spark/sql/AuronQueryTest.scala
@@ -93,4 +93,8 @@ abstract class AuronQueryTest
true
case _ => false
}
+
+ protected def waitUntilListenerBusEmpty(): Unit = {
+ spark.sparkContext.listenerBus.waitUntilEmpty()
+ }
}
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeCollectLimitBase.scala b/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeCollectLimitBase.scala
index d33661030..4a226fc65 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeCollectLimitBase.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeCollectLimitBase.scala
@@ -17,7 +17,6 @@
package org.apache.spark.sql.execution.auron.plan
import scala.collection.mutable
-import scala.collection.mutable.ArrayBuffer
import org.apache.spark.OneToOneDependency
import org.apache.spark.sql.auron.{NativeHelper, NativeRDD, NativeSupports, Shims}
@@ -43,15 +42,7 @@ abstract class NativeCollectLimitBase(limit: Int, offset: Int, override val chil
override def executeCollect(): Array[InternalRow] = {
val partial = Shims.get.createNativeLocalLimitExec(limit, child)
- val buf = new ArrayBuffer[InternalRow]
-
- // collect rows partition-by-partition up to 'limit', avoiding full-partition collect.
- val it = partial.execute().toLocalIterator
- while (buf.size < limit && it.hasNext) {
- val row = it.next().copy()
- buf += row
- }
- val rows = buf.toArray
+ val rows = partial.executeTake(limit)
if (offset > 0) rows.drop(offset) else rows
}