diff --git a/src/coordinator/distributed.rs b/src/coordinator/distributed.rs index a1e123f2..06904fab 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 the in-crate shm/embedder consumer, which files decoded worker metric 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..cf1fbc6e 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 the in-crate shm/embedder consumer, which files decoded worker metric frames into + // the executed plan's store before the per-task EXPLAIN rewrite reads it. + 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/lib.rs b/src/lib.rs index 8c5c33b0..0d2b37e5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -18,7 +18,9 @@ mod worker_resolver; #[cfg(feature = "grpc")] pub use arrow_ipc::CompressionType; -pub use coordinator::DistributedExec; +// `MetricsStore` is re-exposed for the in-crate shm/embedder consumer, which files worker metric +// 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, PartitionRoute, SessionStateBuilderExt, @@ -52,7 +54,7 @@ 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, diff --git a/src/protocol/metrics_proto.rs b/src/protocol/metrics_proto.rs index 27b715f2..ea95295e 100644 --- a/src/protocol/metrics_proto.rs +++ b/src/protocol/metrics_proto.rs @@ -58,6 +58,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"; diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index a08a3eb4..d21fd4b2 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -13,6 +13,7 @@ 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,