diff --git a/benchmarks/cdk/bin/worker.rs b/benchmarks/cdk/bin/worker.rs index 4026b501..cc7974bf 100644 --- a/benchmarks/cdk/bin/worker.rs +++ b/benchmarks/cdk/bin/worker.rs @@ -12,7 +12,8 @@ use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::physical_plan::execute_stream; use datafusion::prelude::SessionContext; use datafusion_distributed::test_utils::work_unit_file_scan::{ - WorkUnitFileScanCodec, WorkUnitFileScanConfig, WorkUnitFileScanTaskEstimator, + WorkUnitFileScanCodec, WorkUnitFileScanConfig, work_unit_file_scan_desired_task_count, + work_unit_file_scan_scale_up_leaf_node, }; use datafusion_distributed::{ ChannelResolver, DistributedExt, DistributedMetricsFormat, NetworkBoundaryExt, @@ -121,7 +122,8 @@ async fn main() -> Result<(), Box> { // Uncomment for enabling WorkUnitFileScans. // .with_physical_optimizer_rule(Arc::new(WorkUnitFileScanRule)) .with_distributed_user_codec(WorkUnitFileScanCodec) - .with_distributed_task_estimator(WorkUnitFileScanTaskEstimator) + .with_distributed_event_handler(work_unit_file_scan_desired_task_count) + .with_distributed_event_handler(work_unit_file_scan_scale_up_leaf_node) .with_distributed_work_unit_feed(|dse: &DataSourceExec| { dse.data_source() .downcast_ref::() diff --git a/benchmarks/src/run.rs b/benchmarks/src/run.rs index 62a81644..bc1585ec 100644 --- a/benchmarks/src/run.rs +++ b/benchmarks/src/run.rs @@ -30,7 +30,7 @@ use datafusion::prelude::*; use datafusion_distributed::test_utils::localhost::LocalHostWorkerResolver; use datafusion_distributed::test_utils::work_unit_file_scan::{ WorkUnitFileScanCodec, WorkUnitFileScanConfig, WorkUnitFileScanRule, - WorkUnitFileScanTaskEstimator, + work_unit_file_scan_desired_task_count, work_unit_file_scan_scale_up_leaf_node, }; use datafusion_distributed::{ DistributedExt, DistributedMetricsFormat, NetworkBoundaryExt, SessionStateBuilderExt, Worker, @@ -216,7 +216,8 @@ impl RunOpt { )? .with_distributed_max_tasks_per_stage(self.max_tasks_per_stage)? .with_distributed_user_codec(WorkUnitFileScanCodec) - .with_distributed_task_estimator(WorkUnitFileScanTaskEstimator) + .with_distributed_event_handler(work_unit_file_scan_desired_task_count) + .with_distributed_event_handler(work_unit_file_scan_scale_up_leaf_node) .with_distributed_work_unit_feed(|dse: &DataSourceExec| { dse.data_source() .downcast_ref::() diff --git a/examples/custom_execution_plan.rs b/examples/custom_execution_plan.rs index e322139c..69cc2c95 100644 --- a/examples/custom_execution_plan.rs +++ b/examples/custom_execution_plan.rs @@ -39,8 +39,9 @@ use datafusion_distributed::test_utils::in_memory_channel_resolver::{ InMemoryChannelResolver, InMemoryWorkerResolver, }; use datafusion_distributed::{ - DistributedExt, DistributedTaskContext, SessionStateBuilderExt, TaskEstimation, TaskEstimator, - WorkerQueryContext, display_plan_ascii, + DesiredTaskCountEvent, DesiredTaskCountEventResponse, DistributedExt, DistributedTaskContext, + ScaleUpLeafNodeEvent, ScaleUpLeafNodeEventResponse, SessionStateBuilderExt, WorkerQueryContext, + display_plan_ascii, }; use datafusion_proto::physical_plan::PhysicalExtensionCodec; use datafusion_proto::protobuf; @@ -297,47 +298,36 @@ impl ConfigExtension for NumbersConfig { const PREFIX: &'static str = "numbers"; } -/// Custom TaskEstimator that tells the planner how to distribute NumbersExec. -#[derive(Debug)] -struct NumbersTaskEstimator; - -impl TaskEstimator for NumbersTaskEstimator { - fn task_estimation( - &self, - plan: &Arc, - cfg: &datafusion::config::ConfigOptions, - ) -> Option { - let plan = plan.downcast_ref::()?; - let cfg: &NumbersConfig = cfg.extensions.get()?; - let task_count = (plan.ranges_per_task[0].end - plan.ranges_per_task[0].start) as f64 - / cfg.numbers_per_task as f64; - - Some(TaskEstimation::desired(task_count.ceil() as usize)) - } - - fn scale_up_leaf_node( - &self, - plan: &Arc, - task_count: usize, - _cfg: &datafusion::config::ConfigOptions, - ) -> Result>> { - let Some(plan) = plan.downcast_ref::() else { - return Ok(None); - }; - let range = &plan.ranges_per_task[0]; - let chunk_size = ((range.end - range.start) as f64 / task_count as f64).ceil() as i64; - - let ranges_per_task = (0..task_count).map(|i| { - let start = range.start + (i as i64 * chunk_size); - let end = (start + chunk_size).min(range.end); - start..end - }); +fn numbers_desired_task_count_handler( + ev: DesiredTaskCountEvent, +) -> Option { + let cfg = ev.session_config; + let plan = ev.plan.downcast_ref::()?; + let cfg: &NumbersConfig = cfg.options().extensions.get()?; + let task_count = (plan.ranges_per_task[0].end - plan.ranges_per_task[0].start) as f64 + / cfg.numbers_per_task as f64; + + Some(DesiredTaskCountEventResponse::desired( + task_count.ceil() as usize + )) +} - Ok(Some(Arc::new(NumbersExec::new( - ranges_per_task, - plan.schema(), - )))) - } +fn numbers_scale_up_leaf_node_handler( + ev: ScaleUpLeafNodeEvent, +) -> Option> { + let plan = ev.plan.downcast_ref::()?; + let range = &plan.ranges_per_task[0]; + let chunk_size = ((range.end - range.start) as f64 / ev.task_count as f64).ceil() as i64; + + let ranges_per_task = (0..ev.task_count).map(|i| { + let start = range.start + (i as i64 * chunk_size); + let end = (start + chunk_size).min(range.end); + start..end + }); + + Some(Ok(ScaleUpLeafNodeEventResponse::new(Arc::new( + NumbersExec::new(ranges_per_task, plan.schema()), + )))) } #[derive(StructOpt)] @@ -381,7 +371,8 @@ async fn main() -> Result<(), Box> { .with_distributed_channel_resolver(channel_resolver) .with_distributed_planner() .with_distributed_user_codec(NumbersExecCodec) - .with_distributed_task_estimator(NumbersTaskEstimator) + .with_distributed_event_handler(numbers_desired_task_count_handler) + .with_distributed_event_handler(numbers_scale_up_leaf_node_handler) .build(); let ctx = SessionContext::from(state); diff --git a/examples/custom_worker_url_routing.rs b/examples/custom_worker_url_routing.rs index fed8e7b5..7742a9e4 100644 --- a/examples/custom_worker_url_routing.rs +++ b/examples/custom_worker_url_routing.rs @@ -6,9 +6,9 @@ //! without changes to the table registration. //! //! Routing is a two-step pipeline: -//! - [`CachedFileScanConfigTaskEstimator::scale_up_leaf_node`] assigns each file to a task slot by +//! - `cached_file_scan_scale_up_leaf_node_handler` assigns each file to a task slot by //! hashing its path (mod task_count), so the same file always lands in the same slot. -//! - [`CachedFileScanConfigTaskEstimator::route_tasks`] maps slot `i` to `sorted_urls[i % n]`, +//! - `cached_file_scan_route_tasks_handler` maps slot `i` to `sorted_urls[i % n]`, //! so each slot always reaches the same worker URL. //! //! Together these guarantee that each worker consistently reads the same set of files and its @@ -41,8 +41,9 @@ use datafusion_distributed::test_utils::localhost::{ LocalHostWorkerResolver, spawn_worker_service, }; use datafusion_distributed::{ - DistributedExt, DistributedGetterExt, DistributedLeafExec, SessionStateBuilderExt, - TaskEstimation, TaskEstimator, TaskRoutingContext, WorkerQueryContext, display_plan_ascii, + DesiredTaskCountEvent, DesiredTaskCountEventResponse, DistributedExt, DistributedGetterExt, + DistributedLeafExec, RouteTasksEvent, RouteTasksEventResponse, ScaleUpLeafNodeEvent, + ScaleUpLeafNodeEventResponse, SessionStateBuilderExt, WorkerQueryContext, display_plan_ascii, }; use datafusion_proto::physical_plan::PhysicalExtensionCodec; use datafusion_proto::protobuf; @@ -55,7 +56,6 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use structopt::StructOpt; use tokio::net::TcpListener; -use url::Url; /// Worker-level cache shared across all task invocations on the same worker, keyed by a stable /// hash of the file group being scanned. @@ -157,68 +157,56 @@ fn hash_key(file_group: &FileGroup) -> usize { hasher.finish() as usize } -/// Assigns each parquet file to a task slot (by hashing its path) and pins each slot to a worker. -#[derive(Debug)] -struct CachedFileScanConfigTaskEstimator; +fn cached_file_scan_desired_task_count_handler( + ev: DesiredTaskCountEvent, +) -> Option { + ev.plan.downcast_ref::()?; + Some(DesiredTaskCountEventResponse::desired(usize::MAX)) +} -impl TaskEstimator for CachedFileScanConfigTaskEstimator { - fn task_estimation( - &self, - plan: &Arc, - _: &ConfigOptions, - ) -> Option { - plan.downcast_ref::()?; - Some(TaskEstimation::desired(usize::MAX)) +fn cached_file_scan_scale_up_leaf_node_handler( + ev: ScaleUpLeafNodeEvent, +) -> Option> { + let cfg = ev.session_config; + let cache_exec = ev.plan.downcast_ref::()?; + let dse = cache_exec.child.downcast_ref::()?; + let fsc = dse.data_source().downcast_ref::()?; + + // Hash each file to a slot so that the same file always lands in the same variant + // regardless of the original file_groups layout. + let mut per_task_files: Vec> = vec![vec![]; ev.task_count]; + for file in fsc.file_groups.iter().flat_map(|g| g.files()) { + let idx = hash_key(&FileGroup::new(vec![file.clone()])) % ev.task_count; + per_task_files[idx].push(file.clone()); } - fn scale_up_leaf_node( - &self, - plan: &Arc, - task_count: usize, - cfg: &ConfigOptions, - ) -> Result>> { - let Some(cache_exec) = plan.downcast_ref::() else { - return Ok(None); - }; - let Some(dse) = cache_exec.child.downcast_ref::() else { - return Ok(None); - }; - let Some(fsc) = dse.data_source().downcast_ref::() else { - return Ok(None); - }; - - // Hash each file to a slot so that the same file always lands in the same variant - // regardless of the original file_groups layout. - let mut per_task_files: Vec> = vec![vec![]; task_count]; - for file in fsc.file_groups.iter().flat_map(|g| g.files()) { - let idx = hash_key(&FileGroup::new(vec![file.clone()])) % task_count; - per_task_files[idx].push(file.clone()); - } + let target_partitions = cfg.target_partitions(); + let variants = (0..ev.task_count) + .map(|i| { + let files = std::mem::take(&mut per_task_files[i]); + // Spread files across up to `target_partitions` FileGroups so DataFusion can + // read them in parallel within the task. + let n_groups = files.len().clamp(1, target_partitions); + let mut groups: Vec> = vec![vec![]; n_groups]; + for (j, file) in files.into_iter().enumerate() { + groups[j % n_groups].push(file); + } + let mut new_fsc = fsc.clone(); + new_fsc.file_groups = groups.into_iter().map(FileGroup::new).collect(); + CacheExec::new(DataSourceExec::from_data_source(new_fsc)) as Arc + }) + .collect::>(); - let target_partitions = cfg.execution.target_partitions; - let variants = (0..task_count) - .map(|i| { - let files = std::mem::take(&mut per_task_files[i]); - // Spread files across up to `target_partitions` FileGroups so DataFusion can - // read them in parallel within the task. - let n_groups = files.len().clamp(1, target_partitions); - let mut groups: Vec> = vec![vec![]; n_groups]; - for (j, file) in files.into_iter().enumerate() { - groups[j % n_groups].push(file); - } - let mut new_fsc = fsc.clone(); - new_fsc.file_groups = groups.into_iter().map(FileGroup::new).collect(); - CacheExec::new(DataSourceExec::from_data_source(new_fsc)) as Arc - }) - .collect::>(); - - Ok(Some(Arc::new(DistributedLeafExec::try_new( - Arc::clone(plan), - variants, - )?))) - } + Some( + DistributedLeafExec::try_new(Arc::clone(ev.plan), variants) + .map(|plan| ScaleUpLeafNodeEventResponse::new(Arc::new(plan))), + ) +} - fn route_tasks(&self, ctx: &TaskRoutingContext<'_>) -> Result>> { +fn cached_file_scan_route_tasks_handler( + ctx: RouteTasksEvent, +) -> Option> { + (|| -> Result> { let available_urls = ctx .task_ctx .session_config() @@ -242,8 +230,9 @@ impl TaskEstimator for CachedFileScanConfigTaskEstimator { } Ok(TreeNodeRecursion::Continue) })?; - Ok(routed) - } + Ok(routed.map(RouteTasksEventResponse::new)) + })() + .transpose() } /// Codec for [`CacheExec`]. The child (`DataSourceExec(FileScanConfig)`) is encoded by the @@ -354,7 +343,9 @@ async fn main() -> Result<(), Box> { .with_distributed_worker_resolver(LocalHostWorkerResolver::new(ports)) .with_distributed_planner() .with_distributed_user_codec(CachedFileScanCodec) - .with_distributed_task_estimator(CachedFileScanConfigTaskEstimator) + .with_distributed_event_handler(cached_file_scan_desired_task_count_handler) + .with_distributed_event_handler(cached_file_scan_scale_up_leaf_node_handler) + .with_distributed_event_handler(cached_file_scan_route_tasks_handler) .build(); state .config_mut() diff --git a/examples/work_unit_feed.rs b/examples/work_unit_feed.rs index 790c32d4..b45e199c 100644 --- a/examples/work_unit_feed.rs +++ b/examples/work_unit_feed.rs @@ -30,8 +30,9 @@ use datafusion_distributed::test_utils::in_memory_channel_resolver::{ InMemoryChannelResolver, InMemoryWorkerResolver, }; use datafusion_distributed::{ - DistributedExt, DistributedTaskContext, SessionStateBuilderExt, TaskEstimation, TaskEstimator, - WorkUnitFeed, WorkUnitFeedProto, WorkUnitFeedProvider, WorkerQueryContext, display_plan_ascii, + DesiredTaskCountEvent, DesiredTaskCountEventResponse, DistributedExt, DistributedTaskContext, + ScaleUpLeafNodeEvent, ScaleUpLeafNodeEventResponse, SessionStateBuilderExt, WorkUnitFeed, + WorkUnitFeedProto, WorkUnitFeedProvider, WorkerQueryContext, display_plan_ascii, }; use datafusion_proto::physical_plan::PhysicalExtensionCodec; use datafusion_proto::protobuf::proto_error; @@ -242,41 +243,30 @@ impl PhysicalExtensionCodec for RemoteScanExecCodec { } } -/// Tells the planner how many tasks the leaf stage gets, and rebuilds the leaf so each task -/// advertises its share of the partitions. -#[derive(Debug)] -struct RemoteScanTaskEstimator; - -impl TaskEstimator for RemoteScanTaskEstimator { - fn task_estimation( - &self, - plan: &Arc, - _: &datafusion::config::ConfigOptions, - ) -> Option { - let task_count = plan - .downcast_ref::()? - .feed - .inner()? - .task_count; - Some(TaskEstimation::desired(task_count)) - } +fn remote_scan_desired_task_count_handler( + ev: DesiredTaskCountEvent, +) -> Option { + let task_count = ev + .plan + .downcast_ref::()? + .feed + .inner()? + .task_count; + Some(DesiredTaskCountEventResponse::desired(task_count)) +} - fn scale_up_leaf_node( - &self, - plan: &Arc, - task_count: usize, - _: &datafusion::config::ConfigOptions, - ) -> Result>> { - let Some(exec) = plan.downcast_ref::() else { - return Ok(None); - }; - let partitions_per_task = exec.feed.try_inner()?.per_partition_chunks.len() / task_count; - Ok(Some(Arc::new(RemoteScanExec::new( +fn remote_scan_scale_up_leaf_node_handler( + ev: ScaleUpLeafNodeEvent, +) -> Option> { + let exec = ev.plan.downcast_ref::()?; + Some(exec.feed.try_inner().map(|feed| { + let partitions_per_task = feed.per_partition_chunks.len() / ev.task_count; + ScaleUpLeafNodeEventResponse::new(Arc::new(RemoteScanExec::new( exec.feed.clone(), partitions_per_task, exec.projection.clone(), - )))) - } + ))) + })) } /// `scan(task_count, 'chunks_p0', 'chunks_p1', ...)` — `task_count` tasks, with one @@ -392,7 +382,8 @@ async fn main() -> Result<(), Box> { .with_distributed_channel_resolver(channel_resolver) .with_distributed_planner() .with_distributed_user_codec(RemoteScanExecCodec) - .with_distributed_task_estimator(RemoteScanTaskEstimator) + .with_distributed_event_handler(remote_scan_desired_task_count_handler) + .with_distributed_event_handler(remote_scan_scale_up_leaf_node_handler) // For every `RemoteScanExec`, hand the planner the feed it must drive from the coordinator. .with_distributed_work_unit_feed(|exec: &RemoteScanExec| Some(&exec.feed)) .build(); diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index a1e123f2..47ca7bc9 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -20,7 +20,7 @@ use std::sync::{Arc, Mutex}; /// [ExecutionPlan] that executes the inner plan in distributed mode. /// Before executing it, two modifications are lazily performed on the plan: /// 1. Assigns worker URLs to all the stages. Unless explicitly set in -/// [crate::TaskEstimator::route_tasks], a random set of URLs are sampled from the +/// [`crate::EventHandler`], a random set of URLs are sampled from the /// channel resolver and assigned to each task in each stage. /// 2. Encodes all the plans in protobuf format so that network boundary nodes can send them /// over the wire. diff --git a/src/coordinator/prepare_dynamic_plan.rs b/src/coordinator/prepare_dynamic_plan.rs index d72d0af6..d0ebc5d2 100644 --- a/src/coordinator/prepare_dynamic_plan.rs +++ b/src/coordinator/prepare_dynamic_plan.rs @@ -1,4 +1,3 @@ -use crate::TaskCountAnnotation::{Desired, Maximum}; use crate::common::{TreeNodeExt, element_wise_sum, vec_avg_reduce, vec_div, vec_mul}; use crate::coordinator::distributed::PreparedPlan; use crate::coordinator::query_coordinator::QueryCoordinator; @@ -6,6 +5,7 @@ use crate::distributed_planner::{ InjectNetworkBoundaryContext, NetworkBoundaryBuilderResult, ProducerHead, calculate_cost, inject_network_boundaries, }; +use crate::events::TaskCountAnnotation::{Desired, Maximum}; use crate::execution_plans::SamplerExec; use crate::stage::{LocalStage, RemoteStage}; use crate::{ @@ -70,7 +70,7 @@ pub(super) async fn prepare_dynamic_plan( .merge(Desired(compute_based_task_count)); // Propagate the final task_count inferred based on runtime statistics and compute cost. - // Here is where leaf nodes are scaled up by TaskEstimator::scale_up_leaf_node, and the + // Here is where leaf nodes are scaled up by EventHandler, and the // plan is finally left ready for distribution. input_stage.plan = nb_ctx .propagate_task_count_until_network_boundaries(&input_stage.plan, task_count)?; diff --git a/src/coordinator/prepare_static_plan.rs b/src/coordinator/prepare_static_plan.rs index 3bfc2d2f..a9e36708 100644 --- a/src/coordinator/prepare_static_plan.rs +++ b/src/coordinator/prepare_static_plan.rs @@ -9,7 +9,7 @@ use std::sync::Arc; /// Prepares the distributed plan for execution, which implies: /// 1. Perform some worker URL assignation, choosing either: -/// - The URLs set by the user with [crate::TaskEstimator::route_tasks]. +/// - The URLs set by the user with [`crate::EventHandler`]. /// - Randomly otherwise /// 2. Sending the sliced subplans to the assigned URLs. For each URL assigned to a task, a /// network call feeding the subplan is necessary. diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index f317f2a6..2b71e04b 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -2,7 +2,7 @@ use crate::common::{TreeNodeExt, now_ns, task_ctx_with_extension}; use crate::config_extension_ext::get_config_extension_propagation_headers; use crate::coordinator::MetricsStore; use crate::coordinator::latency_metric::LatencyMetric; -use crate::distributed_planner::CombinedTaskEstimator; +use crate::events::{RouteTasksEvent, RouteTasksHandler}; use crate::execution_plans::{ChildrenIsolatorUnionExec, DistributedLeafExec}; use crate::passthrough_headers::get_passthrough_headers; use crate::stage::LocalStage; @@ -12,9 +12,9 @@ use crate::worker::LocalWorkerContext; use crate::{ BytesCounterMetric, BytesMetricExt, CoordinatorToWorkerMsg, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedCodec, DistributedTaskContext, - DistributedWorkUnitFeedContext, LoadInfo, NetworkBoundaryExt, SetPlanRequest, Stage, - TaskEstimator, TaskKey, TaskRoutingContext, WorkUnitFeedDeclaration, WorkerToCoordinatorMsg, - get_distributed_channel_resolver, get_distributed_worker_resolver, + DistributedWorkUnitFeedContext, LoadInfo, NetworkBoundaryExt, SetPlanRequest, Stage, TaskKey, + WorkUnitFeedDeclaration, WorkerToCoordinatorMsg, get_distributed_channel_resolver, + get_distributed_worker_resolver, }; use datafusion::common::DataFusionError; use datafusion::common::instant::Instant; @@ -375,28 +375,27 @@ impl<'a> StageCoordinator<'a> { /// is managing. These URLs can be: /// - assigned randomly, if the user did not provide any custom routing. /// - chosen by the user, if they provided an implementation for the - /// [TaskEstimator::route_tasks] method. + /// [`crate::EventHandler`] implementation. pub(super) fn routed_urls(&self) -> Result> { - let session_config = self.task_ctx.session_config(); - let worker_resolver = get_distributed_worker_resolver(session_config)?; - let task_estimator = CombinedTaskEstimator::from_session_config(session_config); - - let routed_urls = match task_estimator.route_tasks(&TaskRoutingContext { + let ev = RouteTasksEvent { task_ctx: Arc::clone(self.task_ctx), plan: self.plan, task_count: self.task_count, - }) { - Ok(Some(routed_urls)) => routed_urls, + }; + let routed_urls = match RouteTasksHandler::handle(ev) { + Some(Ok(response)) => response.urls, // If the user has not defined custom routing with a `route_tasks` implementation, we // default to round-robin task assignation from a randomized starting point. - Ok(None) => { + None => { + let session_config = self.task_ctx.session_config(); + let worker_resolver = get_distributed_worker_resolver(session_config)?; let available_urls = worker_resolver.get_urls()?; let start_idx = rand::rng().random_range(0..available_urls.len()); (0..self.task_count) .map(|i| available_urls[(start_idx + i) % available_urls.len()].clone()) .collect() } - Err(e) => return exec_err!("error routing tasks to workers: {e}"), + Some(Err(e)) => return exec_err!("error routing tasks to workers: {e}"), }; if routed_urls.len() != self.task_count { diff --git a/src/distributed_ext.rs b/src/distributed_ext.rs index 49e1f1af..20b453fc 100644 --- a/src/distributed_ext.rs +++ b/src/distributed_ext.rs @@ -2,14 +2,14 @@ use crate::codec::{set_distributed_user_codec, set_distributed_user_codec_arc}; use crate::config_extension_ext::{ set_distributed_option_extension, set_distributed_option_extension_from_headers, }; -use crate::distributed_planner::set_distributed_task_estimator; +use crate::events::{Event, EventHandler, EventHandlerChain}; use crate::passthrough_headers::set_passthrough_headers; use crate::protocol::set_distributed_channel_resolver; use crate::work_unit_feed::set_distributed_work_unit_feed; use crate::worker_resolver::set_distributed_worker_resolver; use crate::{ - ChannelResolver, DistributedConfig, LocalWorkerContext, TaskEstimator, WorkUnitFeed, - WorkUnitFeedProvider, WorkerResolver, get_distributed_worker_resolver, + ChannelResolver, DistributedConfig, LocalWorkerContext, WorkUnitFeed, WorkUnitFeedProvider, + WorkerResolver, get_distributed_worker_resolver, }; use datafusion::common::DataFusionError; use datafusion::config::ConfigExtension; @@ -278,46 +278,6 @@ pub trait DistributedExt: Sized { resolver: T, ); - /// Adds a distributed task count estimator. [TaskEstimator]s are executed on each node - /// sequentially until one returns an estimation on the number of tasks that should be - /// used for the stage containing that node. - /// - /// Many nodes might decide to provide an estimation, so a reconciliation between all of them - /// is performed internally during planning. - /// - /// ```text - /// ┌───────────────────────┐ - /// │SortPreservingMergeExec│ - /// └───────────────────────┘ - /// ▲ - /// ┌ ─ ─ ─ ─ ─ ─ ─ ┼ ─ ─ ─ ─ ─ ─ ─ ─ Stage 2 - /// ┌───────────┴───────────┐ │ - /// │ │ SortExec │ - /// └───────────────────────┘ │ - /// │ ┌───────────────────────┐ - /// │ AggregateExec │ │ - /// │ └───────────────────────┘ - /// ─ ─ ─ ─ ─ ─ ─ ─▲─ ─ ─ ─ ─ ─ ─ ─ ┘ - /// ┌ ─ ─ ─ ─ ─ ─ ─ ┴ ─ ─ ─ ─ ─ ─ ─ ─ Stage 1 - /// ┌───────────────────────┐ │ - /// │ │ FilterExec │ - /// └───────────────────────┘ │ - /// │ ┌───────────────────────┐ a TaskEstimator estimates the amount of tasks - /// │ SomeExec │◀───┼── based on how much data will be pulled. - /// │ └───────────────────────┘ - /// ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘ - /// ``` - fn with_distributed_task_estimator( - self, - estimator: T, - ) -> Self; - - /// Same as [DistributedExt::with_distributed_task_estimator] but with an in-place mutation. - fn set_distributed_task_estimator( - &mut self, - estimator: T, - ); - /// Sets the number of bytes each partition in a stage with a FileScanConfig node is /// expected to scan. A task runs `target_partitions` partitions, so the task count is /// roughly `total_scan_bytes / bytes_per_partition / target_partitions` (capped at the @@ -611,12 +571,129 @@ pub trait DistributedExt: Sized { /// Same as [DistributedExt::with_distributed_local_worker_context] but with an /// in-place mutation. fn set_distributed_local_worker_context(&mut self, local_worker_context: LocalWorkerContext); + + /// Registers a distributed-planning event handler. + /// + /// The event type is inferred from the handler. Three event types are supported: + /// + /// 1. [`DesiredTaskCountHandler`](crate::DesiredTaskCountHandler) supplies desired or maximum task-count + /// hints while stages are being planned. + /// 2. [`ScaleUpLeafNodeHandler`](crate::ScaleUpLeafNodeHandler) replaces leaf nodes with variants specialized + /// for the final number of stage tasks. + /// 3. [`RouteTasksHandler`](crate::RouteTasksHandler) assigns each task in a stage to a worker URL. + /// + /// Custom handlers are called in registration order before built-in handlers. Returning `None` + /// means that the handler does not accept the event and dispatch should continue. The first + /// `Some` response stops dispatch. + /// + /// # 1. Desired task count + /// + /// A desired-task-count handler is called for plan nodes while the distributed planner is + /// choosing each stage's task count. It returns a desired or maximum task-count hint, or + /// `None` when it does not handle the node. + /// + /// A function with the following signature can be provided as argument: + /// + /// ```rust + /// # use datafusion::execution::SessionStateBuilder; + /// # use datafusion::physical_plan::empty::EmptyExec; + /// # use datafusion_distributed::{DistributedExt, DesiredTaskCountEvent, DesiredTaskCountEventResponse}; + /// + /// fn handle_custom_desired_task_count(event: DesiredTaskCountEvent) -> Option { + /// let _exec = event.plan.downcast_ref::()?; + /// Some(DesiredTaskCountEventResponse::desired(3)) + /// } + /// + /// SessionStateBuilder::new() + /// .with_distributed_event_handler(handle_custom_desired_task_count); + /// ``` + /// + /// ```text + /// ┌────────────────┐ + /// │CustomDataSource│──────────▶ 3 desired tasks + /// └────────────────┘ + /// ``` + /// + /// # 2. Scale up leaf node + /// + /// A scale-up handler is called after the distributed planner has chosen the final task count + /// for a leaf's stage. It can replace the leaf with a distributed variant specialized for + /// that task count. + /// + /// A function with the following signature can be provided as argument: + /// + /// ```rust + /// # use std::sync::Arc; + /// # use datafusion::error::Result; + /// # use datafusion::execution::SessionStateBuilder; + /// # use datafusion::physical_plan::empty::EmptyExec; + /// # use datafusion_distributed::{DistributedExt, DistributedLeafExec, ScaleUpLeafNodeEvent, ScaleUpLeafNodeEventResponse}; + /// + /// fn handle_custom_scale_up_leaf_node(event: ScaleUpLeafNodeEvent) -> Option> { + /// let _exec = event.plan.downcast_ref::()?; + /// Some( + /// DistributedLeafExec::try_new( + /// Arc::clone(event.plan), + /// vec![Arc::clone(event.plan); event.task_count], + /// ) + /// .map(|exec| ScaleUpLeafNodeEventResponse::new(Arc::new(exec))), + /// ) + /// } + /// + /// SessionStateBuilder::new() + /// .with_distributed_event_handler(handle_custom_scale_up_leaf_node); + /// ``` + /// + /// ```text + /// ┌────────────────────────────────────────────────────────────┐ + /// │ DistributedLeafExec │ + /// ┌────────────────┐ │┌──────────────────┐┌──────────────────┐┌──────────────────┐│ + /// │CustomDataSource│─3 tasks─▶││ CustomDataSource ││ CustomDataSource ││ CustomDataSource ││ + /// └────────────────┘ ││ (1/3) ││ (2/3) ││ (3/3) ││ + /// │└──────────────────┘└──────────────────┘└──────────────────┘│ + /// └────────────────────────────────────────────────────────────┘ + /// ``` + /// + /// # 3. Route tasks + /// + /// A routing handler is called before a stage starts execution. It maps the stage's task slots + /// to worker URLs; a successful response must contain one URL per task, in task-index order. + /// + /// A function with the following signature can be provided as argument: + /// + /// ```rust + /// # use datafusion::error::Result; + /// # use datafusion::execution::SessionStateBuilder; + /// # use datafusion_distributed::{DistributedExt, DistributedGetterExt, RouteTasksEvent, RouteTasksEventResponse}; + /// + /// fn handle_custom_route_tasks(event: RouteTasksEvent) -> Option> { + /// let routing = event.task_ctx.session_config() + /// .get_distributed_worker_resolver() + /// .and_then(|resolver| resolver.get_urls()) + /// .map(|workers| RouteTasksEventResponse::new( + /// workers.into_iter().cycle().take(event.task_count).collect() + /// )); + /// Some(routing) + /// } + /// + /// SessionStateBuilder::new() + /// .with_distributed_event_handler(handle_custom_route_tasks); + /// ``` + /// + /// ```text + /// task 0 ──► http://worker1 + /// task 1 ──► http://worker2 EventHandler + /// task 2 ──► http://worker3 + /// ``` + fn with_distributed_event_handler>(self, handler: H) -> Self; + + /// Same as [DistributedExt::with_distributed_event_handler] but with an in-place mutation. + fn set_distributed_event_handler>(&mut self, handler: H); } /// Trait to have a unified interface for getting structs & properties from SessionConfig that are used in distributed context. pub trait DistributedGetterExt: Sized { - /// Gets the [WorkerResolver] from this session's config. Typically called inside - /// [TaskEstimator::route_tasks] to resolve available worker URLs. + /// Gets the [WorkerResolver] from this session's config. fn get_distributed_worker_resolver(&self) -> Result, DataFusionError>; } @@ -652,13 +729,6 @@ impl DistributedExt for SessionConfig { set_distributed_channel_resolver(self, resolver); } - fn set_distributed_task_estimator( - &mut self, - estimator: T, - ) { - set_distributed_task_estimator(self, estimator) - } - fn set_distributed_file_scan_config_bytes_per_partition( &mut self, bytes_per_partition: usize, @@ -783,6 +853,10 @@ impl DistributedExt for SessionConfig { self.set_extension(Arc::new(local_worker_context)); } + fn set_distributed_event_handler>(&mut self, h: H) { + EventHandlerChain::::push_custom(self, Arc::new(h)); + } + delegate! { to self { #[call(set_distributed_option_extension)] @@ -809,10 +883,6 @@ impl DistributedExt for SessionConfig { #[expr($;self)] fn with_distributed_channel_resolver(mut self, resolver: T) -> Self; - #[call(set_distributed_task_estimator)] - #[expr($;self)] - fn with_distributed_task_estimator(mut self, estimator: T) -> Self; - #[call(set_distributed_file_scan_config_bytes_per_partition)] #[expr($?;Ok(self))] fn with_distributed_file_scan_config_bytes_per_partition(mut self, bytes_per_partition: usize) -> Result; @@ -878,9 +948,14 @@ impl DistributedExt for SessionConfig { #[call(set_distributed_local_worker_context)] #[expr($;self)] fn with_distributed_local_worker_context(mut self, local_worker_context: LocalWorkerContext) -> Self; + + #[call(set_distributed_event_handler)] + #[expr($;self)] + fn with_distributed_event_handler>(mut self, h: H) -> Self; } } } + impl DistributedGetterExt for SessionConfig { fn get_distributed_worker_resolver(&self) -> Result, DataFusionError> { get_distributed_worker_resolver(self) @@ -920,11 +995,6 @@ impl DistributedExt for SessionStateBuilder { #[expr($;self)] fn with_distributed_channel_resolver(mut self, resolver: T) -> Self; - fn set_distributed_task_estimator(&mut self, estimator: T); - #[call(set_distributed_task_estimator)] - #[expr($;self)] - fn with_distributed_task_estimator(mut self, estimator: T) -> Self; - fn set_distributed_file_scan_config_bytes_per_partition(&mut self, bytes_per_partition: usize) -> Result<(), DataFusionError>; #[call(set_distributed_file_scan_config_bytes_per_partition)] #[expr($?;Ok(self))] @@ -1011,6 +1081,11 @@ impl DistributedExt for SessionStateBuilder { #[call(set_distributed_local_worker_context)] #[expr($;self)] fn with_distributed_local_worker_context(mut self, local_worker_context: LocalWorkerContext) -> Self; + + fn set_distributed_event_handler>(&mut self, h: H); + #[call(set_distributed_event_handler)] + #[expr($;self)] + fn with_distributed_event_handler>(mut self, h: H) -> Self; } } } @@ -1055,11 +1130,6 @@ impl DistributedExt for SessionState { #[expr($;self)] fn with_distributed_channel_resolver(mut self, resolver: T) -> Self; - fn set_distributed_task_estimator(&mut self, estimator: T); - #[call(set_distributed_task_estimator)] - #[expr($;self)] - fn with_distributed_task_estimator(mut self, estimator: T) -> Self; - fn set_distributed_file_scan_config_bytes_per_partition(&mut self, bytes_per_partition: usize) -> Result<(), DataFusionError>; #[call(set_distributed_file_scan_config_bytes_per_partition)] #[expr($?;Ok(self))] @@ -1146,6 +1216,11 @@ impl DistributedExt for SessionState { #[call(set_distributed_local_worker_context)] #[expr($;self)] fn with_distributed_local_worker_context(mut self, local_worker_context: LocalWorkerContext) -> Self; + + fn set_distributed_event_handler>(&mut self, h: H); + #[call(set_distributed_event_handler)] + #[expr($;self)] + fn with_distributed_event_handler>(mut self, h: H) -> Self; } } } @@ -1183,11 +1258,6 @@ impl DistributedExt for SessionContext { #[expr($;self)] fn with_distributed_channel_resolver(self, resolver: T) -> Self; - fn set_distributed_task_estimator(&mut self, estimator: T); - #[call(set_distributed_task_estimator)] - #[expr($;self)] - fn with_distributed_task_estimator(self, estimator: T) -> Self; - fn set_distributed_file_scan_config_bytes_per_partition(&mut self, bytes_per_partition: usize) -> Result<(), DataFusionError>; #[call(set_distributed_file_scan_config_bytes_per_partition)] #[expr($?;Ok(self))] @@ -1274,6 +1344,11 @@ impl DistributedExt for SessionContext { #[call(set_distributed_local_worker_context)] #[expr($;self)] fn with_distributed_local_worker_context(self, local_worker_context: LocalWorkerContext) -> Self; + + fn set_distributed_event_handler>(&mut self, h: H); + #[call(set_distributed_event_handler)] + #[expr($;self)] + fn with_distributed_event_handler>(self, h: H) -> Self; } } } diff --git a/src/distributed_planner/distributed_query_planner.rs b/src/distributed_planner/distributed_query_planner.rs index 4b6fa46d..cb9d5d66 100644 --- a/src/distributed_planner/distributed_query_planner.rs +++ b/src/distributed_planner/distributed_query_planner.rs @@ -1,5 +1,4 @@ use crate::common::TreeNodeExt; -use crate::distributed_planner::CombinedTaskEstimator; use crate::distributed_planner::inject_network_boundaries::{ CardinalityBasedNetworkBoundaryBuilder, inject_network_boundaries, }; @@ -7,7 +6,8 @@ use crate::distributed_planner::insert_broadcast::insert_broadcast_execs; use crate::distributed_planner::partial_reduce_below_network_shuffles::partial_reduce_below_network_shuffles; use crate::distributed_planner::prepare_network_boundaries::prepare_network_boundaries; use crate::distributed_planner::push_fetch_into_network_coalesce::push_fetch_into_network_coalesce; -use crate::{DistributedConfig, DistributedExec, NetworkBoundaryExt, TaskEstimator}; +use crate::events::{ScaleUpLeafNodeEvent, ScaleUpLeafNodeHandler}; +use crate::{DistributedConfig, DistributedExec, NetworkBoundaryExt}; use async_trait::async_trait; use datafusion::common::tree_node::{Transformed, TreeNode}; use datafusion::execution::SessionState; @@ -83,10 +83,14 @@ impl QueryPlanner for DistributedQueryPlanner { if !plan.children().is_empty() { return Ok(Transformed::no(plan)); } - let task_estimator = CombinedTaskEstimator::from_session_config(session_cfg); - match task_estimator.scale_up_leaf_node(&plan, task_count, cfg)? { + let ev = ScaleUpLeafNodeEvent { + plan: &plan, + task_count, + session_config: session_cfg, + }; + match ScaleUpLeafNodeHandler::handle(ev) { None => Ok(Transformed::no(plan)), - Some(scaled) => Ok(Transformed::yes(scaled)), + Some(response) => Ok(Transformed::yes(response?.plan)), } })?; // Ensure the stages in the plan have nice unique identifiers. @@ -138,7 +142,7 @@ impl QueryPlanner for DistributedQueryPlanner { #[cfg(test)] mod tests { use crate::assert_snapshot; - use crate::test_utils::plans::{BuildSideOneTaskEstimator, TestPlanBuilder}; + use crate::test_utils::plans::{TestPlanBuilder, build_side_one_desired_task_count_handler}; /* schema for the "weather" table MinTemp [type=DOUBLE] [repetitiontype=OPTIONAL] @@ -921,7 +925,7 @@ mod tests { "#; let physical_plan_ascii = TestPlanBuilder::default() .broadcast_joins(true) - .distributed_task_estimator(BuildSideOneTaskEstimator) + .desired_task_count_handler(build_side_one_desired_task_count_handler) .physical_plan_as_ascii(query, false) .await; assert_snapshot!(physical_plan_ascii, @r" diff --git a/src/distributed_planner/inject_network_boundaries.rs b/src/distributed_planner/inject_network_boundaries.rs index 17a69a45..05b47446 100644 --- a/src/distributed_planner/inject_network_boundaries.rs +++ b/src/distributed_planner/inject_network_boundaries.rs @@ -1,16 +1,18 @@ -use crate::TaskCountAnnotation::{Desired, Maximum}; -use crate::distributed_planner::{CombinedTaskEstimator, TaskEstimator}; +use crate::events::TaskCountAnnotation::{Desired, Maximum}; +use crate::events::{ + DesiredTaskCountEvent, DesiredTaskCountHandler, ScaleUpLeafNodeEvent, ScaleUpLeafNodeHandler, + TaskCountAnnotation, +}; use crate::execution_plans::{ChildWeight, ChildrenIsolatorUnionExec}; use crate::stage::LocalStage; use crate::worker_resolver::WorkerResolverExtension; use crate::{ BroadcastExec, DistributedConfig, NetworkBoundaryExt, NetworkBroadcastExec, - NetworkCoalesceExec, NetworkShuffleExec, Stage, TaskCountAnnotation, + NetworkCoalesceExec, NetworkShuffleExec, Stage, }; use async_trait::async_trait; use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion::common::{HashMap, Result, plan_err}; -use datafusion::config::ConfigOptions; use datafusion::physical_expr::Partitioning; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion::physical_plan::execution_plan::CardinalityEffect; @@ -141,12 +143,10 @@ pub(crate) async fn inject_network_boundaries( nb_builder: impl NetworkBoundaryBuilder + Send + Sync, session_cfg: &SessionConfig, ) -> Result> { - let cfg = session_cfg.options(); let ctx = InjectNetworkBoundaryContext { - cfg, - d_cfg: DistributedConfig::from_config_options(cfg)?, + cfg: session_cfg, + d_cfg: DistributedConfig::from_session_config(session_cfg)?, worker_resolver: WorkerResolverExtension::from_session_config(session_cfg), - task_estimator: CombinedTaskEstimator::from_session_config(session_cfg), nb_builder: &nb_builder, task_counts: &Mutex::new(HashMap::new()), query_id: Uuid::new_v4(), @@ -160,9 +160,8 @@ pub(crate) async fn inject_network_boundaries( pub(crate) struct InjectNetworkBoundaryContext<'a> { pub(crate) d_cfg: &'a DistributedConfig, - cfg: &'a ConfigOptions, + cfg: &'a SessionConfig, worker_resolver: Arc, - task_estimator: Arc, nb_builder: &'a (dyn NetworkBoundaryBuilder + Send + Sync), task_counts: &'a Mutex>, query_id: Uuid, @@ -230,13 +229,16 @@ async fn _inject_network_boundaries( nb_ctx: &InjectNetworkBoundaryContext<'_>, ) -> Result> { let broadcast_joins_enabled = nb_ctx.d_cfg.broadcast_joins; - let estimator = nb_ctx.task_estimator.as_ref(); if plan.children().is_empty() { // This is a leaf node, maybe a DataSourceExec, or maybe something else custom from the // user. We need to estimate how many tasks are needed for this leaf node, and we'll take // this decision into account when deciding how many tasks will be actually used. - return if let Some(estimate) = estimator.task_estimation(&plan, nb_ctx.cfg) { + let ev = DesiredTaskCountEvent { + plan: &plan, + session_config: nb_ctx.cfg, + }; + return if let Some(estimate) = DesiredTaskCountHandler::handle(ev) { Ok(nb_ctx.plan_with_task_count(plan, estimate.task_count.limit(nb_ctx.max_tasks()?))) } else { // We could not determine how many tasks this leaf node should run on, so @@ -256,9 +258,11 @@ async fn _inject_network_boundaries( } let processed_children = futures::future::try_join_all(futures).await?; - let mut task_count = estimator - .task_estimation(&plan, nb_ctx.cfg) - .map_or(Desired(1), |v| v.task_count); + let ev = DesiredTaskCountEvent { + plan: &plan, + session_config: nb_ctx.cfg, + }; + let mut task_count = DesiredTaskCountHandler::handle(ev).map_or(Desired(1), |v| v.task_count); if nb_ctx.d_cfg.children_isolator_unions && plan.is::() { // Unions have the chance to decide how many tasks they should run on. If there's a union // with a bunch of children, the user might want to increase parallelism and increase the @@ -418,15 +422,16 @@ impl InjectNetworkBoundaryContext<'_> { ) -> Result> { // Handle leaf nodes. if plan.children().is_empty() { - let scaled_up = self.task_estimator.as_ref().scale_up_leaf_node( + let ev = ScaleUpLeafNodeEvent { plan, - task_count.as_usize(), - self.cfg, - )?; - match scaled_up { + task_count: task_count.as_usize(), + session_config: self.cfg, + }; + match ScaleUpLeafNodeHandler::handle(ev) { None => Ok(self.plan_with_task_count(Arc::clone(plan), task_count)), - Some(scaled_up) => { + Some(response) => { // The scaled up subtree may contain more than 1 node. + let scaled_up = response?.plan; scaled_up.apply(|plan| { self.set_task_count(plan, task_count); Ok(TreeNodeRecursion::Continue) @@ -629,9 +634,8 @@ impl NetworkBoundaryBuilder for CardinalityBasedNetworkBoundaryBuilder { mod tests { use super::*; use crate::distributed_planner::insert_broadcast::insert_broadcast_execs; - use crate::test_utils::plans::{BuildSideOneTaskEstimator, TestPlanBuilder}; - use crate::{TaskEstimation, TaskEstimator, assert_snapshot}; - use datafusion::config::ConfigOptions; + use crate::test_utils::plans::{TestPlanBuilder, build_side_one_desired_task_count_handler}; + use crate::{DesiredTaskCountEvent, DesiredTaskCountEventResponse, assert_snapshot}; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; /* schema for the "weather" table @@ -924,17 +928,13 @@ mod tests { let query = r#" SELECT DISTINCT "RainToday" FROM weather "#; - let task_estimator: Arc = - Arc::new(CallbackEstimator::new(|_: &RepartitionExec| { - Some(TaskEstimation::maximum(1)) - })); let test_plan_builder = TestPlanBuilder::new() .target_partitions(4) .num_workers(4) // annotate_test_plan wants this as false so its s a single node plan .distributed_planner(false) .broadcast_joins(false) - .distributed_task_estimator(task_estimator); + .desired_task_count_handler(repartition_max_one_desired_task_count_handler); let annotated = annotate_test_plan(test_plan_builder, query).await; assert_snapshot!(annotated, @r" AggregateExec: task_count=Desired(1) @@ -952,17 +952,13 @@ mod tests { UNION ALL SELECT "MaxTemp" FROM weather WHERE "RainToday" = 'no' "#; - let task_estimator: Arc = - Arc::new(CallbackEstimator::new(|_: &RepartitionExec| { - Some(TaskEstimation::maximum(1)) - })); let test_plan_builder = TestPlanBuilder::new() .target_partitions(4) .num_workers(4) // annotate_test_plan wants this as false so its s a single node plan .distributed_planner(false) .broadcast_joins(false) - .distributed_task_estimator(task_estimator); + .desired_task_count_handler(repartition_max_one_desired_task_count_handler); let annotated = annotate_test_plan(test_plan_builder, query).await; assert_snapshot!(annotated, @r" ChildrenIsolatorUnionExec: task_count=Desired(2) @@ -1054,7 +1050,7 @@ mod tests { // annotate_test_plan wants this as false so its s a single node plan .distributed_planner(false) .broadcast_joins(true) - .distributed_task_estimator(BuildSideOneTaskEstimator); + .desired_task_count_handler(build_side_one_desired_task_count_handler); let annotated = annotate_test_plan(test_plan_builder, query).await; assert_snapshot!(annotated, @r" HashJoinExec: task_count=Desired(3) @@ -1079,7 +1075,7 @@ mod tests { // annotate_test_plan wants this as false so its s a single node plan .distributed_planner(false) .broadcast_joins(true) - .distributed_task_estimator(BroadcastBuildCoalesceMaxEstimator); + .desired_task_count_handler(broadcast_build_coalesce_max_desired_task_count_handler); let annotated = annotate_test_plan(test_plan_builder, query).await; assert_snapshot!(annotated, @r" HashJoinExec: task_count=Maximum(1) @@ -1193,70 +1189,22 @@ mod tests { "); } - #[allow(clippy::type_complexity)] - struct CallbackEstimator { - f: Arc Option + Send + Sync>, - } - - impl CallbackEstimator { - fn new( - f: impl Fn(&T) -> Option + Send + Sync + 'static, - ) -> Self { - let f = Arc::new(move |plan: &dyn ExecutionPlan| -> Option { - if let Some(plan) = plan.downcast_ref::() { - f(plan) - } else { - None - } - }); - Self { f } - } + fn repartition_max_one_desired_task_count_handler( + ev: DesiredTaskCountEvent, + ) -> Option { + ev.plan + .is::() + .then(|| DesiredTaskCountEventResponse::maximum(1)) } - impl TaskEstimator for CallbackEstimator { - fn task_estimation( - &self, - plan: &Arc, - _: &ConfigOptions, - ) -> Option { - (self.f)(plan.as_ref()) - } - - fn scale_up_leaf_node( - &self, - _: &Arc, - _: usize, - _: &ConfigOptions, - ) -> Result>> { - Ok(None) - } - } - - #[derive(Debug)] - struct BroadcastBuildCoalesceMaxEstimator; - - impl TaskEstimator for BroadcastBuildCoalesceMaxEstimator { - fn task_estimation( - &self, - plan: &Arc, - _: &ConfigOptions, - ) -> Option { - let coalesce = plan.downcast_ref::()?; - if coalesce.input().is::() { - Some(TaskEstimation::maximum(1)) - } else { - None - } - } - - fn scale_up_leaf_node( - &self, - _: &Arc, - _: usize, - _: &ConfigOptions, - ) -> Result>> { - Ok(None) - } + fn broadcast_build_coalesce_max_desired_task_count_handler( + ev: DesiredTaskCountEvent, + ) -> Option { + ev.plan + .downcast_ref::()? + .input() + .is::() + .then(|| DesiredTaskCountEventResponse::maximum(1)) } async fn annotate_test_plan(test_plan_builder: TestPlanBuilder, query: &str) -> String { @@ -1267,10 +1215,9 @@ mod tests { let plan_w_broadcast = insert_broadcast_execs(plan, session_config.options()) .expect("failed to insert broadcasts"); let network_boundaries_ctx = InjectNetworkBoundaryContext { - cfg: session_config.options(), + cfg: &session_config, d_cfg: DistributedConfig::from_config_options(session_config.options()).unwrap(), worker_resolver: WorkerResolverExtension::from_session_config(&session_config), - task_estimator: CombinedTaskEstimator::from_session_config(&session_config), task_counts: &Mutex::new(HashMap::new()), query_id: Uuid::new_v4(), stage_id: &AtomicUsize::new(1), diff --git a/src/distributed_planner/mod.rs b/src/distributed_planner/mod.rs index 96baae96..a033e1f0 100644 --- a/src/distributed_planner/mod.rs +++ b/src/distributed_planner/mod.rs @@ -8,7 +8,6 @@ mod prepare_network_boundaries; mod push_fetch_into_network_coalesce; mod session_state_builder_ext; mod statistics; -mod task_estimator; pub use distributed_config::DistributedConfig; pub(crate) use inject_network_boundaries::{ @@ -18,5 +17,3 @@ pub(crate) use network_boundary::ProducerHead; pub use network_boundary::{NetworkBoundary, NetworkBoundaryExt}; pub use session_state_builder_ext::SessionStateBuilderExt; pub(crate) use statistics::calculate_cost; -pub(crate) use task_estimator::{CombinedTaskEstimator, set_distributed_task_estimator}; -pub use task_estimator::{TaskCountAnnotation, TaskEstimation, TaskEstimator, TaskRoutingContext}; diff --git a/src/distributed_planner/session_state_builder_ext.rs b/src/distributed_planner/session_state_builder_ext.rs index 174bf7f8..01dcc49b 100644 --- a/src/distributed_planner/session_state_builder_ext.rs +++ b/src/distributed_planner/session_state_builder_ext.rs @@ -1,5 +1,9 @@ use crate::distributed_planner::DistributedConfig; use crate::distributed_planner::distributed_query_planner::DistributedQueryPlanner; +use crate::events::{ + DesiredTaskCountHandler, EventHandlerChain, ScaleUpLeafNodeHandler, + file_scan_config_desired_task_count, file_scan_config_scale_up_leaf_node, +}; use datafusion::execution::SessionStateBuilder; use std::sync::Arc; @@ -17,13 +21,20 @@ pub trait SessionStateBuilderExt { impl SessionStateBuilderExt for SessionStateBuilder { fn with_distributed_planner(mut self) -> Self { - DistributedConfig::ensure_in_config(self.config().get_or_insert_default()); - self.config() - .get_or_insert_default() - .options_mut() + let cfg = self.config().get_or_insert_default(); + DistributedConfig::ensure_in_config(cfg); + cfg.options_mut() .optimizer .enable_physical_uncorrelated_scalar_subquery = false; + EventHandlerChain::::push_builtin( + cfg, + Arc::new(file_scan_config_desired_task_count), + ); + EventHandlerChain::::push_builtin( + cfg, + Arc::new(file_scan_config_scale_up_leaf_node), + ); let prev = std::mem::take(self.query_planner()); self.with_query_planner(Arc::new(DistributedQueryPlanner { prev })) } diff --git a/src/distributed_planner/task_estimator.rs b/src/distributed_planner/task_estimator.rs deleted file mode 100644 index 7f6d0319..00000000 --- a/src/distributed_planner/task_estimator.rs +++ /dev/null @@ -1,517 +0,0 @@ -use crate::DistributedConfig; -use crate::execution_plans::DistributedLeafExec; -use TaskCountAnnotation::*; -use datafusion::catalog::memory::DataSourceExec; -use datafusion::config::ConfigOptions; -use datafusion::datasource::physical_plan::{FileGroup, FileGroupPartitioner, FileScanConfig}; -use datafusion::error::Result; -use datafusion::execution::TaskContext; -use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; -use datafusion::prelude::SessionConfig; -use delegate::delegate; -use std::fmt::Debug; -use std::sync::Arc; -use url::Url; - -/// Annotation attached to a single [ExecutionPlan] that determines how many distributed tasks -/// it should run on. -#[derive(Debug, Clone, Copy)] -pub enum TaskCountAnnotation { - /// The desired number of distributed tasks for this node. The final task count for the - /// annotated node might not be exactly this number, it is more like a hint, so depending - /// on the desired task count of adjacent nodes, the final task count might change. - Desired(usize), - /// Sets a maximum number of distributed tasks for this node. Typically used with the inner - /// value of 1, stating that this node cannot be executed in a distributed fashion. - Maximum(usize), -} - -impl From for usize { - fn from(annotation: TaskCountAnnotation) -> Self { - annotation.as_usize() - } -} - -impl TaskCountAnnotation { - pub fn as_usize(&self) -> usize { - match self { - Desired(desired) => *desired, - Maximum(maximum) => *maximum, - } - } - - pub(crate) fn limit(self, limit: usize) -> Self { - match self { - Desired(desired) => Desired(desired.min(limit)), - Maximum(maximum) => Maximum(maximum.min(limit)), - } - } - - pub(crate) fn merge(self, other: TaskCountAnnotation) -> Self { - match (self, other) { - (Desired(a), Desired(b)) => Desired(std::cmp::max(a, b)), - (Desired(_), Maximum(b)) => Maximum(b), - (Maximum(a), Desired(_)) => Maximum(a), - (Maximum(a), Maximum(b)) => Maximum(std::cmp::min(a, b)), - } - } -} - -/// Result of running a [TaskEstimator] on a leaf node. It tells the distributed planner hints -/// about how many tasks should be used in [Stage]s that contain leaf nodes. -pub struct TaskEstimation { - /// The number of tasks that should be used in the [Stage] containing the leaf node. - /// - /// Even if implementations get to decide this number, there are situations where it can - /// get overridden: - /// - If a [Stage] contains multiple leaf nodes, the one that declares the biggest - /// task_count wins. - /// - If there are less available workers than this number, the number of available workers - /// is chosen. - pub task_count: TaskCountAnnotation, -} - -impl TaskEstimation { - /// Tells the distributed planner that the evaluated stage can have **at maximum** the provided - /// number of tasks, setting a hard upper limit. - /// - /// Returning `TaskEstimation::maximum(1)` tells the distributed planner that the evaluated - /// stage cannot be distributed. - /// - /// Even if a `TaskEstimation::maximum(N)` is provided, any other node in the same stage - /// providing a value of `TaskEstimation::maximum(M)` where `M` < `N` will have preference. - pub fn maximum(value: usize) -> Self { - TaskEstimation { - task_count: TaskCountAnnotation::Maximum(value), - } - } - - /// Tells the distributed planner that the evaluated can **optimally** have the provided - /// number of tasks, setting a soft task count hint that can be overridden by others. - /// - /// The provided `TaskEstimation::desired(N)` can be overridden by: - /// - Other nodes providing a `TaskEstimation::desired(M)` where `M` > `N`. - /// - Any other node providing a `TaskEstimation::maximum(M)` where `M` can be anything. - pub fn desired(value: usize) -> Self { - TaskEstimation { - task_count: TaskCountAnnotation::Desired(value), - } - } -} - -/// Given a leaf node, provides an estimation about how many tasks should be used in the -/// stage containing it, and if the leaf node should be replaced by some other. -/// -/// The distributed planner will try many [TaskEstimator]s in order until one provides an -/// estimation for a specific leaf node. Once that's done, upper stages will get their task -/// count calculated based on whether lower stages are reducing the cardinality of the data -/// or increasing it. -pub trait TaskEstimator { - /// Function applied to each node that returns a [TaskEstimation] hinting how many - /// tasks should be used in the [Stage] containing that node. - /// - /// All the [TaskEstimator] registered in the session will be applied to the node - /// until one returns an estimation. - /// - /// - /// If no estimation is returned from any of the registered [TaskEstimator]s, then: - /// - If the node is a leaf node,`Maximum(1)` is assumed, hinting the distributed planner - /// that the leaf node cannot be distributed across tasks. - /// - If the node is a normal node in the plan, then the maximum task count from its children - /// is inherited. - fn task_estimation( - &self, - plan: &Arc, - cfg: &ConfigOptions, - ) -> Option; - - /// After a final task_count is decided, taking into account all the leaf nodes in the [Stage], - /// this allows performing a transformation in the leaf nodes for accounting for the fact that - /// they are going to run in multiple tasks. - fn scale_up_leaf_node( - &self, - plan: &Arc, - task_count: usize, - cfg: &ConfigOptions, - ) -> Result>>; - - /// Optionally defines a custom protocol for routing tasks to specific worker URLs. Receives - /// routing context including task count and a list of available URLs, and returns a vector - /// of routed URLs, in order of task assignment. - /// - /// If Ok(Some(Vec)) is returned, tasks are sent in order to the URLs specified in the - /// returned vector. If Ok(None) is returned, execution defaults to round-robin routing. - fn route_tasks(&self, _routing_ctx: &TaskRoutingContext<'_>) -> Result>> { - Ok(None) - } -} - -/// Context usable for routing tasks to worker URLs. -pub struct TaskRoutingContext<'a> { - /// The task context active at routing time. - pub task_ctx: Arc, - /// The head execution plan of the stage being routed. - pub plan: &'a Arc, - /// The number of tasks to be assigned. - pub task_count: usize, -} - -impl TaskEstimator for usize { - fn task_estimation( - &self, - inputs: &Arc, - _: &ConfigOptions, - ) -> Option { - if inputs.children().is_empty() { - Some(TaskEstimation { - task_count: TaskCountAnnotation::Desired(*self), - }) - } else { - None - } - } - - fn scale_up_leaf_node( - &self, - _: &Arc, - _: usize, - _: &ConfigOptions, - ) -> Result>> { - Ok(None) - } -} - -impl TaskEstimator for Arc { - delegate! { - to self.as_ref() { - fn task_estimation(&self, plan: &Arc, cfg: &ConfigOptions) -> Option; - fn scale_up_leaf_node(&self, plan: &Arc, task_count: usize, cfg: &ConfigOptions) -> Result>>; - fn route_tasks(&self, routing_ctx: &TaskRoutingContext<'_>) -> Result>>; - } - } -} - -impl TaskEstimator for Arc { - delegate! { - to self.as_ref() { - fn task_estimation(&self, plan: &Arc, cfg: &ConfigOptions) -> Option; - fn scale_up_leaf_node(&self, plan: &Arc, task_count: usize, cfg: &ConfigOptions) -> Result>>; - fn route_tasks(&self, routing_ctx: &TaskRoutingContext<'_>) -> Result>>; - } - } -} - -pub(crate) fn set_distributed_task_estimator( - cfg: &mut SessionConfig, - estimator: impl TaskEstimator + Send + Sync + 'static, -) { - let mut combined = cfg - .get_extension::() - .map(|existing| existing.as_ref().clone()) - .unwrap_or_default(); - combined.user_provided.push(Arc::new(estimator)); - cfg.set_extension(Arc::new(combined)); -} - -/// [TaskEstimator] implementation that acts on [DataSourceExec] nodes that contain -/// [FileScanConfig]s data sources (e.g., Parquet or CSV files). It reads the -/// [DistributedConfig].`file_scan_config_bytes_per_partition` field and assigns as many tasks as -/// needed so that no partition scans more than the configured number of bytes. -#[derive(Debug)] -pub(crate) struct FileScanConfigTaskEstimator; - -impl TaskEstimator for FileScanConfigTaskEstimator { - fn task_estimation( - &self, - plan: &Arc, - cfg: &ConfigOptions, - ) -> Option { - let dse: &DataSourceExec = plan.downcast_ref()?; - let file_scan: &FileScanConfig = dse.data_source().downcast_ref()?; - - let d_cfg = cfg.extensions.get::()?; - - let mut total_bytes = 0; - for file_group in &file_scan.file_groups { - for file in file_group.files() { - total_bytes += file.effective_size() as usize - } - } - - let task_count = total_bytes - .div_ceil(d_cfg.file_scan_config_bytes_per_partition) - .div_ceil(cfg.execution.target_partitions); - - Some(TaskEstimation::desired(task_count)) - } - - fn scale_up_leaf_node( - &self, - plan: &Arc, - task_count: usize, - _cfg: &ConfigOptions, - ) -> Result>> { - let Some(dse) = plan.downcast_ref::() else { - return Ok(None); - }; - let Some(file_scan) = dse.data_source().downcast_ref::() else { - return Ok(None); - }; - let partition_count = plan.output_partitioning().partition_count(); - - let rebalanced = if file_scan.partitioned_by_file_group { - let all_partitioned_files = file_scan - .file_groups - .iter() - .flat_map(|file_group| file_group.iter().cloned()) - .collect::>(); - rebalance_round_robin(all_partitioned_files, partition_count * task_count) - .into_iter() - .map(FileGroup::new) - .collect::>() - } else { - FileGroupPartitioner::new() - .with_target_partitions(partition_count * task_count) - // Allow repartitioning beyond normal limits, putting the limit in - // `partition_count * task_count` target partitions, and not in the - // resulting size. - .with_repartition_file_min_size(0) - .with_preserve_order_within_groups(!file_scan.output_ordering.is_empty()) - .repartition_file_groups(&file_scan.file_groups) - .unwrap_or_else(|| file_scan.file_groups.clone()) - .into_iter() - .collect() - }; - - let mut file_scan_template = file_scan.clone(); - file_scan_template.file_groups.clear(); - let mut file_scans = vec![file_scan_template; task_count]; - for (i, file_group) in rebalanced.into_iter().enumerate() { - file_scans[i % task_count].file_groups.push(file_group); - } - - let dle = DistributedLeafExec::try_new( - Arc::clone(plan), - file_scans - .into_iter() - .map(|file_scan| DataSourceExec::from_data_source(file_scan) as _), - )?; - - Ok(Some(Arc::new(dle))) - } -} - -fn rebalance_round_robin(items: Vec, target_groups: usize) -> Vec> { - let mut groups = (0..target_groups) - .map(|_| Vec::new()) - .collect::>>(); - for (idx, item) in items.into_iter().enumerate() { - groups[idx % target_groups].push(item); - } - groups -} - -/// Tries multiple user-provided [TaskEstimator]s until one returns an estimation. If none -/// returns an estimation, a set of default [TaskEstimation] implementations is tried. Right -/// now the only default [TaskEstimation] is [FileScanConfigTaskEstimator]. -#[derive(Clone, Default)] -pub(crate) struct CombinedTaskEstimator { - pub(crate) user_provided: Vec>, -} - -impl CombinedTaskEstimator { - pub(crate) fn from_session_config(cfg: &SessionConfig) -> Arc { - cfg.get_extension::() - .unwrap_or_else(|| Arc::new(Self::default())) - } -} - -impl TaskEstimator for CombinedTaskEstimator { - fn task_estimation( - &self, - plan: &Arc, - cfg: &ConfigOptions, - ) -> Option { - for estimator in &self.user_provided { - if let Some(result) = estimator.task_estimation(plan, cfg) { - return Some(result); - } - } - // We want to execute the default estimators last so that the user-provided ones have - // a chance of providing an estimation. - // If none of the user-provided returned an estimation, the default ones are used. - for default_estimator in [&FileScanConfigTaskEstimator as &dyn TaskEstimator] { - if let Some(result) = default_estimator.task_estimation(plan, cfg) { - return Some(result); - } - } - None - } - - fn scale_up_leaf_node( - &self, - plan: &Arc, - task_count: usize, - cfg: &ConfigOptions, - ) -> Result>> { - for estimator in &self.user_provided { - if let Some(result) = estimator.scale_up_leaf_node(plan, task_count, cfg)? { - return Ok(Some(result)); - } - } - // We want to execute the default estimators last so that the user-provided ones have - // a chance of providing an estimation. - // If none of the user-provided returned an estimation, the default ones are used. - for default_estimator in [&FileScanConfigTaskEstimator as &dyn TaskEstimator] { - if let Some(result) = default_estimator.scale_up_leaf_node(plan, task_count, cfg)? { - return Ok(Some(result)); - } - } - Ok(None) - } - - fn route_tasks(&self, routing_ctx: &TaskRoutingContext<'_>) -> Result>> { - for estimator in &self.user_provided { - if let Some(result) = estimator.route_tasks(routing_ctx)? { - return Ok(Some(result)); - } - } - Ok(None) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::test_utils::parquet::register_parquet_tables; - use datafusion::error::DataFusionError; - use datafusion::prelude::SessionContext; - - #[tokio::test] - async fn test_first_user_estimator_wins() -> Result<(), DataFusionError> { - let mut combined = CombinedTaskEstimator::default(); - combined.push(10); - combined.push(20); - - let node = make_data_source_exec().await?; - assert_eq!(combined.task_count(node, |cfg| cfg), 10); - Ok(()) - } - - #[tokio::test] - async fn test_continues_until_some() -> Result<(), DataFusionError> { - let mut combined = CombinedTaskEstimator::default(); - combined.push(|_: &Arc, _: &ConfigOptions| None); - combined.push(30); - - let node = make_data_source_exec().await?; - assert_eq!(combined.task_count(node, |cfg| cfg), 30); - Ok(()) - } - - #[tokio::test] - async fn test_defaults_to_file_scan_config_task_estimator() -> Result<(), DataFusionError> { - let mut combined = CombinedTaskEstimator::default(); - combined.push(|_: &Arc, _: &ConfigOptions| None); - - // No user estimator returns a value, so the default FileScanConfigTaskEstimator kicks in. - // Size the per-partition budget (with target_partitions pinned to 1) so the scan splits - // into exactly 3 partitions. - let node = make_data_source_exec().await?; - let bytes_per_partition = total_scan_bytes(&node).div_ceil(3); - let task_count = combined.task_count(node, |mut cfg| { - cfg.file_scan_config_bytes_per_partition = bytes_per_partition; - cfg - }); - assert_eq!(task_count, 3); - Ok(()) - } - - fn total_scan_bytes(node: &Arc) -> usize { - let dse = node.downcast_ref::().unwrap(); - let file_scan = dse.data_source().downcast_ref::().unwrap(); - file_scan - .file_groups - .iter() - .flat_map(|file_group| file_group.files()) - .map(|file| file.effective_size() as usize) - .sum() - } - - #[test] - fn test_rebalance_round_robin_fixes_group_boundary_skew() { - let items = (0..8).collect::>(); - let groups = rebalance_round_robin(items, 5); - let sizes = groups.iter().map(Vec::len).collect::>(); - assert_eq!(sizes, vec![2, 2, 2, 1, 1]); - } - - #[test] - fn test_rebalance_round_robin_pads_with_empty_groups() { - // With fewer items than target groups, the extra groups are kept empty rather than - // dropped. This guarantees every task ends up with the same number of partitions. - let items = vec![10, 20, 30]; - let groups = rebalance_round_robin(items, 5); - let sizes = groups.iter().map(Vec::len).collect::>(); - assert_eq!(sizes, vec![1, 1, 1, 0, 0]); - } - - impl CombinedTaskEstimator { - fn push(&mut self, value: impl TaskEstimator + Send + Sync + 'static) { - self.user_provided.push(Arc::new(value)); - } - - fn task_count( - &self, - node: Arc, - f: impl FnOnce(DistributedConfig) -> DistributedConfig, - ) -> usize { - let mut cfg = ConfigOptions::default(); - // Pin target_partitions so the byte-based estimation is deterministic regardless of - // the host's core count. - cfg.execution.target_partitions = 1; - let d_cfg = DistributedConfig { - file_scan_config_bytes_per_partition: 1, - ..Default::default() - }; - cfg.extensions.insert(f(d_cfg)); - self.task_estimation(&node, &cfg) - .unwrap() - .task_count - .as_usize() - } - } - - async fn make_data_source_exec() -> Result, DataFusionError> { - let ctx = SessionContext::new(); - register_parquet_tables(&ctx).await?; - let mut plan = ctx - .sql("SELECT * FROM weather") - .await? - .create_physical_plan() - .await?; - while !plan.children().is_empty() { - plan = Arc::clone(plan.children()[0]) - } - Ok(plan) - } - - impl, &ConfigOptions) -> Option> TaskEstimator for F { - fn task_estimation( - &self, - plan: &Arc, - cfg: &ConfigOptions, - ) -> Option { - self(plan, cfg) - } - - fn scale_up_leaf_node( - &self, - _plan: &Arc, - _task_count: usize, - _cfg: &ConfigOptions, - ) -> Result>> { - Ok(None) - } - } -} diff --git a/src/events/common.rs b/src/events/common.rs new file mode 100644 index 00000000..53d7c6c8 --- /dev/null +++ b/src/events/common.rs @@ -0,0 +1,86 @@ +use datafusion::execution::config::SessionConfig; +use std::sync::Arc; + +/// An event which occurs during distributed planning (which may occur at planning time or execution +/// time if using adaptive query execution). +pub trait Event: Send + Sync + 'static { + /// Data contained in the event. + type Data<'a>: Clone; + /// Result returned when a handler accepts the event. + type Response; +} + +/// Handles [`Event`] data. +pub trait EventHandler: Send + Sync + 'static { + /// Returns `None` when this handler does not accept the event, allowing the next handler in + /// the chain to try. Returns `Some` to stop dispatch and select that response. + fn handle(&self, ev: E::Data<'_>) -> Option; +} + +impl EventHandler for Arc +where + E: Event, + H: EventHandler + ?Sized, +{ + fn handle(&self, ev: E::Data<'_>) -> Option { + self.as_ref().handle(ev) + } +} + +pub(crate) struct EventHandlerChain { + pub(crate) builtin: Vec>>, + pub(crate) custom: Vec>>, +} + +impl Default for EventHandlerChain { + fn default() -> Self { + Self { + builtin: Vec::new(), + custom: Vec::new(), + } + } +} + +impl Clone for EventHandlerChain { + fn clone(&self) -> Self { + Self { + builtin: self.builtin.clone(), + custom: self.custom.clone(), + } + } +} + +impl EventHandlerChain { + pub(crate) fn handle(&self, ev: E::Data<'_>) -> Option { + // Give priority to custom handlers registered by users. + if let Some(res) = self + .custom + .iter() + .find_map(|handler| handler.handle(ev.clone())) + { + return Some(res); + } + // If no user handler handled the event, use the built ins. + self.builtin + .iter() + .find_map(|handler| handler.handle(ev.clone())) + } + + pub(crate) fn push_builtin(cfg: &mut SessionConfig, handler: Arc>) { + let mut handlers = cfg + .get_extension::() + .map(|v| v.as_ref().clone()) + .unwrap_or_default(); + handlers.builtin.push(handler); + cfg.set_extension(Arc::new(handlers)); + } + + pub(crate) fn push_custom(cfg: &mut SessionConfig, handler: Arc>) { + let mut handlers = cfg + .get_extension::() + .map(|v| v.as_ref().clone()) + .unwrap_or_default(); + handlers.custom.push(handler); + cfg.set_extension(Arc::new(handlers)); + } +} diff --git a/src/events/defaults/file_scan_config.rs b/src/events/defaults/file_scan_config.rs new file mode 100644 index 00000000..787b93fd --- /dev/null +++ b/src/events/defaults/file_scan_config.rs @@ -0,0 +1,216 @@ +use crate::DistributedConfig; +#[cfg(test)] +use crate::events::DesiredTaskCountHandler; +use crate::events::{ + DesiredTaskCountEvent, DesiredTaskCountEventResponse, ScaleUpLeafNodeEvent, + ScaleUpLeafNodeEventResponse, +}; +use crate::execution_plans::DistributedLeafExec; +use datafusion::catalog::memory::DataSourceExec; +use datafusion::datasource::physical_plan::{FileGroup, FileGroupPartitioner, FileScanConfig}; +use datafusion::error::Result; +use datafusion::physical_plan::ExecutionPlanProperties; +use std::sync::Arc; + +pub(crate) fn file_scan_config_desired_task_count( + ev: DesiredTaskCountEvent, +) -> Option { + let cfg = ev.session_config; + let dse: &DataSourceExec = ev.plan.downcast_ref()?; + let file_scan: &FileScanConfig = dse.data_source().downcast_ref()?; + + let d_cfg = DistributedConfig::from_session_config(cfg).ok()?; + + let mut total_bytes = 0; + for file_group in &file_scan.file_groups { + for file in file_group.files() { + total_bytes += file.effective_size() as usize + } + } + + let task_count = total_bytes + .div_ceil(d_cfg.file_scan_config_bytes_per_partition) + .div_ceil(cfg.target_partitions()); + + Some(DesiredTaskCountEventResponse::desired(task_count)) +} + +pub(crate) fn file_scan_config_scale_up_leaf_node( + ev: ScaleUpLeafNodeEvent, +) -> Option> { + let dse = ev.plan.downcast_ref::()?; + let file_scan = dse.data_source().downcast_ref::()?; + let partition_count = ev.plan.output_partitioning().partition_count(); + + let rebalanced = if file_scan.partitioned_by_file_group { + let all_partitioned_files = file_scan + .file_groups + .iter() + .flat_map(|file_group| file_group.iter().cloned()) + .collect::>(); + rebalance_round_robin(all_partitioned_files, partition_count * ev.task_count) + .into_iter() + .map(FileGroup::new) + .collect::>() + } else { + FileGroupPartitioner::new() + .with_target_partitions(partition_count * ev.task_count) + .with_repartition_file_min_size(0) + .with_preserve_order_within_groups(!file_scan.output_ordering.is_empty()) + .repartition_file_groups(&file_scan.file_groups) + .unwrap_or_else(|| file_scan.file_groups.clone()) + .into_iter() + .collect() + }; + + let mut file_scan_template = file_scan.clone(); + file_scan_template.file_groups.clear(); + let mut file_scans = vec![file_scan_template; ev.task_count]; + for (i, file_group) in rebalanced.into_iter().enumerate() { + file_scans[i % ev.task_count].file_groups.push(file_group); + } + + let distributed_leaf_result = DistributedLeafExec::try_new( + Arc::clone(ev.plan), + file_scans + .into_iter() + .map(|file_scan| DataSourceExec::from_data_source(file_scan) as _), + ); + let distributed_leaf = match distributed_leaf_result { + Ok(distributed_leaf) => distributed_leaf, + Err(e) => return Some(Err(e)), + }; + + Some(Ok(ScaleUpLeafNodeEventResponse::new(Arc::new( + distributed_leaf, + )))) +} + +fn rebalance_round_robin(items: Vec, target_groups: usize) -> Vec> { + let mut groups = (0..target_groups) + .map(|_| Vec::new()) + .collect::>>(); + for (idx, item) in items.into_iter().enumerate() { + groups[idx % target_groups].push(item); + } + groups +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::DistributedExt; + use crate::test_utils::parquet::register_parquet_tables; + use datafusion::error::DataFusionError; + use datafusion::physical_plan::ExecutionPlan; + use datafusion::prelude::{SessionConfig, SessionContext}; + + #[tokio::test] + async fn test_first_desired_task_count_handler_wins() -> Result<(), DataFusionError> { + let cfg = SessionConfig::new() + .with_distributed_event_handler(desired_ten) + .with_distributed_event_handler(desired_twenty); + + let plan = make_data_source_exec().await?; + let response = DesiredTaskCountHandler::handle(DesiredTaskCountEvent { + plan: &plan, + session_config: &cfg, + }) + .expect("a handler should respond"); + assert_eq!(response.task_count.as_usize(), 10); + Ok(()) + } + + #[tokio::test] + async fn test_desired_task_count_handlers_continue_until_some() -> Result<(), DataFusionError> { + let cfg = SessionConfig::new() + .with_distributed_event_handler(no_desired_task_count) + .with_distributed_event_handler(desired_thirty); + + let plan = make_data_source_exec().await?; + let response = DesiredTaskCountHandler::handle(DesiredTaskCountEvent { + plan: &plan, + session_config: &cfg, + }) + .expect("a handler should respond"); + assert_eq!(response.task_count.as_usize(), 30); + Ok(()) + } + + #[tokio::test] + async fn test_file_scan_config_desired_task_count_handler() -> Result<(), DataFusionError> { + let plan = make_data_source_exec().await?; + let bytes_per_partition = total_scan_bytes(&plan).div_ceil(3); + let mut cfg = SessionConfig::new(); + cfg.options_mut().execution.target_partitions = 1; + cfg.set_distributed_option_extension(DistributedConfig::default()); + cfg.set_distributed_file_scan_config_bytes_per_partition(bytes_per_partition)?; + + let response = file_scan_config_desired_task_count(DesiredTaskCountEvent { + plan: &plan, + session_config: &cfg, + }) + .expect("a file scan should be recognized"); + assert_eq!(response.task_count.as_usize(), 3); + Ok(()) + } + + #[test] + fn test_rebalance_round_robin_fixes_group_boundary_skew() { + let groups = rebalance_round_robin((0..8).collect(), 5); + assert_eq!( + groups.iter().map(Vec::len).collect::>(), + vec![2, 2, 2, 1, 1] + ); + } + + #[test] + fn test_rebalance_round_robin_pads_with_empty_groups() { + let groups = rebalance_round_robin(vec![10, 20, 30], 5); + assert_eq!( + groups.iter().map(Vec::len).collect::>(), + vec![1, 1, 1, 0, 0] + ); + } + + fn total_scan_bytes(plan: &Arc) -> usize { + let dse = plan.downcast_ref::().unwrap(); + let file_scan = dse.data_source().downcast_ref::().unwrap(); + file_scan + .file_groups + .iter() + .flat_map(|file_group| file_group.files()) + .map(|file| file.effective_size() as usize) + .sum() + } + + async fn make_data_source_exec() -> Result, DataFusionError> { + let ctx = SessionContext::new(); + register_parquet_tables(&ctx).await?; + let mut plan = ctx + .sql("SELECT * FROM weather") + .await? + .create_physical_plan() + .await?; + while !plan.children().is_empty() { + plan = Arc::clone(plan.children()[0]); + } + Ok(plan) + } + + fn desired_ten(_: DesiredTaskCountEvent) -> Option { + Some(DesiredTaskCountEventResponse::desired(10)) + } + + fn desired_twenty(_: DesiredTaskCountEvent) -> Option { + Some(DesiredTaskCountEventResponse::desired(20)) + } + + fn no_desired_task_count(_: DesiredTaskCountEvent) -> Option { + None + } + + fn desired_thirty(_: DesiredTaskCountEvent) -> Option { + Some(DesiredTaskCountEventResponse::desired(30)) + } +} diff --git a/src/events/defaults/mod.rs b/src/events/defaults/mod.rs new file mode 100644 index 00000000..c0d6f911 --- /dev/null +++ b/src/events/defaults/mod.rs @@ -0,0 +1,5 @@ +mod file_scan_config; + +pub(crate) use file_scan_config::{ + file_scan_config_desired_task_count, file_scan_config_scale_up_leaf_node, +}; diff --git a/src/events/desired_task_count.rs b/src/events/desired_task_count.rs new file mode 100644 index 00000000..9162c749 --- /dev/null +++ b/src/events/desired_task_count.rs @@ -0,0 +1,152 @@ +use super::TaskCountAnnotation::{Desired, Maximum}; +use super::common::{Event, EventHandler, EventHandlerChain}; +use datafusion::execution::config::SessionConfig; +use datafusion::physical_plan::ExecutionPlan; +use std::sync::Arc; + +/// Annotation attached to a single [ExecutionPlan] that determines how many distributed tasks +/// it should run on. +#[derive(Debug, Clone, Copy)] +pub enum TaskCountAnnotation { + /// The desired number of distributed tasks for this node. The final task count for the + /// annotated node might not be exactly this number, it is more like a hint, so depending + /// on the desired task count of adjacent nodes, the final task count might change. + Desired(usize), + /// Sets a maximum number of distributed tasks for this node. Typically used with the inner + /// value of 1, stating that this node cannot be executed in a distributed fashion. + Maximum(usize), +} + +/// Information supplied when the planner asks a handler for a node's desired task count. +/// +/// Handlers may return a [`DesiredTaskCountEventResponse`] for any execution-plan node. The +/// planner reconciles responses from the nodes in the same stage into its final task count. +#[derive(Clone, Copy)] +pub struct DesiredTaskCountEvent<'a> { + /// The execution-plan node being evaluated. + pub plan: &'a Arc, + /// The session configuration that registered the handlers and holds query options. + pub session_config: &'a SessionConfig, +} + +/// Result of running a [TaskEstimator] on a leaf node. It tells the distributed planner hints +/// about how many tasks should be used in [Stage]s that contain leaf nodes. +pub struct DesiredTaskCountEventResponse { + /// The number of tasks that should be used in the [Stage] containing the leaf node. + /// + /// Even if implementations get to decide this number, there are situations where it can + /// get overridden: + /// - If a [Stage] contains multiple leaf nodes, the one that declares the biggest + /// task_count wins. + /// - If there are less available workers than this number, the number of available workers + /// is chosen. + pub task_count: TaskCountAnnotation, +} + +impl DesiredTaskCountEventResponse { + /// Tells the distributed planner that the evaluated stage can have **at maximum** the provided + /// number of tasks, setting a hard upper limit. + /// + /// Returning `DesiredTaskCountEventResponse::maximum(1)` tells the distributed planner that the + /// evaluated stage cannot be distributed. + /// + /// Even if a `DesiredTaskCountEventResponse::maximum(N)` is provided, any other node in the + /// same stage providing a value of `DesiredTaskCountEventResponse::maximum(M)` where `M` < `N` + /// will have preference. + pub fn maximum(value: usize) -> Self { + DesiredTaskCountEventResponse { + task_count: Maximum(value), + } + } + + /// Tells the distributed planner that the evaluated can **optimally** have the provided + /// number of tasks, setting a soft task count hint that can be overridden by others. + /// + /// The provided `DesiredTaskCountEventResponse::desired(N)` can be overridden by: + /// - Other nodes providing a `DesiredTaskCountEventResponse::desired(M)` where `M` > `N`. + /// - Any other node providing a `DesiredTaskCountEventResponse::maximum(M)` where `M` can be + /// anything. + pub fn desired(value: usize) -> Self { + DesiredTaskCountEventResponse { + task_count: Desired(value), + } + } +} + +impl From for usize { + fn from(annotation: TaskCountAnnotation) -> Self { + annotation.as_usize() + } +} + +impl TaskCountAnnotation { + pub fn as_usize(&self) -> usize { + match self { + Desired(desired) => *desired, + Maximum(maximum) => *maximum, + } + } + + pub(crate) fn limit(self, limit: usize) -> Self { + match self { + Desired(desired) => Desired(desired.min(limit)), + Maximum(maximum) => Maximum(maximum.min(limit)), + } + } + + pub(crate) fn merge(self, other: TaskCountAnnotation) -> Self { + match (self, other) { + (Desired(a), Desired(b)) => Desired(std::cmp::max(a, b)), + (Desired(_), Maximum(b)) => Maximum(b), + (Maximum(a), Desired(_)) => Maximum(a), + (Maximum(a), Maximum(b)) => Maximum(std::cmp::min(a, b)), + } + } +} + +/// Function applied to each node that returns a [DesiredTaskCountEventResponse] hinting how +/// many tasks should be used in the [Stage] containing that node. +/// +/// All the [TaskEstimator] registered in the session will be applied to the node +/// until one returns an estimation. +/// +/// +/// If no estimation is returned from any of the registered [TaskEstimator]s, then: +/// - If the node is a leaf node,`Maximum(1)` is assumed, hinting the distributed planner +/// that the leaf node cannot be distributed across tasks. +/// - If the node is a normal node in the plan, then the maximum task count from its children +/// is inherited. +pub struct DesiredTaskCountHandler; + +impl DesiredTaskCountHandler { + pub(crate) fn handle(ev: DesiredTaskCountEvent<'_>) -> Option { + let handlers = ev + .session_config + .get_extension::>()?; + handlers.handle(ev) + } +} + +impl Event for DesiredTaskCountHandler { + type Data<'a> = DesiredTaskCountEvent<'a>; + type Response = DesiredTaskCountEventResponse; +} + +impl EventHandler for F +where + F: Send + Sync + 'static, + F: for<'a> Fn(DesiredTaskCountEvent<'a>) -> Option, +{ + fn handle(&self, ev: DesiredTaskCountEvent<'_>) -> Option { + self(ev) + } +} + +impl EventHandler for usize { + fn handle(&self, ev: DesiredTaskCountEvent<'_>) -> Option { + ev.plan + .children() + .is_empty() + .then(|| DesiredTaskCountEventResponse::desired(*self)) + } +} diff --git a/src/events/mod.rs b/src/events/mod.rs new file mode 100644 index 00000000..6427fa8b --- /dev/null +++ b/src/events/mod.rs @@ -0,0 +1,19 @@ +mod common; +mod defaults; +mod desired_task_count; +mod route_tasks; +mod scale_up_leaf_node; + +pub(crate) use common::EventHandlerChain; +pub use common::{Event, EventHandler}; +pub(crate) use defaults::{ + file_scan_config_desired_task_count, file_scan_config_scale_up_leaf_node, +}; +pub use desired_task_count::{ + DesiredTaskCountEvent, DesiredTaskCountEventResponse, DesiredTaskCountHandler, + TaskCountAnnotation, +}; +pub use route_tasks::{RouteTasksEvent, RouteTasksEventResponse, RouteTasksHandler}; +pub use scale_up_leaf_node::{ + ScaleUpLeafNodeEvent, ScaleUpLeafNodeEventResponse, ScaleUpLeafNodeHandler, +}; diff --git a/src/events/route_tasks.rs b/src/events/route_tasks.rs new file mode 100644 index 00000000..b22c773f --- /dev/null +++ b/src/events/route_tasks.rs @@ -0,0 +1,76 @@ +use super::common::{Event, EventHandler, EventHandlerChain}; +use datafusion::error::Result; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::ExecutionPlan; +use std::sync::Arc; +use url::Url; + +/// Information supplied when the coordinator assigns a stage's tasks to workers. +/// +/// A routing response contains one worker URL per task, in task-index order. The coordinator +/// validates that a response has exactly [`Self::handle`] URLs before it starts the stage. +#[derive(Clone)] +pub struct RouteTasksEvent<'a> { + /// The task context active for the query being coordinated. + pub task_ctx: Arc, + /// The head execution plan of the stage whose tasks are being routed. + /// WARNING: this is never going to be a custom leaf node, this is the head node of the fragment + /// of the plan that contains the custom leaf node. + pub plan: &'a Arc, + /// The number of task slots that need a worker assignment. + pub task_count: usize, +} + +/// Worker assignments returned by a [`EventHandler`]. +pub struct RouteTasksEventResponse { + /// One worker URL per task, ordered by task index. + pub urls: Vec, +} + +impl RouteTasksEventResponse { + /// Returns a response assigning tasks to `urls` in order. + /// + /// The coordinator rejects the response unless `urls.len()` equals + /// [`RouteTasksEvent::handle`]. + pub fn new(urls: Vec) -> Self { + Self { urls } + } +} + +/// Optionally assigns a stage's task slots to worker URLs. +/// +/// Custom handlers are evaluated in registration order, followed by built-in handlers. Return +/// `Some(Ok(_))` to select a complete routing response and stop dispatch, or `None` to let the next +/// handler try. If every handler returns `None`, the coordinator assigns the tasks round-robin, +/// from a randomized worker offset. Returning `Some(Err(_))` aborts execution before tasks are +/// submitted. +pub struct RouteTasksHandler; + +impl RouteTasksHandler { + pub(crate) fn handle(ev: RouteTasksEvent<'_>) -> Option> { + let handlers = ev + .task_ctx + .session_config() + .get_extension::>()?; + handlers.handle(ev) + } +} + +impl Event for RouteTasksHandler { + type Data<'a> = RouteTasksEvent<'a>; + type Response = Result; +} + +impl EventHandler for F +where + F: Send + Sync + 'static, + F: for<'a> Fn(RouteTasksEvent<'a>) -> Option>, +{ + /// Optionally assigns the tasks described by `ev` to specific worker URLs. + /// + /// Return `None` when this handler does not apply. A successful response must provide exactly + /// one URL for every task, in task-index order. + fn handle(&self, ev: RouteTasksEvent<'_>) -> Option> { + self(ev) + } +} diff --git a/src/events/scale_up_leaf_node.rs b/src/events/scale_up_leaf_node.rs new file mode 100644 index 00000000..b9ee40a7 --- /dev/null +++ b/src/events/scale_up_leaf_node.rs @@ -0,0 +1,76 @@ +use super::common::{Event, EventHandler, EventHandlerChain}; +use datafusion::error::Result; +use datafusion::execution::config::SessionConfig; +use datafusion::physical_plan::ExecutionPlan; +use std::sync::Arc; + +/// Information supplied when a leaf has been assigned its final stage task count. +/// +/// This event runs after desired task counts from all of the stage's leaves have been reconciled, +/// so [`Self::handle`] is the count that the stage will actually use. A handler can use it to +/// replace a leaf with a task-specialized plan, such as a [`crate::DistributedLeafExec`] that +/// selects a different input for every task. +#[derive(Clone, Copy)] +pub struct ScaleUpLeafNodeEvent<'a> { + /// The leaf execution plan to replace, if this handler recognizes it. + pub plan: &'a Arc, + /// The final number of tasks that will execute the leaf's stage. + pub task_count: usize, + /// The session configuration that registered the event handlers and holds query options. + pub session_config: &'a SessionConfig, +} + +/// A replacement plan returned by a [`EventHandler`]. +/// +/// The distributed planner annotates every node in this plan with the final stage task count, so +/// the replacement may be a small subtree rather than only a single leaf node. +pub struct ScaleUpLeafNodeEventResponse { + /// The replacement for [`ScaleUpLeafNodeEvent::plan`]. + pub plan: Arc, +} + +impl ScaleUpLeafNodeEventResponse { + /// Returns a response that replaces the event's leaf with `plan`. + pub fn new(plan: Arc) -> Self { + Self { plan } + } +} + +/// Handles optional leaf rewrites after a stage's task count is final. +/// +/// Custom handlers are evaluated in registration order, followed by built-in handlers. Return +/// `Some(Ok(_))` to select a replacement and stop dispatch, or `None` to let the next handler try +/// the same leaf. If all handlers return `None`, the original leaf is left unchanged. Returning +/// `Some(Err(_))` aborts planning. +pub struct ScaleUpLeafNodeHandler; + +impl ScaleUpLeafNodeHandler { + pub(crate) fn handle( + ev: ScaleUpLeafNodeEvent<'_>, + ) -> Option> { + let handlers = ev + .session_config + .get_extension::>()?; + handlers.handle(ev) + } +} + +impl Event for ScaleUpLeafNodeHandler { + type Data<'a> = ScaleUpLeafNodeEvent<'a>; + type Response = Result; +} + +impl EventHandler for F +where + F: Send + Sync + 'static, + F: for<'a> Fn(ScaleUpLeafNodeEvent<'a>) -> Option>, +{ + /// Optionally replaces the leaf described by `ev`. + /// + /// `ev.task_count` already accounts for the constraints and desired counts of every leaf in + /// the stage. Implementations should use it when creating the per-task variants of the + /// replacement plan. + fn handle(&self, ev: ScaleUpLeafNodeEvent<'_>) -> Option> { + self(ev) + } +} diff --git a/src/lib.rs b/src/lib.rs index c7032742..89f30ae4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -21,7 +21,11 @@ pub use coordinator::DistributedExec; pub use distributed_ext::{DistributedExt, DistributedGetterExt}; pub use distributed_planner::{ DistributedConfig, NetworkBoundary, NetworkBoundaryExt, SessionStateBuilderExt, - TaskCountAnnotation, TaskEstimation, TaskEstimator, TaskRoutingContext, +}; +pub use events::{ + DesiredTaskCountEvent, DesiredTaskCountEventResponse, DesiredTaskCountHandler, EventHandler, + RouteTasksEvent, RouteTasksEventResponse, RouteTasksHandler, ScaleUpLeafNodeEvent, + ScaleUpLeafNodeEventResponse, ScaleUpLeafNodeHandler, TaskCountAnnotation, }; pub use execution_plans::{ BroadcastExec, DistributedLeafExec, NetworkBroadcastExec, NetworkCoalesceExec, @@ -35,8 +39,10 @@ pub use metrics::{ }; pub use worker::LocalWorkerContext; +mod events; #[cfg(any(feature = "integration", test))] pub mod test_utils; + #[cfg(feature = "grpc")] pub use protocol::grpc; diff --git a/src/metrics/task_metrics_collector.rs b/src/metrics/task_metrics_collector.rs index 2d78998d..77b3741a 100644 --- a/src/metrics/task_metrics_collector.rs +++ b/src/metrics/task_metrics_collector.rs @@ -59,7 +59,7 @@ mod tests { .with_distributed_worker_resolver(InMemoryWorkerResolver::new(10)) .with_distributed_channel_resolver(InMemoryChannelResolver::default()) .with_distributed_planner() - .with_distributed_task_estimator(2) + .with_distributed_event_handler(2) .with_distributed_metrics_collection(true) .unwrap() .build(); diff --git a/src/metrics/task_metrics_rewriter.rs b/src/metrics/task_metrics_rewriter.rs index 8913ff87..c7f68ce8 100644 --- a/src/metrics/task_metrics_rewriter.rs +++ b/src/metrics/task_metrics_rewriter.rs @@ -336,7 +336,7 @@ mod tests { .with_distributed_metrics_collection(true) .unwrap() .with_distributed_planner() - .with_distributed_task_estimator(2) + .with_distributed_event_handler(2) } let state = builder.build(); diff --git a/src/test_utils/plans.rs b/src/test_utils/plans.rs index aa56338e..efbadfd7 100644 --- a/src/test_utils/plans.rs +++ b/src/test_utils/plans.rs @@ -1,25 +1,29 @@ #[cfg(test)] use super::parquet::register_parquet_tables; use crate::coordinator::DistributedExec; +#[cfg(test)] +use crate::events::{ + EventHandlerChain, file_scan_config_desired_task_count, file_scan_config_scale_up_leaf_node, +}; use crate::stage::Stage; #[cfg(test)] use crate::{ - DistributedConfig, DistributedExt, SessionStateBuilderExt, TaskEstimation, TaskEstimator, - display_plan_ascii, test_utils::in_memory_channel_resolver::InMemoryWorkerResolver, + DesiredTaskCountEvent, DesiredTaskCountEventResponse, DesiredTaskCountHandler, + DistributedConfig, DistributedExt, EventHandler, ScaleUpLeafNodeHandler, + SessionStateBuilderExt, display_plan_ascii, + test_utils::in_memory_channel_resolver::InMemoryWorkerResolver, }; use crate::{NetworkBoundaryExt, TaskKey}; +use datafusion::{ + common::{HashMap, HashSet}, + physical_plan::ExecutionPlan, +}; #[cfg(test)] use datafusion::{ - common::Result, - config::ConfigOptions, execution::{SessionState, context::SessionContext, session_state::SessionStateBuilder}, physical_plan::displayable, prelude::SessionConfig, }; -use datafusion::{ - common::{HashMap, HashSet}, - physical_plan::ExecutionPlan, -}; use std::sync::Arc; /// count_plan_nodes counts the number of execution plan nodes in a plan using BFS traversal. @@ -133,7 +137,7 @@ pub(crate) struct TestPlanBuilder { distributed_file_scan_config_bytes_per_partition: Option, information_schema: Option, broadcast_joins: bool, - distributed_task_estimator: Option>, + desired_task_count_handler: Option>>, distributed_partial_reduce: Option, distributed_children_isolator_unions: Option, distributed_max_tasks_per_stage: Option, @@ -153,7 +157,7 @@ impl TestPlanBuilder { distributed_file_scan_config_bytes_per_partition: Some(1), information_schema: None, broadcast_joins: false, - distributed_task_estimator: None, + desired_task_count_handler: None, distributed_partial_reduce: None, distributed_children_isolator_unions: None, distributed_max_tasks_per_stage: None, @@ -216,11 +220,11 @@ impl TestPlanBuilder { self } - pub fn distributed_task_estimator( + pub fn desired_task_count_handler( mut self, - task_estimator: impl TaskEstimator + Send + Sync + 'static, + handler: impl EventHandler, ) -> Self { - self.distributed_task_estimator = Some(Arc::new(task_estimator)); + self.desired_task_count_handler = Some(Arc::new(handler)); self } @@ -296,9 +300,19 @@ impl TestPlanBuilder { } if self.distributed_planner { state = state.with_distributed_planner(); + } else { + let cfg = state.config().get_or_insert_default(); + EventHandlerChain::::push_builtin( + cfg, + Arc::new(file_scan_config_desired_task_count), + ); + EventHandlerChain::::push_builtin( + cfg, + Arc::new(file_scan_config_scale_up_leaf_node), + ); } - if let Some(t) = self.distributed_task_estimator.clone() { - state = state.with_distributed_task_estimator(t); + if let Some(handler) = self.desired_task_count_handler.clone() { + state = state.with_distributed_event_handler(handler); } state.build() } @@ -340,7 +354,7 @@ impl Default for TestPlanBuilder { distributed_file_scan_config_bytes_per_partition: Some(1), information_schema: Some(false), broadcast_joins: false, - distributed_task_estimator: None, + desired_task_count_handler: None, distributed_partial_reduce: None, distributed_children_isolator_unions: None, distributed_max_tasks_per_stage: None, @@ -352,35 +366,19 @@ impl Default for TestPlanBuilder { } #[cfg(test)] -#[derive(Debug)] -pub(crate) struct BuildSideOneTaskEstimator; - -#[cfg(test)] -impl TaskEstimator for BuildSideOneTaskEstimator { - fn task_estimation( - &self, - plan: &Arc, - _: &ConfigOptions, - ) -> Option { - if !plan.children().is_empty() { - return None; - } - let schema = plan.schema(); - let has_min_temp = schema.fields().iter().any(|f| f.name() == "MinTemp"); - let has_max_temp = schema.fields().iter().any(|f| f.name() == "MaxTemp"); - if has_min_temp && !has_max_temp { - Some(TaskEstimation::maximum(1)) - } else { - None - } +pub(crate) fn build_side_one_desired_task_count_handler( + ev: DesiredTaskCountEvent, +) -> Option { + let plan = ev.plan; + if !plan.children().is_empty() { + return None; } - - fn scale_up_leaf_node( - &self, - _: &Arc, - _: usize, - _: &ConfigOptions, - ) -> Result>> { - Ok(None) + let schema = plan.schema(); + let has_min_temp = schema.fields().iter().any(|f| f.name() == "MinTemp"); + let has_max_temp = schema.fields().iter().any(|f| f.name() == "MaxTemp"); + if has_min_temp && !has_max_temp { + Some(DesiredTaskCountEventResponse::maximum(1)) + } else { + None } } diff --git a/src/test_utils/routing.rs b/src/test_utils/routing.rs index 40c080b4..15fdabc6 100644 --- a/src/test_utils/routing.rs +++ b/src/test_utils/routing.rs @@ -20,11 +20,13 @@ use futures::stream; use prost::Message; use std::{fmt::Formatter, sync::Arc}; use tonic::async_trait; -use url::Url; use crate::execution_plans::DistributedLeafExec; use crate::worker::LocalWorkerContext; -use crate::{DistributedTaskContext, TaskEstimation, TaskEstimator, WorkerResolver}; +use crate::{ + DesiredTaskCountEvent, DesiredTaskCountEventResponse, DistributedTaskContext, RouteTasksEvent, + RouteTasksEventResponse, ScaleUpLeafNodeEvent, ScaleUpLeafNodeEventResponse, WorkerResolver, +}; use crate::distributed_ext::DistributedGetterExt; // Table function that creates a `URLEmitterExec` for testing task routing. @@ -243,54 +245,47 @@ fn url_emitter_schema() -> SchemaRef { ])) } -#[derive(Clone)] -pub struct URLEmitterTaskEstimator; - -impl TaskEstimator for URLEmitterTaskEstimator { - fn task_estimation( - &self, - plan: &std::sync::Arc, - _cfg: &datafusion::config::ConfigOptions, - ) -> Option { - plan.downcast_ref::() - .map(|exec| TaskEstimation::desired(exec.task_count)) - } +pub fn url_emitter_desired_task_count( + ev: DesiredTaskCountEvent, +) -> Option { + ev.plan + .downcast_ref::() + .map(|exec| DesiredTaskCountEventResponse::desired(exec.task_count)) +} - fn scale_up_leaf_node( - &self, - plan: &Arc, - task_count: usize, - _cfg: &datafusion::config::ConfigOptions, - ) -> Result>> { - let Some(exec) = plan.downcast_ref::() else { - return Ok(None); - }; - let p = exec.properties.partitioning.partition_count(); - // Expose ceil(p / task_count) partitions per task so the network boundary - // computes a consistent output partition count. - let visible = p.div_ceil(task_count).max(1); - let template = Arc::new(exec.clone().with_partitions(visible, visible)); - - // Distribute p partitions across task_count tasks using the floor/remainder algorithm: - // the first (p % task_count) tasks get ceil(p/task_count) effective partitions, the rest - // get floor — using the floor/remainder distribution algorithm. - let q = p / task_count; - let r = p % task_count; - let per_task: Vec> = (0..task_count) - .map(|task_idx| { - let effective = q + if task_idx < r { 1 } else { 0 }; - Arc::new(exec.clone().with_partitions(visible, effective)) as _ - }) - .collect(); - - Ok(Some(Arc::new(DistributedLeafExec::try_new( - template as _, - per_task, - )?))) - } +pub fn url_emitter_scale_up_leaf_node( + ev: ScaleUpLeafNodeEvent, +) -> Option> { + let exec = ev.plan.downcast_ref::()?; + let p = exec.properties.partitioning.partition_count(); + // Expose ceil(p / task_count) partitions per task so the network boundary + // computes a consistent output partition count. + let visible = p.div_ceil(ev.task_count).max(1); + let template = Arc::new(exec.clone().with_partitions(visible, visible)); + + // Distribute p partitions across task_count tasks using the floor/remainder algorithm: + // the first (p % task_count) tasks get ceil(p/task_count) effective partitions, the rest + // get floor — using the floor/remainder distribution algorithm. + let q = p / ev.task_count; + let r = p % ev.task_count; + let per_task: Vec> = (0..ev.task_count) + .map(|task_idx| { + let effective = q + if task_idx < r { 1 } else { 0 }; + Arc::new(exec.clone().with_partitions(visible, effective)) as _ + }) + .collect(); + + Some(Ok(ScaleUpLeafNodeEventResponse::new( + match DistributedLeafExec::try_new(template as _, per_task) { + Ok(exec) => Arc::new(exec), + Err(err) => return Some(Err(err)), + }, + ))) +} - fn route_tasks(&self, routing_ctx: &crate::TaskRoutingContext<'_>) -> Result>> { - let mut routed_urls = routing_ctx +pub fn url_emitter_route_tasks(ev: RouteTasksEvent) -> Option> { + Some((|| { + let mut routed_urls = ev .task_ctx .session_config() .get_distributed_worker_resolver()? @@ -298,9 +293,9 @@ impl TaskEstimator for URLEmitterTaskEstimator { // Trivial routing policy: Assign tasks to URLs in reverse order. routed_urls.reverse(); - routed_urls.truncate(routing_ctx.task_count); - Ok(Some(routed_urls)) - } + routed_urls.truncate(ev.task_count); + Ok(RouteTasksEventResponse::new(routed_urls)) + })()) } #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/src/test_utils/test_work_unit_feed.rs b/src/test_utils/test_work_unit_feed.rs index ce82e5c3..06ed4cbf 100644 --- a/src/test_utils/test_work_unit_feed.rs +++ b/src/test_utils/test_work_unit_feed.rs @@ -1,5 +1,6 @@ use crate::{ - DistributedTaskContext, TaskEstimation, TaskEstimator, WorkUnitFeed, WorkUnitFeedProto, + DesiredTaskCountEvent, DesiredTaskCountEventResponse, DistributedTaskContext, + ScaleUpLeafNodeEvent, ScaleUpLeafNodeEventResponse, WorkUnitFeed, WorkUnitFeedProto, WorkUnitFeedProvider, }; use async_trait::async_trait; @@ -10,7 +11,6 @@ use datafusion::catalog::{Session, TableFunctionImpl}; use datafusion::common::stats::Precision; use datafusion::common::tree_node::{Transformed, TreeNode}; use datafusion::common::{Result, ScalarValue, Statistics, internal_err, plan_err}; -use datafusion::config::ConfigOptions; use datafusion::datasource::{TableProvider, TableType}; use datafusion::error::DataFusionError; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; @@ -319,49 +319,35 @@ impl TableProvider for TestWorkUnitFeedTableProvider { } } -pub struct TestWorkUnitFeedTaskEstimator; - -impl TaskEstimator for TestWorkUnitFeedTaskEstimator { - fn task_estimation( - &self, - plan: &Arc, - _cfg: &ConfigOptions, - ) -> Option { - let exec = plan.downcast_ref::()?; - let provider = exec.feed.clone().try_into_inner().ok()?; - Some(TaskEstimation::desired(provider.task_count)) - } +pub fn row_generator_desired_task_count_handler( + ev: DesiredTaskCountEvent, +) -> Option { + let exec = ev.plan.downcast_ref::()?; + let provider = exec.feed.clone().try_into_inner().ok()?; + Some(DesiredTaskCountEventResponse::desired(provider.task_count)) +} - fn scale_up_leaf_node( - &self, - plan: &Arc, - task_count: usize, - _cfg: &ConfigOptions, - ) -> Result>> { - let Some(exec) = plan.downcast_ref::() else { - return Ok(None); +pub fn row_generator_scale_up_leaf_node_handler( + ev: ScaleUpLeafNodeEvent, +) -> Option> { + let exec = ev.plan.downcast_ref::()?; + let provider = exec.feed.clone().try_into_inner().ok()?; + let partitions_per_task = provider.per_partition_ops.len() / ev.task_count; + + let transformed = Arc::clone(ev.plan).transform_down(|plan| { + if let Some(exec) = plan.downcast_ref::() { + return Ok(Transformed::yes(Arc::new(RowGeneratorExec::new( + exec.feed.clone(), + exec.tag.clone(), + partitions_per_task, + exec.projection.clone(), + exec.total_rows, + )))); }; - let Some(provider) = exec.feed.clone().try_into_inner().ok() else { - return Ok(None); - }; - let partitions_per_task = provider.per_partition_ops.len() / task_count; - - // Rebuild the exec with the decided task count so its partition count matches. - let transformed = Arc::clone(plan).transform_down(|plan| { - if let Some(exec) = plan.downcast_ref::() { - return Ok(Transformed::yes(Arc::new(RowGeneratorExec::new( - exec.feed.clone(), - exec.tag.clone(), - partitions_per_task, - exec.projection.clone(), - exec.total_rows, - )))); - }; - Ok(Transformed::no(plan)) - }); + Ok(Transformed::no(plan)) + }); - Ok(Some(transformed?.data)) - } + Some(transformed.map(|transformed| ScaleUpLeafNodeEventResponse::new(transformed.data))) } impl DisplayAs for RowGeneratorExec { diff --git a/src/test_utils/work_unit_file_scan.rs b/src/test_utils/work_unit_file_scan.rs index df3d8de3..5e843159 100644 --- a/src/test_utils/work_unit_file_scan.rs +++ b/src/test_utils/work_unit_file_scan.rs @@ -5,7 +5,10 @@ //! the latency impact of routing file scan inputs through the work unit //! pipeline as compared to the regular [`FileScanConfig`] path. -use crate::{DistributedConfig, TaskCountAnnotation, TaskEstimation, TaskEstimator}; +use crate::{ + DesiredTaskCountEvent, DesiredTaskCountEventResponse, DistributedConfig, ScaleUpLeafNodeEvent, + ScaleUpLeafNodeEventResponse, +}; use crate::{WorkUnitFeed, WorkUnitFeedProto, WorkUnitFeedProvider}; use datafusion::catalog::memory::DataSourceExec; use datafusion::common::tree_node::{Transformed, TreeNode}; @@ -344,78 +347,61 @@ impl PhysicalOptimizerRule for WorkUnitFileScanRule { } } -/// [`TaskEstimator`] for [`WorkUnitFileScanConfig`] leaves that delegates to -/// the built-in [`FileScanConfigTaskEstimator`]: we synthesize a regular -/// `DataSourceExec(FileScanConfig)` carrying the same files currently stored -/// in the work-unit feed, hand it to the underlying estimator, and then -/// re-wrap the result back into our work-unit-flavored data source. +/// Computes a soft task-count hint from the total size of the files held by a +/// [`WorkUnitFileScanConfig`]'s feed. /// -/// `FileScanConfigTaskEstimator::scale_up_leaf_node` returns a -/// `PartitionIsolatorExec(DataSourceExec(FileScanConfig))`. We unwrap that -/// here: the `PartitionIsolatorExec` itself must not appear in the final plan -/// because the per-task feed routing already handles per-task isolation. We -/// only keep the inner `FileScanConfig`'s file groups, flatten them back into -/// one `PartitionedFile` per feed slot, and feed them into a freshly built -/// `WorkUnitFileScanConfig`. -#[derive(Debug, Default)] -pub struct WorkUnitFileScanTaskEstimator; - -impl TaskEstimator for WorkUnitFileScanTaskEstimator { - fn task_estimation( - &self, - plan: &Arc, - cfg: &ConfigOptions, - ) -> Option { - let dse = plan.downcast_ref::()?; - let wfs = dse.data_source().downcast_ref::()?; - - // Same as FileScanConfigTaskEstimator.task_estimation. - let d_cfg = cfg.extensions.get::()?; - - let mut total_bytes = 0; - for file_group in &wfs.feed.inner()?.file_groups { - for file in file_group.files() { - total_bytes += file.effective_size() as usize; - } +/// This follows the same byte-budget calculation as the default file-scan +/// handler: each task may scan up to +/// `file_scan_config_bytes_per_partition * target_partitions` bytes. +pub fn work_unit_file_scan_desired_task_count( + ev: DesiredTaskCountEvent, +) -> Option { + let cfg = ev.session_config; + let dse = ev.plan.downcast_ref::()?; + let wfs = dse.data_source().downcast_ref::()?; + + let d_cfg = DistributedConfig::from_session_config(cfg).ok()?; + + let mut total_bytes = 0; + for file_group in &wfs.feed.inner()?.file_groups { + for file in file_group.files() { + total_bytes += file.effective_size() as usize; } - - let task_count = total_bytes - .div_ceil(d_cfg.file_scan_config_bytes_per_partition) - .div_ceil(cfg.execution.target_partitions); - - Some(TaskEstimation { - task_count: TaskCountAnnotation::Desired(task_count), - }) } - fn scale_up_leaf_node( - &self, - plan: &Arc, - task_count: usize, - _cfg: &ConfigOptions, - ) -> Result>> { - let Some(dse) = plan.downcast_ref::() else { - return Ok(None); - }; - let Some(wfs) = dse.data_source().downcast_ref::() else { - return Ok(None); - }; - - let wuf_provider = wfs.feed.try_inner()?; + let task_count = total_bytes + .div_ceil(d_cfg.file_scan_config_bytes_per_partition) + .div_ceil(cfg.target_partitions()); - // Same as FileScanConfigTaskEstimator.scale_up_leaf_node - let mut new_file_groups = vec![]; - for file_group in wuf_provider.file_groups.clone() { - new_file_groups.extend(file_group.split_files(task_count)); - } + Some(DesiredTaskCountEventResponse::desired(task_count)) +} - let new_provider = FileScanWorkUnitProvider::new(new_file_groups); - Ok(Some( - DataSourceExec::from_data_source(WorkUnitFileScanConfig { - feed: WorkUnitFeed::new(new_provider), - fsc: wfs.fsc.clone(), - partitions: wfs.partitions, - }) as Arc, - )) +/// Rebuilds a work-unit file scan after the stage task count is finalized. +/// +/// The work-unit feed owns task-local file assignment, so this handler splits +/// the feed's file groups across the selected tasks and builds a fresh +/// [`WorkUnitFileScanConfig`]. It deliberately does not introduce the default +/// file-scan partition isolator: that wrapper would be redundant, because the +/// work-unit feed already isolates each task's files. +pub fn work_unit_file_scan_scale_up_leaf_node( + ev: ScaleUpLeafNodeEvent, +) -> Option> { + let dse = ev.plan.downcast_ref::()?; + let wfs = dse.data_source().downcast_ref::()?; + + let wuf_provider = wfs.feed.inner()?; + + let mut new_file_groups = vec![]; + for file_group in wuf_provider.file_groups.clone() { + new_file_groups.extend(file_group.split_files(ev.task_count)); } + + let new_provider = FileScanWorkUnitProvider::new(new_file_groups); + Some(Ok(ScaleUpLeafNodeEventResponse::new( + DataSourceExec::from_data_source(WorkUnitFileScanConfig { + feed: WorkUnitFeed::new(new_provider), + fsc: wfs.fsc.clone(), + partitions: wfs.partitions, + }) as Arc, + ))) } diff --git a/src/worker_resolver.rs b/src/worker_resolver.rs index 89ea1e3f..e4a5ec50 100644 --- a/src/worker_resolver.rs +++ b/src/worker_resolver.rs @@ -29,7 +29,8 @@ pub(crate) fn set_distributed_worker_resolver( } /// Gets the [WorkerResolver] from the [SessionConfig]'s extensions. Typically called inside -/// [TaskEstimator::route_tasks] to resolve the worker URLs available for distributed tasks. +/// [`crate::EventHandler`] to resolve the worker URLs available for distributed +/// tasks. pub fn get_distributed_worker_resolver( cfg: &SessionConfig, ) -> Result, DataFusionError> { diff --git a/tests/metrics_collection.rs b/tests/metrics_collection.rs index b7e6a4e7..9341ae9a 100644 --- a/tests/metrics_collection.rs +++ b/tests/metrics_collection.rs @@ -13,7 +13,7 @@ mod tests { use datafusion_distributed::test_utils::parquet::register_parquet_tables; use datafusion_distributed::test_utils::test_work_unit_feed::{ RowGeneratorExec, TestWorkUnitFeedExecCodec, TestWorkUnitFeedFunction, - TestWorkUnitFeedTaskEstimator, + row_generator_desired_task_count_handler, row_generator_scale_up_leaf_node_handler, }; use datafusion_distributed::{ DefaultSessionBuilder, DistributedExt, DistributedLeafExec, DistributedMetricsFormat, @@ -320,7 +320,8 @@ mod tests { let (mut ctx, _guard, _) = start_localhost_context(3, build_state).await; ctx.set_distributed_work_unit_feed(|p: &RowGeneratorExec| Some(&p.feed)); ctx.set_distributed_user_codec(TestWorkUnitFeedExecCodec); - ctx.set_distributed_task_estimator(TestWorkUnitFeedTaskEstimator); + ctx.set_distributed_event_handler(row_generator_desired_task_count_handler); + ctx.set_distributed_event_handler(row_generator_scale_up_leaf_node_handler); ctx.register_udtf("test_work_unit", Arc::new(TestWorkUnitFeedFunction)); // Two tasks × two partitions × comma-separated row counts. Total work units sent: diff --git a/tests/task_estimator_test.rs b/tests/task_estimator_test.rs index 6e2fd195..7945648d 100644 --- a/tests/task_estimator_test.rs +++ b/tests/task_estimator_test.rs @@ -8,7 +8,10 @@ mod tests { DistributedExt, WorkerQueryContext, assert_snapshot, display_plan_ascii, test_utils::{ in_memory_channel_resolver::start_in_memory_context, - routing::{URLEmitterExtensionCodec, URLEmitterFunction, URLEmitterTaskEstimator}, + routing::{ + URLEmitterExtensionCodec, URLEmitterFunction, url_emitter_desired_task_count, + url_emitter_route_tasks, url_emitter_scale_up_leaf_node, + }, }, }; @@ -382,7 +385,9 @@ mod tests { async fn run_query(sql: &str) -> Result<(String, String), DataFusionError> { let mut ctx = start_in_memory_context(NUM_WORKERS, build_state).await; - ctx.set_distributed_task_estimator(URLEmitterTaskEstimator); + ctx.set_distributed_event_handler(url_emitter_desired_task_count); + ctx.set_distributed_event_handler(url_emitter_scale_up_leaf_node); + ctx.set_distributed_event_handler(url_emitter_route_tasks); ctx.set_distributed_user_codec(URLEmitterExtensionCodec); ctx.register_udtf("url_emitter", Arc::new(URLEmitterFunction)); ctx.state_ref() diff --git a/tests/work_unit_feed.rs b/tests/work_unit_feed.rs index 36cd0c37..a1d7c004 100644 --- a/tests/work_unit_feed.rs +++ b/tests/work_unit_feed.rs @@ -7,7 +7,7 @@ mod tests { use datafusion_distributed::test_utils::localhost::start_localhost_context; use datafusion_distributed::test_utils::test_work_unit_feed::{ RowGeneratorExec, TestWorkUnitFeedExecCodec, TestWorkUnitFeedFunction, - TestWorkUnitFeedTaskEstimator, + row_generator_desired_task_count_handler, row_generator_scale_up_leaf_node_handler, }; use datafusion_distributed::{ DistributedExt, WorkerQueryContext, assert_snapshot, display_plan_ascii, @@ -863,7 +863,8 @@ mod tests { let (mut ctx, _guard, _) = start_localhost_context(3, build_state).await; ctx.set_distributed_work_unit_feed(|p: &RowGeneratorExec| Some(&p.feed)); ctx.set_distributed_user_codec(TestWorkUnitFeedExecCodec); - ctx.set_distributed_task_estimator(TestWorkUnitFeedTaskEstimator); + ctx.set_distributed_event_handler(row_generator_desired_task_count_handler); + ctx.set_distributed_event_handler(row_generator_scale_up_leaf_node_handler); ctx.register_udtf("test_work_unit", Arc::new(TestWorkUnitFeedFunction)); let mut df_opt = None;