[EPIC] Comet Native Executors: run Spark queries on native executors built on DataFusion Ballista #7288

What is the problem the feature request solves?

Comet accelerates Spark by running native operators inside Spark's JVM executors. #4796 proposed an additional, optional deployment mode with a fully native data plane: a Spark driver (a Spark Connect server, or a regular application's driver) stays the control plane, and queries run on native processes that reuse Comet's Spark-compatible operators and expressions. The draft PoC in #4800 showed this can work: TPC-H Q1's aggregate ran across a Ballista shuffle on a JVM-less executor and matched Spark row for row.

This issue replaces #4796 with a concrete design and tracks the work. It builds on what the PoC learned and on what has changed since: Comet main and Ballista 55.0.0 are now both on DataFusion 55. The mode is called Comet Native Executors. It is purely additive; the existing accelerator is unchanged.

Describe the potential solution

Overview

 PySpark / Spark Connect client, or a Spark application
                    │
                    ▼
 ┌────────────────────────────────────────┐  CONTROL PLANE: Spark driver (JVM)
 │ Catalyst: analyze, optimize            │  • unchanged Spark front end
 │ Comet rules + serde → Comet protobuf   │  • strict-native check
 │ CometNativeExecutorsExec (one leaf)    │  • no Ballista, no datafusion-proto
 └──────────────────┬─────────────────────┘
                    │ Arrow Flight: GetFlightInfo(cmd = Comet query protobuf)
                    ▼
 ┌────────────────────────────────────────┐  DATA PLANE: native, no JVM
 │ comet-scheduler                        │  • Comet front end: validate blocks
 │   Comet Flight front end               │    with Comet's planner, then stage
 │   Ballista scheduler (library)         │    graph → Ballista physical plan
 ├────────────────────────────────────────┤
 │ comet-executor × N                     │  • Ballista executor (library) with
 │   Ballista executor (library)          │    Comet's planner, operators,
 │                                        │    expressions and scans
 └──────────────────┬─────────────────────┘
                    │ Arrow Flight: DoGet(ticket), one per result partition
                    ▼
            Spark driver → client

A Comet Native Executor cluster is one comet-scheduler and N comet-executor processes, built on the Apache DataFusion Ballista crates (used as libraries from crates.io) with Comet's native crates linked in. Ballista provides stage scheduling, the shuffle between stages and task retry; Comet provides the Spark semantics.

Key decisions

1. Comet protobuf is the only contract between the driver and the cluster. The driver serializes a stage graph of Comet operator protos and submits it over Arrow Flight. It never links Ballista or datafusion-proto, so the main Comet build (JAR and libcomet) has no Ballista dependency: Comet can upgrade DataFusion without waiting for a Ballista release, and the cluster can upgrade Ballista without touching drivers. Results come back as Arrow IPC over Flight, which doesn't depend on either version.
2. The cluster is a separate build. It gets its own Cargo workspace and lockfile, depending on Comet's native crates by path and on Ballista from crates.io, with its own CI job. It is built from Comet crates at some revision plus the Ballista release on that revision's DataFusion major. When Comet moves to a DataFusion major that Ballista hasn't released, the cluster build lags until Ballista catches up; a Ballista git dependency can't go into a source release.
3. Static linking first, with a seam narrow enough for FFI later. Inside the executor, Comet's crates and Ballista share one DataFusion version. The fragment seam passes only proto bytes, a partition index, input streams and config, and returns an output stream and metrics. If we later want the executor's Comet library on a different DataFusion version than Ballista, that seam can move behind the Arrow C Stream interface (a spec-level stable ABI that Comet already uses with the JVM) without changing the driver, the proto or the scheduler. datafusion-ffi is an alternative, but its cross-version support is a soft promise: its README still recommends the same version on both sides.
4. Strict native, all or nothing. A query either runs entirely on the native executors or fails with an error naming the unsupported node and its fallback reasons. There are no Spark executors to fall back to, and interleaving the two engines within one query is out of scope.
5. Upstream Ballista, no fork. The Ballista enhancements this needs are filed and contributed upstream (list below).

Driver side (JVM)

• Configs: spark.comet.nativeExecutors.enabled (default false) and spark.comet.nativeExecutors.scheduler.url.
• Hook: the last step of CometRule replaces the converted plan with a single row-based leaf, CometNativeExecutorsExec. It runs in the query-stage-prep instance when AQE applies and in the columnar instance otherwise. Comet's query-stage-prep rule already converts and serializes the whole plan, so this works with AQE on: AQE then finds no exchanges and runs one result stage. (The PoC overrode executeCollect, which Spark Connect never calls. Connect calls executedPlan.execute() and sparkContext.submitJob, so under Connect the PoC's offload silently never ran.)
• Stage graph: walk the native blocks (serializedPlanOpt) separated by CometShuffleExchangeExec and CometBroadcastExchangeExec. Each block becomes a stage, and each exchange becomes a typed edge that reuses Comet's existing Partitioning proto (hash, single, range, round-robin) plus broadcast. Leaf stages carry per-partition scan data from PlanDataInjector.findAllPlanData, so a scan keeps its N partitions (the PoC collapsed each scan to one).
• Strict-native check: built on ExtendedExplainInfo's traversal with an allow-list of the node types the stage-graph builder can translate, plus a scan of each block's proto for expressions that call back into the JVM: JvmScalarUdf (which ScalaUDF codegen dispatch emits), Subquery unless its value has been inlined, and ArrowPythonUdf. Columnar (JVM) shuffle, row-to-columnar transitions and writes are rejected.
• Results: doExecute() returns an RDD with one partition per result partition. Its tasks run on the driver (a Connect server runs local[N]), fetch Arrow IPC with Flight DoGet, and return rows. executeCollect, executeTake and executeToIterator are overridden for classic sessions and subqueries.
• Subqueries: each scalar subquery runs as its own native-executor query, and its value is inlined into the main plan's proto as a literal before submission. DPP and runtime bloom filters are disabled in this mode at first; dropping them never changes results.
• Cancellation: waiting on the remote job happens inside Spark tasks, so when a task is killed (Connect interrupts by job tag) the job is cancelled through Flight CancelFlightInfo. No driver thread blocks indefinitely.
• UI and metrics: at first the SQL tab shows only the leaf. Later, post the inner plan the way AdaptiveSparkPlanExec does, and map the executors' per-operator metrics back by plan_id.

Wire API

A new proto file in native/proto, compiled for both Rust and the JVM:

• QueryRequest: the Comet version and a fingerprint of the proto schema, a config map (the settings native execution reads), and the stages in topological order, root last.
• Stage: the block's Operator proto, its output schema, its inputs in the order of the block's input leaves, and, for leaf stages, the per-partition injected plans.
• StageInput: the producer stage, the edge's Partitioning, and whether it is a broadcast.

Transport is Arrow Flight on the scheduler, following the Flight SQL pattern: GetFlightInfo submits a query and returns one endpoint per result partition, each with an executor location and an opaque ticket; DoGet fetches a partition from an executor; PollFlightInfo handles long-running queries; and CancelFlightInfo cancels. The driver uses arrow-flight, part of arrow-rs, from libcomet at Comet's own Arrow version, and hands batches to the JVM through the existing Arrow C Data path.

Versioning. The boundary between driver and cluster becomes a cross-process API. Today the JAR and libcomet always come from the same build; with a separate cluster, versions can drift, and protobuf only protects the wire format: an older executor silently ignores a field it doesn't know. At first the scheduler requires an exact match of Comet version and proto fingerprint and rejects anything else with a clear error. If mixed versions matter later, we can adopt an additive-only proto policy.

Scheduler: comet-scheduler

Our own process: a Comet Flight front end that wraps Ballista's scheduler, used as a library (create_scheduler returns a SchedulerServer with submit_physical_plan and cancel_job). Executors register with the embedded Ballista scheduler as usual; drivers only talk to the Flight front end. For each query, the front end:

1. checks the Comet version and proto fingerprint;
2. runs Comet's native planner over each stage's block to validate it and check it against the output schema the driver declares for the stage, so a plan the executors can't build fails at submission with the planner's own error, not inside a task;
3. turns the stage graph into a Ballista physical plan (table below) and submits it in-process;
4. returns the result partitions' locations as Flight endpoints.

Fragments still travel to the executors as Comet protos, because each executor rebuilds its block per partition from that partition's scan data and live input streams. Planning a leaf stage also builds its object stores, so validation runs in a mode that skips that, which keeps storage credentials off the scheduler. None of Ballista's client protocol (version header, job status, fetch tickets) leaks to the driver.

How the stage graph maps onto Ballista 55's static planner:

Bucketed scans and storage-partitioned joins are disabled in this mode, because their unshuffled side is partitioned by Spark's Murmur3 hash.

Executor: comet-executor

Ballista's executor used as a library (ExecutorProcessConfig with Comet's codecs), with no JVM.

• CometFragmentExec is a plain DataFusion ExecutionPlan in Comet's core crate, with no Ballista dependency. It holds a block's Operator proto and its declared output schema; its children are a partition descriptor (below) followed by one child per input leaf. execute(p) builds the block with Comet's PhysicalPlanner and session setup on a blocking thread (planning can block on Comet's runtime), feeds the block's input leaves directly from the child plans, and adapts output batches to the declared schema.
• Partition descriptor. Ballista's scheduler slices each stage plan per task, but only for a fixed list of leaf types, and operators only see task-local partition indices. So each fragment gets a MemorySourceConfig child with one single-row batch per partition. The batch carries the global partition index, which seeds rand, monotonically_increasing_id, spark_partition_id and similar expressions, and, for leaf stages, the partition's injected plan. Ballista already slices and serializes that child. A Ballista hook that lets extension nodes slice themselves would make the descriptor unnecessary (Let custom ExecutionPlan nodes restrict themselves to a task's partitions datafusion-ballista#2679).
• Planner seam: the planner's input leaves (Scan, ShuffleScan) take a native child plan instead of a JVM iterator. The PoC instead pumped a ScanExec with block_on, which blocked a Tokio worker on a network stream and no longer fits main.
• Memory: Ballista's per-task FairSpillPool replaces Comet's Spark-backed pool, and spills go to the executor's work dir.

What a JVM-free executor needs, from an audit of main:

• Build without hdfs-opendal, which pulls in hdfs-sys, the only crate that links libjvm, and gate the macOS hdrs dependency behind it. On Linux that alone removes libjvm.
• Make the JVM accessors (JVMClasses::get, with_env) checked instead of get_unchecked, so a JVM-only path that slips through fails cleanly instead of being undefined behavior.
• Reject, or replace with native code, everything that calls back into the JVM: ScalaUDF codegen dispatch and other JVM UDFs; Subquery (inlined by the driver); Parquet encryption key unwrapping; the configurable S3 credential provider callback (the default chain and IRSA are already native); case-insensitive matching of non-ASCII field names; HDFS; and object-store writes, which go through a JVM output stream.

Build and packaging

• Main workspace: the wire proto; CometFragmentExec and the planner seam in core; and the Flight client in libcomet, behind a Cargo feature while this is experimental. No Ballista.
• New workspace (native/native-executors, name open): the codecs, stage-graph assembly, the scheduler front end, and the comet-scheduler and comet-executor binaries.
• Later: Docker images and deployment docs.

Ballista enhancements

Filed upstream:

• Run submitted physical plans through the adaptive planner datafusion-ballista#2678: run submitted physical plans through the adaptive planner. Today only logical plans get AQE.
• Let custom ExecutionPlan nodes restrict themselves to a task's partitions datafusion-ballista#2679: let custom nodes restrict themselves to a task's partitions. PartitionSliceable is a hand-maintained list, which is why the prototype gives every fragment a descriptor child. Inputs that every partition reads whole, such as a broadcast build side, already work by declaring SinglePartition.
• Static planner drops a RepartitionExec with range partitioning instead of planning a range shuffle datafusion-ballista#2680: plan a range shuffle for Partitioning::Range instead of dropping it.
• Let hash shuffles assign rows to partitions with a client's own hash function datafusion-ballista#2681: let hash shuffles place rows with a client's own hash function, such as Spark's Murmur3.

Not filed: Ballista encodes the expressions in its shuffle nodes with DataFusion's default codec, so hash keys must be plain columns. That doesn't block anything, because the driver can project keys into columns in the producing stage.

Out of scope at first

Writes (later: files produced on the executors with the commit on the driver, like the split Iceberg writer); Spark's AQE, which is lost in this mode until Ballista's AQE covers physical plans; fallback to Spark execution; Spark UI parity; and a highly available scheduler.

Proposed sub-tasks

• Prototype, as a draft PR: feat: prototype Comet Native Executors built on DataFusion Ballista #7289. It runs scans, filters, projections, aggregates, shuffled joins, sorts (in a single task) and TPC-H q1 end to end.
• The native side: the proto, CometFragmentExec and the planner seam, comet-scheduler and comet-executor, their Flight tests, and a CI job for the cluster workspace (feat: add comet-scheduler and comet-executor for Comet Native Executors #7299).
• The driver side: the Flight client, the configs and the driver hook (feat: prototype Comet Native Executors built on DataFusion Ballista #7289, rebased after feat: add comet-scheduler and comet-executor for Comet Native Executors #7299). With it, releasing a query's results once the driver has read them, which the scheduler otherwise deletes when their retention runs out.
• Spark Connect end to end; multi-partition scans at scale.
• Inlined scalar subqueries; broadcast joins; DPP.
• JVM-free executor hardening (checked JVM access, the up-front rejection list). That includes native UDFs: planning one loads its library from a path in the plan, on the scheduler while it checks the block as well as on the executors, under a policy the client's conf sets. Until the cluster has a policy of its own, it should refuse them, as it does JVM UDFs.
• TPC-H, then TPC-DS: correctness and benchmarks.
• The Ballista enhancements above: Run submitted physical plans through the adaptive planner datafusion-ballista#2678, Let custom ExecutionPlan nodes restrict themselves to a task's partitions datafusion-ballista#2679, Static planner drops a RepartitionExec with range partitioning instead of planning a range shuffle datafusion-ballista#2680, Let hash shuffles assign rows to partitions with a client's own hash function datafusion-ballista#2681.
• Packaging and deployment docs.

Additional context

• Supersedes Add optional fully-native deployment mode (Ballista-backed) #4796. The PoC is feat: Native execution via Ballista [experimental] #4800, which also led to feat: Allow submitting a pre-built physical plan to the scheduler datafusion-ballista#1924 (physical-plan submission).
• Ballista 55.0.0's dependency ranges match Comet main (DataFusion 55, arrow-flight 59.2, object_store 0.13, prost and tonic 0.14), so the cluster workspace resolves to a single DataFusion today.
• The Comet protobuf and Flight front end also fit the "connect gateway" idea raised on Add optional fully-native deployment mode (Ballista-backed) #4796: a gateway could route a query to Spark, to Comet Native Executors or to DataFusion.