Skip to content

fix: bound and validate requests to the Comet Native Executors cluster - #7546

Draft
andygrove wants to merge 2 commits into
apache:mainfrom
andygrove:native-executors-hardening
Draft

andygrove wants to merge 2 commits into
apache:mainfrom
andygrove:native-executors-hardening

Conversation

@andygrove

Copy link
Copy Markdown
Member

Which issue does this PR close?

Part of #7288.

Rationale for this change

The Comet Native Executors scheduler and executors accept a query request from the network. The counts, sizes and operators in it must be bounded and checked before the planner allocates for them or runs them. Today they are not:

  • A hash exchange with num_partitions of -1 became usize::MAX partitions, and the scheduler tried to allocate for each one.
  • A stage with 2^31 partitions, or a spark.comet.batchSize of 0 or usize::MAX, was accepted and sized every buffer from it.
  • A block with fewer input leaves than stage.inputs was accepted, although the block is fed from the stage's inputs in order.
  • A block containing a Parquet writer, an Iceberg write or a shuffle writer was accepted, so a request could ask the processes that serve every client to write files.
  • The Flight service accepted a request of any size (usize::MAX), and a query ran for as long as it took, even after the driver that submitted it disconnected.

What changes are included in this PR?

  • native/core/src/execution/native_executors: MAX_PARTITIONS (65536) and a MAX_BATCH_SIZE (1048576). partition_descriptor refuses a stage with no partitions or more than MAX_PARTITIONS, and fragment_session refuses a spark.comet.batchSize outside 1..=MAX_BATCH_SIZE.
  • native/native-executors/src/assembly.rs: validates a hash exchange's num_partitions (negative, zero and over the cap) and a stage's partition count. A new check_block requires a stage's block, and each entry of partition_plans, to read as many inputs as the stage declares, and refuses a ParquetWriter, IcebergWrite or ShuffleWriter operator.
  • comet-scheduler: COMET_SCHEDULER_FLIGHT_MAX_REQUEST_BYTES limits the size of a request (512 MiB by default, replacing no limit). COMET_SCHEDULER_QUERY_TIMEOUT_SECONDS fails and cancels a query that runs longer (default 0: no limit). A query is also cancelled when the driver that submitted it disconnects.
  • native/native-executors/README.md and the binaries' documentation: the cluster has no authentication or TLS and must run on a trusted network, the new settings, the refusals above, and the wording on scalar subqueries.

How are these changes tested?

Run locally:

  • cargo fmt --check, cargo clippy --all-targets -- -D warnings and cargo check --all-targets in native/native-executors.
  • cargo test --test assembly --lib in native/native-executors: 20 tests in tests/assembly.rs and 2 unit tests of the query timeout in src/frontend.rs. The new tests cover bounded hash exchange and stage partition counts (including i32::MIN, -1, 0, one over the cap and u32::MAX), the batch size bound, a block and a partition plan that read the wrong number of inputs, and the three refused writers at the root of a block, below another operator and in a partition's own plan.
  • cargo clippy -p datafusion-comet --all-targets -- -D warnings and cargo test -p datafusion-comet --lib native_executors in native.
  • prettier --check on the README.

Not run: tests/cluster.rs, which starts the scheduler and executor binaries and needs sockets, which the environment I developed in does not allow. It compiles (cargo check --all-targets). The new refuses_a_query_larger_than_the_request_limit test in it, and the cancellation of a job when a driver disconnects or a query times out, which no test exercises, have not been run.

This is split out of #7289 so that it can be reviewed on its own against the published Ballista 55.0.0 crates, which this workspace builds on. #7289 carries the limits for pipelined mode.

… before planning them

A query request names the number of partitions of each stage and exchange, the
batch size, and the operators of each stage's block. The scheduler allocated for
each partition and planned each block as given, so a count of -1 became a
usize::MAX allocation, 2^31 partitions exhausted memory, and a block with
fewer input leaves than its stage declares failed only when it ran.

Bound the partition count of a stage and of a hash exchange to between 1 and
65536, bound spark.comet.batchSize to between 1 and 1048576, require a block
and each of its partition plans to read as many inputs as the stage declares,
and refuse a block that writes files (a Parquet writer, an Iceberg write or a
shuffle writer).
Limit the size of a query the Flight service accepts with
COMET_SCHEDULER_FLIGHT_MAX_REQUEST_BYTES (512 MiB by default, where there was
no limit), and add COMET_SCHEDULER_QUERY_TIMEOUT_SECONDS to fail and cancel a
query that runs too long. The scheduler also cancels a job when the driver
that submitted it disconnects, instead of letting it run on.

State in the README and the binaries' documentation that the cluster has no
authentication or TLS and must run on a trusted network, and document the new
settings and the refusals.
@github-actions github-actions Bot added the bug Something isn't working label Nov 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant