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

[lyne7-sc]

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?

• Yes
• No

If yes, include: Generated-by: GPT-5

ASF guidance: https://www.apache.org/legal/generative-tooling.html

[weiqingy]

Thanks for taking this on. A few questions inline.

[weiqingy]

spark.sql.limit.initialNumPartitions was added in Spark 3.4, with a default of 1, and SparkPlan.executeTake only reads it from 3.4 on. So on the spark-3.0 to 3.3 profiles this line is set but never read.

On 3.4 and later it is read, and 100 is larger than the partition count, so the first runJob takes every partition. That run is a single job with no batching. Only the older profiles go through the scale-up loop.

The regression is caught either way, since jobCount < numPartitions holds in both shapes. Is there a reason to pin the value here? Without it the 3.4+ default of 1 applies, and every profile would exercise the batching this test is named for.

[lyne7-sc]

Good catch! I changed initialNumPartitions to 1 so the test exercises incremental partition scanning and distinguishes executeTake from the previous toLocalIterator implementation.

[weiqingy]

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]

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