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/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/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/distributed.rs b/src/coordinator/distributed.rs index fe1bbff3..6377d6d9 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,62 @@ 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 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, + ) -> 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 +257,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 +303,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..34129bb2 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -1,8 +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(crate) use metrics_store::MetricsStore; +pub use latency_metric::LatencyMetric; +pub use metrics_store::MetricsStore; +pub use plan_encoding::{EncodedTaskPlan, encode_task_plan}; +#[cfg(feature = "flight")] +pub(crate) use query_coordinator::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/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..76988dcf 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -1,36 +1,32 @@ -use crate::common::{TreeNodeExt, now_ns, serialize_uuid, task_ctx_with_extension}; +//! 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; 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::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::tree_node::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; -use datafusion_proto::physical_plan::AsExecutionPlan; -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 +40,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 +108,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> { @@ -139,13 +126,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), @@ -154,10 +137,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, })), @@ -187,7 +170,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 +210,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,130 +296,15 @@ 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(()) }); 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)) - } - - /// 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 { - pub(super) plan_bytes_sent: BytesCounterMetric, - pub(super) plan_send_latency: Arc, - pub(super) instantiation_time: u64, -} - -impl CoordinatorToWorkerMetrics { - pub(super) 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 17852a04..da53a20a 100644 --- a/src/distributed_ext.rs +++ b/src/distributed_ext.rs @@ -1,14 +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}; +#[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, + DistributedConfig, TaskEstimator, WorkUnitFeed, WorkUnitFeedProvider, WorkerResolver, + WorkerTransport, }; use arrow_ipc::CompressionType; use datafusion::common::DataFusionError; @@ -224,10 +228,19 @@ 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 /// only one per URL, but users can decide to override that behavior in favor of their own solution. + #[cfg(feature = "flight")] /// /// Example: /// @@ -273,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, @@ -604,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, @@ -611,6 +626,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, @@ -744,10 +763,15 @@ 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; + #[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; @@ -836,11 +860,18 @@ 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; + 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)] @@ -947,11 +978,18 @@ 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; + 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)] @@ -1058,11 +1096,18 @@ 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; + 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/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/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..a1e41a12 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -20,11 +20,18 @@ 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::{ + CoordinatorToWorkerMetrics, DistributedExec, EncodedTaskPlan, LatencyMetric, MetricsStore, + encode_task_plan, +}; 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, @@ -36,33 +43,57 @@ pub use metrics::{ MaxLatencyMetric, MinLatencyMetric, P50LatencyMetric, P75LatencyMetric, P95LatencyMetric, P99LatencyMetric, rewrite_distributed_plan_with_metrics, }; +#[cfg(feature = "flight")] +pub use networking::{ + BoxCloneSyncChannel, ChannelResolver, DefaultChannelResolver, create_worker_client, + get_distributed_channel_resolver, +}; pub use networking::{ - BoxCloneSyncChannel, ChannelResolver, DefaultChannelResolver, WorkerResolver, - create_worker_client, get_distributed_channel_resolver, get_distributed_worker_resolver, + WorkerResolver, 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, + DistributedWorkUnitFeedContext, RemoteWorkUnitFeedRegistry, RemoteWorkUnitFeedRxs, + 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, MappedWorkerSessionBuilder, MappedWorkerSessionBuilderExt, TaskData, - Worker, WorkerQueryContext, WorkerSessionBuilder, + 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 6bab9ae0..1a2c3880 100644 --- a/src/networking/mod.rs +++ b/src/networking/mod.rs @@ -1,11 +1,18 @@ 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}; + +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..de075ae3 --- /dev/null +++ b/src/networking/worker_transport.rs @@ -0,0 +1,62 @@ +use crate::DistributedConfig; +use crate::config_extension_ext::set_distributed_option_extension; +#[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, +) { + 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(default_worker_transport) +} + +#[derive(Clone)] +pub(crate) struct WorkerTransportExtension(pub(crate) Arc); + +impl Default for WorkerTransportExtension { + fn default() -> Self { + 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/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/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/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..078d3a8d 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,14 @@ 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, +}; +#[cfg(feature = "flight")] 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 f914f228..0f1e8b43 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))); @@ -47,6 +47,25 @@ 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, + } +} + +#[cfg(feature = "flight")] pub(crate) fn build_work_unit_batch_msg( id: &Uuid, work_unit_batch: Vec<(usize, Result>)>, @@ -56,23 +75,25 @@ 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(); +} + +#[cfg(feature = "flight")] pub(crate) fn set_work_unit_send_time( mut msg: pb::CoordinatorToWorkerMsg, ) -> pb::CoordinatorToWorkerMsg { @@ -87,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/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/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_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_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/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..c1ce0fc4 --- /dev/null +++ b/src/worker/in_memory.rs @@ -0,0 +1,432 @@ +//! 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), + } + } + + /// 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 + } +} + +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(()) + } + + // 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 e89921fc..dbec1ba9 100644 --- a/src/worker/mod.rs +++ b/src/worker/mod.rs @@ -1,21 +1,37 @@ -pub(crate) mod generated; +#[cfg(feature = "flight")] +mod flight; +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"))] 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}; +// 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}; pub use session_builder::{ DefaultSessionBuilder, MappedWorkerSessionBuilder, MappedWorkerSessionBuilderExt, 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/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..bc166721 100644 --- a/src/worker/worker_connection_pool.rs +++ b/src/worker/worker_connection_pool.rs @@ -1,48 +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; +#[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; +#[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 @@ -51,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, @@ -65,7 +110,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 +138,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,14 +147,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( + // 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, @@ -117,21 +158,7 @@ impl WorkerConnectionPool { 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( - input_stage, - target_partitions, - target_task, - producer_head, - ctx, - &self.metrics, - ) - .map(|v| Box::new(v) as Box<_>) - .map_err(Arc::new) - } }); match conn { @@ -141,16 +168,9 @@ impl WorkerConnectionPool { } } +#[cfg(feature = "flight")] 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 +182,8 @@ 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 { +#[cfg(feature = "flight")] +pub(crate) struct RemoteWorkerConnection { task: Arc>, not_consumed_streams: Arc, cancel_token: CancellationToken, @@ -177,8 +198,9 @@ struct RemoteWorkerConnection { elapsed_compute: Time, } +#[cfg(feature = "flight")] impl RemoteWorkerConnection { - fn init( + pub(crate) fn init( input_stage: &RemoteStage, target_partition_range: Range, target_task: usize, @@ -398,6 +420,7 @@ impl RemoteWorkerConnection { } } +#[cfg(feature = "flight")] impl WorkerConnection for RemoteWorkerConnection { /// Streams the provided `partition` from the remote worker. /// @@ -456,13 +479,15 @@ 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 { - fn init( + pub(crate) fn init( input_stage: &RemoteStage, target_partition_range: Range, target_task: usize, @@ -528,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 { @@ -549,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())); @@ -569,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 { @@ -593,6 +624,7 @@ impl> ElapsedComputeFutureExt for F { } } +#[cfg(feature = "flight")] impl> ElapsedComputeStreamExt for S { fn with_elapsed_compute(self, elapsed_compute: Time) -> ElapsedComputeStream { ElapsedComputeStream { @@ -603,6 +635,7 @@ impl> ElapsedComputeStreamExt for S { } } +#[cfg(feature = "flight")] #[pin_project(PinnedDrop)] struct ElapsedComputeStream { #[pin] @@ -613,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>) { @@ -625,6 +659,7 @@ impl PinnedDrop for ElapsedComputeStream { } } +#[cfg(feature = "flight")] impl> Stream for ElapsedComputeStream { type Item = O; @@ -643,6 +678,7 @@ impl> Stream for ElapsedComputeStream { } } +#[cfg(feature = "flight")] #[pin_project(PinnedDrop)] struct ElapsedComputeFuture { #[pin] @@ -653,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>) { @@ -665,6 +702,7 @@ impl PinnedDrop for ElapsedComputeFuture { } } +#[cfg(feature = "flight")] impl> Future for ElapsedComputeFuture { type Output = O; @@ -683,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;