[viirya]

Which issue does this PR close?

Part of #6499.

Rationale for this change

This is the first part of the split proposed in #6499 for local execution: running a whole query as one DataFusion graph when the application has a local master. It adds the execution lifecycle, the JNI bridge and result delivery, admitting only native blocks whose leaf is a range, so the mechanism can be reviewed before admission grows. It follows the direction agreed in #6499: the graph is built from the operator tree CometExecRule produces and lowered with PhysicalPlanner::create_plan, and the session and memory pool come from create_execution_context.

What changes are included in this PR?

• spark.comet.exec.local.enabled (internal, default false). CometRule offers the converted plan to CometLocalRule, which admits it only as one native block with no JVM input whose only leaf is a CometRangeExec, with deterministic expressions and no subquery or JVM callback. A root CollectLimitExec is kept. These plans have no exchange, so AQE does not wrap them and they are admitted with AQE enabled; there is no Spark version gate.
• The result task creates one graph for all partitions. create_execution_context provides the session and memory pool, so the pool acquires memory through the task's CometTaskMemoryManager and spills to Spark's local directories. Each partition's operators are created by the planner for that partition, so range slices and partition-dependent state match ordinary Comet execution, and RangeScan's mayOverflow keeps overflowing ranges on Spark.
• native/local runs the partitions concurrently and delivers them in partition order, so results keep Spark's row order and a limit returns the same rows as Spark. Closing a query drops its streams on the calling thread, as BatchProducer::stop does.
• CometLocalExec reports UnknownPartitioning(1), so a query over a cached local result still gets its required exchange. CometLocalResultExec hands collect and take results to the driver within the JVM and enforces spark.driver.maxResultSize while copying.
• ExecutionContext::planner now builds the per-partition planner for both execute_plan and local queries.

How are these changes tested?

• native/local tests cover partition-order delivery, early failure from a later partition, cancellation while every Tokio worker is busy, the one-batch handoff and runtime shutdown.
• CometLocalExecutionSuite (17 tests, passing on Spark 3.4, 3.5, 4.0, 4.1 and 4.2) compares results and order with Spark, including ranges near the long boundaries and ranges Spark overflows, and covers declined expressions, AQE on and off, limits, a query over a cached local result, the result size limit, job and task cancellation and native errors. After each test it checks that no native query, imported Arrow memory or result slot is left.

[sunchao]

Reviewed the complete 29-file diff from cf6673808dc92d0be7ed331431fe9a2e97067a7d to cab738b3f24e2b01b0e5c8ca2bc75919cb09a78d. The PR is not a draft. The snapshot and live discussion contained no existing human findings. Found three introduced P2 issues.

Summary

• Prior state and problem: Local applications used separate Spark tasks for native partitions and Spark’s serialized result delivery. This PR introduces a query-wide native execution lifecycle.
• Design approach: An internal, default-off option admits exchange-free native blocks rooted in a range. One Spark task runs their partitions concurrently and receives results in partition order.
• Correctness: Two regressions reproduced. Cached partition-local sorting produces incorrect window ranks, and a limited query fails on an error in a partition Spark never requests. The existing 17-test suite passes but does not cover these cases.
• Compatibility analysis: Checked relevant Spark sources for 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0. Their ordering and window distribution contracts expose the metadata problem. Compared with branch-1.1, local graph execution is new and remains disabled by default.
• Key design decisions: Numeric handles, shared task memory, partition-ordered delivery and UnknownPartitioning(1) are sensible choices. However, concatenating partitions does not preserve their individual sort ordering, and immediate cross-partition error delivery changes take semantics.
• Implementation sketch: CometLocalRule selects a native block. LocalQueryRDD creates its JNI handle, the shared planner builds each partition, and native/local feeds Arrow batches to the result task and driver handoff.
• Performance: The handoff avoids Spark’s result encoding and compression. Each source partition still receives its own plan and up to two buffered batches. No benchmark establishes the net improvement, and no separate measured P1/P2 performance defect was identified.
• Design: Restricting admission and reusing the existing planner/session keeps this first implementation understandable. Ordering guarantees and limited-query failure handling need correction before this execution path is reliable.
• Abstraction & complexity: Separating native lifecycle management from JNI is useful. Extracting ExecutionContext::planner avoids divergent planner setup. No additional abstraction-related issue met the P1/P2 bar.
• Behavioral changes worth calling out: Opted-in queries use one Spark result partition, expose delivery metrics instead of native operator metrics, and enforce spark.driver.maxResultSize using uncompressed row bytes. Required Rust CI also fails on the newly added deprecated atomic call.
• Suggested improvements: Correct the advertised ordering, preserve Spark’s limited-query error behavior, and replace or compatibly handle the deprecated atomic operation. Add the two reproduced semantic cases as regression tests.

Routed skills: review-comet-pr, review-comet-expression-pr, review-comet-ffi-pr and review-comet-memory-pr. Read repository AGENTS.md and the relevant contributor documentation.

Exact-head CI: Linux execution, expression, scan and shuffle suites, TPC-H/TPC-DS verification, native build, and Spark 4.1 SQL suites passed. ubuntu-latest/rust-test and aggregate Required Checks failed because Clippy rejects fetch_update. Rust tests consequently did not run in that job. Other Spark SQL versions and macOS were skipped.

Validation: Compiled the pinned JVM sources with JDK 17. Verified the native CI artifact against GitHub’s SHA-256 digest for this head, then ran CometLocalExecutionSuite: 17/17 passed. Both additional reproductions used these exact-head components on Spark 4.1.3 and compared against Spark and ordinary Comet. Handles, imported Arrow memory and result slots returned to zero. CI-routing checks passed. Local Rust tests could not start because the configured dependency mirror lacks locked DataFusion 55.2.0. Cross-version runtime, spill-stress and performance testing were not performed. The tracked checkout remains unchanged.

[sunchao]

[P2] Clear partition-local ordering when concatenating multiple partitions. With local and range execution enabled, caching spark.range(0, 20, 1, 2).sortWithinPartitions(desc("id")) stores 9..0, 19..10 in one partition, but this method advertises that partition as descending-sorted. InMemoryTableScanExec preserves that ordering, and UnknownPartitioning(1) satisfies an unpartitioned window’s AllTuples requirement. Consequently, row_number().over(Window.orderBy(desc("id"))) skips the necessary sort and assigns id=19 rank 11 instead of 1. This reproduces even with Comet disabled for the consuming window. Return no ordering for multi-partition concatenation unless ordering across partition boundaries is proven, or merge the streams by their sort keys.

Evidence: Reproduced with exact-head JVM classes and verified native CI artifact on Spark 4.1.3/JDK 17, local[2], AQE disabled. Materialized the descending range cache, disabled Comet, then selected id and row_number().over(Window.orderBy(desc("id"))). Spark and ordinary Comet returned (19,1) and (9,11). Local execution returned (19,11) and (9,1). The failing consumer plan was Window -> InMemoryTableScan, with no sort. Probe and output: /tmp/comet6883-review-current-r4qtAM/Review6883Probe.scala and probe.log.

[viirya]

Fixed in 49dbdf8. You're right that Spark's outputOrdering only holds within each partition, so concatenating the partitions can't keep it in general. CometLocalExec now reports the block's ordering only when it holds across the concatenated partitions: with one partition, or when the block's partitioning is a RangePartitioning whose ordering is a prefix of the block's ordering, as a range's slices are. sortWithinPartitions over a multi-slice range now reports no ordering. Your reproduction is now a test: the window over the cached result returns Spark's ranks.

[sunchao]

[P2] Do not surface speculative later-partition failures before a limited query needs those partitions. With ANSI mode, range/local execution enabled and batch size 17, spark.range(0, 200000, 1, 2).filter("id >= 99999").selectExpr("100L DIV (100000L - id) AS v").limit(1).collect() returns [100] in Spark and ordinary Comet. Local execution starts partition 1 eagerly, encounters division by zero at id=100000, and this shared failure check aborts the query before partition 0 delivers its valid row. Spark’s executeTake stops after partition 0 and never evaluates the failing partition. Make error delivery aware of limited consumption, or retain ordinary execution for these plans until that behavior can be preserved.

Evidence: The exact-head reproduction returned Right(ArraySeq(100)) with Spark, the same result with ordinary Comet, and SparkArithmeticException [DIVIDE_BY_ZERO] with local execution. The local plan was CollectLimit 1 -> CometLocalResult -> ColumnarToRow -> CometLocal, containing the projection, filter and two-slice range. Spark’s supported-version SparkPlan.executeTake sources request partitions incrementally. Reproduction and output are recorded in /tmp/comet6883-review-current-r4qtAM/Review6883Probe.scala and probe.log.

[viirya]

Agreed, reporting any partition's failure at once broke executeTake's semantics. Each partition's error now travels through its own channel and is reported when that partition's turn comes, so a limit satisfied by earlier partitions never sees a later partition's failure, while a full collect still fails once it reaches it. That also removed the shared failure slot. Your ANSI DIV example is now a test, together with the full collect that still fails with DIVIDE_BY_ZERO, and the native tests check that a later failure is reported only in its turn.

[sunchao]

[P2] Avoid the deprecated atomic operation in the required Rust build. The stable toolchain used by CI deprecates AtomicI64::fetch_update, so the new call fails cargo clippy --color=never --all-targets --workspace -- -D warnings. This prevents the required Rust job from reaching its tests even when local execution is disabled. Use an equivalent checked compare/exchange loop or a narrowly scoped compatibility allowance that respects the declared Rust 1.95 minimum.

Evidence: Exact-head job https://github.com/apache/datafusion-comet/actions/runs/38217614661/job/114708727490 fails at core/src/local.rs:175:10: use of deprecated method ... fetch_update: renamed to try_update for consistency. The log identifies -D deprecated as implied by -D warnings, exits 101, and skips subsequent Rust test steps. Aggregate Required Checks is also failed.

[viirya]

Replaced with fetch_add and a check that the counter has not wrapped, which works on the 1.95 minimum. cargo clippy --all-targets --workspace -- -D warnings is clean with Rust 1.99.0, the toolchain CI uses.

[sunchao]

Reviewed commit 49dbdf859da3a8cc9022fb75c0633bb7da7c183e. New P1/P2 findings are detailed inline.

[sunchao]

[P2] Preserve bounded delivery for toLocalIterator. This wrapper specializes collect and take but inherits SparkPlan.executeToIterator, which serializes an entire Spark partition before yielding rows. Since LocalQueryRDD combines every source partition into one, iterator consumption now buffers the whole query. With local[2], Comet range/local execution enabled and spark.driver.maxResultSize=64k, consuming spark.range(0, 4000, 1, 8).selectExpr("id", "sha2(cast(id as string), 256) AS v").toLocalIterator() fails with a 267.9 KiB task result. Spark and ordinary Comet both stream all 4,000 rows successfully. This makes previously streamable results fail the driver limit and increases buffering from one source partition to the whole query. Please provide bounded executeToIterator delivery or retain ordinary execution for iterator actions.

Evidence: Reproduced on Spark 4.1.3/JDK 21 with JVM sources compiled at head 49dbdf8 and the matching native CI artifact, SHA-256 4631fc089894cae2019bc78ac2cfdcb6cd11ca1bbbdd458dab5162e6edc67dbd. Spark and ordinary Comet returned rows=4000. Local execution failed: Total size of serialized results of 1 tasks (267.9 KiB) is bigger than spark.driver.maxResultSize (64.0 KiB). The local plan contains CometLocalResult -> ColumnarToRow -> CometLocal, with eight range slices. Spark 3.4.3 through 4.2.0 have identical executeToIterator implementations calling getByteArrayRdd().map(_._2).toLocalIterator. Reproduction and output: /tmp/comet6883-review-active-9jbswbx6/ReviewIteratorProbe.scala and /tmp/comet6883-review-active-9jbswbx6/iterator.log.

[viirya]

Agreed, a single result partition defeats executeToIterator, which fetches one partition at a time to bound what the driver holds. Rather than a bounded handoff from a task that keeps running behind the iterator, which would also need cancellation when the iterator is abandoned, CometLocalResultExec.executeToIterator now runs the admitted native block as ordinary Comet execution, one Spark task per partition (89f4ada). toLocalIterator and the Thrift server's incremental collect then behave as without local execution, including how spark.driver.maxResultSize applies, at the cost of local execution's speedup for these iterator actions. A test checks that the iterator runs one job per partition and matches Spark.

PySpark's toLocalIterator goes through javaToPython and the result RDD instead, so it still receives the whole result as one partition. The docs now list that as a limitation; addressing it would need the local result RDD to expose more than one partition.

[andygrove]

Thanks @viirya, this is a clean first step and it follows what we agreed on #6499. The fixes for sunchao's findings all look right, and the per-partition planning through create_plan keeps range slices and mayOverflow exactly as ordinary execution has them.

Two things:

1. spark-local/ compiles into the spark module through build-helper, which needs the extra scalastyle execution, the spotless include and four compute-changes.py filters. Is it meant to become its own Maven module? If not, spark/src/main/scala/org/apache/comet/local would avoid that plumbing.
2. Could you add a measurement of the result handoff, such as collecting a range with a wide projection at local[8] with the flag on and off? The case for #6499 is performance, and this PR is where the handoff lands, so it's the cheapest point to confirm it pays off.

[sunchao]

Reviewed commit 89f4adae9ee259fdce81b05b9cfc2301585f30ec. New P1/P2 findings are detailed inline.

[sunchao]

[P2] Preserve reverse-partition execution for PySpark tail. With ANSI mode, local[2], and Comet range/local execution enabled, spark.range(0, 20, 1, 2).selectExpr("100L DIV id AS v").tail(1) should return Row(v=5). PySpark calls tailToPython, which invokes the inherited executeTail. That method now reads the single concatenated partition and encounters division by zero at id=0, whereas Spark and ordinary Comet read only the final source partition. Could executeTail delegate to the original block, like executeToIterator, or otherwise retain reverse-partition consumption?

Evidence: Reproduced on Spark 4.1.3/JDK 21 using current-head JVM classes and the checksum-verified, source-identical native CI artifact. The probe calls df.queryExecution.executedPlan.executeTail(1), exactly the execution method used by Dataset.tailToPython. Spark and ordinary Comet returned ArraySeq(5). Local execution threw SparkArithmeticException [DIVIDE_BY_ZERO]. Its plan was CometLocalResult -> ColumnarToRow -> CometLocal(partitions=2) -> CometProject -> CometRange. Spark sources from 3.4.3 through 4.2.0 confirm the same Python route and reverse partition selection. Reproduction and output: /tmp/comet6883-oct19-review-_1702lxj/ReviewTailProbe.scala and /tmp/comet6883-oct19-review-_1702lxj/probe.log.

[viirya]

Thanks, confirmed: tailToPython calls executeTail, which scans partitions from the last one, and the single local partition made it evaluate the whole query. CometLocalResultExec.executeTail now runs the admitted native block as ordinary Comet execution, as executeToIterator does (b82f387). Your ANSI example is now a test that compares executeTail with Spark.

I also went through the other entry points that read a result partition by partition: executeCollect and executeTake, including Dataset.isEmpty, use the result handoff; executeToIterator covers toLocalIterator, the Thrift server's incremental collect and SQL cursors; and Scala's tail plans a CollectTailExec, which is not admitted.

[viirya]

Thanks @andygrove.

1. No, it isn't meant to become its own module. The sources are now in spark/src/main/scala/org/apache/comet/local and the suite in spark/src/test/scala/org/apache/comet/local (b82f387), which drops the build-helper source roots, the extra scalastyle execution, the spotless include and the compute-changes.py and check-ci-config.py entries.

2. I measured collect() of an 8-slice range with a wide projection at local[8] on an M4 Max, each mode in its own JVM, median of 5 runs after 2 warmups, in forward and reverse mode order (ms):

   Query	Spark	Comet	Comet local
   5M rows × 8 longs	718 / 675	660 / 721	441 / 379
   1M rows × 32 longs	429 / 434	455 / 425	296 / 289

   Local execution is about 35–45% faster than ordinary Comet here, which is the encoding, compression and decoding of task results that the handoff skips.

[sunchao]

Reviewed commit b82f3878384044d6607c9c4fdf6cf09c77ddcbcd. New P1/P2 findings are detailed inline.

[sunchao]

[P2] Bound active partitions before initializing their memory consumers. With local[2], 256 MiB off-heap memory, the default fair_unified pool and local/range execution enabled, spark.range(0, 32000, 1, 32).selectExpr("id", "sha2(cast(id as string), 256) AS k").sortWithinPartitions("k").collect() fails with Not enough memory to continue external sort. Spark and ordinary Comet return all 32,000 rows. Calling execute for every partition registers every sort and merge consumer in the same pool, shrinking their fair shares below the sort’s 10 MiB minimum reservation. The two-batch queue does not prevent this. Initialize only a bounded number of active partitions, or fall back for these plans until their reservation requirements can be supported.

Evidence: Reproduced on Spark 4.1.3/JDK 21 using current-head JVM sources and checksum-verified, source-identical native code. Spark and ordinary Comet returned rows=32000, sum=511984000. Local execution failed acquiring 10485760 bytes with a fair limit of 7064090 bytes and 38 registered consumers. The same query with two partitions succeeded in all three modes. DataFusion 55.2.0 ExternalSorter::new registers two consumers per sort, and reserve_memory_for_merge requires the reservation before processing input. Probe and output: /tmp/comet6883-independent-oiwwar3j/ReviewMemoryProbe.scala and memory.log.

[viirya]

Thanks, confirmed: each sort registers its two consumers when its partition starts, and the fair pool divides by every registered consumer, so starting all 32 partitions at once left each consumer less than the 10 MiB merge reservation. Delivery now runs at most spark.executor.cores partitions at once, which local[N] resolves to N, and starts the next one when the first is delivered (61f9540), so no more partitions' operators exist at once than ordinary Comet runs in N tasks. Your query is now a test in the suite, which already runs with local[2] and 256 MiB off-heap, and a native test checks that the later partitions start only as earlier ones are delivered.

[sunchao]

[P2] Use callback syntax accepted by the required syntactic parser. Now that this suite lives under spark/, CI’s scalafix:0.14.6 --check --syntactic command includes it and rejects { _: Iterator[Int] => with `}` expected but `=>` found. This fails the required Scala lint job even though Scala compilation succeeds. Parenthesize the typed parameter, for example { (_: Iterator[Int]) => ... }, so the suite passes the configured parser.

Evidence: Exact-head check PR Build (Linux) / Lint Scala (syntactic) failed at CometLocalExecutionSuite.scala:116:28 and exited 2: https://github.com/apache/datafusion-comet/actions/runs/39289196110/job/118188258091. Retrieved the job log and verified its check-run head SHA is b82f3878384044d6607c9c4fdf6cf09c77ddcbcd. Saved log: /tmp/comet6883-independent-oiwwar3j/scala-lint.log.

[viirya]

Fixed in 61f9540 by parenthesizing the parameter. The new sources have no other typed lambda parameters.

[sunchao]

Reviewed commit 61f954082349ac6fc213ab9991a0f65fc39986c6. New P1/P2 findings are detailed inline.

[sunchao]

[P2] Preserve failures from the partitions Spark requests together in executeTake. With local[2], ANSI mode and Comet range/local execution enabled, spark.range(0, 6, 1, 3).filter("id % 2 = 0").selectExpr("100L DIV (4L - id) AS v").limit(2).collect() throws DIVIDE_BY_ZERO in Spark and ordinary Comet, but returns [25, 50] locally. After partition 0 supplies one row, Spark requests partitions 1 and 2 together and fails on partition 2. The local handoff stops after partition 1 supplies the second row, suppressing that failure and turning a failed ANSI query into success. Please preserve Spark's partition-request/error behavior, or delegate executeTake to ordinary execution until that behavior is supported.

Evidence: Reproduced on Spark 4.1.3/JDK 21 using the exact-head native build and freshly compiled current local JVM sources. actions.log records DIVIDE_BY_ZERO for Spark and ordinary Comet and ArraySeq([25], [50]) for local execution. The local plan is CollectLimit 2 -> CometLocalResult -> ColumnarToRow -> CometLocal(partitions=3). SparkPlan.scala in Spark 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0 computes the next request with ceil(1.5 * left * partsScanned / buf.length) and completes runJob for all requested partitions before trimming rows. Probe and output: /tmp/comet6883-oct20-review-_l0rntl0/ReviewActionProbe.scala and /tmp/comet6883-oct20-review-_l0rntl0/actions.log.

[viirya]

Thanks, confirmed: Spark's executeTake requests partitions in growing batches and runs each batch to completion, so with one row from partition 0 it runs partitions 1 and 2 together, and partition 2 fails even though partition 1 has the missing row. Rather than reproducing that batching in the handoff, CometLocalResultExec.executeTake now runs the admitted native block as ordinary Comet execution (b456e6d), as executeToIterator and executeTail already do. take, head, show, limit(n).collect() and isEmpty therefore evaluate exactly the partitions Spark evaluates, and the result handoff now serves only a full collect, where it pays off. Your example is now a test that fails with DIVIDE_BY_ZERO, as Spark does.

[viirya]

@sunchao @andygrove Can you take another look? Thanks!

[andygrove]

Thanks @viirya. Both of my earlier points are addressed: the sources live under spark/ now, and the benchmark makes the case for the handoff. The fixes for sunchao's rounds look right to me, and moving take, toLocalIterator and tail to ordinary execution keeps their semantics simple. I have one question inline about how many partitions actually make progress at once, now that only a full collect uses the handoff.

[andygrove]

Since b456e6d03 the handoff only serves a full collect, so every partition is read to the end anyway. But only the front partition's channel is drained, and the others stop after PARTITION_BUFFER batches until their turn. For a streaming block where the native work dominates, such as spark.range(0, 80000000, 1, 8).filter("sha2(cast(id as string), 256) like 'ff%'").collect() at local[8], I think that leaves about one partition working at a time, where ordinary Comet runs all eight tasks in parallel. Could you measure a query like that with the flag on and off? If local is slower, could started partitions run to completion, with what they buffer bounded by spark.driver.maxResultSize rather than two batches? The bound on how many partitions start would still protect the memory pool. The "each up to two batches ahead" sentence in local-execution.md would change with it.

[viirya]

Thanks, confirmed. Your query at local[8] (median of 4 runs after a warmup, each mode in its own JVM, forward and reverse order, ms):

	Spark	Comet	Comet local
before	2668 / 3190	891 / 873	4511 / 4337
after (b55a395)	3496 / 2864	1143 / 929	905 / 866

The filter coalesces its output, so each partition yields only a handful of batches, and a partition waiting for its turn stopped after its first two.

As you suggested, the started partitions now run ahead while the batches they hold together total less than spark.driver.maxResultSize (no limit when it is 0); past that, each holds at most two batches until the earlier ones are delivered, so the first partition always progresses. The bound on how many partitions start is unchanged. Two native tests cover it: a later partition runs to its end while the first delivers nothing, and with a zero limit it stops at two batches. local-execution.md describes the new behavior.

The wide-projection collect from before is unaffected (5M rows × 8 longs: Comet 1111 / 979, local 547 / 496).

[andygrove]

This is an AI-assisted review and I have not had time to go through it carefully myself, so please feel free to push back on any of it.

Thanks @viirya. The fixes for the earlier rounds look right to me, and the benchmark for the filter case answers my question about partition progress. I compared against branch-1.2 by reading the diff. With the flag off nothing changes, and with it on I only found intended differences. I have three questions and one docs ask.

Read-ahead memory (native/local/src/lib.rs, Partition::spawn). The batches that started partitions hold ahead of delivery sit in an unbounded channel as native Arrow buffers, and I don't see them reserved in the pool. With the default spark.driver.maxResultSize that is up to 1 GiB of native memory the pool can't see, plus two more batches per partition past that. The JVM also copies the same rows into UnsafeRows as they arrive. Ordinary Comet holds about one batch per task before serialization. Could the held bytes be reserved through the task's pool, or could the default bound be smaller? Could you also share peak RSS for a collect of a few hundred MiB with the flag on and off? If you keep the current design, a sentence in local-execution.md and memory_management.md saying this memory is unaccounted would help the next reader.

Pool shares for non-spillable consumers (native/core/src/local.rs). All N partitions' operators now register in one task's pool, so fair_unified sees one running task instead of N. For spillable consumers that works out, and the 32-partition sort test shows it. For consumers that can't spill, such as a hash join build or a merge, I think the only bound is the task total against pool_size, so N partitions can together take what N tasks would have shared. Is that intended? If so, could local-execution.md say it? If not, would a test with a non-spillable consumer under a small pool show whether it holds?

Single result partition (CometLocalResultExec.doExecute). Because doExecute returns the one-partition RDD, I think df.rdd, foreachPartition, queryExecution.toRdd, cache() and localCheckpoint() also see a single partition when the flag is on. The cache test asserts the cached result is one partition. That means every later query over a cached range is scanned by one task, which gives up the parallelism that ordinary Comet gives it. Is this intended for the first part? If so, could the docs list it next to the PySpark toLocalIterator limitation, and could a test pin down getNumPartitions for df.rdd? If it isn't intended, should admission skip those cases?

Docs (ffi.md). The "Native → JVM Data Flow" section says the JVM pulls batches through CometExecIterator. With this PR LocalQueryIterator also pulls through NativeLocal.nextBatch, using the same export_batch and move_to_spark, so offsets are zeroed correctly. Could that path be added to the section and the ownership table?

CI is green for everything that ran. The new suite is registered in both the Linux and macOS workflows. It runs only on Spark 4.1 in CI, and the flag is off in the Spark SQL tests, so the 3.4 to 4.2 results rest on your local runs.

[comphead]

sunchao's seven inline points and andygrove's earlier ones are fixed in code at this head, and the deterministic gate keeps spark_partition_id, monotonically_increasing_id and rand out. andygrove's review from 2026-10-27 (read-ahead memory outside the pool, pool shares for consumers that cannot spill, one result partition for df.rdd, cache and localCheckpoint, and ffi.md) has no reply or commit after it yet, so I have not repeated those points. Three comments. git merge-tree also shows a conflict with current main in jni_api.rs.

[comphead]

Since b456e6d, CometLocalResultExec.executeTake runs the block as ordinary Comet, so take, head, show and limit(n).collect(), which all reach this through CollectLimitExec.executeCollect, no longer use the local path, and the doc above still says they do. Admitting the limit now only changes CollectLimitExec.doExecute, which serves Spark Connect's collect, Arrow toPandas, limit(n).toLocalIterator() and limit(n).rdd. With one child partition, doExecute skips the per-partition take and the shuffle and returns the first n rows of the concatenated result, while Spark runs take(n) in every child partition. With ANSI on, I expect spark.range(0, 6, 1, 3).selectExpr("100L DIV (4L - id) AS v").limit(1).toLocalIterator() to fail with DIVIDE_BY_ZERO in Spark and to return 25 here. I have not run it. When the block is not admitted, line 83 also converts the whole plan a second time. Could a root CollectLimitExec be left to ordinary conversion?

[comphead]

The only non-test caller, native/core/src/local.rs, always sets in_partition_order, so the execute_stream branch never runs in Comet. Most of tests/execution.rs and several tests/bridge.rs tests build LocalQuery::new(..).execute() without it, so they mainly exercise DataFusion's RepartitionExec and GlobalLimitExec behind execute_stream. Could partition order be the only mode, with max_running and read_ahead passed to LocalQuery::new, and the coalescing mode and its tests land with the part that admits exchanges?

[comphead]

QueryHandle repeats the lifecycle of BatchProducer in jni_api.rs: a Tokio task polls the stream from an Arc<Mutex<Option<_>>>, a flag tells completion from cancellation, catch_unwind turns a panic into an error, and stop drops the stream on the calling thread. Partition::spawn in lib.rs is a third copy with a read-ahead budget, so each batch crosses two channels and two tasks before the JNI thread. Could the JNI pull poll LocalQueryStream directly with a bounded wait, as executePlan already does with block_on for JVM-fed plans? If the separate producer is needed so close can run during a pull, moving BatchProducer into native/local would leave one implementation. I have not prototyped either.

[viirya]

Thanks @andygrove. Addressed in 227e263.

Read-ahead memory is not in the pool. That's right, and I kept it that way on purpose. A reservation for held batches can't be handed back when an operator later needs that memory, so reserving them would fail queries (non-spillable consumers in particular) that ordinary execution runs fine. Spark's own collect also keeps task results outside the execution pool, on the driver. The bound is spark.driver.maxResultSize, the most a collect may bring back anyway, and past it each running partition holds at most two batches. As for the JVM copying the same rows, a delivered batch is closed right after its rows are copied, so only one batch overlaps with the copy.

To size the difference, I measured peak RSS for a 5M-row, 8-long-column collect (about 360 MiB of UnsafeRow) with a fixed 4g heap:

	run 1	run 2
Spark	5.07 G	5.05 G
Comet	5.13 G	5.13 G
Local execution	5.36 G	5.38 G

That's roughly 240 MB more native memory. With a 3g heap, Spark failed with an OOM and Comet with TaskResultLost, while local execution finished (4.27 G / 4.34 G), because it doesn't encode and decode the whole result on the driver. I added this to the accounting gaps in memory_management.md and explained it in local-execution.md.

Non-spillable consumers sharing one task's pool. The total is bounded the same way as with ordinary execution, since all the tasks' consumers already share the executor's pool. What changes is the per-partition cap: one partition's merge, for example, can take more than one task's share and leave less for the others. That's intended. The existing many-partitions sort test in a small pool exercises the merge path. local-execution.md now spells this out.

The result was a single partition for every consumer. Good catch. CometLocalResultExec now executes as the ordinary native block and reports the block's partitioning and ordering. Only collect runs the whole query as one graph. Cache, df.rdd, foreachPartition, localCheckpoint, and PySpark's collect / toLocalIterator (which go through execute()) now see the block's N partitions. The cache test now checks that the cached relation has the ordinary partition count, and a new test checks rdd, toRdd, outputPartitioning, per-partition sizes, and localCheckpoint.

ffi.md. It now describes the local path: NativeLocal.nextBatch exports through the same prepare_output / move_to_spark, returns pending while no batch is ready, and LocalQueryIterator closes imported batches the same way.

[sunchao]

Reviewed 227e263ec against the updated implementation. The new wrapper restores partition counts for RDD consumers, but it still changes the execution path for a bare range. The inline comments include reproductions for that fallback, classloader isolation, task-slot accounting, reservation cleanup, and result-size accounting.

The 28 tests in native/local passed in a standalone run using the current sources, DataFusion 55.2.0 and cached Tokio 1.53.1. The locked Cargo build could not resolve uuid 1.29.0 from the configured registry, so I did not run the full JNI suite. dev/ci/check-suites.py and dev/ci/check-ci-config.py passed. GitHub reports a merge conflict and only the label check on this commit.

[sunchao]

Handle separate driver and executor classloaders. slots is a Scala singleton. A task in the same JVM can still load a second copy when the plugin JAR is shipped via spark.jars and spark.executor.userClassPathFirst=true. I reproduced this on Spark 4.1.3 with the handoff from this commit and one row. Ordinary collection succeeded under both classloading modes. The handoff succeeded with parent-first loading, but under child-first loading this lookup missed, the task returned successfully, and the driver threw Local result not delivered. Could the handoff use a registry shared across those loaders, or fall back to ordinary collection before scheduling in this configuration?

[sunchao]

Use Spark’s concurrent task count here. With local[32] and spark.task.cpus=16, Spark schedules two tasks at once (LocalSchedulerBackend.maxNumConcurrentTasks divides total cores by CPUs per task). This line starts 32 partition streams. In a native probe using this scheduler, DataFusion 55.2.0 sorts and a 256 MiB CometFairMemoryPool, 32 concurrent partitions failed a 10 MiB reservation. Changing only the limit to two returned all 557,056 rows. Could we divide by the task_cpus already passed to createQuery, clamp to at least one, and cover a non-default spark.task.cpus?

[sunchao]

Wait for outstanding plan reservations when closing. Cancelling the query drops its streams, but a sort can own a worker that releases its reservation only when it next yields. Ordinary release_plan drops the context and then waits on plan_memory before Spark can reclaim the task’s memory. This path skips that wait. In a controlled probe using this QueryHandle and PlanMemoryPool, a worker still held 4,096 bytes after cancellation and handle drop, then released them when it resumed. Could both teardown paths share the bounded reservation wait, including early exit on cancellation or the result-size limit?

[sunchao]

Preserve ordinary plan preparation for RDD actions. Constructing ColumnarToRowExec here happens after EliminateRedundantTransitions has run. For a bare range, that rule normally restores Spark’s RangeExec and avoids a native round trip. On Spark 4.1.3 with this commit’s Scala and the same ordinary Comet JNI library in both runs, spark.range(0, 16000000, 1, 8).queryExecution.toRdd.count() returned the same count in both modes, but created zero native iterators with the local flag off and eight with it on. This also bypasses the rule’s columnar-to-row selection for projections. Could the fallback execute the normally prepared block and extend the RDD test to check the bare-range behavior covered by CometRangeExecSuite?

[sunchao]

Keep the output validation switches effective for local collects. Ordinary Comet passes exec_context.debug_native here so spark.comet.debug.enabled runs Arrow’s validate_full before export. The hardcoded false skips it. On import, LocalQueryIterator also omits NativeUtil.checkTimestampLabels, so spark.comet.testing.strictTimestampLabels.enabled misses local output, including nested timestamps. Could the local path forward the debug setting and retain the expected Spark output types needed for the JVM label check?

[sunchao]

Either path can reach the result-size limit first. Spark counts serialized task results, which include per-task overhead. With an 8 KiB cap on Spark 4.1.3, ordinary collection of 128 one-long rows split across 128 partitions failed after 26 task results. Coalescing the same Dataset’s RDD into this commit’s single-partition handoff returned all 128 (2,048 bytes of UnsafeRow payload). Please describe the two accounting rules here and in LocalResultHandoff.collect’s Scaladoc without promising the local limit is always stricter.

[sunchao]

Build the JNI library in the test recipe. cargo test -p datafusion-comet-local builds the Rust scheduling crate and its tests. The NativeLocal JNI entry points live in native/core, so this sequence leaves a clean checkout without the libcomet required by the JVM suite. Adding make core from the repository root before the Maven command would make these steps reproducible.

[viirya]

Thanks @comphead. andygrove's 2026-10-27 points are answered above and in 227e263, and this round is in 28d5702, after merging main (the jni_api.rs conflict was only the new with_output_stream in the planner construction, now in ExecutionContext::planner).

The root CollectLimitExec. You're right. The DIVIDE_BY_ZERO case was real at 64f7257: with one child partition, CollectLimitExec.doExecute skips the per-partition take, so limit(1).toLocalIterator() returned 25 where Spark fails. 227e263 already made CometLocalResultExec.doExecute run the block's N partitions, which fixed that. But it also left admitting the limit with no benefit, since executeCollect goes through executeTake, which is ordinary anyway. So a root limit is no longer admitted: CometLocalRule converts the plan once, and the CollectLimitExec case is gone from isLocal. The limit tests now assert that it isn't admitted and that it matches Spark, and your toLocalIterator example is among them. The docs are updated.

Partition order as the only mode. Done. LocalQuery::new(plan, context, max_running, read_ahead), and the execute_stream branch is removed. I dropped the tests that only exercised RepartitionExec / GlobalLimitExec behind it (round-robin, hash, nested exchanges, pending-exchange cancel, global limit), and moved the rest (empty results, startup errors, a missing runtime, drop and first-error teardown, independent cancellation) to partition order.

QueryHandle duplicating BatchProducer. I took the first option. The JNI pull now polls LocalQueryStream on its own thread, so the producer task, the std channel, the semaphore and the completed flag are gone, and a batch crosses one channel and one task before the pull. Runtime shutdown is still an error rather than EOF, through the per-partition ended flag that PartitionOrder already checks. Two details differ from a plain block_on with a timeout:

• The bounded wait parks the thread (park_timeout, with a waker that unparks it) rather than using a Tokio timer. In the multi-thread runtime, timers are driven by workers, so when every worker is waiting in acquireMemory, a tokio::time timeout wouldn't fire and the task couldn't check for interruption.
• cancel still drops the stream on its own thread. It unparks a waiting reader first, so it waits only for a poll the reader is already in, and PartitionOrder::poll_next only receives from channels and starts partitions.

Partition::spawn stays separate from BatchProducer, because its backpressure is the read-ahead byte bound shared across partitions rather than a fixed per-stream channel capacity.