andygrove opened a new pull request, #7546: URL: https://github.com/apache/datafusion-comet/pull/7546
## 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. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
