diff --git a/native/core/src/execution/native_executors/fragment.rs b/native/core/src/execution/native_executors/fragment.rs index 4e4aec790e0..d926daecfd8 100644 --- a/native/core/src/execution/native_executors/fragment.rs +++ b/native/core/src/execution/native_executors/fragment.rs @@ -22,7 +22,7 @@ use std::sync::Arc; use arrow::array::{Array, AsArray, BinaryArray, RecordBatch, RecordBatchOptions, UInt32Array}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef, UInt32Type}; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{internal_datafusion_err, DataFusionError, Result}; +use datafusion::common::{internal_datafusion_err, plan_datafusion_err, DataFusionError, Result}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::execution::disk_manager::DiskManager; use datafusion::execution::memory_pool::{MemoryPool, UnboundedMemoryPool}; @@ -52,6 +52,14 @@ use crate::execution::spark_config::{ const DEFAULT_BATCH_SIZE: usize = 8192; +/// The most partitions a stage or an exchange may have. A request names the counts, and the +/// scheduler allocates for each partition, so a request can't ask for more than this. Spark's +/// own default is 200 shuffle partitions, and large queries use a few thousand. +pub const MAX_PARTITIONS: usize = 1 << 16; + +/// The largest `spark.comet.batchSize` a request may set. Each operator sizes its buffers by it. +const MAX_BATCH_SIZE: usize = 1 << 20; + fn descriptor_schema() -> SchemaRef { Arc::new(Schema::new(vec![ Field::new("partition", DataType::UInt32, false), @@ -71,6 +79,11 @@ pub fn partition_descriptor( num_partitions: usize, plans: &[Vec], ) -> Result> { + if num_partitions == 0 || num_partitions > MAX_PARTITIONS { + return Err(plan_datafusion_err!( + "A stage must have between 1 and {MAX_PARTITIONS} partitions, not {num_partitions}" + )); + } let schema = descriptor_schema(); let partitions = (0..num_partitions) .map(|index| { @@ -350,14 +363,14 @@ fn fragment_session( ) -> Result> { let mut runtime = comet_runtime_env(memory_pool); runtime.disk_manager = disk_manager; - let session = comet_session_context( - runtime, - conf.get_usize(COMET_BATCH_SIZE, DEFAULT_BATCH_SIZE), - 1, - conf, - spark_plan, - ) - .map_err(|e| DataFusionError::External(Box::new(e)))?; + let batch_size = conf.get_usize(COMET_BATCH_SIZE, DEFAULT_BATCH_SIZE); + if batch_size == 0 || batch_size > MAX_BATCH_SIZE { + return Err(plan_datafusion_err!( + "{COMET_BATCH_SIZE} must be between 1 and {MAX_BATCH_SIZE}, not {batch_size}" + )); + } + let session = comet_session_context(runtime, batch_size, 1, conf, spark_plan) + .map_err(|e| DataFusionError::External(Box::new(e)))?; Ok(Arc::new(session)) } diff --git a/native/core/src/execution/native_executors/mod.rs b/native/core/src/execution/native_executors/mod.rs index fc0b9d6b2f4..eab720f6a03 100644 --- a/native/core/src/execution/native_executors/mod.rs +++ b/native/core/src/execution/native_executors/mod.rs @@ -26,7 +26,7 @@ mod fragment; -pub use fragment::{block_schema, partition_descriptor, CometFragmentExec}; +pub use fragment::{block_schema, partition_descriptor, CometFragmentExec, MAX_PARTITIONS}; /// The version of this Comet build. A Comet Native Executor cluster requires a driver's to match /// its own, together with [`datafusion_comet_proto::PROTO_FINGERPRINT`]. diff --git a/native/native-executors/README.md b/native/native-executors/README.md index a271ba863f9..f8f60a2b1b4 100644 --- a/native/native-executors/README.md +++ b/native/native-executors/README.md @@ -29,6 +29,9 @@ This directory holds the cluster. The Spark driver side is in progress in [#7289](https://github.com/apache/datafusion-comet/pull/7289). It is its own Cargo workspace, so the main Comet build never depends on Ballista. +**Security.** The scheduler and the executors have no authentication and no TLS, and a query names +the files and object stores the executors read. Run a cluster only on a trusted network. + ## The scheduler's Flight service A client runs a query through the scheduler's Arrow Flight service: @@ -39,12 +42,17 @@ A client runs a query through the scheduler's Arrow Flight service: from another Comet build 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 (`src/assembly.rs`), runs it on the Ballista scheduler it - wraps, and answers with one endpoint per result partition. + wraps, and answers with one endpoint per result partition. It refuses a request with more + partitions or a larger batch size than a cluster plausibly needs, a stage whose block reads a + different number of inputs than the stage declares, and a block that writes files. 2. `DoGet` with each endpoint's ticket, at the endpoint's location: the executor that holds the partition. A query's results stay there for `COMET_SCHEDULER_RESULT_RETENTION_SECONDS` after the query ends (five minutes by default), whether or not a client has read them, and each endpoint's expiration time says until when. +The scheduler cancels a job whose client disconnects, or that runs longer than +`COMET_SCHEDULER_QUERY_TIMEOUT_SECONDS`. + A stage that several other stages read runs once for each of them, since Ballista's planner doesn't share a stage between consumers yet. @@ -60,8 +68,11 @@ for each partition. Ballista moves the data between stages: A stage that fails to plan is refused, not run wrongly or loaded with code it names: -- A stage can't call into the JVM: scalar subqueries and JVM UDFs fail to plan, and a scan whose - object store needs a JVM credential provider fails when it runs. +- A stage can't call into the JVM: JVM UDFs and a scalar subquery the driver has not replaced with + its value fail to plan, and a scan whose object store needs a JVM credential provider fails when + it runs. +- A stage can't write files: a Parquet writer, an Iceberg write or a shuffle writer in a block is + refused. The cluster moves data between stages itself. - A stage can't load a native UDF library. These processes serve every client, and a plan names the library path, so a client would choose what the processes load. - A stage that reads files can't also read a hash exchange. Spark places the rows of a bucketed @@ -77,7 +88,13 @@ native/target/debug/comet-executor & Sharing `native/target` with the main workspace saves building the common dependencies twice. Both binaries read their settings from environment variables, listed at the top of -`src/bin/comet-scheduler.rs` and `src/bin/comet-executor.rs`. +`src/bin/comet-scheduler.rs` and `src/bin/comet-executor.rs`. The scheduler's include: + +| Variable | Meaning | +| ------------------------------------------ | -------------------------------------------------------------------- | +| `COMET_SCHEDULER_RESULT_RETENTION_SECONDS` | How long a query's results stay readable after it ends (default 300) | +| `COMET_SCHEDULER_FLIGHT_MAX_REQUEST_BYTES` | The largest query the Flight service accepts (default 512 MiB) | +| `COMET_SCHEDULER_QUERY_TIMEOUT_SECONDS` | The longest a query may take (default 0: no limit) | ## Testing diff --git a/native/native-executors/src/assembly.rs b/native/native-executors/src/assembly.rs index 19cca64beae..87bd6726149 100644 --- a/native/native-executors/src/assembly.rs +++ b/native/native-executors/src/assembly.rs @@ -21,7 +21,7 @@ use std::collections::HashMap; use std::sync::Arc; use comet::execution::native_executors::{ - block_schema, partition_descriptor, CometFragmentExec, COMET_VERSION, + block_schema, partition_descriptor, CometFragmentExec, COMET_VERSION, MAX_PARTITIONS, }; use datafusion::common::{plan_datafusion_err, plan_err, Result}; use datafusion::physical_expr::expressions::Column; @@ -32,6 +32,8 @@ use datafusion::physical_plan::ExecutionPlanProperties; use datafusion::physical_plan::{ExecutionPlan, Partitioning}; use datafusion_comet_proto::spark_expression::expr::ExprStruct; use datafusion_comet_proto::spark_native_executors::{Fragment, QueryRequest, Stage, StageInput}; +use datafusion_comet_proto::spark_operator::operator::OpStruct; +use datafusion_comet_proto::spark_operator::Operator; use datafusion_comet_proto::spark_partitioning::partitioning::PartitioningStruct; use datafusion_comet_proto::PROTO_FINGERPRINT; use prost::Message; @@ -74,11 +76,20 @@ fn build_stage( producers: &[Arc], conf: &HashMap, ) -> Result> { - let block = stage + let plan = stage .plan .as_ref() - .ok_or_else(|| plan_datafusion_err!("The stage has no plan"))? - .encode_to_vec(); + .ok_or_else(|| plan_datafusion_err!("The stage has no plan"))?; + // The block's input leaves are fed from the stage's inputs, in order, so the two must agree. + check_block(plan, stage.inputs.len())?; + for (index, partition_plan) in stage.partition_plans.iter().enumerate() { + let partition_plan = Operator::decode(partition_plan.as_slice()).map_err(|e| { + plan_datafusion_err!("Cannot decode the plan of partition {index}: {e}") + })?; + check_block(&partition_plan, stage.inputs.len()) + .map_err(|e| e.context(format!("Plan of partition {index}")))?; + } + let block = plan.encode_to_vec(); // Spark places the rows of a bucketed scan by Murmur3 and shuffles only the other side of a // join to match. The hash exchange below would place the other side by DataFusion's hash, so // partition `p` of one side would no longer meet the rows of bucket `p` of the other, and the @@ -122,6 +133,11 @@ fn build_stage( if num_partitions == 0 { return plan_err!("The stage has no partitions"); } + if num_partitions > MAX_PARTITIONS { + return plan_err!( + "The stage has {num_partitions} partitions, more than the {MAX_PARTITIONS} allowed" + ); + } if !stage.partition_plans.is_empty() && stage.partition_plans.len() != num_partitions { return plan_err!( "The stage has {num_partitions} partitions but {} partition plans", @@ -161,6 +177,40 @@ fn build_stage( )?)) } +/// Checks a stage's block, `block`, before it is planned: it must read as many inputs as the +/// stage declares, and it must not write files. Like a native UDF library or a call into the JVM, +/// writing is something a driver's stage never does on the cluster, so a request that asks for it +/// is refused rather than run by processes that serve every client. +fn check_block(block: &Operator, inputs: usize) -> Result<()> { + let mut leaves = 0; + let mut pending = vec![block]; + while let Some(op) = pending.pop() { + let writer = match op.op_struct { + Some(OpStruct::ParquetWriter(_)) => Some("ParquetWriter"), + Some(OpStruct::IcebergWrite(_)) => Some("IcebergWrite"), + Some(OpStruct::ShuffleWriter(_)) => Some("ShuffleWriter"), + Some(OpStruct::Scan(_)) | Some(OpStruct::ShuffleScan(_)) => { + leaves += 1; + None + } + _ => None, + }; + if let Some(writer) = writer { + return plan_err!( + "A stage can't contain a {writer} operator: Comet Native Executors don't write \ + files, and the cluster moves data between stages itself" + ); + } + pending.extend(&op.children); + } + if leaves != inputs { + return plan_err!( + "The stage's block reads {leaves} inputs but the stage declares {inputs}" + ); + } + Ok(()) +} + /// The exchange that carries `producer`'s output to a stage, as a node Ballista's planner turns /// into a stage boundary. /// @@ -181,6 +231,16 @@ fn exchange( .and_then(|p| p.partitioning_struct.as_ref()); match partitioning { Some(PartitioningStruct::HashPartition(hash)) => { + let num_partitions = usize::try_from(hash.num_partitions) + .ok() + .filter(|count| (1..=MAX_PARTITIONS).contains(count)) + .ok_or_else(|| { + plan_datafusion_err!( + "A hash exchange must have between 1 and {MAX_PARTITIONS} partitions, \ + not {}", + hash.num_partitions + ) + })?; let schema = producer.schema(); let keys = hash .hash_expression @@ -198,7 +258,7 @@ fn exchange( .collect::>>()?; Ok(Arc::new(RepartitionExec::try_new( producer, - Partitioning::Hash(keys, hash.num_partitions as usize), + Partitioning::Hash(keys, num_partitions), )?)) } // Ballista can't plan a range shuffle for a physical plan yet diff --git a/native/native-executors/src/bin/comet-executor.rs b/native/native-executors/src/bin/comet-executor.rs index 9f0cfc8f78f..4416bf344ea 100644 --- a/native/native-executors/src/bin/comet-executor.rs +++ b/native/native-executors/src/bin/comet-executor.rs @@ -38,6 +38,10 @@ //! memory, or no limit when neither can be found. Memory outside the pools, such as batches in //! flight between stages, comes on top of it. //! - `COMET_EXECUTOR_WORK_DIR`: where shuffle and spill files go (default: a temporary directory) +//! +//! The executor's Flight service and its gRPC connection with the scheduler don't authenticate +//! their callers or use TLS, and a query names the files and object stores the executor reads. +//! Run a cluster only on a trusted network. use std::sync::Arc; diff --git a/native/native-executors/src/bin/comet-scheduler.rs b/native/native-executors/src/bin/comet-scheduler.rs index ba68c19a1d6..0f4a4d9bd44 100644 --- a/native/native-executors/src/bin/comet-scheduler.rs +++ b/native/native-executors/src/bin/comet-scheduler.rs @@ -29,9 +29,14 @@ //! - `COMET_GRPC_MAX_MESSAGE_SIZE`: the largest message, in bytes, of Ballista's gRPC service //! between the scheduler and its executors, such as a task with many partitions of a scan //! (default: Ballista's 16 MiB). Set it on the executors as well, for their side of the same -//! messages. The Flight service for drivers has no limit. Shuffle reads between executors keep -//! the limit of the query's `ballista.client.grpc_max_message_size` (16 MiB by default), which -//! this does not change. +//! messages. Shuffle reads between executors keep the limit of the query's +//! `ballista.client.grpc_max_message_size` (16 MiB by default), which this does not change. +//! - `COMET_SCHEDULER_FLIGHT_MAX_REQUEST_BYTES`: the largest query, in bytes, the Flight service +//! accepts from a driver. A query carries every stage's plan and the files of every scan +//! partition, so a large one is big (default `536870912`, 512 MiB). +//! - `COMET_SCHEDULER_QUERY_TIMEOUT_SECONDS`: the longest a query may take. The scheduler fails +//! and cancels a query that runs longer. A query is also cancelled when the driver that +//! submitted it disconnects. (default `0`: no limit) //! - `COMET_SCHEDULER_EVENT_LOG_DIR`: a directory to write one event log per query to, as //! Ballista's scheduler does (default: no event logs). The scheduler creates it if needed and //! doesn't start if it can't write there. A log has the query's plan and a record for each @@ -46,17 +51,22 @@ //! results, not to logs. //! //! Ballista's REST API, which lists and cancels jobs without authentication, is not served. +//! +//! Neither the Flight service nor Ballista's gRPC service authenticates its callers or uses TLS, +//! and a query names the files and object stores executors read. Run a cluster only on a trusted +//! network. use std::net::SocketAddr; use std::path::Path; use std::sync::Arc; +use std::time::Duration; use arrow_flight::flight_service_server::FlightServiceServer; use ballista_scheduler::cluster::BallistaCluster; use ballista_scheduler::config::SchedulerConfig; use ballista_scheduler::scheduler_process::{create_scheduler, start_grpc_service}; use comet_native_executors::codec::CometPhysicalCodec; -use comet_native_executors::frontend::CometFlightService; +use comet_native_executors::frontend::{CometFlightService, Limits}; use comet_native_executors::{env_var, init_logging}; use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode}; use log::info; @@ -98,6 +108,12 @@ async fn main() -> Result<(), Box> { })?; config.event_log_dir = Some(dir); } + let max_request_bytes: usize = + env_var("COMET_SCHEDULER_FLIGHT_MAX_REQUEST_BYTES")?.unwrap_or(512 << 20); + let mut limits = Limits::default(); + if let Some(seconds) = env_var::("COMET_SCHEDULER_QUERY_TIMEOUT_SECONDS")? { + limits.query_timeout = (seconds > 0).then(|| Duration::from_secs(seconds)); + } let flight_port: u16 = env_var("COMET_SCHEDULER_FLIGHT_PORT")?.unwrap_or(50060); let grpc_addr: SocketAddr = format!("{}:{}", config.bind_host, config.bind_port).parse()?; let flight_addr = SocketAddr::new(grpc_addr.ip(), flight_port); @@ -112,8 +128,8 @@ async fn main() -> Result<(), Box> { let grpc = tokio::spawn(start_grpc_service(grpc_addr, scheduler.clone())); let flight = tonic::transport::Server::builder() .add_service( - FlightServiceServer::new(CometFlightService::new(scheduler)) - .max_decoding_message_size(usize::MAX) + FlightServiceServer::new(CometFlightService::new(scheduler).with_limits(limits)) + .max_decoding_message_size(max_request_bytes) .max_encoding_message_size(usize::MAX), ) .serve(flight_addr); diff --git a/native/native-executors/src/frontend.rs b/native/native-executors/src/frontend.rs index b6077669819..14c00bd422e 100644 --- a/native/native-executors/src/frontend.rs +++ b/native/native-executors/src/frontend.rs @@ -23,6 +23,7 @@ //! also says when the partition goes: Ballista deletes a query's results a set time after the //! query ends, whether or not they have been read (see the comet-scheduler binary). +use std::future::Future; use std::sync::Arc; use std::time::{Duration, UNIX_EPOCH}; @@ -34,13 +35,14 @@ use arrow_flight::{ use ballista_core::extension::SessionConfigExt; use ballista_core::serde::protobuf::{job_status, PartitionLocation, SuccessfulJob}; use ballista_core::serde::scheduler::{Action as BallistaAction, ShuffleFileKind}; +use ballista_core::JobId; use ballista_scheduler::scheduler_server::SchedulerServer; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; use datafusion_comet_proto::spark_native_executors::QueryRequest; use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode}; use futures::stream::BoxStream; -use log::info; +use log::{info, warn}; use prost::Message; use prost_types::Timestamp; use tonic::{Request, Response, Status, Streaming}; @@ -51,15 +53,33 @@ use crate::assembly::{build_plan, check_version}; /// executors cache a runtime per session, so sharing one lets queries share those runtimes. const SESSION_ID: &str = "comet"; +/// How long a query may run. +#[derive(Clone, Debug, Default)] +pub struct Limits { + /// The longest a query may take, from submission to its end. The scheduler fails and cancels a + /// query that runs longer. `None` waits for as long as it takes. + pub query_timeout: Option, +} + /// The Flight service of a Comet scheduler, which runs queries on the Ballista scheduler it /// wraps. pub struct CometFlightService { scheduler: SchedulerServer, + limits: Limits, } impl CometFlightService { pub fn new(scheduler: SchedulerServer) -> Self { - Self { scheduler } + Self { + scheduler, + limits: Limits::default(), + } + } + + /// The service with `limits` instead of the defaults. + pub fn with_limits(mut self, limits: Limits) -> Self { + self.limits = limits; + self } /// Runs `plan` to completion and returns its result: where its output partitions are, in @@ -79,27 +99,37 @@ impl CometFlightService { .await .map_err(|e| Status::internal(format!("Cannot submit the query: {e}")))?; info!("Submitted query as job {job_id}"); - while let Some(status) = receiver.recv().await { - match status.status { - Some(job_status::Status::Successful(mut job)) => { - job.partition_location.sort_by_key(|location| { - let partition = location.partition_id.as_ref().map(|p| p.partition_id); - (partition, location.map_partition_id) - }); - return Ok(job); - } - Some(job_status::Status::Failed(job)) => { - return Err(Status::internal(format!( - "Job {job_id} failed: {}", - job.error - ))) + let mut guard = JobGuard { + scheduler: self.scheduler.clone(), + job_id: Some(job_id.clone()), + }; + let wait = async { + while let Some(status) = receiver.recv().await { + match status.status { + Some(job_status::Status::Successful(mut job)) => { + guard.finished(); + job.partition_location.sort_by_key(|location| { + let partition = location.partition_id.as_ref().map(|p| p.partition_id); + (partition, location.map_partition_id) + }); + return Ok(job); + } + Some(job_status::Status::Failed(job)) => { + guard.finished(); + return Err(Status::internal(format!( + "Job {job_id} failed: {}", + job.error + ))); + } + _ => {} } - _ => {} } - } - Err(Status::internal(format!( - "Job {job_id} ended without a final status" - ))) + guard.finished(); + Err(Status::internal(format!( + "Job {job_id} ended without a final status" + ))) + }; + within(self.limits.query_timeout, &job_id, wait).await } /// When the executors delete the results of a query that ended at `ended_at`, in milliseconds @@ -122,6 +152,58 @@ impl CometFlightService { } } +/// Cancels its job when dropped before [`finished`](Self::finished), as when the driver +/// disconnects and the request's future is dropped, or a wait times out. +struct JobGuard { + scheduler: SchedulerServer, + job_id: Option, +} + +impl JobGuard { + /// The job has ended, so there is nothing to cancel. + fn finished(&mut self) { + self.job_id = None; + } +} + +impl Drop for JobGuard { + fn drop(&mut self) { + let Some(job_id) = self.job_id.take() else { + return; + }; + let Ok(runtime) = tokio::runtime::Handle::try_current() else { + return; + }; + let scheduler = self.scheduler.clone(); + runtime.spawn(async move { + info!("Cancelling job {job_id}, which nothing is waiting for"); + if let Err(e) = scheduler.cancel_job(job_id.clone()).await { + warn!("Cannot cancel job {job_id}: {e}"); + } + }); + } +} + +/// Waits for `wait`, the end of job `job_id`, for at most `timeout`. +async fn within( + timeout: Option, + job_id: &JobId, + wait: impl Future>, +) -> Result { + let Some(timeout) = timeout else { + return wait.await; + }; + tokio::time::timeout(timeout, wait) + .await + .unwrap_or_else(|_| { + Err(Status::deadline_exceeded(format!( + "Job {job_id} did not finish within {} seconds \ + (COMET_SCHEDULER_QUERY_TIMEOUT_SECONDS) and was cancelled", + timeout.as_secs() + ))) + }) +} + /// The Flight endpoint for a result partition: the executor's location, a ticket the executor's /// Flight service reads as a Ballista partition fetch, and when the partition is deleted. fn endpoint( @@ -257,3 +339,36 @@ impl FlightService for CometFlightService { unimplemented("ListActions") } } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn a_wait_past_the_timeout_fails_as_deadline_exceeded() { + let job_id = JobId::new("job"); + let status = within( + Some(Duration::from_millis(10)), + &job_id, + std::future::pending::>(), + ) + .await + .unwrap_err(); + assert_eq!(status.code(), tonic::Code::DeadlineExceeded); + assert!( + status + .message() + .contains("COMET_SCHEDULER_QUERY_TIMEOUT_SECONDS"), + "{status}" + ); + } + + #[tokio::test] + async fn a_wait_within_the_timeout_or_without_one_returns_its_result() { + let job_id = JobId::new("job"); + for timeout in [None, Some(Duration::from_secs(60))] { + let result = within(timeout, &job_id, async { Ok::<_, Status>(7) }).await; + assert_eq!(result.unwrap(), 7); + } + } +} diff --git a/native/native-executors/tests/assembly.rs b/native/native-executors/tests/assembly.rs index 1776fc6a655..cc15883b719 100644 --- a/native/native-executors/tests/assembly.rs +++ b/native/native-executors/tests/assembly.rs @@ -20,6 +20,7 @@ mod common; +use comet::execution::native_executors::MAX_PARTITIONS; use comet_native_executors::assembly::build_plan; use common::*; use datafusion::arrow::datatypes::{DataType, Int32Type, Int64Type}; @@ -28,10 +29,10 @@ use datafusion::prelude::SessionContext; use datafusion_comet_proto::spark_expression::{ expr::ExprStruct, literal::Value, Expr, IfExpr, Literal, Subquery, }; -use datafusion_comet_proto::spark_native_executors::{QueryRequest, Stage}; +use datafusion_comet_proto::spark_native_executors::{QueryRequest, Stage, StageInput}; use datafusion_comet_proto::spark_operator::{ - operator::OpStruct, NativeScan, NativeScanCommon, Operator, Projection, SparkFilePartition, - SparkPartitionedFile, + operator::OpStruct, IcebergWrite, NativeScan, NativeScanCommon, Operator, ParquetWriter, + Projection, Scan, ShuffleWriter, SparkFilePartition, SparkPartitionedFile, }; use prost::Message; use std::sync::Arc; @@ -278,3 +279,178 @@ async fn file_stage_is_checked_without_its_object_store() { .unwrap_err(); assert!(error.to_string().contains("no-such-store"), "{error}"); } + +#[tokio::test(flavor = "multi_thread")] +async fn hash_exchange_partition_counts_are_bounded() { + for count in [i32::MIN, -1, 0, MAX_PARTITIONS as i32 + 1, i32::MAX] { + let request = query(vec![ + range_stage(100, 4), + reading_stage(0, hash_on_first_column(count), 3), + ]); + let error = assemble(request).await.unwrap_err().to_string(); + assert!( + error.contains(&format!( + "between 1 and {MAX_PARTITIONS} partitions, not {count}" + )), + "{error}" + ); + } + let request = query(vec![ + range_stage(100, 4), + reading_stage(0, hash_on_first_column(MAX_PARTITIONS as i32), 3), + ]); + let plan = assemble(request).await.unwrap(); + assert_eq!(plan.output_partitioning().partition_count(), MAX_PARTITIONS); +} + +#[tokio::test(flavor = "multi_thread")] +async fn stage_partition_counts_are_bounded() { + for count in [MAX_PARTITIONS as u32 + 1, 1 << 31, u32::MAX] { + let mut stage = range_stage(100, 4); + stage.num_partitions = count; + let error = assemble(query(vec![stage])).await.unwrap_err().to_string(); + assert!( + error.contains(&format!("more than the {MAX_PARTITIONS} allowed")), + "{error}" + ); + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn batch_size_is_bounded() { + for batch_size in ["0", "2000000", "18446744073709551615"] { + let mut request = query(vec![range_stage(100, 4)]); + request + .conf + .insert("spark.comet.batchSize".to_string(), batch_size.to_string()); + let error = assemble(request).await.unwrap_err().to_string(); + assert!( + error.contains("spark.comet.batchSize must be between 1 and"), + "{error}" + ); + } + let mut request = query(vec![range_stage(100, 4)]); + request + .conf + .insert("spark.comet.batchSize".to_string(), "4096".to_string()); + assemble(request).await.unwrap(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn block_must_read_as_many_inputs_as_the_stage_declares() { + let leaf_less_input = StageInput { + producer: 0, + partitioning: Some(single_partition()), + broadcast: false, + }; + // A block with no input leaves, and an input. + let mut stage = range_stage(100, 4); + stage.inputs = vec![leaf_less_input.clone()]; + let error = assemble(query(vec![range_stage(100, 4), stage])) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("reads 0 inputs but the stage declares 1"), + "{error}" + ); + // A block with one input leaf, and two inputs. + let mut stage = reading_stage(0, single_partition(), 1); + stage.inputs.push(leaf_less_input); + let error = assemble(query(vec![range_stage(100, 4), stage])) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("reads 1 inputs but the stage declares 2"), + "{error}" + ); + // A join's two input leaves, and one input. + let mut stage = joining_stage(0, 1, 3); + stage.inputs.pop(); + let error = assemble(query(vec![range_stage(100, 4), range_stage(100, 2), stage])) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("reads 2 inputs but the stage declares 1"), + "{error}" + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn every_partition_plan_must_read_as_many_inputs_as_the_stage_declares() { + let range = range_stage(100, 2); + let block = range.plan.as_ref().unwrap().encode_to_vec(); + let with_leaf = Operator { + op_struct: Some(OpStruct::Scan(Scan { + fields: vec![long_type()], + source: "test".to_string(), + })), + ..Default::default() + } + .encode_to_vec(); + let mut stage = range; + stage.partition_plans = vec![block, with_leaf]; + let error = assemble(query(vec![stage])).await.unwrap_err(); + assert!( + error.to_string().contains("Plan of partition 1") + && error + .to_string() + .contains("reads 1 inputs but the stage declares 0"), + "{error}" + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn writer_operators_are_refused() { + let writers = [ + ( + "ParquetWriter", + OpStruct::ParquetWriter(ParquetWriter::default()), + ), + ( + "IcebergWrite", + OpStruct::IcebergWrite(IcebergWrite::default()), + ), + ( + "ShuffleWriter", + OpStruct::ShuffleWriter(ShuffleWriter::default()), + ), + ]; + for (name, writer) in writers { + let writing = |child: Operator| Operator { + op_struct: Some(writer.clone()), + children: vec![child], + ..Default::default() + }; + let expected = format!("can't contain a {name} operator"); + let range = range_stage(100, 1); + let block = range.plan.clone().unwrap(); + + // At the root of the block. + let mut stage = range.clone(); + stage.plan = Some(writing(block.clone())); + let error = assemble(query(vec![stage])).await.unwrap_err(); + assert!(error.to_string().contains(&expected), "{error}"); + + // Below other operators. + let mut stage = range.clone(); + stage.plan = Some(Operator { + op_struct: Some(OpStruct::Projection(Projection::default())), + children: vec![writing(block.clone())], + ..Default::default() + }); + let error = assemble(query(vec![stage])).await.unwrap_err(); + assert!(error.to_string().contains(&expected), "{error}"); + + // In a partition's own copy of the block. + let mut stage = range; + stage.partition_plans = vec![writing(block).encode_to_vec()]; + let error = assemble(query(vec![stage])).await.unwrap_err(); + assert!(error.to_string().contains(&expected), "{error}"); + } +} diff --git a/native/native-executors/tests/cluster.rs b/native/native-executors/tests/cluster.rs index 2ed61c2cebf..f396645261c 100644 --- a/native/native-executors/tests/cluster.rs +++ b/native/native-executors/tests/cluster.rs @@ -232,6 +232,24 @@ async fn runs_a_query_across_a_range_exchange() { ); } +#[tokio::test(flavor = "multi_thread")] +async fn refuses_a_query_larger_than_the_request_limit() { + let cluster = Cluster::start_with(&[( + "COMET_SCHEDULER_FLIGHT_MAX_REQUEST_BYTES", + "4096".to_string(), + )]); + cluster.run(&query(vec![range_stage(10, 1)])).await; + let mut large = query(vec![range_stage(10, 1)]); + large.conf.insert("padding".to_string(), "x".repeat(8192)); + match cluster.get_flight_info(&large).await { + Err(FlightError::Tonic(status)) => assert!( + matches!(status.code(), Code::OutOfRange | Code::ResourceExhausted), + "{status}" + ), + other => panic!("Expected the query to be refused, got {other:?}"), + } +} + #[tokio::test(flavor = "multi_thread")] async fn refuses_a_query_from_another_build() { let cluster = Cluster::start(); diff --git a/native/native-executors/tests/common/mod.rs b/native/native-executors/tests/common/mod.rs index 5dc6b530cf7..f80781eef4d 100644 --- a/native/native-executors/tests/common/mod.rs +++ b/native/native-executors/tests/common/mod.rs @@ -28,7 +28,8 @@ use datafusion_comet_proto::spark_expression::data_type::DataTypeId; use datafusion_comet_proto::spark_expression::{self, expr::ExprStruct, BoundReference, Expr}; use datafusion_comet_proto::spark_native_executors::{QueryRequest, Stage, StageInput}; use datafusion_comet_proto::spark_operator::{ - operator::OpStruct, Operator, RangeScan, Scan, SparkStructField, + operator::OpStruct, BuildSide, HashJoin, JoinType, Operator, RangeScan, Scan, ShuffleScan, + SparkStructField, }; use datafusion_comet_proto::spark_partitioning::{ partitioning::PartitioningStruct, HashPartition, Partitioning, RangePartition, SinglePartition, @@ -162,3 +163,38 @@ where values.sort_unstable(); values } + +/// A stage that joins what stages `probe` and `build` send it through hash exchanges of +/// `num_partitions` partitions on their first column, as an inner hash join built on the build +/// side, and returns both sides' columns. +pub fn joining_stage(probe: u32, build: u32, num_partitions: i32) -> Stage { + let leaf = || Operator { + op_struct: Some(OpStruct::ShuffleScan(ShuffleScan { + fields: vec![long_type()], + source: "test".to_string(), + })), + ..Default::default() + }; + let input = |producer| StageInput { + producer, + partitioning: Some(hash_on_first_column(num_partitions)), + broadcast: false, + }; + Stage { + plan: Some(Operator { + children: vec![leaf(), leaf()], + op_struct: Some(OpStruct::HashJoin(HashJoin { + left_join_keys: vec![bound(0, long_type())], + right_join_keys: vec![bound(0, long_type())], + join_type: JoinType::Inner as i32, + build_side: BuildSide::BuildRight as i32, + ..Default::default() + })), + ..Default::default() + }), + inputs: vec![input(probe), input(build)], + num_partitions: num_partitions as u32, + output_schema: vec![field(long_type()), field(long_type())], + ..Default::default() + } +}