diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index a1e123f2..17fe31ab 100644 --- a/src/coordinator/distributed.rs +++ b/src/coordinator/distributed.rs @@ -64,6 +64,13 @@ impl DistributedExec { } } + /// The store where worker task metrics land at runtime, if metrics collection is enabled. + /// Exposed for a driver whose transport returns metrics out-of-band; it files decoded frames + /// here before the per-task EXPLAIN rewrite. + pub fn metrics_store(&self) -> Option> { + self.metrics_store.clone() + } + /// Enables task metrics collection from remote workers. pub fn with_metrics_collection(mut self, enabled: bool) -> Self { self.metrics_store = match enabled { diff --git a/src/coordinator/metrics_store.rs b/src/coordinator/metrics_store.rs index ed55db2c..b38d62a0 100644 --- a/src/coordinator/metrics_store.rs +++ b/src/coordinator/metrics_store.rs @@ -17,7 +17,9 @@ impl MetricsStore { Self { tx, rx } } - pub(crate) fn insert(&self, key: TaskKey, metrics: TaskMetrics) { + // Public for a driver whose transport returns worker metrics out-of-band; it files the + // decoded frames here before the per-task EXPLAIN rewrite reads them. + pub fn insert(&self, key: TaskKey, metrics: TaskMetrics) { self.tx.send_modify(|map| { map.insert(key, metrics); }); diff --git a/src/coordinator/mod.rs b/src/coordinator/mod.rs index c1a8a8dd..be069485 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -6,4 +6,4 @@ mod prepare_static_plan; mod query_coordinator; pub use distributed::DistributedExec; -pub(crate) use metrics_store::MetricsStore; +pub use metrics_store::MetricsStore; diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 887a109f..9309e541 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -11,7 +11,8 @@ use crate::{ DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedCodec, DistributedConfig, DistributedTaskContext, DistributedWorkUnitFeedContext, LoadInfo, NetworkBoundaryExt, SetPlanRequest, Stage, TaskEstimator, TaskKey, TaskRoutingContext, WorkUnitFeedDeclaration, - WorkerToCoordinatorMsg, get_distributed_channel_resolver, get_distributed_worker_resolver, + WorkerToCoordinatorMsg, get_distributed_channel_resolver, get_distributed_dispatch_plan_source, + get_distributed_worker_resolver, }; use datafusion::common::DataFusionError; use datafusion::common::instant::Instant; @@ -140,19 +141,27 @@ 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(); - + // An embedder can serialize the dispatch bytes for this stage itself (e.g. with a codec + // the config's extension point cannot express) instead of the coordinator encoding the + // plan. Either way the bytes describe `specialized`, the ready-to-run per-task plan. let task_key = TaskKey { query_id: self.query_id, stage_id: self.stage_id, task_number: task_i, }; + let plan_proto = match get_distributed_dispatch_plan_source(session_config) + .and_then(|source| source.dispatch_plan_proto(&task_key, &specialized)) + { + Some(bytes) => bytes?, + None => { + let codec = DistributedCodec::new_combined_with_user(session_config); + PhysicalPlanNode::try_from_physical_plan(specialized, &codec)?.encode_to_vec() + } + }; + let plan_size = plan_proto.len(); let msg = CoordinatorToWorkerMsg::SetPlanRequest(SetPlanRequest { task_key, diff --git a/src/dispatch_plan_source.rs b/src/dispatch_plan_source.rs new file mode 100644 index 00000000..d7f84994 --- /dev/null +++ b/src/dispatch_plan_source.rs @@ -0,0 +1,241 @@ +use crate::TaskKey; +use datafusion::common::Result; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::SessionConfig; +use std::sync::Arc; + +/// Serializes the stage subplan the coordinator dispatches, instead of the coordinator encoding +/// it with its own codec. +/// +/// An embedder registers one via +/// [`crate::DistributedExt::with_distributed_dispatch_plan_source`] when the coordinator's codec +/// cannot represent its plan nodes, or when its serialization needs embedder-side handling the +/// codec extension point cannot express (the shm embedder's UDF definitions, for example). The +/// coordinator hands over `specialized`, the same ready-to-run per-task plan it would encode: +/// task-specialized, with nested stages already converted to `Remote`, so a worker executes the +/// decoded bytes as-is. `task` carries the query id, so a source registered on a session that +/// runs concurrent queries can tell them apart. +/// +/// Returning `None` for a task lets the coordinator fall back to encoding the plan itself, so a +/// source that only overrides some stages stays correct. +pub trait DispatchPlanSource: Send + Sync { + fn dispatch_plan_proto( + &self, + task: &TaskKey, + specialized: &Arc, + ) -> Option>>; +} + +#[derive(Clone)] +pub(crate) struct DispatchPlanSourceExtension(pub(crate) Arc); + +pub(crate) fn set_distributed_dispatch_plan_source( + cfg: &mut SessionConfig, + source: impl DispatchPlanSource + 'static, +) { + set_distributed_dispatch_plan_source_arc(cfg, Arc::new(source)) +} + +pub(crate) fn set_distributed_dispatch_plan_source_arc( + cfg: &mut SessionConfig, + source: Arc, +) { + cfg.set_extension(Arc::new(DispatchPlanSourceExtension(source))); +} + +/// Returns the [`DispatchPlanSource`] registered on this config, if any. +pub fn get_distributed_dispatch_plan_source( + cfg: &SessionConfig, +) -> Option> { + cfg.get_extension::() + .map(|ext| Arc::clone(&ext.0)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocol::InProcessChannelResolver; + use crate::{ + DistributedExt, NetworkBoundaryExt, SessionStateBuilderExt, Stage, WorkerResolver, + }; + use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; + use datafusion::common::{DataFusionError, Result}; + use datafusion::execution::SessionStateBuilder; + use datafusion::physical_plan::{ExecutionPlan, collect}; + use datafusion::prelude::{CsvReadOptions, SessionConfig, SessionContext}; + use datafusion_proto::physical_plan::AsExecutionPlan; + use datafusion_proto::protobuf::PhysicalPlanNode; + use prost::Message; + use std::io::Write; + use std::sync::Mutex; + use url::Url; + + struct Workers(usize); + + impl WorkerResolver for Workers { + fn get_urls(&self) -> Result> { + (0..self.0) + .map(|i| Url::parse(&format!("http://worker-{i}"))) + .collect::>() + .map_err(|err| DataFusionError::External(Box::new(err))) + } + } + + type Calls = Arc)>>>; + + /// Records what the coordinator hands over and declines, so the coordinator's own encode + /// still runs and the query is unaffected by the recording. + #[derive(Default)] + struct RecordingSource { + calls: Calls, + } + + impl DispatchPlanSource for RecordingSource { + fn dispatch_plan_proto( + &self, + task: &TaskKey, + specialized: &Arc, + ) -> Option>> { + self.calls + .lock() + .unwrap() + .push((*task, Arc::clone(specialized))); + None + } + } + + const QUERY: &str = "SELECT k, COUNT(*) AS c FROM t GROUP BY k ORDER BY k"; + + async fn distributed_ctx( + name: &str, + source: impl DispatchPlanSource + 'static, + ) -> Result<(SessionContext, std::path::PathBuf)> { + let path = std::env::temp_dir().join(format!("dfd_{name}_{}.csv", std::process::id())); + let mut file = + std::fs::File::create(&path).map_err(|e| DataFusionError::External(Box::new(e)))?; + writeln!(file, "k,v").unwrap(); + for i in 0..200 { + writeln!(file, "{},{}", ["a", "b", "c", "d"][i % 4], i).unwrap(); + } + drop(file); + + let state = SessionStateBuilder::new() + .with_default_features() + .with_config(SessionConfig::new().with_target_partitions(4)) + .with_distributed_planner() + .with_distributed_worker_resolver(Workers(4)) + .with_distributed_channel_resolver(InProcessChannelResolver::default()) + .with_distributed_dispatch_plan_source(source) + .with_distributed_file_scan_config_bytes_per_partition(1) + .unwrap() + .build(); + let ctx = SessionContext::from(state); + ctx.register_csv("t", path.to_str().unwrap(), CsvReadOptions::new()) + .await?; + Ok((ctx, path)) + } + + /// The contract a serializing source relies on: it is consulted once per dispatched task, + /// and the plan it is handed is ready to run, with nested stages already converted to + /// `Remote`. + #[tokio::test] + async fn source_is_consulted_with_the_ready_to_run_plan() -> Result<()> { + let source = RecordingSource::default(); + let calls = Arc::clone(&source.calls); + let (ctx, path) = distributed_ctx("recording", source).await?; + + let physical = ctx.sql(QUERY).await?.create_physical_plan().await?; + collect(physical, ctx.task_ctx()).await?; + + let calls = calls.lock().unwrap(); + assert!( + !calls.is_empty(), + "the coordinator never consulted the source" + ); + let mut seen: datafusion::common::HashSet = Default::default(); + for (task, specialized) in calls.iter() { + assert!(seen.insert(*task), "consulted twice for {task:?}"); + specialized.apply(|node| { + if let Some(nb) = node.as_ref().as_network_boundary() { + assert!( + matches!(nb.input_stage(), Stage::Remote(_)), + "{task:?} carries a Local nested stage; the handed-over plan must be \ + ready to run" + ); + } + Ok(TreeNodeRecursion::Continue) + })?; + } + std::fs::remove_file(&path).ok(); + Ok(()) + } + + /// Serializes with the same codec the coordinator's fallback uses, so the worker decodes + /// source-provided bytes exactly as it decodes coordinator-encoded ones. + struct SerializingSource; + + impl DispatchPlanSource for SerializingSource { + fn dispatch_plan_proto( + &self, + _task: &TaskKey, + specialized: &Arc, + ) -> Option>> { + let codec = crate::DistributedCodec::new_combined_with_user(&SessionConfig::new()); + Some( + PhysicalPlanNode::try_from_physical_plan(Arc::clone(specialized), &codec) + .map(|node| node.encode_to_vec()), + ) + } + } + + /// The bytes the source returns are what the workers run. + #[tokio::test] + async fn source_provided_bytes_run_the_query() -> Result<()> { + let (ctx, path) = distributed_ctx("serializing", SerializingSource).await?; + let got = ctx.sql(QUERY).await?.collect().await?; + let got = datafusion::arrow::util::pretty::pretty_format_batches(&got)?.to_string(); + + let serial = SessionContext::new(); + serial + .register_csv("t", path.to_str().unwrap(), CsvReadOptions::new()) + .await?; + let expected = serial.sql(QUERY).await?.collect().await?; + let expected = + datafusion::arrow::util::pretty::pretty_format_batches(&expected)?.to_string(); + + assert_eq!(got, expected, "source-encoded dispatch != serial"); + std::fs::remove_file(&path).ok(); + Ok(()) + } + + struct FailingSource; + + impl DispatchPlanSource for FailingSource { + fn dispatch_plan_proto( + &self, + _task: &TaskKey, + _specialized: &Arc, + ) -> Option>> { + Some(Err(DataFusionError::Internal( + "the embedder could not serialize this stage".into(), + ))) + } + } + + /// A source failure fails the dispatch instead of falling back to bytes the embedder said it + /// could not produce. + #[tokio::test] + async fn source_error_fails_the_query() -> Result<()> { + let (ctx, path) = distributed_ctx("failing", FailingSource).await?; + let result = ctx.sql(QUERY).await?.collect().await; + let err = result + .expect_err("a failing source must fail the query") + .to_string(); + assert!( + err.contains("could not serialize"), + "unexpected error: {err}" + ); + std::fs::remove_file(&path).ok(); + Ok(()) + } +} diff --git a/src/distributed_ext.rs b/src/distributed_ext.rs index 9efdb420..6f71f7a4 100644 --- a/src/distributed_ext.rs +++ b/src/distributed_ext.rs @@ -2,14 +2,15 @@ use crate::codec::{set_distributed_user_codec, set_distributed_user_codec_arc}; use crate::config_extension_ext::{ set_distributed_option_extension, set_distributed_option_extension_from_headers, }; +use crate::dispatch_plan_source::set_distributed_dispatch_plan_source; use crate::distributed_planner::set_distributed_task_estimator; use crate::passthrough_headers::set_passthrough_headers; use crate::protocol::set_distributed_channel_resolver; use crate::work_unit_feed::set_distributed_work_unit_feed; use crate::worker_resolver::set_distributed_worker_resolver; use crate::{ - ChannelResolver, DistributedConfig, TaskEstimator, WorkUnitFeed, WorkUnitFeedProvider, - WorkerResolver, + ChannelResolver, DispatchPlanSource, DistributedConfig, TaskEstimator, WorkUnitFeed, + WorkUnitFeedProvider, WorkerResolver, }; use datafusion::common::DataFusionError; use datafusion::config::ConfigExtension; @@ -278,6 +279,16 @@ pub trait DistributedExt: Sized { resolver: T, ); + /// Registers a [DispatchPlanSource] the coordinator consults for each stage's dispatch bytes + /// instead of encoding the plan it holds. See [DispatchPlanSource] for when this is needed. + fn with_distributed_dispatch_plan_source( + self, + source: T, + ) -> Self; + + /// Same as [DistributedExt::with_distributed_dispatch_plan_source] but with an in-place mutation. + fn set_distributed_dispatch_plan_source(&mut self, source: T); + /// Adds a distributed task count estimator. [TaskEstimator]s are executed on each node /// sequentially until one returns an estimation on the number of tasks that should be /// used for the stage containing that node. @@ -634,6 +645,10 @@ impl DistributedExt for SessionConfig { set_distributed_channel_resolver(self, resolver); } + fn set_distributed_dispatch_plan_source(&mut self, source: T) { + set_distributed_dispatch_plan_source(self, source); + } + fn set_distributed_task_estimator( &mut self, estimator: T, @@ -787,6 +802,10 @@ impl DistributedExt for SessionConfig { #[expr($;self)] fn with_distributed_channel_resolver(mut self, resolver: T) -> Self; + #[call(set_distributed_dispatch_plan_source)] + #[expr($;self)] + fn with_distributed_dispatch_plan_source(mut self, source: T) -> Self; + #[call(set_distributed_task_estimator)] #[expr($;self)] fn with_distributed_task_estimator(mut self, estimator: T) -> Self; @@ -889,6 +908,11 @@ impl DistributedExt for SessionStateBuilder { #[expr($;self)] fn with_distributed_channel_resolver(mut self, resolver: T) -> Self; + fn set_distributed_dispatch_plan_source(&mut self, source: T); + #[call(set_distributed_dispatch_plan_source)] + #[expr($;self)] + fn with_distributed_dispatch_plan_source(mut self, source: T) -> Self; + fn set_distributed_task_estimator(&mut self, estimator: T); #[call(set_distributed_task_estimator)] #[expr($;self)] @@ -1012,6 +1036,11 @@ impl DistributedExt for SessionState { #[expr($;self)] fn with_distributed_channel_resolver(mut self, resolver: T) -> Self; + fn set_distributed_dispatch_plan_source(&mut self, source: T); + #[call(set_distributed_dispatch_plan_source)] + #[expr($;self)] + fn with_distributed_dispatch_plan_source(mut self, source: T) -> Self; + fn set_distributed_task_estimator(&mut self, estimator: T); #[call(set_distributed_task_estimator)] #[expr($;self)] @@ -1135,6 +1164,11 @@ impl DistributedExt for SessionContext { #[expr($;self)] fn with_distributed_channel_resolver(self, resolver: T) -> Self; + fn set_distributed_dispatch_plan_source(&mut self, source: T); + #[call(set_distributed_dispatch_plan_source)] + #[expr($;self)] + fn with_distributed_dispatch_plan_source(self, source: T) -> Self; + fn set_distributed_task_estimator(&mut self, estimator: T); #[call(set_distributed_task_estimator)] #[expr($;self)] diff --git a/src/distributed_planner/mod.rs b/src/distributed_planner/mod.rs index 18cfe6e0..176bd319 100644 --- a/src/distributed_planner/mod.rs +++ b/src/distributed_planner/mod.rs @@ -15,7 +15,7 @@ pub(crate) use inject_network_boundaries::{ InjectNetworkBoundaryContext, NetworkBoundaryBuilderResult, inject_network_boundaries, }; pub(crate) use network_boundary::ProducerHead; -pub use network_boundary::{NetworkBoundary, NetworkBoundaryExt}; +pub use network_boundary::{NetworkBoundary, NetworkBoundaryExt, PartitionRoute}; pub use session_state_builder_ext::SessionStateBuilderExt; pub(crate) use statistics::calculate_cost; pub(crate) use task_estimator::set_distributed_task_estimator; diff --git a/src/distributed_planner/network_boundary.rs b/src/distributed_planner/network_boundary.rs index 3b7bd3b1..05a00ad6 100644 --- a/src/distributed_planner/network_boundary.rs +++ b/src/distributed_planner/network_boundary.rs @@ -5,7 +5,7 @@ use crate::{ Stage, }; use datafusion::arrow::datatypes::SchemaRef; -use datafusion::common::Result; +use datafusion::common::{Result, internal_err}; use datafusion::execution::TaskContext; use datafusion::physical_expr::Partitioning; use datafusion::physical_plan::repartition::RepartitionExec; @@ -19,6 +19,14 @@ use datafusion_proto::protobuf::proto_error; use prost::Message; 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. @@ -36,6 +44,31 @@ 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; + + /// Maps a producer output partition to the consumer task and the local partition within that + /// task that reads it, for the sliced layout shuffle and broadcast reads use + /// (`global = P_c * consumer_task + local`, where `P_c` is the boundary's own per-task output + /// partition count). A pull-based transport never needs this: its consumers compute their own + /// slice inside the boundary's `execute`. A push-based transport places every produced + /// partition before any consumer asks, so it reads the layout here instead of re-deriving it + /// and drifting when the layout changes. + /// + /// 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.properties().partitioning.partition_count(); + 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 @@ -155,3 +188,98 @@ impl ProducerHead { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::protocol::InProcessChannelResolver; + use crate::{DistributedExt, NetworkBoundaryExt, SessionStateBuilderExt, WorkerResolver}; + use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; + use datafusion::error::DataFusionError; + use datafusion::execution::SessionStateBuilder; + use datafusion::prelude::{CsvReadOptions, SessionConfig, SessionContext}; + use std::io::Write; + use url::Url; + + struct Workers(usize); + + impl WorkerResolver for Workers { + fn get_urls(&self) -> Result> { + (0..self.0) + .map(|i| Url::parse(&format!("http://worker-{i}"))) + .collect::>() + .map_err(|err| DataFusionError::External(Box::new(err))) + } + } + + /// Pins the sliced routing (`global = P_c * consumer_task + local`) on boundaries the + /// planner actually built, with `P_c` read off the boundary's own properties, so an override + /// or a change in what `properties()` reports fails here first. The data-level guarantee + /// that the slicing matches what consumers read comes from the in-process transport's + /// end-to-end test. `NetworkCoalesceExec` must refuse to route: its consumers read whole + /// per-producer-task groups, and the sliced formula would misroute them. + #[tokio::test] + async fn route_partition_matches_the_consumer_slicing() -> Result<()> { + let path = std::env::temp_dir().join(format!("dfd_routing_{}.csv", std::process::id())); + let mut file = + std::fs::File::create(&path).map_err(|e| DataFusionError::External(Box::new(e)))?; + writeln!(file, "k,v").unwrap(); + for i in 0..200 { + writeln!(file, "{},{}", ["a", "b", "c", "d"][i % 4], i).unwrap(); + } + drop(file); + + let state = SessionStateBuilder::new() + .with_default_features() + .with_config(SessionConfig::new().with_target_partitions(4)) + .with_distributed_planner() + .with_distributed_worker_resolver(Workers(4)) + .with_distributed_channel_resolver(InProcessChannelResolver::default()) + .with_distributed_file_scan_config_bytes_per_partition(1) + .unwrap() + .build(); + let ctx = SessionContext::from(state); + ctx.register_csv("t", path.to_str().unwrap(), CsvReadOptions::new()) + .await?; + let physical = ctx + .sql("SELECT k, COUNT(*) AS c FROM t GROUP BY k ORDER BY k") + .await? + .create_physical_plan() + .await?; + + let mut sliced = 0usize; + let mut grouped = 0usize; + physical.apply(|node| { + let Some(nb) = node.as_ref().as_network_boundary() else { + return Ok(TreeNodeRecursion::Continue); + }; + if node + .as_ref() + .downcast_ref::() + .is_some() + { + assert!( + nb.route_partition(0).is_err(), + "a per-task-group boundary must refuse the sliced routing" + ); + grouped += 1; + return Ok(TreeNodeRecursion::Continue); + } + let p_c = nb.properties().partitioning.partition_count(); + assert!(p_c > 0); + for consumer_task in 0..3 { + for local in 0..p_c { + let route = nb.route_partition(p_c * consumer_task + local)?; + assert_eq!(route.consumer_task, consumer_task); + assert_eq!(route.consumer_partition, local); + } + } + sliced += 1; + Ok(TreeNodeRecursion::Continue) + })?; + assert!(sliced > 0, "the plan grew no sliced-layout boundary"); + assert!(grouped > 0, "the plan grew no per-task-group boundary"); + std::fs::remove_file(&path).ok(); + Ok(()) + } +} diff --git a/src/execution_plans/network_coalesce.rs b/src/execution_plans/network_coalesce.rs index 703843cc..4379910a 100644 --- a/src/execution_plans/network_coalesce.rs +++ b/src/execution_plans/network_coalesce.rs @@ -182,6 +182,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 32176a88..3f4f1a80 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,6 +4,7 @@ mod codec; mod common; mod config_extension_ext; mod coordinator; +mod dispatch_plan_source; mod distributed_ext; mod distributed_planner; mod execution_plans; @@ -17,10 +18,12 @@ mod worker_resolver; #[cfg(feature = "grpc")] pub use arrow_ipc::CompressionType; -pub use coordinator::DistributedExec; +// `MetricsStore` is public for a driver whose transport returns worker metrics out-of-band; it +// files the decoded frames into the executed plan's store before the per-task EXPLAIN rewrite. +pub use coordinator::{DistributedExec, MetricsStore}; pub use distributed_ext::DistributedExt; pub use distributed_planner::{ - DistributedConfig, NetworkBoundary, NetworkBoundaryExt, SessionStateBuilderExt, + DistributedConfig, NetworkBoundary, NetworkBoundaryExt, PartitionRoute, SessionStateBuilderExt, TaskCountAnnotation, TaskEstimation, TaskEstimator, TaskRoutingContext, }; pub use execution_plans::{ @@ -45,23 +48,25 @@ pub use protocol::grpc; pub use protocol::generated::worker as proto; pub use codec::DistributedCodec; +pub use dispatch_plan_source::{DispatchPlanSource, get_distributed_dispatch_plan_source}; pub use worker_resolver::{WorkerResolver, get_distributed_worker_resolver}; pub use protocol::{ ChannelResolver, CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, GetWorkerInfoResponse, InProcessChannelResolver, LoadInfo, ProducerHeadSpec, SetPlanRequest, TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, - WorkerToCoordinatorMsg, get_distributed_channel_resolver, + WorkerToCoordinatorMsg, decode_task_metrics, get_distributed_channel_resolver, }; pub use stage::{ DistributedTaskContext, Stage, display_plan_ascii, display_plan_graphviz, explain_analyze, }; pub use work_unit_feed::{ - DistributedWorkUnitFeedContext, WorkUnit, WorkUnitFeed, WorkUnitFeedProto, WorkUnitFeedProvider, + DistributedWorkUnitFeedContext, WorkUnit, WorkUnitFeed, WorkUnitFeedProto, + WorkUnitFeedProvider, set_received_time, }; pub use worker::{ DefaultSessionBuilder, MappedWorkerSessionBuilder, MappedWorkerSessionBuilderExt, TaskData, - Worker, WorkerQueryContext, WorkerSessionBuilder, + Worker, WorkerQueryContext, WorkerSessionBuilder, collect_plan_metrics_protos, }; #[cfg(all(feature = "grpc", any(feature = "integration", test)))] diff --git a/src/protocol/grpc/mod.rs b/src/protocol/grpc/mod.rs index fc367585..cc2f6576 100644 --- a/src/protocol/grpc/mod.rs +++ b/src/protocol/grpc/mod.rs @@ -1,6 +1,5 @@ mod channel_resolver; mod errors; -mod metrics_proto; mod observability; mod on_drop_stream; mod spawn_select_all; diff --git a/src/protocol/grpc/worker_client.rs b/src/protocol/grpc/worker_client.rs index 93291d91..c2866827 100644 --- a/src/protocol/grpc/worker_client.rs +++ b/src/protocol/grpc/worker_client.rs @@ -1,16 +1,16 @@ use super::channel_resolver::BoxCloneSyncChannel; use super::errors::{map_flight_to_datafusion_error, map_status_to_datafusion_error}; -use super::metrics_proto::metrics_set_proto_to_df; use crate::common::serialize_uuid; use crate::grpc::on_drop_stream::on_drop_stream; use crate::protocol::generated::worker as pb; use crate::protocol::generated::worker::FlightAppMetadata; +use crate::protocol::metrics_proto::decode_task_metrics; use crate::{ BytesMetricExt, CoordinatorToWorkerMsg, DistributedConfig, ExecuteTaskRequest, FirstLatencyMetric, GetWorkerInfoRequest, GetWorkerInfoResponse, LatencyMetricExt, LoadInfo, MaxLatencyMetric, MinLatencyMetric, P50LatencyMetric, P95LatencyMetric, ProducerHeadSpec, - SetPlanRequest, TaskKey, TaskMetrics, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, - WorkerChannel, WorkerToCoordinatorMsg, + SetPlanRequest, TaskKey, WorkUnitBatch, WorkUnitFeedDeclaration, WorkUnitMsg, WorkerChannel, + WorkerToCoordinatorMsg, }; use arrow_flight::FlightData; use arrow_flight::decode::FlightRecordBatchStream; @@ -491,21 +491,6 @@ fn decode_worker_to_coordinator_msg( ) } -fn decode_task_metrics(task_metrics: pb::TaskMetrics) -> Result { - Ok(TaskMetrics { - pre_order_plan_metrics: task_metrics - .pre_order_plan_metrics - .into_iter() - .map(|metrics_set| metrics_set_proto_to_df(&metrics_set)) - .collect::>()?, - task_metrics: metrics_set_proto_to_df( - &task_metrics - .task_metrics - .ok_or_else(|| missing("task_metrics"))?, - )?, - }) -} - fn decode_load_info(load_info: pb::LoadInfo) -> LoadInfo { LoadInfo { partition: load_info.partition as usize, diff --git a/src/protocol/grpc/worker_service.rs b/src/protocol/grpc/worker_service.rs index f3472d43..87c1ef4b 100644 --- a/src/protocol/grpc/worker_service.rs +++ b/src/protocol/grpc/worker_service.rs @@ -1,7 +1,7 @@ use super::errors::{datafusion_error_to_tonic_status, map_status_to_datafusion_error}; -use super::metrics_proto::df_metrics_set_to_proto; use super::spawn_select_all::spawn_select_all; use crate::protocol::generated::worker as pb; +use crate::protocol::metrics_proto::df_metrics_set_to_proto; use crate::common::{deserialize_uuid, now_ns}; use crate::protocol::ProducerHeadSpec; diff --git a/src/protocol/grpc/metrics_proto.rs b/src/protocol/metrics_proto.rs similarity index 95% rename from src/protocol/grpc/metrics_proto.rs rename to src/protocol/metrics_proto.rs index 654b6cee..f72e0edb 100644 --- a/src/protocol/grpc/metrics_proto.rs +++ b/src/protocol/metrics_proto.rs @@ -1,4 +1,4 @@ -use crate::protocol::generated::worker as pb; +use super::generated::worker as pb; use chrono::DateTime; use datafusion::common::internal_err; use datafusion::error::DataFusionError; @@ -55,6 +55,27 @@ pub fn metrics_set_proto_to_df( Ok(metrics_set) } +/// Decode a wire [`pb::TaskMetrics`] into the in-memory [`crate::TaskMetrics`]. The no-gRPC push +/// embedder receives metric frames as proto and files them into the plain metrics store, so it +/// needs the decode direction without the gRPC client. +pub fn decode_task_metrics( + task_metrics: pb::TaskMetrics, +) -> Result { + Ok(crate::TaskMetrics { + pre_order_plan_metrics: task_metrics + .pre_order_plan_metrics + .iter() + .map(metrics_set_proto_to_df) + .collect::>()?, + task_metrics: metrics_set_proto_to_df( + task_metrics + .task_metrics + .as_ref() + .ok_or_else(|| DataFusionError::Internal("Missing field 'task_metrics'".into()))?, + )?, + }) +} + /// Custom metrics are not supported in proto conversion. const CUSTOM_METRICS_NOT_SUPPORTED: &str = "custom metrics are not supported in metrics proto conversion"; @@ -1305,3 +1326,50 @@ mod tests { } } } + +#[cfg(test)] +mod decode_tests { + use super::*; + use crate::TaskKey; + use crate::coordinator::MetricsStore; + use datafusion::physical_plan::metrics::{Count, Metric, MetricValue, MetricsSet}; + + /// Pins the driver-facing decode path: a proto metrics frame (the wire form a transport + /// delivers out-of-band) decodes to the plain form and files into the store the per-task + /// EXPLAIN rewrite reads. + #[test] + fn a_metrics_frame_decodes_and_files_into_the_store() { + let count = Count::new(); + count.add(42); + let mut set = MetricsSet::new(); + set.push(Arc::new(Metric::new( + MetricValue::OutputRows(count), + Some(0), + ))); + let node_metrics = df_metrics_set_to_proto(&set).expect("encode"); + + let decoded = decode_task_metrics(pb::TaskMetrics { + pre_order_plan_metrics: vec![node_metrics], + task_metrics: Some(pb::MetricsSet::default()), + }) + .expect("decode"); + assert_eq!(decoded.pre_order_plan_metrics.len(), 1); + assert_eq!( + decoded.pre_order_plan_metrics[0].output_rows(), + Some(42), + "the counter survives the wire round trip" + ); + + let store = MetricsStore::new(); + let key = TaskKey { + query_id: uuid::Uuid::new_v4(), + stage_id: 1, + task_number: 0, + }; + store.insert(key, decoded); + assert!( + store.get(&key).is_some(), + "the driver-filed entry is readable" + ); + } +} diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index dd878f2d..eb7abc57 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -7,11 +7,15 @@ mod channel_resolver; // stack. pub(crate) mod generated; mod in_process; +// The metrics codec sits off gRPC for the same reason as the message types: a transport that +// delivers metrics out-of-band decodes the same frames the gRPC client does. +pub(crate) mod metrics_proto; mod worker_channel; pub use channel_resolver::{ChannelResolver, get_distributed_channel_resolver}; pub(crate) use channel_resolver::{ChannelResolverExtension, set_distributed_channel_resolver}; pub use in_process::InProcessChannelResolver; +pub use metrics_proto::decode_task_metrics; pub use worker_channel::{ CoordinatorToWorkerMsg, ExecuteTaskRequest, GetWorkerInfoRequest, GetWorkerInfoResponse, diff --git a/src/work_unit_feed/mod.rs b/src/work_unit_feed/mod.rs index 1bd55ebf..ec74f52c 100644 --- a/src/work_unit_feed/mod.rs +++ b/src/work_unit_feed/mod.rs @@ -11,6 +11,7 @@ pub(crate) use remote_work_unit_feed::{ }; pub(crate) use work_unit_feed_registry::{WorkUnitFeedRegistry, set_distributed_work_unit_feed}; +pub use remote_work_unit_feed::set_received_time; pub use work_unit::WorkUnit; pub use work_unit_feed::{WorkUnitFeed, WorkUnitFeedProto}; pub use work_unit_feed_provider::{DistributedWorkUnitFeedContext, WorkUnitFeedProvider}; diff --git a/src/work_unit_feed/remote_work_unit_feed.rs b/src/work_unit_feed/remote_work_unit_feed.rs index 18c51ffe..78ecc42a 100644 --- a/src/work_unit_feed/remote_work_unit_feed.rs +++ b/src/work_unit_feed/remote_work_unit_feed.rs @@ -99,6 +99,13 @@ pub(crate) fn set_work_unit_received_time( msg } +/// Stamps the receive time on a bare proto work unit. A transport that hands units across without +/// going through [set_work_unit_received_time] stamps each one as it crosses into the worker, so the +/// worker-side latency math has a delivery timestamp; a missing stamp reads as zero latency. +pub fn set_received_time(work_unit: &mut crate::proto::WorkUnit) { + work_unit.received_timestamp_unix_nanos = now_ns(); +} + /// Remove implementation of a [WorkUnitFeedProvider] that pulls [crate::WorkUnitMsg]s coming over /// the wire from a [RemoteWorkUnitFeedRegistry]. /// diff --git a/src/worker/impl_coordinator_channel.rs b/src/worker/impl_coordinator_channel.rs index cf16e890..9cb5d7bc 100644 --- a/src/worker/impl_coordinator_channel.rs +++ b/src/worker/impl_coordinator_channel.rs @@ -1,5 +1,6 @@ use crate::common::TreeNodeExt; use crate::execution_plans::SamplerExec; +use crate::protocol::metrics_proto::df_metrics_set_to_proto; use crate::work_unit_feed::{RemoteWorkUnitFeedRegistry, set_work_unit_received_time}; use crate::worker::LocalWorkerContext; use crate::worker::task_data::TaskDataMetrics; @@ -231,3 +232,24 @@ fn send_metrics_via_channel( task_metrics: task_data_metrics.to_metrics_set(), }); } + +/// Collects the per-node metrics of `plan` in pre-order traversal order, encoded as protos. A +/// push transport reuses this for its metrics frame so the per-node shape matches what the +/// coordinator channel reports. A node without metrics, or one whose metrics fail to encode, +/// contributes an empty set in place: the consumer matches entries to plan nodes by pre-order +/// position, so dropping an entry would mis-attribute every set after it. +pub fn collect_plan_metrics_protos( + plan: &Arc, + dt_ctx: DistributedTaskContext, +) -> Vec { + let mut pre_order_plan_metrics = vec![]; + let _ = plan.apply_with_dt_ctx(dt_ctx, |node, _| { + pre_order_plan_metrics.push( + node.metrics() + .and_then(|m| df_metrics_set_to_proto(&m).ok()) + .unwrap_or_default(), + ); + Ok(TreeNodeRecursion::Continue) + }); + pre_order_plan_metrics +} diff --git a/src/worker/mod.rs b/src/worker/mod.rs index a0ab8959..759dee43 100644 --- a/src/worker/mod.rs +++ b/src/worker/mod.rs @@ -10,6 +10,7 @@ pub(crate) mod test_utils; mod worker_connection_pool; mod worker_service; +pub use impl_coordinator_channel::collect_plan_metrics_protos; pub(crate) use single_write_multi_read::SingleWriteMultiRead; pub(crate) use worker_connection_pool::{LocalWorkerContext, WorkerConnectionPool};