Skip to content

[AURON #2457] Fix per-partition job submission in NativeCollectLimit #2458

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
slfan1989 merged 4 commits into from
Aug 16, 2026
Merged

[AURON #2457] Fix per-partition job submission in NativeCollectLimit #2458

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
6e42143
fix per-partition job submission in NativeCollectLimit
lyne7-sc Aug 9, 2026
73ccf3b
fix flaky CollectLimit suite on Spark 3.0
lyne7-sc Aug 9, 2026
73d8306
strengthen collect limit job submission suite
lyne7-sc Aug 10, 2026
e97ec25
address review comment
lyne7-sc Aug 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions spark-extension-shims-spark/src/test/scala/org/apache/auron/exec/AuronExecSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

@weiqingy weiqingy Aug 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

getJobIdsForGroup reads from AppStatusStore, which is filled by a listener on Spark's async event bus. So it counts jobs whose events have been processed by now, not jobs that were submitted.

The race falls on the unhelpful side. If a JobStart is still queued, the count comes back low, and a low count passes this assertion. On the old one-job-per-partition code, one late event would be enough to make this test go green.

I have not seen it flake, this is just from reading the path. Would draining the bus before the read make it deterministic? AuronAdaptiveQueryExecSuite.scala:109 uses spark.sparkContext.listenerBus.waitUntilEmpty() for the same reason. It is private[spark], so the caller has to sit under org.apache.spark, and AuronQueryTest already does.

@lyne7-sc lyne7-sc Aug 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point. I added sparkContext.listenerBus.waitUntilEmpty() before checking jobCount.

assert(jobCount > 0 && jobCount < numPartitions)
} finally {
spark.sparkContext.clearJobGroup()
}
}
}
}

test("CollectLimit with offset") {
if (AuronTestUtils.isSparkV34OrGreater) {
withTempView("t1") {
Expand Down
4 changes: 4 additions & 0 deletions spark-extension-shims-spark/src/test/scala/org/apache/spark/sql/AuronQueryTest.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,4 +93,8 @@ abstract class AuronQueryTest
true
case _ => false
}

protected def waitUntilListenerBusEmpty(): Unit = {
spark.sparkContext.listenerBus.waitUntilEmpty()
}
}
11 changes: 1 addition & 10 deletions ...ion/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeCollectLimitBase.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand All @@ -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
}

Expand Down
Loading