From a2778c7c74483b1dcd3a82624af3498e8b36eef3 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sat, 20 Jun 2026 02:18:02 -0700 Subject: [PATCH 1/4] Added a pluggable `WorkerTransport` abstraction. Carries the extension points a transport needs beyond plan delivery and reads: a per-execution cancellation token that `DistributedExec::execute` attaches and fires when the output stream drops, `NetworkBoundary` partition routing (`route_partition` / `partitions_per_consumer_task` / `PartitionRoute`), and `DistributedExec::prepare_in_process_plan` for a transport that drives the head stage itself. The default stays Arrow Flight; in-memory and shared-memory transports plug in through these extension points without further changes to the core. --- src/common/cancellation.rs | 29 ++++ src/common/mod.rs | 7 +- src/common/recursion.rs | 2 +- src/common/uuid.rs | 4 +- src/config_extension_ext.rs | 2 +- src/coordinator/distributed.rs | 122 ++++++++++++-- src/coordinator/latency_metric.rs | 4 +- src/coordinator/metrics_store.rs | 2 +- src/coordinator/mod.rs | 4 +- src/coordinator/prepare_static_plan.rs | 89 +++++++--- src/coordinator/query_coordinator.rs | 152 ++++++------------ src/distributed_ext.rs | 38 ++++- src/distributed_planner/distributed_config.rs | 23 ++- src/distributed_planner/mod.rs | 4 +- src/distributed_planner/network_boundary.rs | 38 ++++- src/execution_plans/network_coalesce.rs | 10 ++ src/lib.rs | 20 ++- src/networking/mod.rs | 4 + src/networking/worker_transport.rs | 45 ++++++ src/passthrough_headers.rs | 2 +- src/work_unit_feed/remote_work_unit_feed.rs | 16 +- src/worker/flight.rs | 62 +++++++ src/worker/mod.rs | 5 + src/worker/transport.rs | 108 +++++++++++++ src/worker/worker_connection_pool.rs | 47 ++---- 25 files changed, 633 insertions(+), 206 deletions(-) create mode 100644 src/common/cancellation.rs create mode 100644 src/networking/worker_transport.rs create mode 100644 src/worker/flight.rs create mode 100644 src/worker/transport.rs diff --git a/src/common/cancellation.rs b/src/common/cancellation.rs new file mode 100644 index 00000000..c8626954 --- /dev/null +++ b/src/common/cancellation.rs @@ -0,0 +1,29 @@ +use datafusion::execution::TaskContext; +use tokio_util::sync::CancellationToken; + +/// Per-execution cancellation token, attached to the [TaskContext] by `DistributedExec` when a +/// plan starts. One token per execution: not on the transport (a shared, process-wide instance +/// cannot own per-execution state) and not per connection (too granular). +#[derive(Clone)] +pub(crate) struct DistributedCancellationToken(pub(crate) CancellationToken); + +/// Returns the per-execution [CancellationToken] attached to `ctx`, or a fresh never-cancelled one +/// if none is set (a context that did not come through `DistributedExec`). A transport's producers +/// and consumers watch this instead of the transport carrying a `cancellation()` method. +/// +/// The token does not replace transport-level stream close. Close propagates hop by hop and +/// surfaces only when the next blocking operation fails; the token reaches every in-process +/// participant the moment the head stream drops, including ones whose channels were never opened +/// or never read. +/// +/// Two caveats for watchers: +/// - The token fires on any drop of the head stream, including after normal exhaustion. Treat it +/// as teardown, not failure. +/// - It lives in a session-config extension, so it is process-local: it does not cross plan +/// serialization to remote workers. Out-of-process producers need their own teardown signal. +pub fn get_distributed_cancellation_token(ctx: &TaskContext) -> CancellationToken { + ctx.session_config() + .get_extension::() + .map(|t| t.0.clone()) + .unwrap_or_default() +} diff --git a/src/common/mod.rs b/src/common/mod.rs index bf9ed549..f272f137 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -1,3 +1,4 @@ +mod cancellation; mod children_helpers; mod on_drop_stream; mod once_lock; @@ -6,10 +7,12 @@ mod task_context_helpers; mod time; mod uuid; +pub(crate) use cancellation::DistributedCancellationToken; +pub use cancellation::get_distributed_cancellation_token; pub(crate) use children_helpers::require_one_child; pub(crate) use on_drop_stream::on_drop_stream; pub(crate) use once_lock::OnceLockResult; -pub(crate) use recursion::TreeNodeExt; +pub use recursion::TreeNodeExt; pub(crate) use task_context_helpers::task_ctx_with_extension; pub(crate) use time::now_ns; -pub(crate) use uuid::{deserialize_uuid, serialize_uuid}; +pub use uuid::{deserialize_uuid, serialize_uuid}; diff --git a/src/common/recursion.rs b/src/common/recursion.rs index 2f2b9463..752af33a 100644 --- a/src/common/recursion.rs +++ b/src/common/recursion.rs @@ -6,7 +6,7 @@ use datafusion::physical_plan::ExecutionPlan; use std::cell::RefCell; use std::sync::Arc; -pub(crate) trait TreeNodeExt { +pub trait TreeNodeExt { /// Applies `f` to the node then each of its children, recursively (a /// top-down, pre-order traversal), propagating the [DistributedTaskContext] correctly /// across nodes that mutate this context, and ignoring nodes that do not belong to diff --git a/src/common/uuid.rs b/src/common/uuid.rs index f0178a3d..a80a384d 100644 --- a/src/common/uuid.rs +++ b/src/common/uuid.rs @@ -2,10 +2,10 @@ use datafusion::common::Result; use datafusion_proto::protobuf::proto_error; use uuid::Uuid; -pub(crate) fn deserialize_uuid(id: &[u8]) -> Result { +pub fn deserialize_uuid(id: &[u8]) -> Result { Uuid::from_slice(id).map_err(|err| proto_error(format!("Invalid Uuid bytes: {err}"))) } -pub(crate) fn serialize_uuid(id: &Uuid) -> Vec { +pub fn serialize_uuid(id: &Uuid) -> Vec { id.as_bytes().to_vec() } diff --git a/src/config_extension_ext.rs b/src/config_extension_ext.rs index 88556ffe..07ac52c3 100644 --- a/src/config_extension_ext.rs +++ b/src/config_extension_ext.rs @@ -80,7 +80,7 @@ struct ConfigExtensionPropagationContext { prefixes: Vec<&'static str>, } -pub(crate) fn get_config_extension_propagation_headers( +pub fn get_config_extension_propagation_headers( cfg: &SessionConfig, ) -> Result { fn parse_err(err: impl Error) -> DataFusionError { diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index fe1bbff3..2f22fa04 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -1,20 +1,28 @@ -use crate::common::{require_one_child, serialize_uuid}; +use crate::DistributedConfig; +use crate::common::{ + DistributedCancellationToken, on_drop_stream, require_one_child, serialize_uuid, + task_ctx_with_extension, +}; use crate::coordinator::metrics_store::MetricsStore; use crate::coordinator::prepare_static_plan::prepare_static_plan; -use crate::coordinator::query_coordinator::QueryCoordinator; use crate::distributed_planner::NetworkBoundaryExt; +use crate::networking::get_distributed_worker_transport; use crate::worker::generated::worker::TaskKey; use datafusion::common::internal_datafusion_err; +use datafusion::common::runtime::JoinSet; use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion::common::{Result, exec_err}; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_expr_common::metrics::MetricsSet; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; -use datafusion::physical_plan::stream::RecordBatchReceiverStreamBuilder; +use datafusion::physical_plan::stream::{ + RecordBatchReceiverStreamBuilder, RecordBatchStreamAdapter, +}; use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties}; use futures::StreamExt; use std::fmt::Formatter; use std::sync::{Arc, Mutex}; +use tokio_util::sync::CancellationToken; /// [ExecutionPlan] that executes the inner plan in distributed mode. /// Before executing it, two modifications are lazily performed on the plan: @@ -72,6 +80,14 @@ impl DistributedExec { self } + /// Exposes the metrics store so a transport that collects worker `TaskMetrics` out-of-band can + /// file them where [`rewrite_distributed_plan_with_metrics`] reads them. + /// + /// [`rewrite_distributed_plan_with_metrics`]: crate::rewrite_distributed_plan_with_metrics + pub fn metrics_store(&self) -> Option> { + self.metrics_store.clone() + } + /// Waits until all worker tasks have reported their metrics back via the coordinator channel. /// /// Metrics are delivered asynchronously after query execution completes, so callers that need @@ -133,6 +149,63 @@ impl DistributedExec { .clone() .ok_or_else(|| internal_datafusion_err!("No head stage found. Was execute() called?")) } + + /// Routes and dispatches every stage through the registered transport, then returns the head + /// stage for the caller to drive synchronously, with no `execute()` background task. + /// + /// This is the extension point the in-crate shared-memory transport (and its in-process test) executes + /// through: it owns the runtime and drives the head stage itself. The transport's dispatch must + /// complete synchronously: a transport that spawns background delivery work, or a plan that + /// declares work-unit feeds (which are pumped by that background work), is rejected rather than + /// left to stall. + pub fn prepare_in_process_plan( + &self, + ctx: &Arc, + ) -> Result> { + let d_cfg = DistributedConfig::from_config_options(ctx.session_config().options())?; + let registry = &d_cfg.__private_work_unit_feed_registry; + let mut has_feeds = false; + self.base_plan.apply(|plan| { + if registry.get_work_unit_feed(plan).is_some() { + has_feeds = true; + return Ok(TreeNodeRecursion::Stop); + } + Ok(TreeNodeRecursion::Continue) + })?; + if has_feeds { + return exec_err!( + "the plan declares work-unit feeds, which are not delivered on the in-process path" + ); + } + + let dispatcher = get_distributed_worker_transport(ctx.session_config()).dispatcher(); + let mut join_set = JoinSet::new(); + let result = prepare_static_plan( + &self.base_plan, + ctx, + dispatcher.as_ref(), + &mut join_set, + &self.metrics, + self.metrics_store.as_ref(), + )?; + // Dropping the join set aborts anything still on it, which would kill an async delivery + // mid-flight and surface as a hang far from the cause. Reject instead. + if !join_set.is_empty() { + return exec_err!( + "the registered transport spawned background delivery work; \ + prepare_in_process_plan requires a transport whose dispatch completes synchronously" + ); + } + self.plan_for_viz + .lock() + .expect("poisoned lock") + .replace(result.plan_for_viz); + self.head_stage + .lock() + .expect("poisoned lock") + .replace(Arc::clone(&result.head_stage)); + Ok(result.head_stage) + } } impl DisplayAs for DistributedExec { @@ -185,20 +258,36 @@ impl ExecutionPlan for DistributedExec { let base_plan = Arc::clone(&self.base_plan); let plan_for_viz = Arc::clone(&self.plan_for_viz); let head_stage = Arc::clone(&self.head_stage); + let metrics = self.metrics.clone(); + let metrics_store = self.metrics_store.clone(); - let query_coordinator = QueryCoordinator::new( - Arc::clone(&context), - &self.metrics, - self.metrics_store.clone(), - ); + // One token per execution. In-process transports watch it to tear down the moment this + // output stream drops, whether the consumer read to the end or abandoned it early (a LIMIT + // gather). It rides the dispatch context so every worker fragment shares it. + let token = CancellationToken::new(); + let context = Arc::new(task_ctx_with_extension( + &context, + DistributedCancellationToken(token.clone()), + )); let mut builder = RecordBatchReceiverStreamBuilder::new(self.schema(), 1); let tx = builder.tx(); builder.spawn(async move { - let _guard = query_coordinator.end_query_guard(); + let dispatcher = + get_distributed_worker_transport(context.session_config()).dispatcher(); + // The query owns the join set: background plan-delivery and metrics tasks spawned by the + // dispatcher run on it, so draining it below is what waits for the workers to finish. + let mut join_set = JoinSet::new(); - let result = prepare_static_plan(&query_coordinator, &base_plan)?; + let result = prepare_static_plan( + &base_plan, + &context, + dispatcher.as_ref(), + &mut join_set, + &metrics, + metrics_store.as_ref(), + )?; plan_for_viz .lock() @@ -215,11 +304,20 @@ impl ExecutionPlan for DistributedExec { } } drop(tx); - query_coordinator.drain_pending_tasks().await?; + for res in join_set.join_all().await { + res?; + } + // Dropping the dispatcher ends the coordinator->worker streams, propagating EOS so the + // workers can clean up. It is held until here so those streams stay open during draining. + drop(dispatcher); Ok(()) }); - Ok(builder.build()) + let stream = builder.build(); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + on_drop_stream(stream, move || token.cancel()), + ))) } fn metrics(&self) -> Option { diff --git a/src/coordinator/latency_metric.rs b/src/coordinator/latency_metric.rs index eea47932..a519d052 100644 --- a/src/coordinator/latency_metric.rs +++ b/src/coordinator/latency_metric.rs @@ -8,7 +8,7 @@ use std::time::Duration; /// DataFusion metrics system is pretty limited from an API standpoint. This intermediate struct /// bridges the gaps that are not satisfied by upstream API for measuring latency. -pub(super) struct LatencyMetric { +pub struct LatencyMetric { max: Time, avg: Time, max_latency_micros: AtomicU64, @@ -53,7 +53,7 @@ impl LatencyMetric { } } - pub(super) fn record(&self, start: &Instant) { + pub fn record(&self, start: &Instant) { let micros = start.elapsed().as_micros() as u64; self.max_latency_micros.fetch_max(micros, Ordering::Relaxed); self.sum_latency_micros.fetch_add(micros, Ordering::Relaxed); diff --git a/src/coordinator/metrics_store.rs b/src/coordinator/metrics_store.rs index 4b9b0ffc..2d4cc596 100644 --- a/src/coordinator/metrics_store.rs +++ b/src/coordinator/metrics_store.rs @@ -18,7 +18,7 @@ impl MetricsStore { Self { tx, rx } } - pub(crate) fn insert(&self, key: TaskKey, metrics: pb::TaskMetrics) { + pub fn insert(&self, key: TaskKey, metrics: pb::TaskMetrics) { self.tx.send_modify(|map| { map.insert(key, metrics); }); diff --git a/src/coordinator/mod.rs b/src/coordinator/mod.rs index 8fe771d3..5302aa6c 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -5,4 +5,6 @@ mod prepare_static_plan; mod query_coordinator; pub use distributed::DistributedExec; -pub(crate) use metrics_store::MetricsStore; +pub use latency_metric::LatencyMetric; +pub use metrics_store::MetricsStore; +pub(crate) use query_coordinator::FlightWorkerDispatch; diff --git a/src/coordinator/prepare_static_plan.rs b/src/coordinator/prepare_static_plan.rs index 65d74276..9cb6b2c7 100644 --- a/src/coordinator/prepare_static_plan.rs +++ b/src/coordinator/prepare_static_plan.rs @@ -1,25 +1,34 @@ +use crate::coordinator::MetricsStore; use crate::coordinator::distributed::PreparedPlan; -use crate::coordinator::query_coordinator::QueryCoordinator; +use crate::networking::get_distributed_worker_resolver; use crate::stage::RemoteStage; -use crate::{NetworkBoundaryExt, Stage}; +use crate::worker::{WorkerDispatch, WorkerDispatchRequest}; +use crate::{DistributedConfig, NetworkBoundaryExt, Stage, TaskEstimator, TaskRoutingContext}; +use datafusion::common::runtime::JoinSet; use datafusion::common::tree_node::{Transformed, TreeNode}; use datafusion::common::{Result, exec_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr_common::metrics::ExecutionPlanMetricsSet; use datafusion::physical_plan::ExecutionPlan; +use rand::Rng; use std::sync::Arc; +use url::Url; /// 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]. /// - 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. -/// 3. In each network boundary, set the input plan to `None`. That way, network boundaries -/// become nodes without children and traversing them will not go further down in. -/// 4. Spawn a background task per worker that waits for the worker to finish and collects -/// its metrics into [DistributedExec::task_metrics] via the coordinator channel. +/// 2. Hand each network boundary's stage to the [WorkerDispatch], which delivers the sliced +/// subplans to the assigned workers and wires up whatever back-channels the transport needs. +/// 3. In each network boundary, set the input stage to its `Remote` form holding the assigned +/// worker URLs. Traversing them will not go further down, as they become leaf-like. pub(super) fn prepare_static_plan( - query_coordinator: &QueryCoordinator, base_plan: &Arc, + task_ctx: &Arc, + dispatcher: &dyn WorkerDispatch, + join_set: &mut JoinSet>, + metrics: &ExecutionPlanMetricsSet, + metrics_store: Option<&Arc>, ) -> Result { let prepared = Arc::clone(base_plan).transform_up(|plan| { // The following logic is just applied on network boundaries. @@ -31,25 +40,22 @@ pub(super) fn prepare_static_plan( return exec_err!("Input stage from network boundary was not in Local state"); }; - let mut stage_coordinator = query_coordinator.stage_coordinator(stage); + let routed_urls = routed_urls(task_ctx, &stage.plan, stage.tasks)?; - let routed_urls = stage_coordinator.routed_urls()?; - - let mut workers = Vec::with_capacity(stage.tasks); - for (i, routed_url) in routed_urls.into_iter().enumerate() { - workers.push(routed_url.clone()); - // Spawn a task that sends the subplan to the chosen URL. - // There will be as many spawned tasks as workers. - let (worker_tx, worker_rx) = stage_coordinator.send_plan_task(i, routed_url)?; - stage_coordinator.worker_to_coordinator_task(i, worker_rx); - stage_coordinator.coordinator_to_worker_task(i, worker_tx)?; - } + dispatcher.dispatch(WorkerDispatchRequest { + stage, + routed_urls: &routed_urls, + task_ctx, + metrics, + metrics_store, + join_set, + })?; Ok(Transformed::yes(plan.with_input_stage(Stage::Remote( RemoteStage { query_id: stage.query_id, num: stage.num, - workers, + workers: routed_urls, }, ))?)) })?; @@ -60,3 +66,42 @@ pub(super) fn prepare_static_plan( plan_for_viz: Arc::clone(base_plan), }) } + +/// Returns as many URLs as the `task_count` for the stage. These URLs can be: +/// - chosen by the user, if they provided an implementation for [TaskEstimator::route_tasks]. +/// - assigned via round-robin from a randomized starting point otherwise. +fn routed_urls( + task_ctx: &Arc, + plan: &Arc, + task_count: usize, +) -> Result> { + let session_config = task_ctx.session_config(); + let d_cfg = DistributedConfig::from_config_options(session_config.options())?; + let worker_resolver = get_distributed_worker_resolver(session_config)?; + let task_estimator = &d_cfg.__private_task_estimator; + + let routed_urls = match task_estimator.route_tasks(&TaskRoutingContext { + task_ctx: Arc::clone(task_ctx), + plan, + task_count, + }) { + Ok(Some(routed_urls)) => routed_urls, + Ok(None) => { + let available_urls = worker_resolver.get_urls()?; + let start_idx = rand::rng().random_range(0..available_urls.len()); + (0..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}"), + }; + + if routed_urls.len() != task_count { + return exec_err!( + "number of tasks ({}) was not equal to number of urls ({}) at execution time", + task_count, + routed_urls.len() + ); + } + Ok(routed_urls) +} diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 3d19c750..3ef3d030 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -5,21 +5,21 @@ use crate::coordinator::latency_metric::LatencyMetric; use crate::execution_plans::{ChildrenIsolatorUnionExec, DistributedLeafExec}; use crate::passthrough_headers::get_passthrough_headers; use crate::protobuf::tonic_status_to_datafusion_error; -use crate::stage::LocalStage; use crate::work_unit_feed::{build_work_unit_batch_msg, set_work_unit_send_time}; use crate::worker::generated::worker as pb; use crate::worker::generated::worker::coordinator_to_worker_msg::Inner; use crate::worker::generated::worker::set_plan_request::WorkUnitFeedDeclaration; +use crate::worker::{WorkerDispatch, WorkerDispatchRequest}; use crate::{ BytesCounterMetric, BytesMetricExt, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedCodec, - DistributedConfig, DistributedTaskContext, DistributedWorkUnitFeedContext, TaskEstimator, - TaskKey, TaskRoutingContext, get_distributed_channel_resolver, get_distributed_worker_resolver, + DistributedConfig, DistributedTaskContext, DistributedWorkUnitFeedContext, TaskKey, + get_distributed_channel_resolver, }; +use datafusion::common::Result; use datafusion::common::instant::Instant; use datafusion::common::runtime::JoinSet; use datafusion::common::tree_node::{Transformed, TreeNodeRecursion}; use datafusion::common::{DataFusionError, exec_datafusion_err}; -use datafusion::common::{Result, exec_err}; use datafusion::execution::TaskContext; use datafusion::physical_expr_common::metrics::{ExecutionPlanMetricsSet, Label, MetricBuilder}; use datafusion::physical_plan::ExecutionPlan; @@ -28,9 +28,7 @@ use datafusion_proto::protobuf::PhysicalPlanNode; use futures::{Stream, StreamExt}; use http::Extensions; use prost::Message; -use rand::Rng; -use std::ops::DerefMut; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, OnceLock}; use tokio::sync::Notify; use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; use tokio_stream::wrappers::UnboundedReceiverStream; @@ -44,65 +42,56 @@ use uuid::Uuid; /// small batches. See [StreamExt::ready_chunks] docs for more details about how chunking works. const WORK_UNIT_FEED_CHUNK_SIZE: usize = 256; -/// Manages communication between coordinator and workers for a single query. -/// -/// The [QueryCoordinator]'s lifetime is scoped to a single query , and will instantiate independent -/// [StageCoordinator] scoped to each individual stage. -pub(super) struct QueryCoordinator { - task_ctx: Arc, +/// The Arrow-Flight plan-delivery side: a [WorkerDispatch] that ships each task's plan over the +/// bidirectional coordinator-to-worker gRPC stream and wires up the work-unit feed and metrics +/// back-channel. Per-query state (plan-send metrics, the keep-alive notifier) is created lazily on +/// the first stage's dispatch and shared across stages. Dropping the dispatcher fires the notifier, +/// which closes the coordinator->worker streams and propagates EOS to the workers so they can clean +/// up; the coordinator holds it until the query's result stream is drained. +#[derive(Default)] +pub(crate) struct FlightWorkerDispatch { + state: OnceLock, +} + +struct FlightDispatchState { coordinator_to_worker_metrics: CoordinatorToWorkerMetrics, - metrics_store: Option>, end_stream_notifier: Arc, - join_set: Mutex>>, } -impl QueryCoordinator { - /// Builds a new [QueryCoordinator] scoped to a query. - pub(super) fn new( - task_ctx: Arc, - metrics_set: &ExecutionPlanMetricsSet, - metrics_store: Option>, - ) -> Self { - Self { - task_ctx, - metrics_store, - coordinator_to_worker_metrics: CoordinatorToWorkerMetrics::new(metrics_set), +impl WorkerDispatch for FlightWorkerDispatch { + fn dispatch(&self, request: WorkerDispatchRequest<'_>) -> Result<()> { + let state = self.state.get_or_init(|| FlightDispatchState { + coordinator_to_worker_metrics: CoordinatorToWorkerMetrics::new(request.metrics), end_stream_notifier: Arc::new(Notify::new()), - join_set: Mutex::new(JoinSet::new()), - } - } - - /// Builds a new [StageCoordinator] that will manage coordinator-worker connections for the given - /// stage. - pub(super) fn stage_coordinator<'a>(&'a self, stage: &'a LocalStage) -> StageCoordinator<'a> { - StageCoordinator { + }); + let stage = request.stage; + let mut stage_coordinator = StageCoordinator { plan: &stage.plan, query_id: stage.query_id, stage_id: stage.num, task_count: stage.tasks, - task_ctx: &self.task_ctx, - metrics: &self.coordinator_to_worker_metrics, - metrics_store: &self.metrics_store, - end_stream_notifier: &self.end_stream_notifier, - join_set: &self.join_set, + task_ctx: request.task_ctx, + metrics: &state.coordinator_to_worker_metrics, + metrics_store: request.metrics_store, + end_stream_notifier: &state.end_stream_notifier, + join_set: request.join_set, + }; + for (task_i, url) in request.routed_urls.iter().enumerate() { + let (worker_tx, worker_rx) = stage_coordinator.send_plan_task(task_i, url.clone())?; + stage_coordinator.worker_to_coordinator_task(task_i, worker_rx); + stage_coordinator.coordinator_to_worker_task(task_i, worker_tx)?; } + Ok(()) } +} - /// returns a guard that, when dropped, it signals all the coordinator->worker connections that - /// the query is finished, ending them, and propagating the EOS to the workers so that they can - /// clean up any remaining state. - pub(super) fn end_query_guard(&self) -> NotifyGuard { - NotifyGuard(Arc::clone(&self.end_stream_notifier)) - } - - /// Blocks until all background tasks have finished (e.g., sending WorkUnit feeds, or collecting - /// metrics) - pub(super) async fn drain_pending_tasks(self) -> Result<()> { - let join_set = std::mem::take(self.join_set.lock().unwrap().deref_mut()); - for res in join_set.join_all().await { - res?; +impl Drop for FlightWorkerDispatch { + fn drop(&mut self) { + if let Some(state) = self.state.get() { + // Signal the coordinator->worker streams that the query is finished, ending them and + // propagating the EOS so workers can clean up any remaining state. + state.end_stream_notifier.notify_waiters(); } - Ok(()) } } @@ -121,9 +110,9 @@ pub(super) struct StageCoordinator<'a> { task_count: usize, task_ctx: &'a Arc, metrics: &'a CoordinatorToWorkerMetrics, - metrics_store: &'a Option>, + metrics_store: Option<&'a Arc>, end_stream_notifier: &'a Arc, - join_set: &'a Mutex>>, + join_set: &'a mut JoinSet>, } impl<'a> StageCoordinator<'a> { @@ -187,7 +176,7 @@ impl<'a> StageCoordinator<'a> { let metrics = self.metrics.clone(); - self.join_set.lock().unwrap().spawn(async move { + self.join_set.spawn(async move { let start = Instant::now(); let mut client = channel_resolver.get_worker_client_for_url(&url).await?; let response = client.coordinator_channel(request).await.map_err(|e| { @@ -227,7 +216,7 @@ impl<'a> StageCoordinator<'a> { stage_id: self.stage_id as u64, task_number: task_i as u64, }; - let task_metrics = self.metrics_store.clone(); + let task_metrics = self.metrics_store.cloned(); // Cannot use self.join_set because that's tied to the lifetime of the query, and the // metrics collection process might outlive the query's lifetime. @@ -313,7 +302,7 @@ impl<'a> StageCoordinator<'a> { } } - self.join_set.lock().unwrap().spawn(async move { + self.join_set.spawn(async move { let _guard = WorkUnitEosOnDrop(tx); futures::future::try_join_all(futures).await?; Ok(()) @@ -362,59 +351,12 @@ impl<'a> StageCoordinator<'a> { })?; Ok((transformed.data, work_unit_feed_declarations)) } - - /// Returns as many URLs as the task count for the stage this [StageCoordinator] - /// 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. - pub(super) fn routed_urls(&self) -> Result> { - let session_config = self.task_ctx.session_config(); - let d_cfg = DistributedConfig::from_config_options(session_config.options())?; - let worker_resolver = get_distributed_worker_resolver(session_config)?; - let task_estimator = &d_cfg.__private_task_estimator; - - let routed_urls = match task_estimator.route_tasks(&TaskRoutingContext { - task_ctx: Arc::clone(self.task_ctx), - plan: self.plan, - task_count: self.task_count, - }) { - Ok(Some(routed_urls)) => routed_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) => { - 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}"), - }; - - if routed_urls.len() != self.task_count { - return exec_err!( - "number of tasks ({}) was not equal to number of urls ({}) at execution time", - self.task_count, - routed_urls.len() - ); - } - Ok(routed_urls) - } } fn keep_stream_alive(notify: Arc) -> impl Stream + 'static { futures::stream::once(notify.notified_owned()).filter_map(|()| futures::future::ready(None)) } -pub(super) struct NotifyGuard(Arc); - -impl Drop for NotifyGuard { - fn drop(&mut self) { - self.0.notify_waiters(); - } -} - /// Metrics that measure network details about communications between [DistributedExec] and a worker. #[derive(Clone)] pub(super) struct CoordinatorToWorkerMetrics { diff --git a/src/distributed_ext.rs b/src/distributed_ext.rs index 17852a04..e44818ad 100644 --- a/src/distributed_ext.rs +++ b/src/distributed_ext.rs @@ -2,13 +2,16 @@ 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::networking::{set_distributed_channel_resolver, set_distributed_worker_resolver}; +use crate::networking::{ + set_distributed_channel_resolver, set_distributed_worker_resolver, + set_distributed_worker_transport, +}; use crate::passthrough_headers::set_passthrough_headers; use crate::protobuf::{set_distributed_user_codec, set_distributed_user_codec_arc}; use crate::work_unit_feed::set_distributed_work_unit_feed; use crate::{ ChannelResolver, DistributedConfig, TaskEstimator, WorkUnitFeed, WorkUnitFeedProvider, - WorkerResolver, + WorkerResolver, WorkerTransport, }; use arrow_ipc::CompressionType; use datafusion::common::DataFusionError; @@ -224,6 +227,14 @@ pub trait DistributedExt: Sized { /// Same as [DistributedExt::with_distributed_channel_resolver] but with an in-place mutation. fn set_distributed_worker_resolver(&mut self, resolver: T); + /// Overrides the [WorkerTransport] used to open worker connections and dispatch plans. Defaults + /// to the Arrow-Flight gRPC transport. Embedders that run workers behind something other than + /// gRPC (e.g. shared-memory queues) plug their own transport in here. + fn with_distributed_worker_transport(self, transport: T) -> Self; + + /// Same as [DistributedExt::with_distributed_worker_transport] but with an in-place mutation. + fn set_distributed_worker_transport(&mut self, transport: T); + /// This is what tells Distributed DataFusion how to build a Worker gRPC client out of a worker URL. /// /// There's a default implementation that caches the Worker client instances so that there's @@ -611,6 +622,10 @@ impl DistributedExt for SessionConfig { set_distributed_channel_resolver(self, resolver); } + fn set_distributed_worker_transport(&mut self, transport: T) { + set_distributed_worker_transport(self, transport); + } + fn set_distributed_task_estimator( &mut self, estimator: T, @@ -748,6 +763,10 @@ impl DistributedExt for SessionConfig { #[expr($;self)] fn with_distributed_channel_resolver(mut self, resolver: T) -> Self; + #[call(set_distributed_worker_transport)] + #[expr($;self)] + fn with_distributed_worker_transport(mut self, transport: T) -> Self; + #[call(set_distributed_task_estimator)] #[expr($;self)] fn with_distributed_task_estimator(mut self, estimator: T) -> Self; @@ -841,6 +860,11 @@ impl DistributedExt for SessionStateBuilder { #[expr($;self)] fn with_distributed_channel_resolver(mut self, resolver: T) -> Self; + fn set_distributed_worker_transport(&mut self, transport: T); + #[call(set_distributed_worker_transport)] + #[expr($;self)] + fn with_distributed_worker_transport(mut self, transport: T) -> Self; + fn set_distributed_task_estimator(&mut self, estimator: T); #[call(set_distributed_task_estimator)] #[expr($;self)] @@ -952,6 +976,11 @@ impl DistributedExt for SessionState { #[expr($;self)] fn with_distributed_channel_resolver(mut self, resolver: T) -> Self; + fn set_distributed_worker_transport(&mut self, transport: T); + #[call(set_distributed_worker_transport)] + #[expr($;self)] + fn with_distributed_worker_transport(mut self, transport: T) -> Self; + fn set_distributed_task_estimator(&mut self, estimator: T); #[call(set_distributed_task_estimator)] #[expr($;self)] @@ -1063,6 +1092,11 @@ impl DistributedExt for SessionContext { #[expr($;self)] fn with_distributed_channel_resolver(self, resolver: T) -> Self; + fn set_distributed_worker_transport(&mut self, transport: T); + #[call(set_distributed_worker_transport)] + #[expr($;self)] + fn with_distributed_worker_transport(self, transport: T) -> Self; + fn set_distributed_task_estimator(&mut self, estimator: T); #[call(set_distributed_task_estimator)] #[expr($;self)] diff --git a/src/distributed_planner/distributed_config.rs b/src/distributed_planner/distributed_config.rs index d6faee39..6c1f43dd 100644 --- a/src/distributed_planner/distributed_config.rs +++ b/src/distributed_planner/distributed_config.rs @@ -1,5 +1,7 @@ use crate::distributed_planner::task_estimator::CombinedTaskEstimator; -use crate::networking::{ChannelResolverExtension, WorkerResolverExtension}; +use crate::networking::{ + ChannelResolverExtension, WorkerResolverExtension, WorkerTransportExtension, +}; use crate::work_unit_feed::WorkUnitFeedRegistry; use crate::{TaskEstimator, WorkerResolver}; use datafusion::common::{DataFusionError, extensions_options, not_impl_err, plan_err}; @@ -75,6 +77,9 @@ extensions_options! { /// [WorkerResolver] implementation that tells the distributed planner information about /// the available workers ready to execute distributed tasks. pub(crate) __private_worker_resolver: WorkerResolverExtension, default = WorkerResolverExtension::not_implemented() + /// [crate::WorkerTransport] used to open worker connections and dispatch plans. Defaults to + /// the Arrow-Flight gRPC transport; embedders can override it with their own transport. + pub(crate) __private_worker_transport: WorkerTransportExtension, default = WorkerTransportExtension::default() /// [WorkUnitFeedRegistry] that contains a set of getters that, applied to each node in a /// plan, will return the [crate::WorkUnitFeed]s present in all nodes. pub(crate) __private_work_unit_feed_registry: WorkUnitFeedRegistry, default = WorkUnitFeedRegistry::default() @@ -179,6 +184,22 @@ impl Debug for WorkerResolverExtension { } } +impl ConfigField for WorkerTransportExtension { + fn visit(&self, _: &mut V, _: &str, _: &'static str) { + // nothing to do. + } + + fn set(&mut self, _: &str, _: &str) -> datafusion::common::Result<()> { + not_impl_err!("Not implemented") + } +} + +impl Debug for WorkerTransportExtension { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "WorkerTransportExtension") + } +} + impl ConfigField for CombinedTaskEstimator { fn visit(&self, _: &mut V, _: &str, _: &'static str) { //nothing to do. diff --git a/src/distributed_planner/mod.rs b/src/distributed_planner/mod.rs index a1a7c9fb..a219c548 100644 --- a/src/distributed_planner/mod.rs +++ b/src/distributed_planner/mod.rs @@ -10,8 +10,8 @@ mod session_state_builder_ext; mod task_estimator; pub use distributed_config::DistributedConfig; -pub use network_boundary::{NetworkBoundary, NetworkBoundaryExt}; -pub(crate) use network_boundary::{ProducerHead, insert_producer_head}; +pub(crate) use network_boundary::insert_producer_head; +pub use network_boundary::{NetworkBoundary, NetworkBoundaryExt, PartitionRoute, ProducerHead}; pub use session_state_builder_ext::SessionStateBuilderExt; pub(crate) use task_estimator::set_distributed_task_estimator; pub use task_estimator::{TaskCountAnnotation, TaskEstimation, TaskEstimator, TaskRoutingContext}; diff --git a/src/distributed_planner/network_boundary.rs b/src/distributed_planner/network_boundary.rs index 858a7e00..cc08ba64 100644 --- a/src/distributed_planner/network_boundary.rs +++ b/src/distributed_planner/network_boundary.rs @@ -1,10 +1,18 @@ use crate::{BroadcastExec, NetworkBroadcastExec, NetworkCoalesceExec, NetworkShuffleExec, Stage}; -use datafusion::common::Result; +use datafusion::common::{Result, internal_err}; use datafusion::physical_expr::Partitioning; use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use std::sync::Arc; +/// Where a producer's output partition should be sent: which consumer task, and the local partition +/// index within that task's slice. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PartitionRoute { + pub consumer_task: usize, + pub consumer_partition: usize, +} + /// This trait represents a node that introduces the necessity of a network boundary in the plan. /// The distributed planner, upon stepping into one of these, will break the plan and build a stage /// out of it. @@ -22,6 +30,34 @@ pub trait NetworkBoundary: ExecutionPlan { /// implementation have. This information is used during planning an executing for ensuring /// the head of a stage has the appropriate shape for consumption. fn producer_head(&self, consumer_tasks: usize) -> ProducerHead; + + /// `P_c`: how many partitions each consumer task reads in the sliced layout + /// (`global = P_c * consumer_task + local`) that shuffle and broadcast reads use. Surfaced so a + /// transport that has to place a produced partition does not re-derive it from node properties. + /// Meaningless for [NetworkCoalesceExec], whose consumers read whole per-producer-task groups. + fn partitions_per_consumer_task(&self) -> usize { + self.properties().partitioning.partition_count() + } + + /// Maps a producer output partition to the consumer task and the local partition within that + /// task that reads it, for the `global = P_c * consumer_task + local` layout. + /// + /// Boundaries whose consumers do not read that layout must override this with an error; the + /// default would silently misroute them. A zero-partition boundary is a planner bug, so it + /// errors instead of routing everything to task `0`. + fn route_partition(&self, output_partition: usize) -> Result { + let p_c = self.partitions_per_consumer_task(); + if p_c == 0 { + return internal_err!( + "cannot route output partition {output_partition}: the boundary reports 0 \ + partitions per consumer task" + ); + } + Ok(PartitionRoute { + consumer_task: output_partition / p_c, + consumer_partition: output_partition % p_c, + }) + } } /// Defines what shape should the head node of a stage have upon getting executed. Depending diff --git a/src/execution_plans/network_coalesce.rs b/src/execution_plans/network_coalesce.rs index 8fb06e74..fcd8202b 100644 --- a/src/execution_plans/network_coalesce.rs +++ b/src/execution_plans/network_coalesce.rs @@ -180,6 +180,16 @@ impl NetworkBoundary for NetworkCoalesceExec { fn producer_head(&self, _consumer_task_count: usize) -> ProducerHead { ProducerHead::None } + + fn route_partition( + &self, + output_partition: usize, + ) -> Result { + internal_err!( + "NetworkCoalesceExec routes by producer task group, not by output partition; \ + partition {output_partition} has no slice-layout route" + ) + } } impl DisplayAs for NetworkCoalesceExec { diff --git a/src/lib.rs b/src/lib.rs index d2a3a463..4928b8c8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -20,11 +20,15 @@ pub mod test_utils; mod work_unit_feed; pub use arrow_ipc::CompressionType; -pub use coordinator::DistributedExec; +pub use common::{ + TreeNodeExt, deserialize_uuid, get_distributed_cancellation_token, serialize_uuid, +}; +pub use config_extension_ext::get_config_extension_propagation_headers; +pub use coordinator::{DistributedExec, LatencyMetric, MetricsStore}; pub use distributed_ext::DistributedExt; pub use distributed_planner::{ - DistributedConfig, NetworkBoundary, NetworkBoundaryExt, SessionStateBuilderExt, - TaskCountAnnotation, TaskEstimation, TaskEstimator, TaskRoutingContext, + DistributedConfig, NetworkBoundary, NetworkBoundaryExt, PartitionRoute, ProducerHead, + SessionStateBuilderExt, TaskCountAnnotation, TaskEstimation, TaskEstimator, TaskRoutingContext, }; pub use execution_plans::{ BroadcastExec, DistributedLeafExec, NetworkBroadcastExec, NetworkCoalesceExec, @@ -39,9 +43,12 @@ pub use metrics::{ pub use networking::{ BoxCloneSyncChannel, ChannelResolver, DefaultChannelResolver, WorkerResolver, create_worker_client, get_distributed_channel_resolver, get_distributed_worker_resolver, + get_distributed_worker_transport, set_distributed_worker_transport, }; +pub use passthrough_headers::get_passthrough_headers; pub use stage::{ - DistributedTaskContext, Stage, display_plan_ascii, display_plan_graphviz, explain_analyze, + DistributedTaskContext, RemoteStage, Stage, display_plan_ascii, display_plan_graphviz, + explain_analyze, }; pub use work_unit_feed::{ DistributedWorkUnitFeedContext, WorkUnit, WorkUnitFeed, WorkUnitFeedProto, WorkUnitFeedProvider, @@ -50,8 +57,9 @@ pub use worker::generated::worker::worker_service_client::WorkerServiceClient; pub use worker::generated::worker::worker_service_server::WorkerServiceServer; pub use worker::generated::worker::{GetWorkerInfoRequest, GetWorkerInfoResponse, TaskKey}; pub use worker::{ - DefaultSessionBuilder, MappedWorkerSessionBuilder, MappedWorkerSessionBuilderExt, TaskData, - Worker, WorkerQueryContext, WorkerSessionBuilder, + DefaultSessionBuilder, FlightWorkerTransport, MappedWorkerSessionBuilder, + MappedWorkerSessionBuilderExt, TaskData, Worker, WorkerConnection, WorkerDispatch, + WorkerDispatchRequest, WorkerQueryContext, WorkerSessionBuilder, WorkerTransport, }; pub use observability::{ diff --git a/src/networking/mod.rs b/src/networking/mod.rs index 6bab9ae0..0c2f56ce 100644 --- a/src/networking/mod.rs +++ b/src/networking/mod.rs @@ -1,5 +1,6 @@ mod channel_resolver; mod worker_resolver; +mod worker_transport; pub use channel_resolver::{ BoxCloneSyncChannel, ChannelResolver, DefaultChannelResolver, create_worker_client, @@ -9,3 +10,6 @@ pub(crate) use channel_resolver::{ChannelResolverExtension, set_distributed_chan pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver}; pub(crate) use worker_resolver::{WorkerResolverExtension, set_distributed_worker_resolver}; + +pub(crate) use worker_transport::WorkerTransportExtension; +pub use worker_transport::{get_distributed_worker_transport, set_distributed_worker_transport}; diff --git a/src/networking/worker_transport.rs b/src/networking/worker_transport.rs new file mode 100644 index 00000000..f38550aa --- /dev/null +++ b/src/networking/worker_transport.rs @@ -0,0 +1,45 @@ +use crate::DistributedConfig; +use crate::config_extension_ext::set_distributed_option_extension; +use crate::worker::{FlightWorkerTransport, WorkerTransport}; +use datafusion::prelude::SessionConfig; +use std::sync::Arc; + +pub fn set_distributed_worker_transport( + cfg: &mut SessionConfig, + worker_transport: impl WorkerTransport + 'static, +) { + let opts = cfg.options_mut(); + let worker_transport = WorkerTransportExtension(Arc::new(worker_transport)); + if let Some(distributed_cfg) = opts.extensions.get_mut::() { + distributed_cfg.__private_worker_transport = worker_transport; + } else { + set_distributed_option_extension( + cfg, + DistributedConfig { + __private_worker_transport: worker_transport, + ..Default::default() + }, + ) + } +} + +/// Returns the [WorkerTransport] in scope, defaulting to the Arrow-Flight gRPC transport. Network +/// boundaries call this at execute time to open connections and dispatch plans, so a custom +/// transport set via [crate::DistributedExt::with_distributed_worker_transport] takes over both the +/// read and write sides. +pub fn get_distributed_worker_transport(cfg: &SessionConfig) -> Arc { + cfg.options() + .extensions + .get::() + .map(|distributed_cfg| Arc::clone(&distributed_cfg.__private_worker_transport.0)) + .unwrap_or_else(|| Arc::new(FlightWorkerTransport)) +} + +#[derive(Clone)] +pub(crate) struct WorkerTransportExtension(pub(crate) Arc); + +impl Default for WorkerTransportExtension { + fn default() -> Self { + Self(Arc::new(FlightWorkerTransport)) + } +} diff --git a/src/passthrough_headers.rs b/src/passthrough_headers.rs index d269e685..6ac085c3 100644 --- a/src/passthrough_headers.rs +++ b/src/passthrough_headers.rs @@ -32,7 +32,7 @@ pub(crate) fn set_passthrough_headers( /// Gets passthrough headers from a SessionConfig extension. /// Returns an empty HeaderMap if none are set. -pub(crate) fn get_passthrough_headers(cfg: &SessionConfig) -> HeaderMap { +pub fn get_passthrough_headers(cfg: &SessionConfig) -> HeaderMap { cfg.get_extension::() .map(|h| h.0.clone()) .unwrap_or_default() diff --git a/src/work_unit_feed/remote_work_unit_feed.rs b/src/work_unit_feed/remote_work_unit_feed.rs index f914f228..b724de73 100644 --- a/src/work_unit_feed/remote_work_unit_feed.rs +++ b/src/work_unit_feed/remote_work_unit_feed.rs @@ -13,10 +13,10 @@ use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; use tokio_stream::wrappers::UnboundedReceiverStream; use uuid::Uuid; -pub(crate) type WorkUnitTx = UnboundedSender>; -pub(crate) type WorkUnitRx = UnboundedReceiver>; -pub(crate) type RemoteWorkUnitFeedRxs = HashMap<(Uuid, usize), Mutex>>; -pub(crate) type RemoteWorkUnitFeedTxs = HashMap<(Uuid, usize), WorkUnitTx>; +pub type WorkUnitTx = UnboundedSender>; +pub type WorkUnitRx = UnboundedReceiver>; +pub type RemoteWorkUnitFeedRxs = HashMap<(Uuid, usize), Mutex>>; +pub type RemoteWorkUnitFeedTxs = HashMap<(Uuid, usize), WorkUnitTx>; /// Bridge between the worker's gRPC layer and the remote-variant /// [`crate::WorkUnitFeed`]s installed in the deserialized plan. @@ -30,15 +30,15 @@ pub(crate) type RemoteWorkUnitFeedTxs = HashMap<(Uuid, usize), WorkUnitTx>; /// concrete `T::WorkUnit` type so the leaf sees the same typed stream as it would in a /// single-node execution. #[derive(Default)] -pub(crate) struct RemoteWorkUnitFeedRegistry { - pub(crate) receivers: RemoteWorkUnitFeedRxs, - pub(crate) senders: RemoteWorkUnitFeedTxs, +pub struct RemoteWorkUnitFeedRegistry { + pub receivers: RemoteWorkUnitFeedRxs, + pub senders: RemoteWorkUnitFeedTxs, } impl RemoteWorkUnitFeedRegistry { /// Creates all the receivers and senders for a specific [WorkUnit] Feed id. One feed per /// partition is created. - pub(crate) fn add(&mut self, id: Uuid, partitions: usize) { + pub fn add(&mut self, id: Uuid, partitions: usize) { for partition in 0..partitions { let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); self.receivers.insert((id, partition), Mutex::new(Some(rx))); diff --git a/src/worker/flight.rs b/src/worker/flight.rs new file mode 100644 index 00000000..7bb8a792 --- /dev/null +++ b/src/worker/flight.rs @@ -0,0 +1,62 @@ +use crate::coordinator::FlightWorkerDispatch; +use crate::distributed_planner::ProducerHead; +use crate::stage::RemoteStage; +use crate::worker::transport::{WorkerConnection, WorkerDispatch, WorkerTransport}; +use crate::worker::worker_connection_pool::{ + LocalWorkerConnection, LocalWorkerContext, RemoteWorkerConnection, +}; +use datafusion::common::{Result, internal_datafusion_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr_common::metrics::ExecutionPlanMetricsSet; +use std::ops::Range; +use std::sync::Arc; + +/// The Arrow-Flight gRPC transport, used by default. The read side opens one bidirectional gRPC +/// stream per worker task and demultiplexes it into per-partition streams; the write side ships +/// plans over the coordinator-to-worker channel via [FlightWorkerDispatch]. +#[derive(Default)] +pub struct FlightWorkerTransport; + +impl WorkerTransport for FlightWorkerTransport { + fn open( + &self, + input_stage: &RemoteStage, + target_partitions: Range, + target_task: usize, + producer_head: ProducerHead, + ctx: &Arc, + metrics: &ExecutionPlanMetricsSet, + ) -> Result> { + let Some(target_url) = input_stage.workers.get(target_task) else { + return Err(internal_datafusion_err!( + "input_stage.workers[{target_task}] out of range." + )); + }; + if let Some(lw_ctx) = ctx.session_config().get_extension::() + && &lw_ctx.self_url == target_url + { + // Reaching ourselves: pull from the local task registry instead of a gRPC round-trip. + Ok(Box::new(LocalWorkerConnection::init( + input_stage, + target_partitions, + target_task, + producer_head, + ctx, + metrics, + )?)) + } else { + Ok(Box::new(RemoteWorkerConnection::init( + input_stage, + target_partitions, + target_task, + producer_head, + ctx, + metrics, + )?)) + } + } + + fn dispatcher(&self) -> Box { + Box::new(FlightWorkerDispatch::default()) + } +} diff --git a/src/worker/mod.rs b/src/worker/mod.rs index e89921fc..f7ac23a2 100644 --- a/src/worker/mod.rs +++ b/src/worker/mod.rs @@ -1,3 +1,4 @@ +mod flight; pub(crate) mod generated; mod impl_coordinator_channel; mod impl_execute_task; @@ -7,12 +8,16 @@ mod spawn_select_all; mod task_data; #[cfg(any(test, feature = "integration"))] pub(crate) mod test_utils; +mod transport; mod worker_connection_pool; mod worker_service; pub(crate) use single_write_multi_read::SingleWriteMultiRead; pub(crate) use worker_connection_pool::{LocalWorkerContext, WorkerConnectionPool}; +pub use flight::FlightWorkerTransport; +pub use transport::{WorkerConnection, WorkerDispatch, WorkerDispatchRequest, WorkerTransport}; + pub use session_builder::{ DefaultSessionBuilder, MappedWorkerSessionBuilder, MappedWorkerSessionBuilderExt, WorkerQueryContext, WorkerSessionBuilder, diff --git a/src/worker/transport.rs b/src/worker/transport.rs new file mode 100644 index 00000000..db76d229 --- /dev/null +++ b/src/worker/transport.rs @@ -0,0 +1,108 @@ +use crate::coordinator::MetricsStore; +use crate::distributed_planner::ProducerHead; +use crate::stage::{LocalStage, RemoteStage}; +use datafusion::arrow::array::RecordBatch; +use datafusion::common::Result; +use datafusion::common::runtime::JoinSet; +use datafusion::execution::TaskContext; +use datafusion::physical_expr_common::metrics::ExecutionPlanMetricsSet; +use futures::stream::BoxStream; +use std::ops::Range; +use std::sync::Arc; +use url::Url; + +/// A live connection to a single worker that demultiplexes the underlying transport into one +/// stream per partition. +/// +/// One connection handles every partition in the `target_partitions` range requested at open +/// time, so the implementation can reuse a single underlying network/IPC stream and fan messages +/// out to per-partition queues. Each partition can be streamed exactly once. +pub trait WorkerConnection: Send + Sync { + /// Streams the given output `partition`. The connection is opened per stage, so it closes over + /// the stage rather than taking it per call. Streaming the same partition twice is an error. + fn execute(&self, partition: usize) -> Result>>; +} + +/// Everything a [WorkerDispatch] needs to deliver one stage's plans to its workers. +/// +/// The coordinator computes the per-task worker assignment (`routed_urls[i]` is the worker for +/// task `i`) and hands the transport this request; the transport delivers each task's plan and +/// wires up whatever per-task back-channels it needs. `join_set` is the query's, so background +/// delivery work spawned onto it propagates failures to the query head. +#[non_exhaustive] +pub struct WorkerDispatchRequest<'a> { + pub stage: &'a LocalStage, + pub routed_urls: &'a [Url], + pub task_ctx: &'a Arc, + pub metrics: &'a ExecutionPlanMetricsSet, + /// Back-channel for task metrics. Only the Flight transport reads it (workers push their + /// metrics back over its gRPC stream); other transports ignore it. + pub metrics_store: Option<&'a Arc>, + pub join_set: &'a mut JoinSet>, +} + +/// The plan-delivery (write) side of a transport, symmetric to [WorkerConnection] (the read side). +/// A dispatcher is a per-query object: [WorkerTransport::dispatcher] creates it before the first +/// stage is dispatched, and every stage of that query goes through the same instance. +/// +/// Flight resolves each worker's URL, encodes the plan, and ships a `SetPlanRequest` over a +/// bidirectional gRPC stream that also carries the work-unit feed and the metrics back-channel. A +/// co-located transport registers the plan in a local table. The coordinator just calls +/// `dispatch`; delivery is no longer a fixed gRPC step it special-cases. +pub trait WorkerDispatch: Send + Sync { + fn dispatch(&self, request: WorkerDispatchRequest<'_>) -> Result<()>; +} + +/// Factory that opens a [WorkerConnection] to a single worker task and delivers plans to workers. +/// +/// The default implementation is the Arrow-Flight gRPC transport baked into this crate +/// (`FlightWorkerTransport`). Custom transports (e.g. shared-memory queues for an embedded +/// execution context) plug in via [crate::DistributedExt::with_distributed_worker_transport]. +pub trait WorkerTransport: Send + Sync { + /// Opens a connection to the worker hosting `target_task` of `input_stage` covering the + /// partitions in `target_partitions`. The returned [WorkerConnection] takes ownership of any + /// background resources (gRPC streams, demux tasks, ...) and cleans them up on drop. Bypassing + /// the network for a worker co-located with the coordinator is the implementation's concern + /// (Flight compares the target URL against its own). + fn open( + &self, + input_stage: &RemoteStage, + target_partitions: Range, + target_task: usize, + producer_head: ProducerHead, + ctx: &Arc, + metrics: &ExecutionPlanMetricsSet, + ) -> Result>; + + /// Creates the plan-delivery side of this transport for one query. + /// + /// Called once per query, before any stage is dispatched, so the returned [WorkerDispatch] can + /// hold per-query state. Flight uses this to share one set of plan-send metrics across every + /// stage's dispatch; a fresh dispatcher per stage would register duplicate metrics. + fn dispatcher(&self) -> Box; +} + +impl WorkerTransport for Arc { + fn open( + &self, + input_stage: &RemoteStage, + target_partitions: Range, + target_task: usize, + producer_head: ProducerHead, + ctx: &Arc, + metrics: &ExecutionPlanMetricsSet, + ) -> Result> { + self.as_ref().open( + input_stage, + target_partitions, + target_task, + producer_head, + ctx, + metrics, + ) + } + + fn dispatcher(&self) -> Box { + self.as_ref().dispatcher() + } +} diff --git a/src/worker/worker_connection_pool.rs b/src/worker/worker_connection_pool.rs index 7acd629c..78313844 100644 --- a/src/worker/worker_connection_pool.rs +++ b/src/worker/worker_connection_pool.rs @@ -1,13 +1,14 @@ use crate::common::{OnceLockResult, on_drop_stream, serialize_uuid}; use crate::distributed_planner::ProducerHead; use crate::metrics::LatencyMetricExt; -use crate::networking::get_distributed_channel_resolver; +use crate::networking::{get_distributed_channel_resolver, get_distributed_worker_transport}; use crate::passthrough_headers::get_passthrough_headers; use crate::protobuf::{datafusion_error_to_tonic_status, map_flight_to_datafusion_error}; use crate::stage::RemoteStage; use crate::worker::generated::worker::FlightAppMetadata; use crate::worker::generated::worker::{ExecuteTaskRequest, TaskKey}; use crate::worker::impl_execute_task::execute_local_task; +use crate::worker::transport::WorkerConnection; use crate::worker::worker_service::TaskDataEntries; use crate::{BytesMetricExt, ChannelResolver, DistributedConfig}; use arrow_flight::FlightData; @@ -65,7 +66,7 @@ pub(crate) struct LocalWorkerContext { /// it will initialize the corresponding position in the vector matching the provided `target_task` /// index. pub(crate) struct WorkerConnectionPool { - connections: Vec>>, + connections: Vec>>, pub(crate) metrics: ExecutionPlanMetricsSet, } @@ -93,7 +94,7 @@ impl WorkerConnectionPool { target_task: usize, producer_head: ProducerHead, ctx: &Arc, - ) -> Result<&(dyn WorkerConnection + Sync + Send)> { + ) -> Result<&dyn WorkerConnection> { let Some(worker_connection) = self.connections.get(target_task) else { return internal_err!( "WorkerConnections: Task index {target_task} not found, only have {} tasks", @@ -102,26 +103,10 @@ impl WorkerConnectionPool { }; let conn = worker_connection.get_or_init(|| { - let Some(target_url) = input_stage.workers.get(target_task) else { - internal_err!("input_stage.workers[{target_task}] out of range.")? - }; - if let Some(lw_ctx) = ctx.session_config().get_extension::() - && &lw_ctx.self_url == target_url - { - // Instead of making a gRPC call to ourselves, better to just use local comms. - LocalWorkerConnection::init( - input_stage, - target_partitions, - target_task, - producer_head, - ctx, - &self.metrics, - ) - .map(|v| Box::new(v) as Box<_>) - .map_err(Arc::new) - } else { - // We are trying to reach a URL different from ours, so use normal gRPC streams. - RemoteWorkerConnection::init( + // Local-vs-remote is the transport's concern now; the pool only caches one connection + // per worker task and hands the routing decision to the configured transport. + get_distributed_worker_transport(ctx.session_config()) + .open( input_stage, target_partitions, target_task, @@ -129,9 +114,7 @@ impl WorkerConnectionPool { ctx, &self.metrics, ) - .map(|v| Box::new(v) as Box<_>) .map_err(Arc::new) - } }); match conn { @@ -143,14 +126,6 @@ impl WorkerConnectionPool { type WorkerMsg = Result<(FlightData, FlightAppMetadata), Status>; -/// Abstraction that allows treating remote and local comms as equal. Network boundaries do not -/// care if the stream comes over the wire or locally. -pub(crate) trait WorkerConnection { - /// Streams the specified partition. Consumers do not care if the implementation pulls data - /// from in-memory or from local comms. - fn execute(&self, partition: usize) -> Result>>; -} - /// Represents a connection to one [Worker]. Network boundaries will use this for streaming /// data from single partitions while the actual network communication is handling all the partitions /// under the hood. @@ -162,7 +137,7 @@ pub(crate) trait WorkerConnection { /// the same underlying TCP connection, there do is some overhead in having one gRPC stream per /// partition VS a single gRPC stream interleaving multiple partitions. The whole serialized plan /// needs to be sent over the wire on every gRPC call, so the less gRPC calls we do the better. -struct RemoteWorkerConnection { +pub(crate) struct RemoteWorkerConnection { task: Arc>, not_consumed_streams: Arc, cancel_token: CancellationToken, @@ -178,7 +153,7 @@ struct RemoteWorkerConnection { } impl RemoteWorkerConnection { - fn init( + pub(crate) fn init( input_stage: &RemoteStage, target_partition_range: Range, target_task: usize, @@ -462,7 +437,7 @@ pub(crate) struct LocalWorkerConnection { } impl LocalWorkerConnection { - fn init( + pub(crate) fn init( input_stage: &RemoteStage, target_partition_range: Range, target_task: usize, From e9b3070c4cbd7abf707c0709f30e3646a529319d Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sat, 20 Jun 2026 03:06:12 -0700 Subject: [PATCH 2/4] Added an in-process `InMemoryWorkerTransport`. This PR adds a `WorkerTransport` that hosts its workers in the current process: plans are delivered with a direct `Worker::set_task_plan` call and partitions are read straight from the local task registry, with no gRPC underneath. It is the reference implementation for the transport extension points and the basis for running distributed plans without the Flight stack. To keep the dispatch paths from drifting, plan delivery is factored into transport-neutral pieces both transports share: `encode_task_plan` (task specialization + codec), the `Worker::set_task_plan` core that the Flight coordinator stream now wraps, and `collect_task_work_unit_feeds`. The in-memory read side runs one `execute_local_task` over the whole partition range and pumps each partition into a buffer, so a consumer that interleaves partition polls of a partitioned join can't leave partitions empty. --- src/coordinator/mod.rs | 4 +- src/coordinator/plan_encoding.rs | 78 ++++ src/coordinator/query_coordinator.rs | 76 +--- src/lib.rs | 15 +- src/work_unit_feed/coordinator.rs | 61 +++ src/work_unit_feed/mod.rs | 9 +- src/work_unit_feed/remote_work_unit_feed.rs | 41 +- src/worker/impl_coordinator_channel.rs | 109 +----- src/worker/impl_set_plan.rs | 127 +++++++ src/worker/in_memory.rs | 402 ++++++++++++++++++++ src/worker/mod.rs | 3 + 11 files changed, 753 insertions(+), 172 deletions(-) create mode 100644 src/coordinator/plan_encoding.rs create mode 100644 src/work_unit_feed/coordinator.rs create mode 100644 src/worker/impl_set_plan.rs create mode 100644 src/worker/in_memory.rs diff --git a/src/coordinator/mod.rs b/src/coordinator/mod.rs index 5302aa6c..3d46f72a 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -1,10 +1,12 @@ mod distributed; mod latency_metric; mod metrics_store; +mod plan_encoding; mod prepare_static_plan; mod query_coordinator; pub use distributed::DistributedExec; pub use latency_metric::LatencyMetric; pub use metrics_store::MetricsStore; -pub(crate) use query_coordinator::FlightWorkerDispatch; +pub use plan_encoding::{EncodedTaskPlan, encode_task_plan}; +pub(crate) use query_coordinator::{CoordinatorToWorkerMetrics, FlightWorkerDispatch}; diff --git a/src/coordinator/plan_encoding.rs b/src/coordinator/plan_encoding.rs new file mode 100644 index 00000000..d6baca1c --- /dev/null +++ b/src/coordinator/plan_encoding.rs @@ -0,0 +1,78 @@ +use crate::common::{TreeNodeExt, serialize_uuid}; +use crate::execution_plans::{ChildrenIsolatorUnionExec, DistributedLeafExec}; +use crate::worker::generated::worker::set_plan_request::WorkUnitFeedDeclaration; +use crate::{DistributedCodec, DistributedConfig, DistributedTaskContext}; +use datafusion::common::Result; +use datafusion::common::tree_node::Transformed; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::SessionConfig; +use datafusion_proto::physical_plan::AsExecutionPlan; +use datafusion_proto::protobuf::PhysicalPlanNode; +use prost::Message; +use std::sync::Arc; + +/// A stage plan specialized for one task and encoded for delivery, together with the work-unit +/// feeds it declares. +pub struct EncodedTaskPlan { + pub plan_proto: Vec, + pub feed_declarations: Vec, + /// Output partitions of the specialized plan. Task-isolated nodes (a union running one child + /// per task) make this differ from the unspecialized stage plan, so a push-based transport + /// sizes its per-task sinks from this. The pull transports ignore it; a push transport (such + /// as the shared-memory one) reads it. + pub partitions: usize, +} + +/// Specializes `plan` for `task_index` and encodes it with the session's combined codec. +/// +/// Shared by every transport's dispatch path so task specialization and codec selection cannot +/// drift between transports: each task must see its own slice of task-isolated nodes (a union +/// executing one child per task, an unexecuted [DistributedLeafExec] variant), and the worker +/// decodes with the same codec stack the coordinator encoded with. +pub fn encode_task_plan( + plan: &Arc, + task_index: usize, + task_count: usize, + cfg: &SessionConfig, +) -> Result { + let d_cfg = DistributedConfig::from_config_options(cfg.options())?; + let wuf_registry = &d_cfg.__private_work_unit_feed_registry; + + let mut feed_declarations = vec![]; + let d_ctx = DistributedTaskContext { + task_index, + task_count, + }; + + let specialized = Arc::clone(plan).transform_down_with_dt_ctx(d_ctx, |plan, d_ctx| { + if let Some(wuf) = wuf_registry.get_work_unit_feed(&plan) { + feed_declarations.push(WorkUnitFeedDeclaration { + id: serialize_uuid(&wuf.id()), + partitions: plan.properties().partitioning.partition_count() as u64, + }); + }; + + if let Some(ciu) = plan.downcast_ref::() { + let ciu = ciu.to_task_specialized(d_ctx.task_index); + return Ok(Transformed::yes(Arc::new(ciu))); + }; + + if let Some(dle) = plan.downcast_ref::() { + let specialized = dle.to_task_specialized(d_ctx.task_index); + return Ok(Transformed::yes(specialized)); + } + + Ok(Transformed::no(plan)) + })?; + + let codec = DistributedCodec::new_combined_with_user(cfg); + let partitions = specialized.data.properties().partitioning.partition_count(); + let plan_proto = + PhysicalPlanNode::try_from_physical_plan(specialized.data, &codec)?.encode_to_vec(); + + Ok(EncodedTaskPlan { + plan_proto, + feed_declarations, + partitions, + }) +} diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 3ef3d030..2e87f77c 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -2,32 +2,28 @@ use crate::common::{TreeNodeExt, now_ns, serialize_uuid, 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::execution_plans::{ChildrenIsolatorUnionExec, DistributedLeafExec}; +use crate::coordinator::plan_encoding::encode_task_plan; use crate::passthrough_headers::get_passthrough_headers; use crate::protobuf::tonic_status_to_datafusion_error; use crate::work_unit_feed::{build_work_unit_batch_msg, set_work_unit_send_time}; use crate::worker::generated::worker as pb; use crate::worker::generated::worker::coordinator_to_worker_msg::Inner; -use crate::worker::generated::worker::set_plan_request::WorkUnitFeedDeclaration; use crate::worker::{WorkerDispatch, WorkerDispatchRequest}; use crate::{ - BytesCounterMetric, BytesMetricExt, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedCodec, - DistributedConfig, DistributedTaskContext, DistributedWorkUnitFeedContext, TaskKey, + BytesCounterMetric, BytesMetricExt, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedConfig, + DistributedTaskContext, DistributedWorkUnitFeedContext, TaskKey, get_distributed_channel_resolver, }; use datafusion::common::Result; use datafusion::common::instant::Instant; use datafusion::common::runtime::JoinSet; -use datafusion::common::tree_node::{Transformed, TreeNodeRecursion}; +use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{DataFusionError, exec_datafusion_err}; use datafusion::execution::TaskContext; use datafusion::physical_expr_common::metrics::{ExecutionPlanMetricsSet, Label, MetricBuilder}; use datafusion::physical_plan::ExecutionPlan; -use datafusion_proto::physical_plan::AsExecutionPlan; -use datafusion_proto::protobuf::PhysicalPlanNode; use futures::{Stream, StreamExt}; use http::Extensions; -use prost::Message; use std::sync::{Arc, OnceLock}; use tokio::sync::Notify; use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; @@ -128,13 +124,9 @@ impl<'a> StageCoordinator<'a> { UnboundedReceiver, )> { let session_config = self.task_ctx.session_config(); - let codec = DistributedCodec::new_combined_with_user(session_config); - let (specialized, work_unit_feed_declarations) = self.task_specialized_plan(task_i)?; - - let plan_proto = - PhysicalPlanNode::try_from_physical_plan(specialized, &codec)?.encode_to_vec(); - let plan_size = plan_proto.len(); + let encoded = encode_task_plan(self.plan, task_i, self.task_count, session_config)?; + let plan_size = encoded.plan_proto.len(); let task_key = TaskKey { query_id: serialize_uuid(&self.query_id), @@ -143,10 +135,10 @@ impl<'a> StageCoordinator<'a> { }; let msg = pb::CoordinatorToWorkerMsg { inner: Some(Inner::SetPlanRequest(pb::SetPlanRequest { - plan_proto, + plan_proto: encoded.plan_proto, task_count: self.task_count as u64, task_key: Some(task_key.clone()), - work_unit_feed_declarations, + work_unit_feed_declarations: encoded.feed_declarations, target_worker_url: url.to_string(), query_start_time_ns: self.metrics.instantiation_time, })), @@ -309,48 +301,6 @@ impl<'a> StageCoordinator<'a> { }); Ok(()) } - - /// Specializes the [Arc] for this stage to provided task index. This implies - /// trimming down any unnecessary information that the specific `task_i` task is not going to - /// need, like unexecuted branches in [ChildrenIsolatorUnionExec], or unexecuted variants of - /// [DistributedLeafExec]. - fn task_specialized_plan( - &self, - task_i: usize, - ) -> Result<(Arc, Vec)> { - let session_config = self.task_ctx.session_config(); - let d_cfg = DistributedConfig::from_config_options(session_config.options())?; - let wuf_registry = &d_cfg.__private_work_unit_feed_registry; - - let mut work_unit_feed_declarations = vec![]; - let d_ctx = DistributedTaskContext { - task_index: task_i, - task_count: self.task_count, - }; - - let plan = Arc::clone(self.plan); - let transformed = plan.transform_down_with_dt_ctx(d_ctx, |plan, d_ctx| { - if let Some(wuf) = wuf_registry.get_work_unit_feed(&plan) { - work_unit_feed_declarations.push(WorkUnitFeedDeclaration { - id: serialize_uuid(&wuf.id()), - partitions: plan.properties().partitioning.partition_count() as u64, - }); - }; - - if let Some(ciu) = plan.downcast_ref::() { - let ciu = ciu.to_task_specialized(d_ctx.task_index); - return Ok(Transformed::yes(Arc::new(ciu))); - }; - - if let Some(dle) = plan.downcast_ref::() { - let specialized = dle.to_task_specialized(d_ctx.task_index); - return Ok(Transformed::yes(specialized)); - } - - Ok(Transformed::no(plan)) - })?; - Ok((transformed.data, work_unit_feed_declarations)) - } } fn keep_stream_alive(notify: Arc) -> impl Stream + 'static { @@ -359,14 +309,14 @@ fn keep_stream_alive(notify: Arc) -> impl Stream + /// Metrics that measure network details about communications between [DistributedExec] and a worker. #[derive(Clone)] -pub(super) struct CoordinatorToWorkerMetrics { - pub(super) plan_bytes_sent: BytesCounterMetric, - pub(super) plan_send_latency: Arc, - pub(super) instantiation_time: u64, +pub(crate) struct CoordinatorToWorkerMetrics { + pub(crate) plan_bytes_sent: BytesCounterMetric, + pub(crate) plan_send_latency: Arc, + pub(crate) instantiation_time: u64, } impl CoordinatorToWorkerMetrics { - pub(super) fn new(metrics: &ExecutionPlanMetricsSet) -> Self { + pub(crate) fn new(metrics: &ExecutionPlanMetricsSet) -> Self { Self { // Metric that measures to total sum of bytes worth of subplans sent. plan_bytes_sent: MetricBuilder::new(metrics) diff --git a/src/lib.rs b/src/lib.rs index 4928b8c8..5c5540e2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -24,7 +24,9 @@ pub use common::{ TreeNodeExt, deserialize_uuid, get_distributed_cancellation_token, serialize_uuid, }; pub use config_extension_ext::get_config_extension_propagation_headers; -pub use coordinator::{DistributedExec, LatencyMetric, MetricsStore}; +pub use coordinator::{ + DistributedExec, EncodedTaskPlan, LatencyMetric, MetricsStore, encode_task_plan, +}; pub use distributed_ext::DistributedExt; pub use distributed_planner::{ DistributedConfig, NetworkBoundary, NetworkBoundaryExt, PartitionRoute, ProducerHead, @@ -51,15 +53,18 @@ pub use stage::{ explain_analyze, }; pub use work_unit_feed::{ - DistributedWorkUnitFeedContext, WorkUnit, WorkUnitFeed, WorkUnitFeedProto, WorkUnitFeedProvider, + DistributedWorkUnitFeedContext, RemoteWorkUnitFeedRegistry, RemoteWorkUnitFeedRxs, + RemoteWorkUnitFeedTxs, WorkUnit, WorkUnitFeed, WorkUnitFeedProto, WorkUnitFeedProvider, + WorkUnitRx, WorkUnitTx, collect_task_work_unit_feeds, set_received_time, set_sent_time, }; pub use worker::generated::worker::worker_service_client::WorkerServiceClient; pub use worker::generated::worker::worker_service_server::WorkerServiceServer; pub use worker::generated::worker::{GetWorkerInfoRequest, GetWorkerInfoResponse, TaskKey}; pub use worker::{ - DefaultSessionBuilder, FlightWorkerTransport, MappedWorkerSessionBuilder, - MappedWorkerSessionBuilderExt, TaskData, Worker, WorkerConnection, WorkerDispatch, - WorkerDispatchRequest, WorkerQueryContext, WorkerSessionBuilder, WorkerTransport, + DefaultSessionBuilder, FlightWorkerTransport, InMemoryWorkerTransport, + MappedWorkerSessionBuilder, MappedWorkerSessionBuilderExt, TaskData, Worker, WorkerConnection, + WorkerDispatch, WorkerDispatchRequest, WorkerQueryContext, WorkerSessionBuilder, + WorkerTransport, }; pub use observability::{ diff --git a/src/work_unit_feed/coordinator.rs b/src/work_unit_feed/coordinator.rs new file mode 100644 index 00000000..e897d10c --- /dev/null +++ b/src/work_unit_feed/coordinator.rs @@ -0,0 +1,61 @@ +use crate::common::{TreeNodeExt, task_ctx_with_extension}; +use crate::work_unit_feed::remote_work_unit_feed::build_work_unit; +use crate::worker::generated::worker as pb; +use crate::{DistributedConfig, DistributedTaskContext, DistributedWorkUnitFeedContext}; +use datafusion::common::Result; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::ExecutionPlan; +use futures::StreamExt; +use futures::stream::BoxStream; +use std::sync::Arc; + +/// Drives every registered [crate::WorkUnitFeed] in `plan` for one task and returns its +/// per-partition work-unit streams, already encoded. +/// +/// Transport-neutral: the Flight dispatch wraps each [pb::WorkUnit] in its envelope and pushes it +/// over the coordinator-to-worker gRPC stream; an in-process transport hands them straight to the +/// worker-side channels. Each unit carries its `(feed id, task-local partition)`; two tasks emit +/// identical pairs, so routing to the right task stays the transport's job, as with Flight's +/// per-task stream. A task owns a non-overlapping window of partition feeds, offset by its task +/// index, and each partition feed can be consumed once. +pub fn collect_task_work_unit_feeds( + plan: &Arc, + ctx: &Arc, + task_index: usize, + task_count: usize, +) -> Result>>> { + let d_cfg = DistributedConfig::from_config_options(ctx.session_config().options())?; + let registry = &d_cfg.__private_work_unit_feed_registry; + + let d_ctx = DistributedTaskContext { + task_index, + task_count, + }; + let mut streams = vec![]; + plan.apply_with_dt_ctx(d_ctx, |plan, d_ctx| { + let Some(wuf) = registry.get_work_unit_feed(plan) else { + return Ok(TreeNodeRecursion::Continue); + }; + + let partitions = plan.properties().partitioning.partition_count(); + let start_partition = partitions * d_ctx.task_index; + let end_partition = start_partition + partitions; + + let dist_feed_ctx = DistributedWorkUnitFeedContext { + fan_out_tasks: d_ctx.task_count, + }; + let t_ctx = Arc::new(task_ctx_with_extension(ctx, dist_feed_ctx)); + let id = wuf.id(); + + for (partition, feed_idx) in (start_partition..end_partition).enumerate() { + let stream = wuf + .feed(feed_idx, Arc::clone(&t_ctx))? + .map(move |res| res.map(|wu| build_work_unit(&id, partition, wu))) + .boxed(); + streams.push(stream); + } + Ok(TreeNodeRecursion::Continue) + })?; + Ok(streams) +} diff --git a/src/work_unit_feed/mod.rs b/src/work_unit_feed/mod.rs index 1bd55ebf..2e24fcfa 100644 --- a/src/work_unit_feed/mod.rs +++ b/src/work_unit_feed/mod.rs @@ -1,3 +1,4 @@ +mod coordinator; mod remote_work_unit_feed; mod work_unit; #[allow(clippy::module_inception)] @@ -5,9 +6,13 @@ mod work_unit_feed; mod work_unit_feed_provider; mod work_unit_feed_registry; +pub use coordinator::collect_task_work_unit_feeds; +pub use remote_work_unit_feed::{ + RemoteWorkUnitFeedRegistry, RemoteWorkUnitFeedRxs, RemoteWorkUnitFeedTxs, WorkUnitRx, + WorkUnitTx, set_received_time, set_sent_time, +}; pub(crate) use remote_work_unit_feed::{ - RemoteWorkUnitFeedRegistry, build_work_unit_batch_msg, set_work_unit_received_time, - set_work_unit_send_time, + build_work_unit_batch_msg, set_work_unit_received_time, set_work_unit_send_time, }; pub(crate) use work_unit_feed_registry::{WorkUnitFeedRegistry, set_distributed_work_unit_feed}; diff --git a/src/work_unit_feed/remote_work_unit_feed.rs b/src/work_unit_feed/remote_work_unit_feed.rs index b724de73..d0d61bb8 100644 --- a/src/work_unit_feed/remote_work_unit_feed.rs +++ b/src/work_unit_feed/remote_work_unit_feed.rs @@ -47,6 +47,24 @@ impl RemoteWorkUnitFeedRegistry { } } +/// Encodes one feed element into a [pb::WorkUnit] for `partition` of feed `id`. The send/receive +/// stamps stay zero for the transport to fill in right before delivery. +pub(crate) fn build_work_unit( + id: &Uuid, + partition: usize, + work_unit: Box, +) -> pb::WorkUnit { + pb::WorkUnit { + id: serialize_uuid(id), + partition: partition as u64, + body: work_unit.encode_to_bytes(), + created_timestamp_unix_nanos: now_ns(), + sent_timestamp_unix_nanos: 0, + received_timestamp_unix_nanos: 0, + processed_timestamp_unix_nanos: 0, + } +} + pub(crate) fn build_work_unit_batch_msg( id: &Uuid, work_unit_batch: Vec<(usize, Result>)>, @@ -56,23 +74,24 @@ pub(crate) fn build_work_unit_batch_msg( pb::WorkUnitBatch { batch: work_unit_batch .into_iter() - .map(|(partition, work_unit)| { - Ok(pb::WorkUnit { - id: serialize_uuid(id), - partition: partition as u64, - body: work_unit?.encode_to_bytes(), - created_timestamp_unix_nanos: now_ns(), - sent_timestamp_unix_nanos: 0, - received_timestamp_unix_nanos: 0, - processed_timestamp_unix_nanos: 0, - }) - }) + .map(|(partition, work_unit)| Ok(build_work_unit(id, partition, work_unit?))) .collect::>()?, }, )), }) } +/// Stamps the send time on a bare unit. Any transport can stamp before delivery; the worker-side +/// latency math treats a missing stamp as zero latency. +pub fn set_sent_time(work_unit: &mut pb::WorkUnit) { + work_unit.sent_timestamp_unix_nanos = now_ns(); +} + +/// Stamps the receive time on a bare unit. See [set_sent_time]. +pub fn set_received_time(work_unit: &mut pb::WorkUnit) { + work_unit.received_timestamp_unix_nanos = now_ns(); +} + pub(crate) fn set_work_unit_send_time( mut msg: pb::CoordinatorToWorkerMsg, ) -> pb::CoordinatorToWorkerMsg { diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index d1c2097f..5bcc4314 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -1,26 +1,15 @@ +use crate::Worker; use crate::common::deserialize_uuid; -use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, set_work_unit_received_time}; +use crate::work_unit_feed::set_work_unit_received_time; use crate::worker::LocalWorkerContext; use crate::worker::generated::worker::coordinator_to_worker_msg::Inner; -use crate::worker::generated::worker::set_plan_request::WorkUnitFeedDeclaration; use crate::worker::generated::worker::worker_service_server::WorkerService; use crate::worker::generated::worker::{ CoordinatorToWorkerMsg, WorkerToCoordinatorMsg, worker_to_coordinator_msg, }; -use crate::worker::task_data::TaskDataMetrics; -use crate::{ - DistributedCodec, DistributedConfig, DistributedExt, DistributedTaskContext, TaskData, Worker, - WorkerQueryContext, -}; use datafusion::common::DataFusionError; -use datafusion::execution::SessionStateBuilder; -use datafusion::prelude::SessionConfig; -use datafusion_proto::physical_plan::AsExecutionPlan; -use datafusion_proto::protobuf::PhysicalPlanNode; use futures::{FutureExt, StreamExt, TryStreamExt}; -use std::sync::atomic::AtomicUsize; -use std::sync::{Arc, OnceLock}; -use tokio::sync::oneshot; +use std::sync::Arc; use tonic::{Request, Response, Status, Streaming}; use url::Url; @@ -40,88 +29,28 @@ impl Worker { "First Coordinator message must be SetPlanRequest", )); }; - let key = request.task_key.ok_or_else(missing("task_key"))?; - - let entry = self - .task_data_entries - .get_with(key.clone(), async { Default::default() }) - .await; - - let mut remote_work_unit_feed_registry = RemoteWorkUnitFeedRegistry::default(); - for WorkUnitFeedDeclaration { id, partitions } in &request.work_unit_feed_declarations { - if let Ok(id) = deserialize_uuid(id) { - remote_work_unit_feed_registry.add(id, *partitions as usize); - } - } + let key = request.task_key.clone().ok_or_else(missing("task_key"))?; + let target_worker_url = request.target_worker_url.clone(); + let headers = grpc_headers.into_headers(); - let (metrics_tx, metrics_rx) = oneshot::channel(); - - let task_data = || async { - let headers = grpc_headers.into_headers(); - - let mut cfg = SessionConfig::default() - .with_extension(Arc::new(remote_work_unit_feed_registry.receivers)) - .with_extension(Arc::new(DistributedTaskContext { - task_index: key.task_number as usize, - task_count: request.task_count as usize, - })) - .with_extension(Arc::new(LocalWorkerContext { - task_data_entries: Arc::clone(&self.task_data_entries), - self_url: Url::parse(&request.target_worker_url) + // Flight's local-bypass read needs the worker's own registry + URL on the session, so a + // co-located target is served from memory instead of dialing back over gRPC. + let task_data_entries = Arc::clone(&self.task_data_entries); + let outcome = self + .set_task_plan(request, headers, move |cfg| { + Ok(cfg.with_extension(Arc::new(LocalWorkerContext { + task_data_entries, + self_url: Url::parse(&target_worker_url) .map_err(|e| DataFusionError::External(Box::new(e)))?, - })) - .with_distributed_option_extension_from_headers::(&headers)?; - - let d_cfg = DistributedConfig::from_config_options(cfg.options())?; - let shuffle_batch_size = d_cfg.shuffle_batch_size; - let collect_metrics = d_cfg.collect_metrics; - if shuffle_batch_size != 0 { - cfg = cfg.with_batch_size(shuffle_batch_size); - } - - let session_state = self - .session_builder - .build_session_state(WorkerQueryContext { - builder: SessionStateBuilder::new() - .with_default_features() - .with_config(cfg) - .with_runtime_env(Arc::clone(&self.runtime)), - headers, - }) - .await?; - - let codec = DistributedCodec::new_combined_with_user(session_state.config()); - let task_ctx = session_state.task_ctx(); - let proto_node = PhysicalPlanNode::try_decode(request.plan_proto.as_ref())?; - let mut plan = proto_node.try_into_physical_plan(&task_ctx, &codec)?; - - for hook in self.hooks.on_plan.iter() { - plan = hook(plan, session_state.config())?; - } - - // Initialize partition count to the number of partitions in the stage - let total_partitions = plan.properties().partitioning.partition_count(); - Ok::<_, DataFusionError>(TaskData { - base_plan: plan, - final_plan: Arc::new(OnceLock::new()), - task_ctx, - num_partitions_remaining: Arc::new(AtomicUsize::new(total_partitions)), - metrics_tx: match collect_metrics { - true => Arc::new(std::sync::Mutex::new(Some(metrics_tx))), - false => Arc::new(std::sync::Mutex::new(None)), - }, - task_data_metrics: Arc::new(TaskDataMetrics::new(request.query_start_time_ns)), + }))) }) - }; + .await + .map_err(|e| Status::internal(e.to_string()))?; - entry.write(task_data().await.map_err(Arc::new)).map_err(|_| { - Status::internal(format!( - "Logic error while setting plan for TaskKey {key:?}: the plan was set twice. This is a bug in datafusion-distributed, please report it." - )) - })?; + let metrics_rx = outcome.metrics_rx; // Continue reading remaining messages (work unit feed data) in the background. - let mut work_unit_senders = Some(remote_work_unit_feed_registry.senders); + let mut work_unit_senders = Some(outcome.work_unit_senders); let task_data_entries = Arc::clone(&self.task_data_entries); #[allow(clippy::disallowed_methods)] tokio::spawn(async move { diff --git a/src/worker/impl_set_plan.rs b/src/worker/impl_set_plan.rs new file mode 100644 index 00000000..edb5a37c --- /dev/null +++ b/src/worker/impl_set_plan.rs @@ -0,0 +1,127 @@ +use crate::common::deserialize_uuid; +use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, RemoteWorkUnitFeedTxs}; +use crate::worker::generated::worker as pb; +use crate::worker::generated::worker::set_plan_request::WorkUnitFeedDeclaration; +use crate::worker::task_data::TaskDataMetrics; +use crate::worker::worker_service::Worker; +use crate::{ + DistributedCodec, DistributedConfig, DistributedExt, DistributedTaskContext, TaskData, + WorkerQueryContext, +}; +use datafusion::common::{DataFusionError, Result, exec_datafusion_err}; +use datafusion::execution::SessionStateBuilder; +use datafusion::prelude::SessionConfig; +use datafusion_proto::physical_plan::AsExecutionPlan; +use datafusion_proto::protobuf::PhysicalPlanNode; +use http::HeaderMap; +use std::sync::atomic::AtomicUsize; +use std::sync::{Arc, OnceLock}; +use tokio::sync::oneshot; + +/// What a transport gets back from [Worker::set_task_plan]: the channels to push work units into, +/// and the one-shot receiver that resolves with the task's metrics once every partition finished +/// or was dropped. +pub(crate) struct SetPlanOutcome { + pub(crate) work_unit_senders: RemoteWorkUnitFeedTxs, + pub(crate) metrics_rx: oneshot::Receiver, +} + +impl Worker { + /// Stores one task's plan so [`super::impl_execute_task::execute_local_task`] can pick it up: + /// builds the worker-side session through this worker's [`crate::WorkerSessionBuilder`], + /// decodes the plan against it, and publishes the resulting [TaskData] under its task key. + /// + /// This is the transport-neutral core of plan delivery. The Flight service wraps it in its + /// coordinator gRPC stream; an in-process transport calls it directly. `customize_cfg` lets the + /// caller attach transport-specific session extensions (Flight adds its local-bypass context) + /// without this method knowing about them. + pub(crate) async fn set_task_plan( + &self, + request: pb::SetPlanRequest, + headers: HeaderMap, + customize_cfg: impl FnOnce(SessionConfig) -> Result, + ) -> Result { + let key = request + .task_key + .clone() + .ok_or_else(|| exec_datafusion_err!("Missing field 'task_key'"))?; + + let entry = self + .task_data_entries + .get_with(key.clone(), async { Default::default() }) + .await; + + let mut remote_work_unit_feed_registry = RemoteWorkUnitFeedRegistry::default(); + for WorkUnitFeedDeclaration { id, partitions } in &request.work_unit_feed_declarations { + if let Ok(id) = deserialize_uuid(id) { + remote_work_unit_feed_registry.add(id, *partitions as usize); + } + } + + let (metrics_tx, metrics_rx) = oneshot::channel(); + + let task_data = async { + let mut cfg = SessionConfig::default() + .with_extension(Arc::new(remote_work_unit_feed_registry.receivers)) + .with_extension(Arc::new(DistributedTaskContext { + task_index: key.task_number as usize, + task_count: request.task_count as usize, + })); + cfg = customize_cfg(cfg)?; + cfg = + cfg.with_distributed_option_extension_from_headers::(&headers)?; + + let d_cfg = DistributedConfig::from_config_options(cfg.options())?; + let shuffle_batch_size = d_cfg.shuffle_batch_size; + let collect_metrics = d_cfg.collect_metrics; + if shuffle_batch_size != 0 { + cfg = cfg.with_batch_size(shuffle_batch_size); + } + + let session_state = self + .session_builder + .build_session_state(WorkerQueryContext { + builder: SessionStateBuilder::new() + .with_default_features() + .with_config(cfg) + .with_runtime_env(Arc::clone(&self.runtime)), + headers, + }) + .await?; + + let codec = DistributedCodec::new_combined_with_user(session_state.config()); + let task_ctx = session_state.task_ctx(); + let proto_node = PhysicalPlanNode::try_decode(request.plan_proto.as_ref())?; + let mut plan = proto_node.try_into_physical_plan(&task_ctx, &codec)?; + + for hook in self.hooks.on_plan.iter() { + plan = hook(plan, session_state.config())?; + } + + // Initialize partition count to the number of partitions in the stage. + let total_partitions = plan.properties().partitioning.partition_count(); + Ok::<_, DataFusionError>(TaskData { + base_plan: plan, + final_plan: Arc::new(OnceLock::new()), + task_ctx, + num_partitions_remaining: Arc::new(AtomicUsize::new(total_partitions)), + metrics_tx: match collect_metrics { + true => Arc::new(std::sync::Mutex::new(Some(metrics_tx))), + false => Arc::new(std::sync::Mutex::new(None)), + }, + task_data_metrics: Arc::new(TaskDataMetrics::new(request.query_start_time_ns)), + }) + }; + + entry.write(task_data.await.map_err(Arc::new)).map_err(|_| { + exec_datafusion_err!( + "Logic error while setting plan for TaskKey {key:?}: the plan was set twice. This is a bug in datafusion-distributed, please report it." + ) + })?; + + Ok(SetPlanOutcome { + work_unit_senders: remote_work_unit_feed_registry.senders, + metrics_rx, + }) + } +} diff --git a/src/worker/in_memory.rs b/src/worker/in_memory.rs new file mode 100644 index 00000000..f50106eb --- /dev/null +++ b/src/worker/in_memory.rs @@ -0,0 +1,402 @@ +//! An in-process [WorkerTransport]: every "worker" is the current process, so plans are delivered +//! with a function call and partitions are read straight from the local task registry. It is the +//! default transport when the `flight` feature is off, and the reference implementation for the +//! transport extension points: it goes through the same plan encode/decode, session building, +//! work-unit feed, and metrics delivery as a remote transport, just without a wire underneath. + +use crate::common::{deserialize_uuid, serialize_uuid}; +use crate::config_extension_ext::get_config_extension_propagation_headers; +use crate::coordinator::{CoordinatorToWorkerMetrics, encode_task_plan}; +use crate::distributed_planner::ProducerHead; +use crate::networking::set_distributed_worker_transport; +use crate::passthrough_headers::get_passthrough_headers; +use crate::stage::RemoteStage; +use crate::work_unit_feed::{collect_task_work_unit_feeds, set_received_time, set_sent_time}; +use crate::worker::WorkerSessionBuilder; +use crate::worker::generated::worker as pb; +use crate::worker::impl_execute_task::execute_local_task; +use crate::worker::transport::{ + WorkerConnection, WorkerDispatch, WorkerDispatchRequest, WorkerTransport, +}; +use crate::worker::worker_service::{TaskDataEntries, Worker}; +use datafusion::arrow::array::RecordBatch; +use datafusion::common::instant::Instant; +use datafusion::common::runtime::SpawnedTask; +use datafusion::common::{DataFusionError, Result, internal_datafusion_err, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr_common::metrics::ExecutionPlanMetricsSet; +use datafusion::physical_plan::metrics::MetricBuilder; +use futures::StreamExt; +use futures::stream::BoxStream; +use std::ops::Range; +use std::sync::{Arc, Mutex, OnceLock}; +use tokio_stream::wrappers::UnboundedReceiverStream; + +/// [WorkerTransport] that hosts its workers in the current process. +/// +/// All tasks share one [Worker] (one task registry, one session builder): task keys carry the +/// query, stage, and task number, so a single registry isolates them. The worker resolver's URLs +/// only size the stages; nothing is dialed. +/// +/// With the `flight` feature off this is the default transport, so distributed plans run out of the +/// box without the Flight stack. Embedders with their own comms register a custom transport. +#[derive(Clone, Default)] +pub struct InMemoryWorkerTransport { + worker: Worker, +} + +impl InMemoryWorkerTransport { + /// Builds the transport around an existing [Worker], sharing its task registry, session + /// builder, and runtime environment. + pub fn new(worker: Worker) -> Self { + Self { worker } + } + + /// Builds the transport with a custom [WorkerSessionBuilder], the same customization hook a + /// remote worker offers. + pub fn from_session_builder( + session_builder: impl WorkerSessionBuilder + Send + Sync + 'static, + ) -> Self { + Self { + worker: Worker::from_session_builder(session_builder), + } + } + + /// The in-process [Worker] backing this transport. + pub fn worker(&self) -> &Worker { + &self.worker + } +} + +impl WorkerTransport for InMemoryWorkerTransport { + fn open( + &self, + input_stage: &RemoteStage, + target_partitions: Range, + target_task: usize, + producer_head: ProducerHead, + ctx: &Arc, + metrics: &ExecutionPlanMetricsSet, + ) -> Result> { + Ok(Box::new(InMemoryWorkerConnection::init( + input_stage, + target_partitions, + target_task, + producer_head, + Arc::clone(&self.worker.task_data_entries), + ctx, + metrics, + )?)) + } + + fn dispatcher(&self) -> Box { + Box::new(InMemoryWorkerDispatcher { + worker: self.worker.clone(), + metrics: OnceLock::new(), + }) + } +} + +/// Per-query plan-delivery state for the in-memory transport. As with Flight, the plan-send +/// metrics and the query start timestamp live for the whole query, not per stage. +struct InMemoryWorkerDispatcher { + worker: Worker, + metrics: OnceLock, +} + +/// Delivery is [Worker::set_task_plan] called directly: the plan still round-trips through the +/// session's codec stack, so each task executes its own decoded instance (plan nodes carry +/// per-execution state, sharing one tree between tasks is not sound) and codec gaps surface the +/// same way they would over a wire. +impl WorkerDispatch for InMemoryWorkerDispatcher { + fn dispatch(&self, request: WorkerDispatchRequest<'_>) -> Result<()> { + let WorkerDispatchRequest { + stage, + routed_urls, + task_ctx, + metrics, + metrics_store, + join_set, + .. + } = request; + let metrics = self + .metrics + .get_or_init(|| CoordinatorToWorkerMetrics::new(metrics)) + .clone(); + + let mut headers = get_config_extension_propagation_headers(task_ctx.session_config())?; + headers.extend(get_passthrough_headers(task_ctx.session_config())); + + for (task_i, url) in routed_urls.iter().enumerate() { + let encoded = + encode_task_plan(&stage.plan, task_i, stage.tasks, task_ctx.session_config())?; + let plan_size = encoded.plan_proto.len(); + + let task_key = pb::TaskKey { + query_id: serialize_uuid(&stage.query_id), + stage_id: stage.num as u64, + task_number: task_i as u64, + }; + let set_plan = pb::SetPlanRequest { + plan_proto: encoded.plan_proto, + task_count: stage.tasks as u64, + task_key: Some(task_key.clone()), + work_unit_feed_declarations: encoded.feed_declarations, + target_worker_url: url.to_string(), + query_start_time_ns: metrics.instantiation_time, + }; + + // Collected before spawning so the providers see the same eager `feed()` timing as + // they do under Flight. + let feed_streams = + collect_task_work_unit_feeds(&stage.plan, task_ctx, task_i, stage.tasks)?; + + let worker = self.worker.clone(); + let transport = InMemoryWorkerTransport::new(self.worker.clone()); + let headers = headers.clone(); + let metrics = metrics.clone(); + let metrics_store = metrics_store.cloned(); + join_set.spawn(async move { + let start = Instant::now(); + let outcome = worker + .set_task_plan(set_plan, headers, move |mut cfg| { + // Child-stage reads inside the decoded worker plan consult the worker + // session for a transport; they must land on this same task registry. + set_distributed_worker_transport(&mut cfg, transport); + Ok(cfg) + }) + .await?; + metrics.plan_send_latency.record(&start); + metrics.plan_bytes_sent.add_bytes(plan_size); + + // Detached like Flight's metrics collection: the receiver resolves only once every + // partition finished or was dropped, and a task whose partitions are never opened + // must not stall query completion. + let metrics_rx = outcome.metrics_rx; + #[allow(clippy::disallowed_methods)] + tokio::spawn(async move { + if let (Ok(task_metrics), Some(store)) = (metrics_rx.await, metrics_store) { + store.insert(task_key, task_metrics); + } + }); + + // Pump the work-unit feeds straight into the worker-side channels. Both hop stamps + // are set here; the latency metrics read as (near) zero, which is what an in-process + // hop is. + let senders = Arc::new(outcome.work_unit_senders); + let mut pumps = vec![]; + for mut stream in feed_streams { + let senders = Arc::clone(&senders); + pumps.push(async move { + while let Some(unit) = stream.next().await { + let mut unit = unit?; + set_sent_time(&mut unit); + set_received_time(&mut unit); + let Ok(id) = deserialize_uuid(&unit.id) else { + continue; + }; + let Some(tx) = senders.get(&(id, unit.partition as usize)) else { + continue; + }; + if tx.send(Ok(unit)).is_err() { + break; // channel closed + } + } + Ok::<_, DataFusionError>(()) + }); + } + futures::future::try_join_all(pumps).await?; + Ok(()) + }); + } + Ok(()) + } +} + +/// [WorkerConnection] over the local task registry: builds the per-partition record batch streams +/// straight from [execute_local_task], with no encoding in between. +pub(crate) struct InMemoryWorkerConnection { + partition_start: usize, + local_streams: Vec>>>>, + /// Drives the single `execute_local_task` call that produces every partition stream. Held so it + /// lives as long as the connection (the pool keeps it for the query), not aborted early. + _driver: SpawnedTask<()>, +} + +impl InMemoryWorkerConnection { + pub(crate) fn init( + input_stage: &RemoteStage, + target_partition_range: Range, + target_task: usize, + producer_head: ProducerHead, + task_data_entries: Arc, + ctx: &Arc, + metrics: &ExecutionPlanMetricsSet, + ) -> Result { + MetricBuilder::new(metrics) + .global_counter("local_connections_used") + .add(1); + + let task_key = pb::TaskKey { + query_id: serialize_uuid(&input_stage.query_id), + stage_id: input_stage.num as u64, + task_number: target_task as u64, + }; + let producer_head = producer_head.to_proto(ctx)?; + + let partition_start = target_partition_range.start; + let n_partitions = target_partition_range.len(); + + // Execute the whole partition range in ONE `execute_local_task` call, then drive each + // partition stream from a dedicated pump into a buffer the consumer reads from. This + // decouples the worker plan's execution from when (or whether) the consumer pulls, exactly + // as Flight gets from its gRPC stream. Letting the consumer drive the worker plan directly + // (its `SortPreservingMergeExec` interleaving the polls of a partitioned `HashJoinExec`) can + // otherwise leave some partitions empty. + let request = pb::ExecuteTaskRequest { + task_key: Some(task_key), + target_partition_start: target_partition_range.start as u64, + target_partition_end: target_partition_range.end as u64, + producer_head: Some(producer_head), + }; + + let mut senders = Vec::with_capacity(n_partitions); + let mut local_streams = Vec::with_capacity(n_partitions); + for _ in 0..n_partitions { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel::>(); + senders.push(tx); + local_streams.push(Mutex::new(Some(UnboundedReceiverStream::new(rx).boxed()))); + } + + let driver = SpawnedTask::spawn(async move { + match execute_local_task(&task_data_entries, request).await { + Ok((streams, _)) => { + let pumps = senders + .into_iter() + .zip(streams) + .map(|(tx, mut stream)| async move { + while let Some(item) = stream.next().await { + if tx.send(item).is_err() { + break; // consumer dropped this partition + } + } + }) + .collect::>(); + futures::future::join_all(pumps).await; + } + Err(err) => { + let err = Arc::new(err); + for tx in senders { + let _ = tx.send(Err(DataFusionError::Shared(Arc::clone(&err)))); + } + } + } + }); + + Ok(Self { + partition_start, + local_streams, + _driver: driver, + }) + } +} + +impl WorkerConnection for InMemoryWorkerConnection { + fn execute(&self, partition: usize) -> Result>> { + let Some(relative_i) = partition.checked_sub(self.partition_start) else { + return internal_err!( + "InMemoryWorkerConnection received an invalid partition {partition}, the starting partition is {}", + self.partition_start + ); + }; + let Some(slot) = self.local_streams.get(relative_i) else { + return internal_err!( + "InMemoryWorkerConnection has no stream for partition {partition}. Was it already consumed?" + ); + }; + slot.lock().unwrap().take().ok_or_else(|| { + internal_datafusion_err!( + "InMemoryWorkerConnection stream for partition {partition} was already consumed" + ) + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::in_memory_channel_resolver::InMemoryWorkerResolver; + use crate::test_utils::session_context::register_temp_parquet_table; + use crate::{DistributedExt, SessionStateBuilderExt, display_plan_ascii}; + use datafusion::arrow::array::{Int32Array, StringArray}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::record_batch::RecordBatch as ArrowRecordBatch; + use datafusion::arrow::util::pretty::pretty_format_batches; + use datafusion::execution::SessionStateBuilder; + use datafusion::physical_plan::execute_stream; + use datafusion::prelude::SessionContext; + use futures::TryStreamExt; + + const QUERY: &str = "SELECT tag, count(*) AS c, sum(val) AS s FROM t GROUP BY tag ORDER BY tag"; + + fn sample_batch() -> ArrowRecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("tag", DataType::Utf8, false), + Field::new("val", DataType::Int32, false), + ])); + let tags: Vec = (0..100).map(|i| format!("tag{}", i % 7)).collect(); + let vals: Vec = (0..100).collect(); + ArrowRecordBatch::try_new( + schema, + vec![ + Arc::new(StringArray::from(tags)), + Arc::new(Int32Array::from(vals)), + ], + ) + .unwrap() + } + + fn distributed_ctx(transport: Option) -> SessionContext { + let mut builder = SessionStateBuilder::new() + .with_default_features() + .with_distributed_planner() + .with_distributed_task_estimator(2) + .with_distributed_worker_resolver(InMemoryWorkerResolver::new(3)); + if let Some(transport) = transport { + builder = builder.with_distributed_worker_transport(transport); + } + let mut state = builder.build(); + state.config_mut().options_mut().execution.target_partitions = 3; + SessionContext::from(state) + } + + async fn run(ctx: &SessionContext) -> Result<(String, String)> { + let plan = ctx.sql(QUERY).await?.create_physical_plan().await?; + let display = display_plan_ascii(plan.as_ref(), false); + let batches: Vec<_> = execute_stream(plan, ctx.task_ctx())?.try_collect().await?; + Ok((display, pretty_format_batches(&batches)?.to_string())) + } + + #[tokio::test] + async fn distributed_query_matches_single_node() -> Result<()> { + let ctx = distributed_ctx(Some(InMemoryWorkerTransport::default())); + let path = + register_temp_parquet_table("t", sample_batch().schema(), vec![sample_batch()], &ctx) + .await?; + + let (display, distributed) = run(&ctx).await?; + assert!( + display.contains("NetworkShuffleExec"), + "the query did not distribute:\n{display}" + ); + + let single = SessionContext::default(); + single + .register_parquet("t", path.to_string_lossy().as_ref(), Default::default()) + .await?; + let (_, expected) = run(&single).await?; + + assert_eq!(distributed, expected); + Ok(()) + } +} diff --git a/src/worker/mod.rs b/src/worker/mod.rs index f7ac23a2..347683db 100644 --- a/src/worker/mod.rs +++ b/src/worker/mod.rs @@ -2,6 +2,8 @@ mod flight; pub(crate) mod generated; mod impl_coordinator_channel; mod impl_execute_task; +mod impl_set_plan; +mod in_memory; mod session_builder; mod single_write_multi_read; mod spawn_select_all; @@ -16,6 +18,7 @@ pub(crate) use single_write_multi_read::SingleWriteMultiRead; pub(crate) use worker_connection_pool::{LocalWorkerContext, WorkerConnectionPool}; pub use flight::FlightWorkerTransport; +pub use in_memory::InMemoryWorkerTransport; pub use transport::{WorkerConnection, WorkerDispatch, WorkerDispatchRequest, WorkerTransport}; pub use session_builder::{ From 53ff81772a08afa8eed4b0ca27a186a31c685bb8 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sat, 20 Jun 2026 11:32:01 -0700 Subject: [PATCH 3/4] Made the Arrow-Flight transport optional behind a `flight` feature. With `flight` off there is no `tonic` / `arrow-flight` and the in-memory transport is the default, so distributed plans still run. The integration suite runs over the in-memory transport in both build configurations: `start_localhost_context` builds an `InMemoryWorkerTransport` cluster instead of a gRPC one, and the gRPC harness moved to `start_localhost_flight_context`. A `unit-test-flight-transport` job sets `DATAFUSION_DISTRIBUTED_TEST_TRANSPORT=flight` to keep full Flight coverage. For the no-flight suite to run, not just build, the `tpch` / `tpcds` / `clickbench` / `stateful_data_cleanup` dataset tests move into the benchmarks crate: as a dev-dependency of the library they re-enabled `flight` on every test build through feature unification. With that gone, `cargo test --no-default-features --features integration --lib --tests` is genuinely Flight-free, and a `unit-test-no-flight` job runs it. The only tests still gated on `flight` are the ones that need a real wire: the `URLEmitter` routing tests, which assert per-URL worker identity, and the network-boundary connection metrics (`bytes_transferred`, latency). The rest, including the worker-hook and metrics tests, register `InMemoryWorkerTransport` directly and run either way. --- .github/workflows/ci.yml | 54 +++++++++-- Cargo.lock | 3 +- Cargo.toml | 21 +++-- benchmarks/Cargo.toml | 11 +++ .../tests}/clickbench_correctness_test.rs | 2 +- .../tests}/clickbench_plans_test.rs | 2 +- .../tests}/stateful_data_cleanup.rs | 2 +- .../tests}/tpcds_correctness_test.rs | 2 +- .../tests}/tpcds_plans_test.rs | 2 +- .../tests}/tpch_correctness_test.rs | 2 +- .../tests}/tpch_plans_test.rs | 2 +- src/coordinator/dispatch_metrics.rs | 33 +++++++ src/coordinator/mod.rs | 6 +- src/coordinator/query_coordinator.rs | 34 +------ src/distributed_ext.rs | 23 +++-- src/execution_plans/benchmarks/fixture.rs | 6 ++ src/execution_plans/benchmarks/mod.rs | 4 + src/lib.rs | 44 ++++++--- src/metrics/task_metrics_collector.rs | 7 +- src/metrics/task_metrics_rewriter.rs | 7 +- src/networking/channel_resolver.rs | 41 ++++++++- src/networking/mod.rs | 5 +- src/networking/worker_transport.rs | 23 ++++- src/observability/generated/mod.rs | 1 + src/observability/generated/observability.rs | 2 + src/observability/mod.rs | 4 + src/protobuf/mod.rs | 2 + src/test_utils/in_memory_channel_resolver.rs | 41 ++++++--- src/test_utils/localhost.rs | 73 +++++++++++++-- src/test_utils/routing.rs | 20 +++-- src/work_unit_feed/mod.rs | 1 + src/work_unit_feed/remote_work_unit_feed.rs | 3 + src/worker/generated/mod.rs | 3 +- src/worker/generated/worker.rs | 2 + src/worker/impl_execute_task.rs | 79 +++++++++++----- src/worker/in_memory.rs | 30 +++++++ src/worker/mod.rs | 16 +++- src/worker/single_write_multi_read.rs | 1 + src/worker/task_data.rs | 2 + src/worker/test_utils/mod.rs | 1 + src/worker/worker_connection_pool.rs | 89 ++++++++++++++++--- src/worker/worker_service.rs | 46 +++++++--- tests/extension_points.rs | 73 +++++++++++++++ tests/metrics_collection.rs | 14 ++- tests/task_estimator_test.rs | 4 +- 45 files changed, 670 insertions(+), 173 deletions(-) rename {tests => benchmarks/tests}/clickbench_correctness_test.rs (99%) rename {tests => benchmarks/tests}/clickbench_plans_test.rs (99%) rename {tests => benchmarks/tests}/stateful_data_cleanup.rs (98%) rename {tests => benchmarks/tests}/tpcds_correctness_test.rs (99%) rename {tests => benchmarks/tests}/tpcds_plans_test.rs (99%) rename {tests => benchmarks/tests}/tpch_correctness_test.rs (99%) rename {tests => benchmarks/tests}/tpch_plans_test.rs (99%) create mode 100644 src/coordinator/dispatch_metrics.rs create mode 100644 tests/extension_points.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index cc217c7b..657e0690 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -32,6 +32,8 @@ jobs: components: clippy - run: cargo clippy --all-targets --all-features -- -D warnings + # The generic suite runs over the in-memory transport (what `start_localhost_context` builds by + # default), so it covers the distributed paths without a gRPC server. unit-test: runs-on: ubuntu-latest steps: @@ -41,12 +43,50 @@ jobs: - uses: ./.github/actions/setup - run: cargo test --features integration + # The same suite over the Arrow-Flight gRPC transport. `start_localhost_context` switches to it + # when `DATAFUSION_DISTRIBUTED_TEST_TRANSPORT=flight`, so both transports get full coverage. + unit-test-flight-transport: + runs-on: ubuntu-latest + env: + DATAFUSION_DISTRIBUTED_TEST_TRANSPORT: flight + steps: + - uses: actions/checkout@v4 + with: + lfs: true + - uses: ./.github/actions/setup + - run: cargo test --features integration + + # Builds the lib and the suite with `flight` off: no `tonic` / `arrow-flight`, in-memory transport + # as the default. The dataset suites moved to the benchmarks crate, so the suite no longer drags + # the benchmarks dev-dependency in (which re-enabled `flight` through feature unification). + # Examples are left out: they demo the Flight cluster, so `--all-features` clippy covers them. + no-flight: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: ./.github/actions/setup + with: + components: clippy + - run: cargo build --no-default-features --features integration + - run: cargo clippy --no-default-features --features integration --lib --tests -- -D warnings + + # The integration suite over the in-memory transport with `flight` off, so the no-flight runtime + # path is run, not just built. Flight-specific tests gate themselves out. + unit-test-no-flight: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + lfs: true + - uses: ./.github/actions/setup + - run: cargo test --no-default-features --features integration --lib --tests + tpch-test: runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 - uses: ./.github/actions/setup - - run: cargo test --features tpch --test 'tpch_*' + - run: cargo test -p datafusion-distributed-benchmarks --features tpch --test 'tpch_*' tpcds-correctness-test: runs-on: ubuntu-latest @@ -59,9 +99,9 @@ jobs: - uses: ./.github/actions/setup - uses: actions/cache@v4 with: - path: testdata/tpcds/main.zip + path: benchmarks/testdata/tpcds/main.zip key: "main.zip" - - run: cargo test --features tpcds --test tpcds_correctness_test shard${{ matrix.shard }} + - run: cargo test -p datafusion-distributed-benchmarks --features tpcds --test tpcds_correctness_test shard${{ matrix.shard }} tpcds-plans-test: runs-on: ubuntu-latest @@ -70,9 +110,9 @@ jobs: - uses: ./.github/actions/setup - uses: actions/cache@v4 with: - path: testdata/tpcds/main.zip + path: benchmarks/testdata/tpcds/main.zip key: "main.zip" - - run: cargo test --features tpcds --test tpcds_plans_test + - run: cargo test -p datafusion-distributed-benchmarks --features tpcds --test tpcds_plans_test clickbench-test: runs-on: ubuntu-latest @@ -81,9 +121,9 @@ jobs: - uses: ./.github/actions/setup - uses: actions/cache@v4 with: - path: testdata/clickbench/ + path: benchmarks/testdata/clickbench/ key: "data" - - run: cargo test --features clickbench --test 'clickbench_*' + - run: cargo test -p datafusion-distributed-benchmarks --features clickbench --test 'clickbench_*' format-check: runs-on: ubuntu-latest diff --git a/Cargo.lock b/Cargo.lock index ebe698d2..ecaf4617 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2199,7 +2199,6 @@ dependencies = [ "crossbeam-queue", "dashmap", "datafusion", - "datafusion-distributed-benchmarks", "datafusion-proto", "delegate", "futures", @@ -2248,11 +2247,13 @@ dependencies = [ "object_store", "openssl", "parquet", + "pretty_assertions", "reqwest", "serde", "serde_json", "structopt", "sysinfo", + "test-case", "tokio", "tonic", "tpchgen", diff --git a/Cargo.toml b/Cargo.toml index 4d6e3e7e..7f7296cd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,13 +23,13 @@ datafusion = { workspace = true, features = [ "datetime_expressions", ] } datafusion-proto = { workspace = true } -arrow-flight = "58" +arrow-flight = { version = "58", optional = true } arrow-select = "58" arrow-ipc = { version = "58", features = ["zstd"] } async-trait = "0.1.89" tokio = { version = "1.48", features = ["full"] } -tonic = { version = "0.14.1", features = ["transport"] } -tower = "0.5.2" +tonic = { version = "0.14.1", features = ["transport"], optional = true } +tower = { version = "0.5.2", optional = true } http = "1.3.1" itertools = "0.14.0" futures = "0.3.31" @@ -49,7 +49,7 @@ crossbeam-queue = "0.3" sysinfo = { version = "0.30", optional = true } sketches-ddsketch = { version = "0.3", features = ["use_serde"] } bincode = "1" -tonic-prost = "0.14.2" +tonic-prost = { version = "0.14.2", optional = true } # integration_tests deps insta = { version = "1.46.0", features = ["filters"], optional = true } @@ -58,19 +58,22 @@ arrow = { version = "58", optional = true, features = ["test_utils"] } hyper-util = { version = "0.1.16", optional = true } [features] +default = ["flight"] +# Arrow Flight gRPC transport. When off, the crate builds with no `tonic` / `arrow-flight` +# and the in-memory transport becomes the default: distributed plans run inside the current +# process. Multi-process execution then needs an embedder-registered transport. +flight = ["dep:arrow-flight", "dep:tonic", "dep:tonic-prost", "dep:tower"] avro = ["datafusion/avro"] +# Independent from `flight` so the integration suite runs against both the gRPC transport +# (`--features integration`) and the in-memory one (`--no-default-features --features +# integration`). The Flight-specific tests gate themselves on `flight`. integration = ["insta", "parquet", "arrow", "hyper-util"] system-metrics = ["sysinfo"] -tpch = ["integration"] -tpcds = ["integration"] -clickbench = ["integration"] -slow-tests = [] sysinfo = ["dep:sysinfo"] [dev-dependencies] -datafusion-distributed-benchmarks = { path = "benchmarks" } structopt = "0.3" insta = { version = "1.46.0", features = ["filters"] } parquet = "58" diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 0a03ead0..f2cc695d 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -33,9 +33,20 @@ aws-sdk-ec2 = "1" openssl = { version = "0.10", features = ["vendored"] } # Keep this. Necessary for the remote benchmarks worker. mimalloc = "0.1" +[features] +# Gates for the dataset test suites under `tests/`. They live here instead of the library because +# they depend on it: as a dev-dependency of the library they re-enable `flight` on every test build +# through feature unification, which makes the no-flight config untestable. +tpch = [] +tpcds = [] +clickbench = [] +slow-tests = [] + [dev-dependencies] criterion = "0.5" sysinfo = "0.30" +pretty_assertions = "1.4" +test-case = "3.3.1" [build-dependencies] built = { version = "0.8", features = ["git2", "chrono"] } diff --git a/tests/clickbench_correctness_test.rs b/benchmarks/tests/clickbench_correctness_test.rs similarity index 99% rename from tests/clickbench_correctness_test.rs rename to benchmarks/tests/clickbench_correctness_test.rs index 9df2a8d3..e7bf868e 100644 --- a/tests/clickbench_correctness_test.rs +++ b/benchmarks/tests/clickbench_correctness_test.rs @@ -1,4 +1,4 @@ -#[cfg(all(feature = "integration", feature = "clickbench", test))] +#[cfg(all(feature = "clickbench", test))] mod tests { use datafusion::arrow::array::RecordBatch; use datafusion::common::plan_err; diff --git a/tests/clickbench_plans_test.rs b/benchmarks/tests/clickbench_plans_test.rs similarity index 99% rename from tests/clickbench_plans_test.rs rename to benchmarks/tests/clickbench_plans_test.rs index 2c202201..910194b1 100644 --- a/tests/clickbench_plans_test.rs +++ b/benchmarks/tests/clickbench_plans_test.rs @@ -1,4 +1,4 @@ -#[cfg(all(feature = "integration", feature = "clickbench", test))] +#[cfg(all(feature = "clickbench", test))] mod tests { use datafusion::error::Result; use datafusion_distributed::test_utils::in_memory_channel_resolver::start_in_memory_context; diff --git a/tests/stateful_data_cleanup.rs b/benchmarks/tests/stateful_data_cleanup.rs similarity index 98% rename from tests/stateful_data_cleanup.rs rename to benchmarks/tests/stateful_data_cleanup.rs index a3fe7bca..d45f3451 100644 --- a/tests/stateful_data_cleanup.rs +++ b/benchmarks/tests/stateful_data_cleanup.rs @@ -1,4 +1,4 @@ -#[cfg(all(feature = "integration", feature = "tpch", test))] +#[cfg(all(feature = "tpch", test))] mod tests { use datafusion::common::instant::Instant; use datafusion::error::Result; diff --git a/tests/tpcds_correctness_test.rs b/benchmarks/tests/tpcds_correctness_test.rs similarity index 99% rename from tests/tpcds_correctness_test.rs rename to benchmarks/tests/tpcds_correctness_test.rs index bb6011a7..5b203e41 100644 --- a/tests/tpcds_correctness_test.rs +++ b/benchmarks/tests/tpcds_correctness_test.rs @@ -1,4 +1,4 @@ -#[cfg(all(feature = "integration", feature = "tpcds", test))] +#[cfg(all(feature = "tpcds", test))] mod tests { use datafusion::arrow::array::RecordBatch; use datafusion::common::plan_err; diff --git a/tests/tpcds_plans_test.rs b/benchmarks/tests/tpcds_plans_test.rs similarity index 99% rename from tests/tpcds_plans_test.rs rename to benchmarks/tests/tpcds_plans_test.rs index 667ed4d1..04a6fd0a 100644 --- a/tests/tpcds_plans_test.rs +++ b/benchmarks/tests/tpcds_plans_test.rs @@ -1,4 +1,4 @@ -#[cfg(all(feature = "integration", feature = "tpcds", test))] +#[cfg(all(feature = "tpcds", test))] mod tests { use datafusion::error::Result; use datafusion_distributed::test_utils::in_memory_channel_resolver::start_in_memory_context; diff --git a/tests/tpch_correctness_test.rs b/benchmarks/tests/tpch_correctness_test.rs similarity index 99% rename from tests/tpch_correctness_test.rs rename to benchmarks/tests/tpch_correctness_test.rs index 3b6bc5a4..6f74d58f 100644 --- a/tests/tpch_correctness_test.rs +++ b/benchmarks/tests/tpch_correctness_test.rs @@ -1,4 +1,4 @@ -#[cfg(all(feature = "integration", feature = "tpch", test))] +#[cfg(all(feature = "tpch", test))] mod tests { use datafusion::physical_plan::execute_stream; use datafusion::prelude::SessionContext; diff --git a/tests/tpch_plans_test.rs b/benchmarks/tests/tpch_plans_test.rs similarity index 99% rename from tests/tpch_plans_test.rs rename to benchmarks/tests/tpch_plans_test.rs index 231ee3b0..3e4a719a 100644 --- a/tests/tpch_plans_test.rs +++ b/benchmarks/tests/tpch_plans_test.rs @@ -1,4 +1,4 @@ -#[cfg(all(feature = "integration", feature = "tpch", test))] +#[cfg(all(feature = "tpch", test))] mod tests { use datafusion_distributed::test_utils::in_memory_channel_resolver::start_in_memory_context; use datafusion_distributed::{ diff --git a/src/coordinator/dispatch_metrics.rs b/src/coordinator/dispatch_metrics.rs new file mode 100644 index 00000000..87fbf605 --- /dev/null +++ b/src/coordinator/dispatch_metrics.rs @@ -0,0 +1,33 @@ +use crate::common::now_ns; +use crate::coordinator::latency_metric::LatencyMetric; +use crate::{BytesCounterMetric, BytesMetricExt, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL}; +use datafusion::physical_expr_common::metrics::{ExecutionPlanMetricsSet, Label, MetricBuilder}; +use std::sync::Arc; + +/// Metrics that measure network details about communications between [crate::DistributedExec] and a +/// worker. Shared by every transport's dispatch path, so the plan-send byte/latency counters read +/// the same regardless of how plans reach the workers. +#[derive(Clone)] +pub struct CoordinatorToWorkerMetrics { + pub plan_bytes_sent: BytesCounterMetric, + pub plan_send_latency: Arc, + pub instantiation_time: u64, +} + +impl CoordinatorToWorkerMetrics { + pub fn new(metrics: &ExecutionPlanMetricsSet) -> Self { + Self { + // Metric that measures to total sum of bytes worth of subplans sent. + plan_bytes_sent: MetricBuilder::new(metrics) + .with_label(Label::new(DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, "0")) + .bytes_counter("plan_bytes_sent"), + // Latency statistics about the network calls issued to the workers for feeding subplans. + plan_send_latency: Arc::new(LatencyMetric::new( + "plan_send_latency", + |b| b.with_label(Label::new(DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, "0")), + metrics, + )), + instantiation_time: now_ns(), + } + } +} diff --git a/src/coordinator/mod.rs b/src/coordinator/mod.rs index 3d46f72a..34129bb2 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -1,12 +1,16 @@ +mod dispatch_metrics; mod distributed; mod latency_metric; mod metrics_store; mod plan_encoding; mod prepare_static_plan; +#[cfg(feature = "flight")] mod query_coordinator; +pub use dispatch_metrics::CoordinatorToWorkerMetrics; pub use distributed::DistributedExec; pub use latency_metric::LatencyMetric; pub use metrics_store::MetricsStore; pub use plan_encoding::{EncodedTaskPlan, encode_task_plan}; -pub(crate) use query_coordinator::{CoordinatorToWorkerMetrics, FlightWorkerDispatch}; +#[cfg(feature = "flight")] +pub(crate) use query_coordinator::FlightWorkerDispatch; diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 2e87f77c..ab26cde5 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -1,7 +1,7 @@ -use crate::common::{TreeNodeExt, now_ns, serialize_uuid, task_ctx_with_extension}; +use crate::common::{TreeNodeExt, serialize_uuid, task_ctx_with_extension}; use crate::config_extension_ext::get_config_extension_propagation_headers; +use crate::coordinator::CoordinatorToWorkerMetrics; use crate::coordinator::MetricsStore; -use crate::coordinator::latency_metric::LatencyMetric; use crate::coordinator::plan_encoding::encode_task_plan; use crate::passthrough_headers::get_passthrough_headers; use crate::protobuf::tonic_status_to_datafusion_error; @@ -10,8 +10,7 @@ use crate::worker::generated::worker as pb; use crate::worker::generated::worker::coordinator_to_worker_msg::Inner; use crate::worker::{WorkerDispatch, WorkerDispatchRequest}; use crate::{ - BytesCounterMetric, BytesMetricExt, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedConfig, - DistributedTaskContext, DistributedWorkUnitFeedContext, TaskKey, + DistributedConfig, DistributedTaskContext, DistributedWorkUnitFeedContext, TaskKey, get_distributed_channel_resolver, }; use datafusion::common::Result; @@ -20,7 +19,6 @@ use datafusion::common::runtime::JoinSet; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{DataFusionError, exec_datafusion_err}; use datafusion::execution::TaskContext; -use datafusion::physical_expr_common::metrics::{ExecutionPlanMetricsSet, Label, MetricBuilder}; use datafusion::physical_plan::ExecutionPlan; use futures::{Stream, StreamExt}; use http::Extensions; @@ -306,29 +304,3 @@ impl<'a> StageCoordinator<'a> { fn keep_stream_alive(notify: Arc) -> impl Stream + 'static { futures::stream::once(notify.notified_owned()).filter_map(|()| futures::future::ready(None)) } - -/// Metrics that measure network details about communications between [DistributedExec] and a worker. -#[derive(Clone)] -pub(crate) struct CoordinatorToWorkerMetrics { - pub(crate) plan_bytes_sent: BytesCounterMetric, - pub(crate) plan_send_latency: Arc, - pub(crate) instantiation_time: u64, -} - -impl CoordinatorToWorkerMetrics { - pub(crate) fn new(metrics: &ExecutionPlanMetricsSet) -> Self { - Self { - // Metric that measures to total sum of bytes worth of subplans sent. - plan_bytes_sent: MetricBuilder::new(metrics) - .with_label(Label::new(DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, "0")) - .bytes_counter("plan_bytes_sent"), - // Latency statistics about the network calls issued to the workers for feeding subplans. - plan_send_latency: Arc::new(LatencyMetric::new( - "plan_send_latency", - |b| b.with_label(Label::new(DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, "0")), - metrics, - )), - instantiation_time: now_ns(), - } - } -} diff --git a/src/distributed_ext.rs b/src/distributed_ext.rs index e44818ad..da53a20a 100644 --- a/src/distributed_ext.rs +++ b/src/distributed_ext.rs @@ -1,17 +1,18 @@ +#[cfg(feature = "flight")] +use crate::ChannelResolver; 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::networking::{ - set_distributed_channel_resolver, set_distributed_worker_resolver, - set_distributed_worker_transport, -}; +#[cfg(feature = "flight")] +use crate::networking::set_distributed_channel_resolver; +use crate::networking::{set_distributed_worker_resolver, set_distributed_worker_transport}; use crate::passthrough_headers::set_passthrough_headers; use crate::protobuf::{set_distributed_user_codec, set_distributed_user_codec_arc}; use crate::work_unit_feed::set_distributed_work_unit_feed; use crate::{ - ChannelResolver, DistributedConfig, TaskEstimator, WorkUnitFeed, WorkUnitFeedProvider, - WorkerResolver, WorkerTransport, + DistributedConfig, TaskEstimator, WorkUnitFeed, WorkUnitFeedProvider, WorkerResolver, + WorkerTransport, }; use arrow_ipc::CompressionType; use datafusion::common::DataFusionError; @@ -239,6 +240,7 @@ pub trait DistributedExt: Sized { /// /// There's a default implementation that caches the Worker client instances so that there's /// only one per URL, but users can decide to override that behavior in favor of their own solution. + #[cfg(feature = "flight")] /// /// Example: /// @@ -284,6 +286,7 @@ pub trait DistributedExt: Sized { ) -> Self; /// Same as [DistributedExt::with_distributed_channel_resolver] but with an in-place mutation. + #[cfg(feature = "flight")] fn set_distributed_channel_resolver( &mut self, resolver: T, @@ -615,6 +618,7 @@ impl DistributedExt for SessionConfig { set_distributed_worker_resolver(self, resolver); } + #[cfg(feature = "flight")] fn set_distributed_channel_resolver( &mut self, resolver: T, @@ -759,6 +763,7 @@ impl DistributedExt for SessionConfig { #[expr($;self)] fn with_distributed_worker_resolver(mut self, resolver: T) -> Self; + #[cfg(feature = "flight")] #[call(set_distributed_channel_resolver)] #[expr($;self)] fn with_distributed_channel_resolver(mut self, resolver: T) -> Self; @@ -855,7 +860,9 @@ impl DistributedExt for SessionStateBuilder { #[expr($;self)] fn with_distributed_worker_resolver(mut self, resolver: T) -> Self; + #[cfg(feature = "flight")] fn set_distributed_channel_resolver(&mut self, resolver: T); + #[cfg(feature = "flight")] #[call(set_distributed_channel_resolver)] #[expr($;self)] fn with_distributed_channel_resolver(mut self, resolver: T) -> Self; @@ -971,7 +978,9 @@ impl DistributedExt for SessionState { #[expr($;self)] fn with_distributed_worker_resolver(mut self, resolver: T) -> Self; + #[cfg(feature = "flight")] fn set_distributed_channel_resolver(&mut self, resolver: T); + #[cfg(feature = "flight")] #[call(set_distributed_channel_resolver)] #[expr($;self)] fn with_distributed_channel_resolver(mut self, resolver: T) -> Self; @@ -1087,7 +1096,9 @@ impl DistributedExt for SessionContext { #[expr($;self)] fn with_distributed_worker_resolver(self, resolver: T) -> Self; + #[cfg(feature = "flight")] fn set_distributed_channel_resolver(&mut self, resolver: T); + #[cfg(feature = "flight")] #[call(set_distributed_channel_resolver)] #[expr($;self)] fn with_distributed_channel_resolver(self, resolver: T) -> Self; diff --git a/src/execution_plans/benchmarks/fixture.rs b/src/execution_plans/benchmarks/fixture.rs index 4423d7d0..4004bd1a 100644 --- a/src/execution_plans/benchmarks/fixture.rs +++ b/src/execution_plans/benchmarks/fixture.rs @@ -1,4 +1,6 @@ +#[cfg(feature = "flight")] use crate::worker::generated::worker::worker_service_client::WorkerServiceClient; +#[cfg(feature = "flight")] use crate::{BoxCloneSyncChannel, ChannelResolver, create_worker_client}; use arrow::datatypes::DataType::{ Boolean, Dictionary, Float64, Int32, Int64, List, Timestamp, UInt8, Utf8, @@ -8,6 +10,7 @@ use arrow::record_batch::RecordBatch; use arrow::util::data_gen::create_random_batch; use datafusion::common::{Result, exec_err}; use std::sync::Arc; +#[cfg(feature = "flight")] use url::Url; pub(super) fn benchmark_schema() -> Arc { @@ -58,6 +61,7 @@ pub(super) fn make_input_partitions( Ok(partitions) } +#[cfg(feature = "flight")] pub(super) fn rows_for_producer( total_rows: usize, producer_tasks: usize, @@ -70,11 +74,13 @@ pub(super) fn rows_for_producer( /// [ChannelResolver] implementation that returns gRPC clients backed by an in-memory /// tokio duplex rather than a TCP connection. +#[cfg(feature = "flight")] #[derive(Clone)] pub(super) struct InMemoryChannelsResolver { pub channels: Vec, } +#[cfg(feature = "flight")] #[async_trait::async_trait] impl ChannelResolver for InMemoryChannelsResolver { async fn get_worker_client_for_url( diff --git a/src/execution_plans/benchmarks/mod.rs b/src/execution_plans/benchmarks/mod.rs index a8753d10..b5021daa 100644 --- a/src/execution_plans/benchmarks/mod.rs +++ b/src/execution_plans/benchmarks/mod.rs @@ -1,10 +1,14 @@ mod fixture; mod local_repartition_bench; +#[cfg(feature = "flight")] mod shuffle_bench; +#[cfg(feature = "flight")] mod transport_bench; pub use local_repartition_bench::{ LocalRepartitionBench, LocalRepartitionFixture, LocalRepartitionMode, }; +#[cfg(feature = "flight")] pub use shuffle_bench::{ShuffleBench, ShuffleFixture}; +#[cfg(feature = "flight")] pub use transport_bench::{TransportBench, TransportBenchMode, TransportFixture}; diff --git a/src/lib.rs b/src/lib.rs index 5c5540e2..a1e41a12 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -25,7 +25,8 @@ pub use common::{ }; pub use config_extension_ext::get_config_extension_propagation_headers; pub use coordinator::{ - DistributedExec, EncodedTaskPlan, LatencyMetric, MetricsStore, encode_task_plan, + CoordinatorToWorkerMetrics, DistributedExec, EncodedTaskPlan, LatencyMetric, MetricsStore, + encode_task_plan, }; pub use distributed_ext::DistributedExt; pub use distributed_planner::{ @@ -42,10 +43,14 @@ pub use metrics::{ MaxLatencyMetric, MinLatencyMetric, P50LatencyMetric, P75LatencyMetric, P95LatencyMetric, P99LatencyMetric, rewrite_distributed_plan_with_metrics, }; +#[cfg(feature = "flight")] pub use networking::{ - BoxCloneSyncChannel, ChannelResolver, DefaultChannelResolver, WorkerResolver, - create_worker_client, get_distributed_channel_resolver, get_distributed_worker_resolver, - get_distributed_worker_transport, set_distributed_worker_transport, + BoxCloneSyncChannel, ChannelResolver, DefaultChannelResolver, create_worker_client, + get_distributed_channel_resolver, +}; +pub use networking::{ + WorkerResolver, get_distributed_worker_resolver, get_distributed_worker_transport, + set_distributed_worker_transport, }; pub use passthrough_headers::get_passthrough_headers; pub use stage::{ @@ -57,25 +62,38 @@ pub use work_unit_feed::{ RemoteWorkUnitFeedTxs, WorkUnit, WorkUnitFeed, WorkUnitFeedProto, WorkUnitFeedProvider, WorkUnitRx, WorkUnitTx, collect_task_work_unit_feeds, set_received_time, set_sent_time, }; +#[cfg(feature = "flight")] +pub use worker::FlightWorkerTransport; +// `protobuf` already names a private module, so the generated worker messages re-export as `proto` +// for an out-of-crate transport to build the same frames the Flight path does. +pub use worker::generated::worker as proto; +#[cfg(feature = "flight")] pub use worker::generated::worker::worker_service_client::WorkerServiceClient; +#[cfg(feature = "flight")] pub use worker::generated::worker::worker_service_server::WorkerServiceServer; pub use worker::generated::worker::{GetWorkerInfoRequest, GetWorkerInfoResponse, TaskKey}; pub use worker::{ - DefaultSessionBuilder, FlightWorkerTransport, InMemoryWorkerTransport, - MappedWorkerSessionBuilder, MappedWorkerSessionBuilderExt, TaskData, Worker, WorkerConnection, - WorkerDispatch, WorkerDispatchRequest, WorkerQueryContext, WorkerSessionBuilder, - WorkerTransport, + DefaultSessionBuilder, InMemoryWorkerTransport, MappedWorkerSessionBuilder, + MappedWorkerSessionBuilderExt, ResultTaskData, SingleWriteMultiRead, TaskData, TaskDataEntries, + Worker, WorkerConnection, WorkerDispatch, WorkerDispatchRequest, WorkerQueryContext, + WorkerSessionBuilder, WorkerTransport, collect_plan_metrics_protos, execute_local_task, }; pub use observability::{ GetClusterWorkersRequest, GetClusterWorkersResponse, GetTaskProgressRequest, - GetTaskProgressResponse, ObservabilityService, ObservabilityServiceClient, - ObservabilityServiceImpl, ObservabilityServiceServer, PingRequest, PingResponse, TaskProgress, - TaskStatus, WorkerMetrics, + GetTaskProgressResponse, PingRequest, PingResponse, TaskProgress, TaskStatus, WorkerMetrics, +}; +#[cfg(feature = "flight")] +pub use observability::{ + ObservabilityService, ObservabilityServiceClient, ObservabilityServiceImpl, + ObservabilityServiceServer, }; #[cfg(any(feature = "integration", test))] pub use execution_plans::benchmarks::{ - LocalRepartitionBench, LocalRepartitionFixture, LocalRepartitionMode, ShuffleBench, - ShuffleFixture, TransportBench, TransportBenchMode, TransportFixture, + LocalRepartitionBench, LocalRepartitionFixture, LocalRepartitionMode, +}; +#[cfg(all(feature = "flight", any(feature = "integration", test)))] +pub use execution_plans::benchmarks::{ + ShuffleBench, ShuffleFixture, TransportBench, TransportBenchMode, TransportFixture, }; diff --git a/src/metrics/task_metrics_collector.rs b/src/metrics/task_metrics_collector.rs index 553db1cc..b4ecca15 100644 --- a/src/metrics/task_metrics_collector.rs +++ b/src/metrics/task_metrics_collector.rs @@ -29,14 +29,13 @@ mod tests { use futures::StreamExt; use crate::coordinator::DistributedExec; - use crate::test_utils::in_memory_channel_resolver::{ - InMemoryChannelResolver, InMemoryWorkerResolver, - }; + use crate::test_utils::in_memory_channel_resolver::InMemoryWorkerResolver; use crate::test_utils::parquet::register_parquet_tables; use crate::test_utils::plans::{ count_plan_nodes_up_to_network_boundary, get_stages_and_task_keys, }; use crate::test_utils::session_context::register_temp_parquet_table; + use crate::worker::InMemoryWorkerTransport; use crate::{DistributedExt, SessionStateBuilderExt}; use datafusion::execution::{SessionStateBuilder, context::SessionContext}; use datafusion::prelude::SessionConfig; @@ -57,7 +56,7 @@ mod tests { .with_default_features() .with_config(config) .with_distributed_worker_resolver(InMemoryWorkerResolver::new(10)) - .with_distributed_channel_resolver(InMemoryChannelResolver::default()) + .with_distributed_worker_transport(InMemoryWorkerTransport::default()) .with_distributed_planner() .with_distributed_task_estimator(2) .with_distributed_metrics_collection(true) diff --git a/src/metrics/task_metrics_rewriter.rs b/src/metrics/task_metrics_rewriter.rs index 3ba5c88c..5c62071c 100644 --- a/src/metrics/task_metrics_rewriter.rs +++ b/src/metrics/task_metrics_rewriter.rs @@ -290,12 +290,11 @@ mod tests { annotate_metrics_set_with_task_id, stage_metrics_rewriter, }; use crate::metrics::{DistributedMetricsFormat, rewrite_distributed_plan_with_metrics}; - use crate::test_utils::in_memory_channel_resolver::{ - InMemoryChannelResolver, InMemoryWorkerResolver, - }; + use crate::test_utils::in_memory_channel_resolver::InMemoryWorkerResolver; use crate::test_utils::metrics::make_test_metrics_set_proto_from_seed; use crate::test_utils::plans::count_plan_nodes_up_to_network_boundary; use crate::test_utils::session_context::register_temp_parquet_table; + use crate::worker::InMemoryWorkerTransport; use crate::worker::generated::worker as pb; use crate::{DistributedExec, SessionStateBuilderExt}; use datafusion::arrow::array::{Int32Array, StringArray}; @@ -339,7 +338,7 @@ mod tests { if distributed { builder = builder .with_distributed_worker_resolver(InMemoryWorkerResolver::new(10)) - .with_distributed_channel_resolver(InMemoryChannelResolver::default()) + .with_distributed_worker_transport(InMemoryWorkerTransport::default()) .with_distributed_metrics_collection(true) .unwrap() .with_distributed_planner() diff --git a/src/networking/channel_resolver.rs b/src/networking/channel_resolver.rs index d15c8bef..9a847897 100644 --- a/src/networking/channel_resolver.rs +++ b/src/networking/channel_resolver.rs @@ -1,18 +1,36 @@ +#[cfg(feature = "flight")] use crate::DistributedConfig; +#[cfg(feature = "flight")] use crate::config_extension_ext::set_distributed_option_extension; +#[cfg(feature = "flight")] use crate::worker::generated::worker::worker_service_client::WorkerServiceClient; +#[cfg(feature = "flight")] use async_trait::async_trait; +#[cfg(feature = "flight")] use datafusion::common::{DataFusionError, config_datafusion_err, exec_datafusion_err}; +#[cfg(feature = "flight")] use datafusion::execution::TaskContext; +#[cfg(feature = "flight")] use datafusion::prelude::SessionConfig; +#[cfg(feature = "flight")] use futures::FutureExt; +#[cfg(feature = "flight")] use futures::future::Shared; -use std::sync::{Arc, LazyLock}; +#[cfg(feature = "flight")] +use std::sync::Arc; +#[cfg(feature = "flight")] +use std::sync::LazyLock; +#[cfg(feature = "flight")] use std::time::Duration; +#[cfg(feature = "flight")] use tonic::body::Body; +#[cfg(feature = "flight")] use tonic::codegen::BoxFuture; +#[cfg(feature = "flight")] use tonic::transport::Channel; +#[cfg(feature = "flight")] use tower::ServiceExt; +#[cfg(feature = "flight")] use url::Url; /// Allows users to customize the way Worker clients are created. A common use case is to @@ -29,6 +47,7 @@ use url::Url; /// [`create_worker_client`] helper function to ensure clients are configured with /// appropriate message size limits for internal communication. This helps avoid message /// size errors when transferring large datasets. +#[cfg(feature = "flight")] #[async_trait] pub trait ChannelResolver { /// For a given URL, get a Worker gRPC client for communicating to it. @@ -46,6 +65,7 @@ pub trait ChannelResolver { ) -> Result, DataFusionError>; } +#[cfg(feature = "flight")] pub(crate) fn set_distributed_channel_resolver( cfg: &mut SessionConfig, channel_resolver: impl ChannelResolver + Send + Sync + 'static, @@ -72,6 +92,7 @@ pub(crate) fn set_distributed_channel_resolver( // The Tonic channels need to be established and reused under a whole RuntimeEnv scope, not a single // TaskContext, which forces us to put the default implementation in a static global variable that // stores and reuses tonic channels per RuntimeEnv's pointer address. +#[cfg(feature = "flight")] static DEFAULT_CHANNEL_RESOLVER_PER_RUNTIME: LazyLock< moka::sync::Cache< /* Arc pointer address */ usize, @@ -79,6 +100,7 @@ static DEFAULT_CHANNEL_RESOLVER_PER_RUNTIME: LazyLock< >, > = LazyLock::new(|| moka::sync::Cache::builder().max_capacity(256).build()); +#[cfg(feature = "flight")] pub fn get_distributed_channel_resolver( task_ctx: &TaskContext, ) -> Arc { @@ -93,27 +115,36 @@ pub fn get_distributed_channel_resolver( .get_with(runtime_addr, || Arc::new(DefaultChannelResolver::default())) } +#[cfg(feature = "flight")] pub type BoxCloneSyncChannel = tower::util::BoxCloneSyncService< http::Request, http::Response, tonic::transport::Error, >; +#[cfg(feature = "flight")] type ChannelCacheValue = Shared>>; +/// Holds the user-provided [ChannelResolver], if any. Always present in [DistributedConfig] (the +/// `extensions_options!` macro can't gate a field), but the inner channel-resolver handle only +/// exists with the `flight` feature; without it the field is empty and nothing dials. #[derive(Clone, Default)] -pub(crate) struct ChannelResolverExtension(Option>); +pub(crate) struct ChannelResolverExtension( + #[cfg(feature = "flight")] Option>, +); /// Default implementation of a [ChannelResolver] that connects to the workers given the URL once /// and stores the connection instance in a TTI cache. /// /// Sane default over which other [ChannelResolver] can be built for better customization of the /// [WorkerServiceClient]s. +#[cfg(feature = "flight")] #[derive(Clone)] pub struct DefaultChannelResolver { cache: Arc>, } +#[cfg(feature = "flight")] impl Default for DefaultChannelResolver { fn default() -> Self { Self { @@ -130,6 +161,7 @@ impl Default for DefaultChannelResolver { } } +#[cfg(feature = "flight")] impl DefaultChannelResolver { /// Gets the cached [BoxCloneSyncChannel] for the given URL, or builds a new one. pub async fn get_channel(&self, url: &Url) -> Result { @@ -170,6 +202,7 @@ impl DefaultChannelResolver { } } +#[cfg(feature = "flight")] #[async_trait] impl ChannelResolver for DefaultChannelResolver { async fn get_worker_client_for_url( @@ -180,6 +213,7 @@ impl ChannelResolver for DefaultChannelResolver { } } +#[cfg(feature = "flight")] #[async_trait] impl ChannelResolver for Arc { async fn get_worker_client_for_url( @@ -217,6 +251,7 @@ impl ChannelResolver for Arc { /// } /// } /// ``` +#[cfg(feature = "flight")] pub fn create_worker_client( channel: BoxCloneSyncChannel, ) -> WorkerServiceClient { @@ -225,7 +260,7 @@ pub fn create_worker_client( .max_encoding_message_size(usize::MAX) } -#[cfg(test)] +#[cfg(all(test, feature = "flight"))] mod tests { use super::*; use crate::Worker; diff --git a/src/networking/mod.rs b/src/networking/mod.rs index 0c2f56ce..1a2c3880 100644 --- a/src/networking/mod.rs +++ b/src/networking/mod.rs @@ -2,11 +2,14 @@ mod channel_resolver; mod worker_resolver; mod worker_transport; +pub(crate) use channel_resolver::ChannelResolverExtension; +#[cfg(feature = "flight")] +pub(crate) use channel_resolver::set_distributed_channel_resolver; +#[cfg(feature = "flight")] pub use channel_resolver::{ BoxCloneSyncChannel, ChannelResolver, DefaultChannelResolver, create_worker_client, get_distributed_channel_resolver, }; -pub(crate) use channel_resolver::{ChannelResolverExtension, set_distributed_channel_resolver}; pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver}; pub(crate) use worker_resolver::{WorkerResolverExtension, set_distributed_worker_resolver}; diff --git a/src/networking/worker_transport.rs b/src/networking/worker_transport.rs index f38550aa..de075ae3 100644 --- a/src/networking/worker_transport.rs +++ b/src/networking/worker_transport.rs @@ -1,9 +1,26 @@ use crate::DistributedConfig; use crate::config_extension_ext::set_distributed_option_extension; -use crate::worker::{FlightWorkerTransport, WorkerTransport}; +#[cfg(feature = "flight")] +use crate::worker::FlightWorkerTransport; +#[cfg(not(feature = "flight"))] +use crate::worker::InMemoryWorkerTransport; +use crate::worker::WorkerTransport; use datafusion::prelude::SessionConfig; use std::sync::Arc; +/// The transport used when none is registered: Arrow-Flight with the `flight` feature on (the +/// default), the in-process transport when it is off so distributed plans still run. +fn default_worker_transport() -> Arc { + #[cfg(feature = "flight")] + { + Arc::new(FlightWorkerTransport) + } + #[cfg(not(feature = "flight"))] + { + Arc::new(InMemoryWorkerTransport::default()) + } +} + pub fn set_distributed_worker_transport( cfg: &mut SessionConfig, worker_transport: impl WorkerTransport + 'static, @@ -32,7 +49,7 @@ pub fn get_distributed_worker_transport(cfg: &SessionConfig) -> Arc() .map(|distributed_cfg| Arc::clone(&distributed_cfg.__private_worker_transport.0)) - .unwrap_or_else(|| Arc::new(FlightWorkerTransport)) + .unwrap_or_else(default_worker_transport) } #[derive(Clone)] @@ -40,6 +57,6 @@ pub(crate) struct WorkerTransportExtension(pub(crate) Arc); impl Default for WorkerTransportExtension { fn default() -> Self { - Self(Arc::new(FlightWorkerTransport)) + Self(default_worker_transport()) } } diff --git a/src/observability/generated/mod.rs b/src/observability/generated/mod.rs index e8cd9143..f40122f1 100644 --- a/src/observability/generated/mod.rs +++ b/src/observability/generated/mod.rs @@ -1,3 +1,4 @@ // This file is @generated by tonic-prost-build +#[cfg_attr(not(feature = "flight"), allow(dead_code))] pub mod observability; diff --git a/src/observability/generated/observability.rs b/src/observability/generated/observability.rs index c8da3372..0aee74f1 100644 --- a/src/observability/generated/observability.rs +++ b/src/observability/generated/observability.rs @@ -71,6 +71,7 @@ impl TaskStatus { } } /// Generated client implementations. +#[cfg(feature = "flight")] pub mod observability_service_client { #![allow( unused_variables, @@ -218,6 +219,7 @@ pub mod observability_service_client { } } /// Generated server implementations. +#[cfg(feature = "flight")] pub mod observability_service_server { #![allow( unused_variables, diff --git a/src/observability/mod.rs b/src/observability/mod.rs index 555c6577..31b0fa0e 100644 --- a/src/observability/mod.rs +++ b/src/observability/mod.rs @@ -1,7 +1,10 @@ mod generated; +#[cfg(feature = "flight")] mod service; +#[cfg(feature = "flight")] pub use generated::observability::observability_service_client::ObservabilityServiceClient; +#[cfg(feature = "flight")] pub use generated::observability::observability_service_server::{ ObservabilityService, ObservabilityServiceServer, }; @@ -10,4 +13,5 @@ pub use generated::observability::{ GetClusterWorkersRequest, GetClusterWorkersResponse, GetTaskProgressRequest, GetTaskProgressResponse, PingRequest, PingResponse, TaskProgress, TaskStatus, WorkerMetrics, }; +#[cfg(feature = "flight")] pub use service::ObservabilityServiceImpl; diff --git a/src/protobuf/mod.rs b/src/protobuf/mod.rs index 47deeb40..9bee040d 100644 --- a/src/protobuf/mod.rs +++ b/src/protobuf/mod.rs @@ -1,9 +1,11 @@ mod distributed_codec; +#[cfg(feature = "flight")] mod errors; mod producer_head; mod user_codec; pub use distributed_codec::DistributedCodec; +#[cfg(feature = "flight")] pub(crate) use errors::{ datafusion_error_to_tonic_status, map_flight_to_datafusion_error, tonic_status_to_datafusion_error, diff --git a/src/test_utils/in_memory_channel_resolver.rs b/src/test_utils/in_memory_channel_resolver.rs index 12ad84bb..e25195a3 100644 --- a/src/test_utils/in_memory_channel_resolver.rs +++ b/src/test_utils/in_memory_channel_resolver.rs @@ -1,27 +1,37 @@ +use crate::worker::InMemoryWorkerTransport; +use crate::{DistributedExt, SessionStateBuilderExt, Worker, WorkerResolver, WorkerSessionBuilder}; +use datafusion::common::DataFusionError; +use datafusion::execution::SessionStateBuilder; +use datafusion::prelude::{SessionConfig, SessionContext}; +use url::Url; + +#[cfg(feature = "flight")] use crate::worker::generated::worker::worker_service_client::WorkerServiceClient; +#[cfg(feature = "flight")] use crate::{ - BoxCloneSyncChannel, ChannelResolver, DefaultSessionBuilder, DistributedExt, - MappedWorkerSessionBuilderExt, SessionStateBuilderExt, Worker, WorkerResolver, - WorkerSessionBuilder, create_worker_client, + BoxCloneSyncChannel, ChannelResolver, DefaultSessionBuilder, MappedWorkerSessionBuilderExt, + create_worker_client, }; +#[cfg(feature = "flight")] use async_trait::async_trait; -use datafusion::common::DataFusionError; -use datafusion::execution::SessionStateBuilder; -use datafusion::prelude::{SessionConfig, SessionContext}; +#[cfg(feature = "flight")] use hyper_util::rt::TokioIo; +#[cfg(feature = "flight")] use tonic::transport::{Endpoint, Server}; -use url::Url; +#[cfg(feature = "flight")] const DUMMY_URL: &str = "http://localhost:50051"; const DUMMY_URL_PREFIX: &str = "http://url-"; /// [ChannelResolver] implementation that returns gRPC clients backed by an in-memory /// tokio duplex rather than a TCP connection. +#[cfg(feature = "flight")] #[derive(Clone)] pub struct InMemoryChannelResolver { channel: WorkerServiceClient, } +#[cfg(feature = "flight")] impl InMemoryChannelResolver { /// Build an [InMemoryChannelResolver] with a custom [WorkerSessionBuilder]. /// This allows you to inject your own DataFusion extensions in the in-memory worker @@ -72,12 +82,14 @@ impl InMemoryChannelResolver { } } +#[cfg(feature = "flight")] impl Default for InMemoryChannelResolver { fn default() -> Self { Self::from_session_builder(DefaultSessionBuilder) } } +#[cfg(feature = "flight")] #[async_trait] impl ChannelResolver for InMemoryChannelResolver { async fn get_worker_client_for_url( @@ -107,8 +119,11 @@ impl WorkerResolver for InMemoryWorkerResolver { } } -/// Creates a distributed session context backed by a single in-memory worker service. -/// The set of produced worker URLs is deterministic, taking the form http://worker-. +#[cfg(feature = "flight")] +/// Creates a distributed session context backed by a single in-memory gRPC worker service. The +/// produced worker URLs are deterministic, taking the form http://url-; routing tests that emit +/// and assert per-URL worker identity need the distinct dialed workers this gives, which the single +/// in-process worker can't represent. pub async fn start_in_memory_context( num_workers: usize, session_builder: impl WorkerSessionBuilder + Send + Sync + 'static, @@ -123,7 +138,7 @@ pub async fn start_in_memory_context( SessionContext::from(state) } -/// Creates a distributed session context backed by a configurable in-memory worker service. +/// Creates a distributed session context backed by a configurable in-process worker. /// /// Like [crate::test_utils::localhost::start_localhost_context], this uses tiny file-scan /// partitions so small test datasets still cross worker boundaries. @@ -132,14 +147,14 @@ pub async fn start_configured_in_memory_context( session_builder: impl WorkerSessionBuilder + Send + Sync + 'static, configure_worker: impl Fn(Worker) -> Worker + Send + Sync + 'static, ) -> SessionContext { - let channel_resolver = - InMemoryChannelResolver::from_configured_worker(session_builder, configure_worker); + let transport = + InMemoryWorkerTransport::from_configured_worker(session_builder, configure_worker); let state = SessionStateBuilder::new() .with_default_features() .with_config(SessionConfig::new().with_target_partitions(num_workers)) .with_distributed_planner() .with_distributed_worker_resolver(InMemoryWorkerResolver::new(num_workers)) - .with_distributed_channel_resolver(channel_resolver) + .with_distributed_worker_transport(transport) .with_distributed_file_scan_config_bytes_per_partition(1) .unwrap() .build(); diff --git a/src/test_utils/localhost.rs b/src/test_utils/localhost.rs index 80682952..4eef1644 100644 --- a/src/test_utils/localhost.rs +++ b/src/test_utils/localhost.rs @@ -1,23 +1,37 @@ -use crate::{DistributedExt, SessionStateBuilderExt, Worker, WorkerResolver, WorkerSessionBuilder}; +#[cfg(feature = "flight")] +use crate::WorkerResolver; +use crate::test_utils::in_memory_channel_resolver::InMemoryWorkerResolver; +use crate::worker::InMemoryWorkerTransport; +use crate::{DistributedExt, SessionStateBuilderExt, Worker, WorkerSessionBuilder}; +#[cfg(feature = "flight")] use async_trait::async_trait; +#[cfg(feature = "flight")] use datafusion::common::DataFusionError; use datafusion::common::runtime::JoinSet; use datafusion::execution::SessionStateBuilder; -use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion::prelude::SessionContext; +#[cfg(feature = "flight")] use std::error::Error; +#[cfg(feature = "flight")] use std::time::Duration; +#[cfg(feature = "flight")] use tokio::net::TcpListener; +#[cfg(feature = "flight")] use tonic::transport::Server; +#[cfg(feature = "flight")] use url::Url; -/// Create workers and context on localhost with a fixed number of target partitions. +/// Create workers and context on localhost with a fixed number of target partitions, behind the +/// Arrow-Flight gRPC transport. For flight-specific tests (network metrics, URL routing); the +/// generic suite runs through [start_localhost_context] instead. /// /// Creates `num_workers` listeners, all bound to a random OS decided port on `127.0.0.1`, then /// attaches a channel resolver that is aware of these addresses to `session_builder` and uses it /// to spawn a flight service behind each listener. /// /// Returns a session context aware of these workers, and a join set of all spawned worker tasks. -pub async fn start_localhost_context( +#[cfg(feature = "flight")] +pub async fn start_localhost_flight_context( num_workers: usize, session_builder: B, ) -> (SessionContext, JoinSet<()>, Vec) @@ -63,26 +77,66 @@ where tokio::time::sleep(Duration::from_millis(100)).await; let worker_resolver = LocalHostWorkerResolver::new(ports); - let state = SessionStateBuilder::new() + let mut state = SessionStateBuilder::new() .with_default_features() - .with_config(SessionConfig::new().with_target_partitions(3)) .with_distributed_planner() .with_distributed_worker_resolver(worker_resolver) // Test datasets are tiny, so budget one byte per partition: the estimator then asks for far - // more partitions than exist, which gets capped at the worker count, fanning every scan out - // across the whole (small) test cluster so the distributed paths are exercised. + // more partitions than exist, capped at the worker count, fanning every scan across the + // whole (small) test cluster so the distributed paths are exercised. .with_distributed_file_scan_config_bytes_per_partition(1) .unwrap() .build(); + state.config_mut().options_mut().execution.target_partitions = 3; (SessionContext::from(state), join_set, workers) } +/// Workers and context with a fixed number of target partitions, hosted in-process by an +/// [InMemoryWorkerTransport] built from `session_builder`. Plans are delivered with a function call +/// and partitions are read straight from the local registry, so this is what the integration suite +/// exercises by default in both build configurations. Nothing listens on localhost; the returned +/// [Worker]s are handles onto the shared in-process task registry. +pub async fn start_localhost_context( + num_workers: usize, + session_builder: B, +) -> (SessionContext, JoinSet<()>, Vec) +where + B: WorkerSessionBuilder + Send + Sync + 'static, + B: Clone, +{ + // CI runs this same suite over Flight as a separate job, so both transports get full coverage. + #[cfg(feature = "flight")] + if std::env::var("DATAFUSION_DISTRIBUTED_TEST_TRANSPORT").as_deref() == Ok("flight") { + return start_localhost_flight_context(num_workers, session_builder).await; + } + let transport = InMemoryWorkerTransport::from_session_builder(session_builder); + let workers = (0..num_workers) + .map(|_| transport.worker().clone()) + .collect(); + + let mut state = SessionStateBuilder::new() + .with_default_features() + .with_distributed_planner() + .with_distributed_worker_resolver(InMemoryWorkerResolver::new(num_workers)) + .with_distributed_worker_transport(transport) + // Tiny test datasets: budget one byte per partition so the estimator fans scans out across + // the whole (small) cluster and the distributed paths actually run. + .with_distributed_file_scan_config_bytes_per_partition(1) + .unwrap() + .build(); + state.config_mut().options_mut().execution.target_partitions = 3; + + (SessionContext::from(state), JoinSet::new(), workers) +} + +#[cfg(feature = "flight")] #[derive(Clone)] pub struct LocalHostWorkerResolver { ports: Vec, } +#[cfg(feature = "flight")] impl LocalHostWorkerResolver { pub fn new, I: IntoIterator>(ports: I) -> Self where @@ -94,6 +148,7 @@ impl LocalHostWorkerResolver { } } +#[cfg(feature = "flight")] #[async_trait] impl WorkerResolver for LocalHostWorkerResolver { fn get_urls(&self) -> Result, DataFusionError> { @@ -105,6 +160,7 @@ impl WorkerResolver for LocalHostWorkerResolver { } } +#[cfg(feature = "flight")] pub async fn spawn_worker_service( session_builder: impl WorkerSessionBuilder + Send + Sync + 'static, incoming: TcpListener, @@ -119,6 +175,7 @@ pub async fn spawn_worker_service( .await?) } +#[cfg(feature = "flight")] fn external_err(err: impl Error + Send + Sync + 'static) -> DataFusionError { DataFusionError::External(Box::new(err)) } diff --git a/src/test_utils/routing.rs b/src/test_utils/routing.rs index 26dc6c10..6f387e4d 100644 --- a/src/test_utils/routing.rs +++ b/src/test_utils/routing.rs @@ -2,6 +2,7 @@ use arrow::{ array::{Int64Array, RecordBatch, StringArray}, datatypes::{DataType, Field, Schema, SchemaRef}, }; +use async_trait::async_trait; use datafusion::{ catalog::{Session, TableFunctionImpl, TableProvider}, common::{ScalarValue, Statistics, internal_err, plan_err}, @@ -19,10 +20,10 @@ use datafusion_proto::{physical_plan::PhysicalExtensionCodec, protobuf::proto_er use futures::stream; use prost::Message; use std::{fmt::Formatter, sync::Arc}; -use tonic::async_trait; use url::Url; use crate::execution_plans::DistributedLeafExec; +#[cfg(feature = "flight")] use crate::worker::LocalWorkerContext; use crate::{ DistributedConfig, DistributedTaskContext, TaskEstimation, TaskEstimator, WorkerResolver, @@ -201,17 +202,24 @@ impl ExecutionPlan for URLEmitterExec { return EmptyExec::new(schema).execute(0, context); } let distributed_ctx = DistributedTaskContext::from_ctx(&context); - let local_worker_ctx = context + // The worker URL each task lands on only exists with the Flight transport; the in-process + // transport runs every task in this process, so it has no per-task self URL to emit. + #[cfg(feature = "flight")] + let self_url = context .session_config() .get_extension::() - .expect("URLEmitterExec requires LocalWorkerContext during distributed execution"); + .expect("URLEmitterExec requires LocalWorkerContext during distributed execution") + .self_url + .as_str() + .trim_end_matches('/') + .to_string(); + #[cfg(not(feature = "flight"))] + let self_url = String::new(); let mut columns: Vec> = vec![ Arc::new(Int64Array::from(vec![distributed_ctx.task_count as i64])), Arc::new(Int64Array::from(vec![distributed_ctx.task_index as i64])), Arc::new(StringArray::from(vec![self.tag.as_str()])), - Arc::new(StringArray::from(vec![ - local_worker_ctx.self_url.as_str().trim_end_matches('/'), - ])), + Arc::new(StringArray::from(vec![self_url.as_str()])), ]; if let Some(indices) = &self.projection { columns = indices.iter().map(|&i| Arc::clone(&columns[i])).collect(); diff --git a/src/work_unit_feed/mod.rs b/src/work_unit_feed/mod.rs index 2e24fcfa..078d3a8d 100644 --- a/src/work_unit_feed/mod.rs +++ b/src/work_unit_feed/mod.rs @@ -11,6 +11,7 @@ pub use remote_work_unit_feed::{ RemoteWorkUnitFeedRegistry, RemoteWorkUnitFeedRxs, RemoteWorkUnitFeedTxs, WorkUnitRx, WorkUnitTx, set_received_time, set_sent_time, }; +#[cfg(feature = "flight")] pub(crate) use remote_work_unit_feed::{ build_work_unit_batch_msg, set_work_unit_received_time, set_work_unit_send_time, }; diff --git a/src/work_unit_feed/remote_work_unit_feed.rs b/src/work_unit_feed/remote_work_unit_feed.rs index d0d61bb8..0f1e8b43 100644 --- a/src/work_unit_feed/remote_work_unit_feed.rs +++ b/src/work_unit_feed/remote_work_unit_feed.rs @@ -65,6 +65,7 @@ pub(crate) fn build_work_unit( } } +#[cfg(feature = "flight")] pub(crate) fn build_work_unit_batch_msg( id: &Uuid, work_unit_batch: Vec<(usize, Result>)>, @@ -92,6 +93,7 @@ pub fn set_received_time(work_unit: &mut pb::WorkUnit) { work_unit.received_timestamp_unix_nanos = now_ns(); } +#[cfg(feature = "flight")] pub(crate) fn set_work_unit_send_time( mut msg: pb::CoordinatorToWorkerMsg, ) -> pb::CoordinatorToWorkerMsg { @@ -106,6 +108,7 @@ pub(crate) fn set_work_unit_send_time( msg } +#[cfg(feature = "flight")] pub(crate) fn set_work_unit_received_time( mut msg: pb::CoordinatorToWorkerMsg, ) -> pb::CoordinatorToWorkerMsg { diff --git a/src/worker/generated/mod.rs b/src/worker/generated/mod.rs index 844c269c..e3ebdfcb 100644 --- a/src/worker/generated/mod.rs +++ b/src/worker/generated/mod.rs @@ -1 +1,2 @@ -pub(crate) mod worker; +#[cfg_attr(not(feature = "flight"), allow(dead_code))] +pub mod worker; diff --git a/src/worker/generated/worker.rs b/src/worker/generated/worker.rs index fe7a0137..e43ebd59 100644 --- a/src/worker/generated/worker.rs +++ b/src/worker/generated/worker.rs @@ -434,6 +434,7 @@ pub struct MaxGauge { pub value: u64, } /// Generated client implementations. +#[cfg(feature = "flight")] pub mod worker_service_client { #![allow( unused_variables, @@ -583,6 +584,7 @@ pub mod worker_service_client { } } /// Generated server implementations. +#[cfg(feature = "flight")] pub mod worker_service_server { #![allow( unused_variables, diff --git a/src/worker/impl_execute_task.rs b/src/worker/impl_execute_task.rs index 9be4e51b..35248472 100644 --- a/src/worker/impl_execute_task.rs +++ b/src/worker/impl_execute_task.rs @@ -1,38 +1,61 @@ -use crate::common::{TreeNodeExt, now_ns, on_drop_stream}; +use crate::common::{TreeNodeExt, on_drop_stream}; use crate::metrics::proto::df_metrics_set_to_proto; -use crate::protobuf::datafusion_error_to_tonic_status; -use crate::worker::generated::worker::{FlightAppMetadata, TaskMetrics}; -use crate::worker::worker_service::{TaskDataEntries, Worker}; +use crate::worker::generated::worker::TaskMetrics; +use crate::worker::worker_service::TaskDataEntries; use crate::{DistributedConfig, DistributedTaskContext}; -use arrow_flight::encode::{DictionaryHandling, FlightDataEncoder, FlightDataEncoderBuilder}; -use arrow_flight::error::FlightError; -use arrow_select::dictionary::garbage_collect_any_dictionary; -use datafusion::arrow::array::{Array, AsArray, RecordBatch, RecordBatchOptions}; +#[cfg(feature = "flight")] +use datafusion::arrow::array::RecordBatch; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{Result, exec_err, internal_err}; use crate::worker::generated::worker::ExecuteTaskRequest; -use crate::worker::generated::worker::worker_service_server::WorkerService; -use crate::worker::spawn_select_all::spawn_select_all; use crate::worker::task_data::TaskDataMetrics; -use datafusion::arrow::ipc::CompressionType; -use datafusion::arrow::ipc::writer::IpcWriteOptions; use datafusion::common::exec_datafusion_err; use datafusion::error::DataFusionError; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; -use futures::TryStreamExt; -use prost::Message; use std::sync::Arc; use std::sync::Mutex; use std::sync::atomic::Ordering; use std::time::Duration; use tokio::sync::oneshot::Sender; + +#[cfg(feature = "flight")] +use crate::common::now_ns; +#[cfg(feature = "flight")] +use crate::protobuf::datafusion_error_to_tonic_status; +#[cfg(feature = "flight")] +use crate::worker::generated::worker::FlightAppMetadata; +#[cfg(feature = "flight")] +use crate::worker::generated::worker::worker_service_server::WorkerService; +#[cfg(feature = "flight")] +use crate::worker::spawn_select_all::spawn_select_all; +#[cfg(feature = "flight")] +use crate::worker::worker_service::Worker; +#[cfg(feature = "flight")] +use arrow_flight::encode::{DictionaryHandling, FlightDataEncoder, FlightDataEncoderBuilder}; +#[cfg(feature = "flight")] +use arrow_flight::error::FlightError; +#[cfg(feature = "flight")] +use arrow_select::dictionary::garbage_collect_any_dictionary; +#[cfg(feature = "flight")] +use datafusion::arrow::array::{Array, AsArray, RecordBatchOptions}; +#[cfg(feature = "flight")] +use datafusion::arrow::ipc::CompressionType; +#[cfg(feature = "flight")] +use datafusion::arrow::ipc::writer::IpcWriteOptions; +#[cfg(feature = "flight")] +use futures::TryStreamExt; +#[cfg(feature = "flight")] +use prost::Message; +#[cfg(feature = "flight")] use tokio_stream::StreamExt; +#[cfg(feature = "flight")] use tonic::{Request, Response, Status}; /// How many record batches to buffer from the plan execution. +#[cfg(feature = "flight")] const RECORD_BATCH_BUFFER_SIZE: usize = 2; const WAIT_PLAN_TIMEOUT_SECS: u64 = 10; @@ -41,7 +64,7 @@ const WAIT_PLAN_TIMEOUT_SECS: u64 = 10; /// /// This method is async mainly for the key retrieval operation from [TaskDataEntries], but it does /// not start polling any stream, it just instantiates them. -pub(crate) async fn execute_local_task( +pub async fn execute_local_task( task_data_entries: &Arc, body: ExecuteTaskRequest, ) -> Result<(Vec, Arc)> { @@ -120,6 +143,7 @@ pub(crate) async fn execute_local_task( /// /// This method eagerly starts streaming data from the task, and communicates via channels the /// produced [RecordBatch]s already encoded as Arrow Flight data. +#[cfg(feature = "flight")] pub(crate) async fn execute_remote_task( task_data_entries: &Arc, request: Request, @@ -170,6 +194,7 @@ pub(crate) async fn execute_remote_task( })))) } +#[cfg(feature = "flight")] fn build_flight_data_stream( stream: SendableRecordBatchStream, compression_type: Option, @@ -208,14 +233,13 @@ fn build_flight_data_stream( Ok(stream) } -/// Collects metrics from the plan in pre-order traversal order and sends them via the -/// coordinator channel oneshot. -fn send_metrics_via_channel( - metrics_tx: &Arc>>>, +/// Collects the per-node metrics of `plan` in pre-order traversal order, encoded as protos. Shared +/// by the gRPC metrics oneshot and a push transport's metrics frame, so both report the same +/// per-node shape. +pub fn collect_plan_metrics_protos( plan: &Arc, dt_ctx: DistributedTaskContext, - task_data_metrics: &Arc, -) { +) -> Vec { let mut pre_order_plan_metrics = vec![]; let _ = plan.apply_with_dt_ctx(dt_ctx, |node, _| { pre_order_plan_metrics.push( @@ -225,6 +249,18 @@ fn send_metrics_via_channel( ); Ok(TreeNodeRecursion::Continue) }); + pre_order_plan_metrics +} + +/// Collects metrics from the plan in pre-order traversal order and sends them via the +/// coordinator channel oneshot. +fn send_metrics_via_channel( + metrics_tx: &Arc>>>, + plan: &Arc, + dt_ctx: DistributedTaskContext, + task_data_metrics: &Arc, +) { + let pre_order_plan_metrics = collect_plan_metrics_protos(plan, dt_ctx); let tx = { let mut guard = match metrics_tx.lock() { @@ -248,6 +284,7 @@ fn send_metrics_via_channel( /// /// Unused values can arise from operations such as filtering, where /// some keys may no longer be referenced in the filtered result. +#[cfg(feature = "flight")] fn garbage_collect_arrays(batch: RecordBatch) -> Result { let (schema, arrays, row_count) = batch.into_parts(); diff --git a/src/worker/in_memory.rs b/src/worker/in_memory.rs index f50106eb..c1ce0fc4 100644 --- a/src/worker/in_memory.rs +++ b/src/worker/in_memory.rs @@ -62,6 +62,18 @@ impl InMemoryWorkerTransport { } } + /// Builds the transport from a custom [WorkerSessionBuilder], then runs `configure_worker` over + /// the hosted [Worker], so a test can install worker-level hooks the session builder can't reach + /// (e.g. `add_on_plan_hook`). + pub fn from_configured_worker( + session_builder: impl WorkerSessionBuilder + Send + Sync + 'static, + configure_worker: impl FnOnce(Worker) -> Worker, + ) -> Self { + Self { + worker: configure_worker(Worker::from_session_builder(session_builder)), + } + } + /// The in-process [Worker] backing this transport. pub fn worker(&self) -> &Worker { &self.worker @@ -399,4 +411,22 @@ mod tests { assert_eq!(distributed, expected); Ok(()) } + + // With `flight` compiled out, no transport is registered here: the query must run through the + // process-wide default, which is the in-memory transport. + #[cfg(not(feature = "flight"))] + #[tokio::test] + async fn no_flight_default_runs_distributed_queries() -> Result<()> { + let ctx = distributed_ctx(None); + register_temp_parquet_table("t", sample_batch().schema(), vec![sample_batch()], &ctx) + .await?; + + let (display, results) = run(&ctx).await?; + assert!( + display.contains("NetworkShuffleExec"), + "the query did not distribute:\n{display}" + ); + assert!(results.contains("tag0")); + Ok(()) + } } diff --git a/src/worker/mod.rs b/src/worker/mod.rs index 347683db..dbec1ba9 100644 --- a/src/worker/mod.rs +++ b/src/worker/mod.rs @@ -1,11 +1,14 @@ +#[cfg(feature = "flight")] mod flight; -pub(crate) mod generated; +pub mod generated; +#[cfg(feature = "flight")] mod impl_coordinator_channel; mod impl_execute_task; mod impl_set_plan; mod in_memory; mod session_builder; mod single_write_multi_read; +#[cfg(feature = "flight")] mod spawn_select_all; mod task_data; #[cfg(any(test, feature = "integration"))] @@ -14,9 +17,14 @@ mod transport; mod worker_connection_pool; mod worker_service; -pub(crate) use single_write_multi_read::SingleWriteMultiRead; -pub(crate) use worker_connection_pool::{LocalWorkerContext, WorkerConnectionPool}; +// Surface an out-of-crate transport executes fragments in-process through. +pub use impl_execute_task::{collect_plan_metrics_protos, execute_local_task}; +pub use single_write_multi_read::SingleWriteMultiRead; +#[cfg(feature = "flight")] +pub(crate) use worker_connection_pool::LocalWorkerContext; +pub(crate) use worker_connection_pool::WorkerConnectionPool; +#[cfg(feature = "flight")] pub use flight::FlightWorkerTransport; pub use in_memory::InMemoryWorkerTransport; pub use transport::{WorkerConnection, WorkerDispatch, WorkerDispatchRequest, WorkerTransport}; @@ -26,4 +34,4 @@ pub use session_builder::{ WorkerQueryContext, WorkerSessionBuilder, }; pub use task_data::TaskData; -pub use worker_service::Worker; +pub use worker_service::{ResultTaskData, TaskDataEntries, Worker}; diff --git a/src/worker/single_write_multi_read.rs b/src/worker/single_write_multi_read.rs index b9632352..622c2bb9 100644 --- a/src/worker/single_write_multi_read.rs +++ b/src/worker/single_write_multi_read.rs @@ -63,6 +63,7 @@ impl SingleWriteMultiRead { } /// Reads the current value, if any, not waiting for it to be set by a writer. + #[cfg(feature = "flight")] pub(crate) fn read_now(&self) -> Option { self.rx.borrow().clone() } diff --git a/src/worker/task_data.rs b/src/worker/task_data.rs index 28f5ca5d..b1d61b59 100644 --- a/src/worker/task_data.rs +++ b/src/worker/task_data.rs @@ -110,11 +110,13 @@ impl TaskDataMetrics { impl TaskData { /// Returns the number of partitions remaining to be processed. + #[cfg(feature = "flight")] pub(crate) fn num_partitions_remaining(&self) -> usize { self.num_partitions_remaining.load(Ordering::SeqCst) } /// Returns the total number of partitions in this task. + #[cfg(feature = "flight")] pub(crate) fn total_partitions(&self) -> usize { match self.final_plan.get() { Some(Ok(plan)) => plan.output_partitioning().partition_count(), diff --git a/src/worker/test_utils/mod.rs b/src/worker/test_utils/mod.rs index ff232c92..6b6cb112 100644 --- a/src/worker/test_utils/mod.rs +++ b/src/worker/test_utils/mod.rs @@ -1 +1,2 @@ +#[cfg(feature = "flight")] pub(crate) mod worker_handles; diff --git a/src/worker/worker_connection_pool.rs b/src/worker/worker_connection_pool.rs index 78313844..bc166721 100644 --- a/src/worker/worker_connection_pool.rs +++ b/src/worker/worker_connection_pool.rs @@ -1,49 +1,92 @@ -use crate::common::{OnceLockResult, on_drop_stream, serialize_uuid}; +use crate::common::OnceLockResult; use crate::distributed_planner::ProducerHead; +use crate::networking::get_distributed_worker_transport; +use crate::stage::RemoteStage; +use crate::worker::transport::WorkerConnection; +use datafusion::common::{DataFusionError, Result, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr_common::metrics::ExecutionPlanMetricsSet; +use std::fmt::{Debug, Formatter}; +use std::ops::Range; +use std::sync::{Arc, OnceLock}; + +#[cfg(feature = "flight")] +use crate::common::{on_drop_stream, serialize_uuid}; +#[cfg(feature = "flight")] use crate::metrics::LatencyMetricExt; -use crate::networking::{get_distributed_channel_resolver, get_distributed_worker_transport}; +#[cfg(feature = "flight")] +use crate::networking::get_distributed_channel_resolver; +#[cfg(feature = "flight")] use crate::passthrough_headers::get_passthrough_headers; +#[cfg(feature = "flight")] use crate::protobuf::{datafusion_error_to_tonic_status, map_flight_to_datafusion_error}; -use crate::stage::RemoteStage; +#[cfg(feature = "flight")] use crate::worker::generated::worker::FlightAppMetadata; +#[cfg(feature = "flight")] use crate::worker::generated::worker::{ExecuteTaskRequest, TaskKey}; +#[cfg(feature = "flight")] use crate::worker::impl_execute_task::execute_local_task; -use crate::worker::transport::WorkerConnection; +#[cfg(feature = "flight")] use crate::worker::worker_service::TaskDataEntries; +#[cfg(feature = "flight")] use crate::{BytesMetricExt, ChannelResolver, DistributedConfig}; +#[cfg(feature = "flight")] use arrow_flight::FlightData; +#[cfg(feature = "flight")] use arrow_flight::decode::FlightRecordBatchStream; +#[cfg(feature = "flight")] use arrow_flight::error::FlightError; +#[cfg(feature = "flight")] use dashmap::DashMap; +#[cfg(feature = "flight")] use datafusion::arrow::array::RecordBatch; +#[cfg(feature = "flight")] use datafusion::common::instant::Instant; +#[cfg(feature = "flight")] use datafusion::common::runtime::SpawnedTask; -use datafusion::common::{ - DataFusionError, Result, exec_err, internal_datafusion_err, internal_err, -}; -use datafusion::execution::TaskContext; +#[cfg(feature = "flight")] +use datafusion::common::{exec_err, internal_datafusion_err}; +#[cfg(feature = "flight")] use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; -use datafusion::physical_expr_common::metrics::{ExecutionPlanMetricsSet, MetricValue}; +#[cfg(feature = "flight")] +use datafusion::physical_expr_common::metrics::MetricValue; +#[cfg(feature = "flight")] use datafusion::physical_plan::metrics::{MetricBuilder, Time}; +#[cfg(feature = "flight")] use futures::stream::BoxStream; +#[cfg(feature = "flight")] use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt}; +#[cfg(feature = "flight")] use http::Extensions; +#[cfg(feature = "flight")] use pin_project::{pin_project, pinned_drop}; +#[cfg(feature = "flight")] use prost::Message; +#[cfg(feature = "flight")] use std::borrow::Cow; -use std::fmt::{Debug, Formatter}; -use std::ops::Range; +#[cfg(feature = "flight")] use std::pin::Pin; +#[cfg(feature = "flight")] +use std::sync::Mutex; +#[cfg(feature = "flight")] use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::{Arc, Mutex, OnceLock}; +#[cfg(feature = "flight")] use std::task::{Context, Poll}; +#[cfg(feature = "flight")] use std::time::{Duration, SystemTime, UNIX_EPOCH}; +#[cfg(feature = "flight")] use tokio::sync::Notify; +#[cfg(feature = "flight")] use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; +#[cfg(feature = "flight")] use tokio_stream::wrappers::UnboundedReceiverStream; +#[cfg(feature = "flight")] use tokio_util::sync::CancellationToken; +#[cfg(feature = "flight")] use tonic::metadata::MetadataMap; +#[cfg(feature = "flight")] use tonic::{Request, Status}; +#[cfg(feature = "flight")] use url::Url; /// Context set by [crate::Worker::coordinator_channel] in DataFusion's @@ -52,6 +95,7 @@ use url::Url; /// /// This information can be used for executing tasks locally bypassing gRPC comms if the tasks that /// needs to be remotely executed happens to be owned by this same worker. +#[cfg(feature = "flight")] pub(crate) struct LocalWorkerContext { /// The registry of in-flight tasks the [crate::Worker] in the current scope owns. pub(crate) task_data_entries: Arc, @@ -124,6 +168,7 @@ impl WorkerConnectionPool { } } +#[cfg(feature = "flight")] type WorkerMsg = Result<(FlightData, FlightAppMetadata), Status>; /// Represents a connection to one [Worker]. Network boundaries will use this for streaming @@ -137,6 +182,7 @@ type WorkerMsg = Result<(FlightData, FlightAppMetadata), Status>; /// the same underlying TCP connection, there do is some overhead in having one gRPC stream per /// partition VS a single gRPC stream interleaving multiple partitions. The whole serialized plan /// needs to be sent over the wire on every gRPC call, so the less gRPC calls we do the better. +#[cfg(feature = "flight")] pub(crate) struct RemoteWorkerConnection { task: Arc>, not_consumed_streams: Arc, @@ -152,6 +198,7 @@ pub(crate) struct RemoteWorkerConnection { elapsed_compute: Time, } +#[cfg(feature = "flight")] impl RemoteWorkerConnection { pub(crate) fn init( input_stage: &RemoteStage, @@ -373,6 +420,7 @@ impl RemoteWorkerConnection { } } +#[cfg(feature = "flight")] impl WorkerConnection for RemoteWorkerConnection { /// Streams the provided `partition` from the remote worker. /// @@ -431,11 +479,13 @@ impl WorkerConnection for RemoteWorkerConnection { /// Equivalent to [RemoteWorkerConnection], but that pulls data from the local registry of tasks /// rather than doing it across a gRPC interface. +#[cfg(feature = "flight")] pub(crate) struct LocalWorkerConnection { partition_start: usize, local_streams: Vec>>>>, } +#[cfg(feature = "flight")] impl LocalWorkerConnection { pub(crate) fn init( input_stage: &RemoteStage, @@ -503,6 +553,7 @@ impl LocalWorkerConnection { } } +#[cfg(feature = "flight")] impl WorkerConnection for LocalWorkerConnection { fn execute(&self, partition: usize) -> Result>> { let Some(relative_i) = partition.checked_sub(self.partition_start) else { @@ -524,6 +575,7 @@ impl WorkerConnection for LocalWorkerConnection { } } +#[cfg(feature = "flight")] fn fanout(o_txs: &[UnboundedSender], err: Status) { for o_tx in o_txs { let _ = o_tx.send(Err(err.clone())); @@ -544,20 +596,24 @@ impl Clone for WorkerConnectionPool { } } +#[cfg(feature = "flight")] impl Debug for RemoteWorkerConnection { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("WorkerConnection").finish() } } +#[cfg(feature = "flight")] trait ElapsedComputeFutureExt: Future + Sized { fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeFuture; } +#[cfg(feature = "flight")] trait ElapsedComputeStreamExt: Stream + Sized { fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeStream; } +#[cfg(feature = "flight")] impl> ElapsedComputeFutureExt for F { fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeFuture { ElapsedComputeFuture { @@ -568,6 +624,7 @@ impl> ElapsedComputeFutureExt for F { } } +#[cfg(feature = "flight")] impl> ElapsedComputeStreamExt for S { fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeStream { ElapsedComputeStream { @@ -578,6 +635,7 @@ impl> ElapsedComputeStreamExt for S { } } +#[cfg(feature = "flight")] #[pin_project(PinnedDrop)] struct ElapsedComputeStream { #[pin] @@ -588,6 +646,7 @@ struct ElapsedComputeStream { /// Drop implementation that ensures that any accumulated time is properly dumped to the metric /// in case the stream gets dropped before completion. +#[cfg(feature = "flight")] #[pinned_drop] impl PinnedDrop for ElapsedComputeStream { fn drop(self: Pin<&mut Self>) { @@ -600,6 +659,7 @@ impl PinnedDrop for ElapsedComputeStream { } } +#[cfg(feature = "flight")] impl> Stream for ElapsedComputeStream { type Item = O; @@ -618,6 +678,7 @@ impl> Stream for ElapsedComputeStream { } } +#[cfg(feature = "flight")] #[pin_project(PinnedDrop)] struct ElapsedComputeFuture { #[pin] @@ -628,6 +689,7 @@ struct ElapsedComputeFuture { /// Drop implementation that ensures that any accumulated time is properly dumped to the metric /// in case the future gets dropped before completion. +#[cfg(feature = "flight")] #[pinned_drop] impl PinnedDrop for ElapsedComputeFuture { fn drop(self: Pin<&mut Self>) { @@ -640,6 +702,7 @@ impl PinnedDrop for ElapsedComputeFuture { } } +#[cfg(feature = "flight")] impl> Future for ElapsedComputeFuture { type Output = O; @@ -658,7 +721,7 @@ impl> Future for ElapsedComputeFuture { } } -#[cfg(test)] +#[cfg(all(test, feature = "flight"))] mod tests { use super::*; use futures::StreamExt; diff --git a/src/worker/worker_service.rs b/src/worker/worker_service.rs index 5d9c375f..253dd7ce 100644 --- a/src/worker/worker_service.rs +++ b/src/worker/worker_service.rs @@ -1,17 +1,8 @@ +use crate::DefaultSessionBuilder; use crate::worker::WorkerSessionBuilder; -use crate::worker::generated::worker::worker_service_server::{WorkerService, WorkerServiceServer}; -use crate::worker::generated::worker::{ - CoordinatorToWorkerMsg, ExecuteTaskRequest, TaskKey, WorkerToCoordinatorMsg, -}; -use crate::worker::impl_execute_task::execute_remote_task; +use crate::worker::generated::worker::TaskKey; use crate::worker::single_write_multi_read::SingleWriteMultiRead; use crate::worker::task_data::TaskData; -use crate::{ - DefaultSessionBuilder, GetWorkerInfoRequest, GetWorkerInfoResponse, ObservabilityServiceImpl, - ObservabilityServiceServer, WorkerResolver, -}; -use arrow_flight::FlightData; -use async_trait::async_trait; use datafusion::common::DataFusionError; use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::physical_plan::ExecutionPlan; @@ -20,7 +11,27 @@ use moka::future::Cache; use std::borrow::Cow; use std::sync::Arc; use std::time::Duration; + +#[cfg(feature = "flight")] +use crate::worker::generated::worker::worker_service_server::{WorkerService, WorkerServiceServer}; +#[cfg(feature = "flight")] +use crate::worker::generated::worker::{ + CoordinatorToWorkerMsg, ExecuteTaskRequest, WorkerToCoordinatorMsg, +}; +#[cfg(feature = "flight")] +use crate::worker::impl_execute_task::execute_remote_task; +#[cfg(feature = "flight")] +use crate::{ + GetWorkerInfoRequest, GetWorkerInfoResponse, ObservabilityServiceImpl, + ObservabilityServiceServer, WorkerResolver, +}; +#[cfg(feature = "flight")] +use arrow_flight::FlightData; +#[cfg(feature = "flight")] +use async_trait::async_trait; +#[cfg(feature = "flight")] use tonic::codegen::BoxStream; +#[cfg(feature = "flight")] use tonic::{Request, Response, Status, Streaming}; const TASK_CACHE_TTI: Duration = Duration::from_mins(10); @@ -35,8 +46,8 @@ pub(super) struct WorkerHooks { pub(super) on_plan: Vec>, } -pub(crate) type ResultTaskData = Result>; -pub(crate) type TaskDataEntries = Cache>>; +pub type ResultTaskData = Result>; +pub type TaskDataEntries = Cache>>; #[derive(Clone)] pub struct Worker { @@ -66,6 +77,12 @@ impl Default for Worker { } impl Worker { + /// The registry of in-flight task plans this worker hosts. An in-process transport reads it to + /// execute fragments locally, the same registry the Flight service writes to. + pub fn task_data_entries(&self) -> &Arc { + &self.task_data_entries + } + /// Builds a [Worker] with a custom [WorkerSessionBuilder]. Use this /// method whenever you need to add custom stuff to the `SessionContext` that executes the query. pub fn from_session_builder( @@ -151,6 +168,7 @@ impl Worker { /// /// # } /// ``` + #[cfg(feature = "flight")] pub fn into_worker_server(self) -> WorkerServiceServer { WorkerServiceServer::new(self) .max_decoding_message_size(usize::MAX) @@ -162,6 +180,7 @@ impl Worker { /// /// The returned server is meant to be added to the same [`tonic::transport::Server`] as the /// Flight service — gRPC multiplexes both services on a single port. + #[cfg(feature = "flight")] pub fn with_observability_service( &self, worker_resolver: Arc, @@ -192,6 +211,7 @@ impl Worker { /// /// The methods are delegated to plan `impl Worker` implementations so that they can be implemented /// in different files. +#[cfg(feature = "flight")] #[async_trait] impl WorkerService for Worker { type CoordinatorChannelStream = BoxStream; diff --git a/tests/extension_points.rs b/tests/extension_points.rs new file mode 100644 index 00000000..a1755125 --- /dev/null +++ b/tests/extension_points.rs @@ -0,0 +1,73 @@ +//! Asserts the extension surface an out-of-crate WorkerTransport needs stays public. + +#![allow(dead_code)] + +// A `tests/*.rs` file links the lib as an external crate, so any item that is still `pub(crate)` +// (not truly `pub`) fails to compile here. Each reference below pins one extension point. + +// Traits: a generic bound rejects a non-`pub` trait regardless of object safety. +fn _assert_tree_node_ext() {} +fn _assert_worker_transport() {} +fn _assert_worker_connection() {} +fn _assert_worker_dispatch() {} +fn _assert_worker_resolver() {} +fn _assert_network_boundary_ext() {} +fn _assert_distributed_ext() {} +fn _assert_session_state_builder_ext() {} + +#[test] +fn extension_points_are_public() { + // Free functions referenced as values. + let _ = datafusion_distributed::execute_local_task; + let _ = datafusion_distributed::collect_plan_metrics_protos; + let _ = datafusion_distributed::encode_task_plan; + let _ = datafusion_distributed::collect_task_work_unit_feeds; + let _ = datafusion_distributed::set_sent_time; + let _ = datafusion_distributed::set_received_time; + let _ = datafusion_distributed::serialize_uuid; + let _ = datafusion_distributed::deserialize_uuid; + let _ = datafusion_distributed::get_config_extension_propagation_headers; + let _ = datafusion_distributed::get_passthrough_headers; + let _: fn(_, datafusion_distributed::InMemoryWorkerTransport) = + datafusion_distributed::set_distributed_worker_transport; + let _ = datafusion_distributed::get_distributed_cancellation_token; + let _ = datafusion_distributed::get_distributed_worker_transport; + let _ = datafusion_distributed::display_plan_ascii; + + // Types referenced in type position. + let _: Option = None; + let _: Option = None; + let _: Option> = None; + let _: Option = None; + let _: Option = None; + let _: Option = None; + let _: Option = None; + let _: Option = None; + let _: Option = None; + let _: Option = None; + let _: Option = None; + let _: Option = None; + let _: Option = None; + let _: Option = None; + let _: Option = None; + let _: Option> = None; + let _: Option = None; + let _: Option = None; + let _: Option = None; + + // The generated proto messages, re-exported as `proto` because `protobuf` already names a + // private module. + datafusion_distributed::proto::TaskKey::default(); + let _: Option = None; + let _: Option = None; + let _: Option = None; + let _: Option = None; + let _: Option = None; + + // `Worker::task_data_entries` is public, referenced as a method value. + let _ = datafusion_distributed::Worker::task_data_entries; + + // An embedder files worker `TaskMetrics` it collected out-of-band through these. + let _ = datafusion_distributed::DistributedExec::metrics_store; + let _ = datafusion_distributed::MetricsStore::insert; +} diff --git a/tests/metrics_collection.rs b/tests/metrics_collection.rs index ccb43a57..360e57f8 100644 --- a/tests/metrics_collection.rs +++ b/tests/metrics_collection.rs @@ -9,6 +9,8 @@ mod tests { use datafusion::physical_plan::{ExecutionPlan, execute_stream}; use datafusion::prelude::SessionContext; use datafusion_distributed::test_utils::localhost::start_localhost_context; + #[cfg(feature = "flight")] + use datafusion_distributed::test_utils::localhost::start_localhost_flight_context; use datafusion_distributed::test_utils::parquet::register_parquet_tables; use datafusion_distributed::test_utils::test_work_unit_feed::{ RowGeneratorExec, TestWorkUnitFeedExecCodec, TestWorkUnitFeedFunction, @@ -16,9 +18,11 @@ mod tests { }; use datafusion_distributed::{ DefaultSessionBuilder, DistributedExt, DistributedLeafExec, DistributedMetricsFormat, - NetworkCoalesceExec, NetworkShuffleExec, WorkerQueryContext, display_plan_ascii, - rewrite_distributed_plan_with_metrics, + WorkerQueryContext, display_plan_ascii, rewrite_distributed_plan_with_metrics, }; + // Only the Flight-gated network-boundary metrics test reads these. + #[cfg(feature = "flight")] + use datafusion_distributed::{NetworkCoalesceExec, NetworkShuffleExec}; use futures::TryStreamExt; use std::sync::Arc; use test_case::test_case; @@ -138,13 +142,17 @@ mod tests { Ok(()) } + // The network-connection metrics asserted here (bytes_transferred, network latency, buffered + // memory) are produced by the Flight gRPC connection; the shared-memory transport has no + // equivalent, so this one runs against the Flight cluster directly. + #[cfg(feature = "flight")] #[test_case(DistributedMetricsFormat::Aggregated ; "aggregated_metrics")] #[test_case(DistributedMetricsFormat::PerTask ; "per_task_metrics")] #[tokio::test] async fn test_metric_collection_network_boundaries( format: DistributedMetricsFormat, ) -> Result<(), Box> { - let (d_ctx, _guard, _) = start_localhost_context(3, DefaultSessionBuilder).await; + let (d_ctx, _guard, _) = start_localhost_flight_context(3, DefaultSessionBuilder).await; let query = r#"SELECT count(*), "RainToday" FROM weather GROUP BY "RainToday" ORDER BY count(*)"#; diff --git a/tests/task_estimator_test.rs b/tests/task_estimator_test.rs index 2cb7d303..a6753426 100644 --- a/tests/task_estimator_test.rs +++ b/tests/task_estimator_test.rs @@ -1,4 +1,6 @@ -#[cfg(all(feature = "integration", test))] +// The `URLEmitter` tests emit and assert per-URL worker identity, so they need the distinct dialed +// workers the gRPC in-memory cluster gives, not the single in-process worker. +#[cfg(all(feature = "integration", feature = "flight", test))] mod tests { use std::sync::Arc; From 6c4373db600261b569b6eb42c66b0c41d4518bf0 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 22 Jun 2026 13:38:17 -0700 Subject: [PATCH 4/4] Clarified the in-process extension point and Flight module docs. The `prepare_in_process_plan` doc referenced an out-of-crate transport that won't be upstreamed; reworded it to stand on its own. The `query_coordinator` module note makes its `flight` gating visible without opening `mod.rs`. --- src/coordinator/distributed.rs | 9 ++++----- src/coordinator/query_coordinator.rs | 4 ++++ 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index 2f22fa04..6377d6d9 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -153,11 +153,10 @@ impl DistributedExec { /// Routes and dispatches every stage through the registered transport, then returns the head /// stage for the caller to drive synchronously, with no `execute()` background task. /// - /// This is the extension point the in-crate shared-memory transport (and its in-process test) executes - /// through: it owns the runtime and drives the head stage itself. The transport's dispatch must - /// complete synchronously: a transport that spawns background delivery work, or a plan that - /// declares work-unit feeds (which are pumped by that background work), is rejected rather than - /// left to stall. + /// This is the extension point for an in-process transport that owns the runtime and drives the + /// head stage itself. The transport's dispatch must complete synchronously: a transport that + /// spawns background delivery work, or a plan that declares work-unit feeds (which are pumped by + /// that background work), is rejected rather than left to stall. pub fn prepare_in_process_plan( &self, ctx: &Arc, diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index ab26cde5..76988dcf 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -1,3 +1,7 @@ +//! Arrow-Flight plan delivery: ships each task's plan to its worker over the coordinator-to-worker +//! gRPC stream. Flight-specific, so `coordinator/mod.rs` gates the whole module on the `flight` +//! feature. + use crate::common::{TreeNodeExt, serialize_uuid, task_ctx_with_extension}; use crate::config_extension_ext::get_config_extension_propagation_headers; use crate::coordinator::CoordinatorToWorkerMetrics;