Skip to content

feat: add comet-scheduler and comet-executor for Comet Native Executors - #7299

Open
andygrove wants to merge 8 commits into
apache:mainfrom
andygrove:native-executors-binaries
Open

andygrove wants to merge 8 commits into
apache:mainfrom
andygrove:native-executors-binaries

Conversation

@andygrove

@andygrove andygrove commented Oct 28, 2026 •

Copy link
Copy Markdown
Member

Which issue does this PR close?

Part of #7288.

Rationale for this change

#7288 designs Comet Native Executors: an optional deployment mode where a Spark driver plans each query and a cluster of native processes built on Apache DataFusion Ballista runs it, with no JVM in the data plane. The draft prototype, #7289, runs TPC-H q1 that way end to end.

This PR lands the native side of that prototype: everything needed to run a Comet stage graph on a cluster, without connecting it to Spark yet. Users see no change, and the main Comet build gains no dependency. The Spark driver side follows in #7289, which will be rebased onto main once this merges.

What changes are included in this PR?

One commit per part; it reviews best commit by commit.

  1. Wire protocol. native_executors.proto defines QueryRequest: a stage graph of Comet native blocks connected by the exchanges between them. Each stage carries its Spark output schema and, for stages that read files, each partition's plan. Fragment carries a stage from the scheduler to its executors. The proto crate exports a fingerprint of the .proto files, so a cluster can refuse a query from another Comet build.

  2. Native blocks outside Spark. CometFragmentExec runs a stage's block for each of its partitions.

    • Inputs: a new planner seam, PhysicalPlanner::with_native_inputs, feeds the block's input leaves from other DataFusion plans instead of from the JVM. An expression that calls into the JVM, such as a scalar subquery or a JVM UDF, fails to plan. PhysicalPlanner::without_data plans a block the same way without its data, to check it: its input leaves are empty, and its file scans build no object store.
    • Output types: each partition's output is cast to the stage's declared Spark types before it leaves the stage, as Comet's shuffle writer does in a Spark task, because where a hash exchange sends a row depends on its keys' types as well as their values. Each batch then carries the stage's schema, with every column nullable: the cast keeps a column nullable wherever the block's output is, even where Spark knows it has no nulls.
    • Memory and spills: the block runs in a Comet session over the host's memory pool and disk manager.
    • Partition descriptor: an extra child, an in-memory source with one row per partition, carries each partition's index and, for a stage that reads files, its own copy of the block. That works around Ballista only slicing a fixed list of leaf types per task (Let custom ExecutionPlan nodes restrict themselves to a task's partitions datafusion-ballista#2679).
    • Main build: nothing in the main build calls any of this yet.
  3. native/native-executors. A new Cargo workspace, with its own lockfile, that builds two binaries on the Ballista 55 crates:

    • comet-scheduler wraps Ballista's scheduler in an Arrow Flight service. GetFlightInfo takes a QueryRequest, checks the client's Comet version and plan format, and plans each block with Comet's planner to check it. That builds no object store for the block's files, so the scheduler uses no storage credentials. It then builds a Ballista physical plan, runs it, and returns one endpoint per result partition. Each endpoint carries an opaque ticket for DoGet on the executor that holds the partition, and the time the partition is deleted: Ballista deletes a query's results COMET_SCHEDULER_RESULT_RETENTION_SECONDS after the query ends (five minutes by default), whether or not they have been read.
    • comet-executor is Ballista's executor with Comet's codec. It links no libjvm.
    • Both take their settings from environment variables, with Ballista's defaults otherwise.

    Exchanges map onto Ballista's planner like this, and a stage that reads exchanges has as many partitions as they deliver:

    Exchange Ballista plan
    Hash RepartitionExec(Hash), partitioned with DataFusion's hash (Let hash shuffles assign rows to partitions with a client's own hash function datafusion-ballista#2681)
    Single partition, broadcast CoalescePartitionsExec
    Range CoalescePartitionsExec, so the sort runs in one task (Static planner drops a RepartitionExec with range partitioning instead of planning a range shuffle datafusion-ballista#2680)
  4. Tests through Flight. tests/cluster.rs starts the two binaries and runs queries the way a driver would.

  5. CI. A new linux-test-native-executors job runs clippy and the workspace's tests, since the existing Rust jobs don't cover a separate workspace. The lint job checks its formatting. The workflow comment and the README say what to do when Comet moves to a DataFusion major that Ballista hasn't released yet: the upgrade can turn the job off until Ballista catches up.

Also: the release process now refreshes the new workspace's lockfile along with the others.

Known limits:

  • A stage that several other stages read runs once for each of them, because Ballista's planner doesn't share a stage between consumers yet.
  • Nothing cancels a query whose client goes away. Cancellation comes with the driver side.
  • A query's results are deleted when their retention runs out, not once the client has read them. Releasing them when read also comes with the driver side.
  • Other Ballista changes this will want are Run submitted physical plans through the adaptive planner datafusion-ballista#2678 (AQE for submitted physical plans) and the issues linked above.

How are these changes tested?

  • Stage graphs on DataFusion alone: tests/assembly.rs (12 tests) runs them with no Ballista. It covers:
    • a RangeScan leaf stage feeding hash, range and single-partition exchanges;
    • the cast of a stage's output to its declared types, and that each batch carries the schema its stage declares, which Ballista's shuffle writer asserts;
    • the partition count of a stage that reads an exchange;
    • a check of a stage that reads files, which builds no object store;
    • rejection of mismatched columns, mismatched partition plans, a stage with no partitions, and a block that calls into the JVM.
  • The binaries through Flight: tests/cluster.rs (4 tests) starts comet-scheduler and comet-executor and runs queries through the scheduler's Flight service, reading every result partition with DoGet. It covers a hash shuffle and a range exchange, checks that a query from another Comet build is refused, and checks, with a short retention, that each endpoint says when its partition goes and that a partition no client read is gone after that. It runs in the new CI job.
  • End to end from Spark: feat: prototype Comet Native Executors built on DataFusion Ballista #7289 holds this PR's commits with the driver side on top, and its Spark suite passes against these binaries. It runs scans, aggregates, shuffled joins and TPC-H q1 on SF1 data from a Spark driver.
  • Existing tests: the core planner and jni_api unit tests pass. Clippy and cargo machete are clean for the main workspace, and clippy is clean for the new one.

A Spark driver submits a query to a Comet Native Executor cluster as a
QueryRequest: a stage graph of Comet native blocks, connected by the
exchanges between them, with each stage's Spark output schema and, for
stages that read files, each partition's plan. A Fragment carries a
stage from the scheduler to its executors.

The proto crate also exports a fingerprint of the .proto files, so that
a cluster can refuse a query from a client that speaks another plan
format. See apache#7288.
@andygrove
andygrove force-pushed the native-executors-binaries branch 3 times, most recently from a1799e5 to 7542b77 Compare October 28, 2026 16:12

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Reviewed 7542b7764300fe0195a9ebf9e54dc5ecfe31d208, including the schema-normalization update, across five independent review scopes. One result-lifetime issue needs addressing before approval; details inline. The documented stage recomputation, single-partition range exchange, and absent cancellation were treated as known limits.

Validation: inspected the exact Ballista 55.0.0 cleanup and Flight read paths. This is a source-confirmed finding, not a timed cluster reproduction. No local assembly or cluster tests ran: the configured offline registry lacked arrow-flight. Current-head checks are still queued, so there is no completed CI verdict for this revision. The cluster tests cover small non-broadcast queries; successful per-partition file plans and broadcast consumption remain unvalidated in this review.

Comment thread native/native-executors/src/bin/comet-scheduler.rs
CometFragmentExec runs a stage's native block for each of its
partitions, with the block's input leaves fed by other DataFusion plans,
such as a distributed engine's shuffle readers, instead of by the JVM.
PhysicalPlanner::with_native_inputs is the seam: an input leaf (Scan or
ShuffleScan) takes the next native input, cast to the types Spark
declared for it, and an expression that calls into the JVM fails to
plan.

Each partition's output is cast to the stage's declared Spark schema
before it leaves the stage, as a Spark task's shuffle writer casts its
input, since a hash exchange places rows by the types of their values
too. The block runs in a Comet session over the host's memory pool and
disk manager.

A partition descriptor child, an in-memory source with one row per
partition, carries each partition's index in its stage and, for a stage
that reads files, the partition's own copy of the block. An engine that
splits a stage between tasks slices it with the stage's other leaves.

block_schema plans a block with empty inputs, so that a scheduler can
check a block before running it. See apache#7288.
@andygrove
andygrove force-pushed the native-executors-binaries branch from 7542b77 to 47db6ec Compare October 28, 2026 16:49
…ista

A new Cargo workspace, native/native-executors, builds the two processes
of a Comet Native Executor cluster on the Ballista 55 crates. It is a
workspace of its own, so the main Comet build never depends on Ballista.

comet-scheduler wraps Ballista's scheduler in an Arrow Flight service.
GetFlightInfo takes a QueryRequest, checks that the client runs the same
Comet build, plans each block with Comet's planner, turns the stage
graph into a Ballista physical plan, runs it, and returns one endpoint
per result partition with an opaque ticket for the executor's DoGet.

Hash exchanges become RepartitionExec(Hash), and single-partition and
broadcast exchanges become CoalescePartitionsExec, which Ballista's
planner turns into stage boundaries. A stage that reads exchanges has
as many partitions as they deliver. Ballista can't plan a range shuffle
for a physical plan yet (apache/datafusion-ballista#2680), so a range
exchange is read as one partition, and the stage that sorts it runs in
a single task.

comet-executor is Ballista's executor with Comet's codec, which carries
each CometFragmentExec's Fragment through Ballista's stage and task
plans. It needs no JVM. Both binaries take their settings from the
environment, with Ballista's defaults otherwise. See apache#7288.
tests/cluster.rs starts the two binaries, submits stage graphs to the
scheduler's Flight service as a driver would, and reads every result
partition from the executors with DoGet. It covers a hash shuffle, a
range exchange, and a query from another Comet build, which the
scheduler refuses.
native/native-executors is a Cargo workspace of its own, so the Rust
jobs don't cover it. The new linux-test-native-executors job runs
clippy and the workspace's tests, including the ones that start a
scheduler and an executor, and the lint job checks its formatting.
@andygrove
andygrove force-pushed the native-executors-binaries branch from 47db6ec to 16f004a Compare October 28, 2026 17:05

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Re-reviewed 16f004ab2ab4500cb72814ee9acd8218bf59b4f6 across five independent scopes, including the schema-only planning change and the latest retention fix.

The result-retention finding is addressed. The scheduler exposes the retention interval, zero disables cleanup, and each Flight endpoint advertises completion time plus that same interval. I checked the pinned Ballista timestamp units and cleanup ordering; the advertised deadline is conservative. The delayed-second-endpoint test covers the reported boundary and reuses the existing cluster harness. No additional actionable finding from this pass.

Validation remains incomplete on this commit: two reported checks are queued and one has succeeded, with no completed native-executors build/test verdict. I attempted the focused cluster regression with --locked --offline, but dependency resolution stopped because the configured registry lacks arrow-flight; no local test ran. I am resolving the original code thread, but withholding approval until current-head validation completes.

@comphead comphead left a comment

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.

Carrying Comet's operator proto inside an opaque CometFragmentExec over Ballista's scheduler and Flight shuffle fits well, and the JNI path is unchanged by the planner seams as far as I can tell. sunchao's retention point is settled. Four comments, on a silent wrong-result shape, the trust model of shared processes, a JVM path that require_jvm misses, and the separate lockfile. None of this was run.

if num_partitions == 0 {
return plan_err!("The stage has no partitions");
}
if !stage.partition_plans.is_empty() && stage.partition_plans.len() != num_partitions {

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.

This accepts a stage that reads files and also reads a hash exchange whenever the counts match. In a Spark plan that shape is a bucketed scan joined with a shuffled side. CometNativeScanExec.outputPartitioning reports the bucketed scan's HashPartitioning, so EnsureRequirements shuffles only the other side, into numBuckets partitions with Murmur3. exchange turns that side into a RepartitionExec with DataFusion's hash, so partition p no longer holds the rows of bucket p, and the join returns wrong rows without an error. The doc on exchange already names this case. Could build_stage reject a stage that has partition_plans and a non-broadcast HashPartition input, so it fails instead? A one-partition file stage zipped with a SinglePartition input stays valid. I traced this by reading and have not run it.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Agreed, that shape returns wrong rows without an error. e38d3d1 makes build_stage reject a stage that has partition plans and a non-broadcast hash input, and the doc on exchange now points at that check. A one-partition file stage zipped with a SinglePartition input still assembles. file_stage_cannot_read_a_hash_exchange and file_stage_can_read_a_single_partition_exchange in tests/assembly.rs cover both.

/// this build doesn't support, or one that calls into the JVM. A problem with the files or their
/// store shows only when the block runs. Planning can block on Comet's own runtime, so don't call
/// this on an async runtime's thread.
pub fn block_schema(block: &[u8], conf: &HashMap<String, String>) -> Result<SchemaRef> {

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.

block_schema plans each block on the scheduler, and planning an ExprStruct::NativeScalarUdf loads call.library_path with libloading. The policy check reads NativeUdfPolicy::from_spark_config over the request's conf, so the client decides it, and the defaults allow any path (spark.comet.nativeUdf.enabled is true and allowedPaths is empty). So any client that reaches the Flight port can have the scheduler, and then each executor, load a library path it names and run its initializers, and the library stays loaded for later queries through the process-wide c_udf::cache. In Spark that code reaches only the submitting user's own executors, while these processes are shared. Could native executors refuse NativeScalarUdf the way require_jvm refuses JVM UDFs, or take the policy from the process environment rather than the request?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Good catch. fcb2312 makes a planner that runs without a JVM (the scheduler check and the executors) refuse NativeScalarUdf, the way require_jvm refuses JVM UDFs, so the library path in a request is never loaded there. I preferred that to a process-owned policy for the first PR. Native UDFs on these processes can come back with an allowlist that comes from the process environment. native_scalar_udf_is_refused_on_native_executors in the planner tests covers both the checking and the running planner.

ObjectStoreBackend::Local,
));
}
let (url, _, backend) = prepare_object_store_with_configs(

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.

require_jvm covers scalar subqueries and JVM UDFs, but a scan can still call into the JVM through its object store. When the options set fs.s3a.comet.credential.provider.class, s3::create_store builds a CometS3CredentialBridge, whose constructor calls JVMClasses::with_env. That runs JAVA_VM.get_unchecked() behind only a debug_assert!, so on an executor with no JVM it is undefined behaviour in a release build. The scheduler check does not catch it because without_data skips the store. encryption_support.rs and name_fold.rs also call with_env. Could with_env return a CometError when JAVA_VM is unset instead of using get_unchecked? That would turn any JVM path missed here into a plain error.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Done in fcb2312: JVMClasses::with_env takes JAVA_VM.get() and returns a CometError::Internal when no JVM was started, so the S3 credential bridge, encryption_support and anything else missed here fail with an error instead of undefined behavior. with_env_fails_without_a_jvm in jni-bridge covers it.

# This is its own workspace, with its own lockfile, so that the main Comet build never depends
# on Ballista. It builds Comet's native crates by path together with the Ballista release on the
# same DataFusion major version.
[workspace]

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.

Nothing keeps this lockfile current. .github/dependabot.yml covers only /native, which excludes this workspace, so the third-party crates here, including the network-facing tonic stack, stay at today's versions. The lock already differs from native/Cargo.lock: arrow-array, arrow-ipc and the other arrow subcrates resolve to 59.3.2 here and 59.3.1 there, so the new job tests the core crate against different arrow builds than libcomet ships with. Could this directory get a dependabot entry with the same ignore rules?

@andygrove andygrove Oct 29, 2026 •

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Added in 53ea7d4: a cargo entry for /native/native-executors with the same ignore and grouping rules as /native. The lock resolves arrow to 59.3.2 while native/Cargo.lock has 59.3.1, so the job currently tests the core crate against a newer arrow than libcomet ships with. I left that alone because the next Dependabot run on /native should close the gap, and the README now says the two lockfiles can differ between updates.

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Re-reviewed 16f004ab2ab4500cb72814ee9acd8218bf59b4f6 with five independent scopes. Retention is addressed, and I verified the current-head native-executors CI log: all 12 assembly and 4 cluster tests passed, including the delayed endpoint test.

The subsequently posted findings still need changes. I independently traced these current source paths:

  • Mixed file/hash stages accept matching partition counts without ensuring Spark bucket placement agrees with DataFusion hash placement.
  • Native-UDF policy comes from the same request that supplies the library path, allowing loading a server-accessible library during scheduler planning rather than enforcing a process-owned policy.
  • Missing JVM initialization is reachable through the S3 credential-provider bridge and reaches JAVA_VM.get_unchecked() with only a debug assertion. The name_fold path already checks for an absent JVM, so the S3 path is the substantiated example here.

The separate-workspace dependency-maintenance request is also still applicable. These are existing threads, so I have not duplicated them. Validation is source-level for those three boundaries, not a bucketed-query reproduction or executed unsafe JNI/library-loading probe. Green CI does not cover these cases. Keeping changes requested.

…Native Executor

A native executor serves every client and has no JVM. JVMClasses::with_env now returns an error when no
JVM was started, instead of dereferencing an unset VM, so a JVM path the planner doesn't screen out
(such as a scan's S3 credential provider) fails cleanly. A native UDF loads the library path its plan
names under a policy taken from the client's own configuration, so planning one is refused on a
native executor.
Spark places the rows of a bucketed scan with Murmur3 and shuffles only the other side of a join to match.
The hash exchange would place that side with DataFusion's hash, so partitions of the two sides would no
longer line up and the join would return wrong rows. build_stage now rejects the shape.
The workspace has its own Cargo.lock, which the /native entry doesn't cover. Use the same ignore and
grouping rules, and document the limits of what a stage can do on a native executor.
@github-actions github-actions Bot added the area:ffi Arrow FFI / JNI boundary label Oct 29, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:ffi Arrow FFI / JNI boundary area:scan Parquet scan / data reading enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants