Propagate collected metrics optionally

Author: gabotechs

src/channel_resolver_ext.rs

@@ -1,4 +1,5 @@
 use crate::DistributedConfig;
+use crate::config_extension_ext::set_distributed_option_extension;
 use arrow_flight::flight_service_client::FlightServiceClient;
 use async_trait::async_trait;
 use datafusion::common::exec_err;
@@ -17,10 +18,10 @@ pub(crate) fn set_distributed_channel_resolver(
     if let Some(distributed_cfg) = opts.extensions.get_mut::<DistributedConfig>() {
         distributed_cfg.__private_channel_resolver = channel_resolver_ext;
     } else {
-        opts.extensions.insert(DistributedConfig {
+        set_distributed_option_extension(cfg, DistributedConfig {
             __private_channel_resolver: channel_resolver_ext,
             ..Default::default()
-        });
+        }).expect("Calling set_distributed_option_extension with a default DistributedConfig should never fail");
     }
 }
 

src/config_extension_ext.rs

@@ -19,6 +19,14 @@ pub(crate) fn set_distributed_option_extension<T: ConfigExtension + Default>(
     let mut meta = HeaderMap::new();
 
     for entry in t.entries() {
+        // assume that fields starting with "__" are private, and are not supposed to be sent
+        // over the wire. This accounts for the fact that we need to send our DistributedConfig
+        // options without setting the __private_task_estimator and __private_channel_resolver.
+        // Ideally those two fields should not even be there on the first place, but until
+        // https://github.com/apache/datafusion/pull/18739 we need to put them there.
+        if entry.key.starts_with("__") {
+            continue;
+        }
         if let Some(value) = entry.value {
             meta.insert(
                 HeaderName::from_str(&format!(
@@ -44,29 +52,51 @@ pub(crate) fn set_distributed_option_extension<T: ConfigExtension + Default>(
     Ok(())
 }
 
-pub(crate) fn set_distributed_option_extension_from_headers<T: ConfigExtension + Default>(
-    cfg: &mut SessionConfig,
+pub(crate) fn set_distributed_option_extension_from_headers<'a, T: ConfigExtension + Default>(
+    cfg: &'a mut SessionConfig,
     headers: &HeaderMap,
-) -> Result<(), DataFusionError> {
-    let mut result = T::default();
-    let mut found_some = false;
+) -> Result<&'a T, DataFusionError> {
+    enum MutOrOwned<'a, T> {
+        Mut(&'a mut T),
+        Owned(T),
+    }
+
+    impl<'a, T> MutOrOwned<'a, T> {
+        fn as_mut(&mut self) -> &mut T {
+            match self {
+                MutOrOwned::Mut(v) => v,
+                MutOrOwned::Owned(v) => v,
+            }
+        }
+    }
+
+    // If the config extension existed before, we want to modify instead of adding a new one from
+    // scratch. If not, we'll start from scratch with a new one.
+    let mut result = match cfg.options_mut().extensions.get_mut::<T>() {
+        Some(v) => MutOrOwned::Mut(v),
+        None => MutOrOwned::Owned(T::default()),
+    };
+
     for (k, v) in headers.iter() {
         let key = k.as_str().trim_start_matches(FLIGHT_METADATA_CONFIG_PREFIX);
         let prefix = format!("{}.", T::PREFIX);
         if key.starts_with(&prefix) {
-            found_some = true;
-            result.set(
+            result.as_mut().set(
                 key.trim_start_matches(&prefix),
                 v.to_str()
                     .map_err(|err| internal_datafusion_err!("Cannot parse header value: {err}"))?,
             )?;
         }
     }
-    if !found_some {
-        return Ok(());
+
+    // Only insert the extension if it is not already there. If this is otherwise MutOrOwned::Mut it
+    // means that the extension was already there, and we already modified it.
+    if let MutOrOwned::Owned(v) = result {
+        cfg.options_mut().extensions.insert(v);
     }
-    cfg.options_mut().extensions.insert(result);
-    Ok(())
+    cfg.options().extensions.get().ok_or_else(|| {
+        internal_datafusion_err!("ProgrammingError: a config option extension was just inserted, but it was not immediately retrievable")
+    })
 }
 
 #[derive(Clone, Debug, Default)]
@@ -190,8 +220,15 @@ mod tests {
             &Default::default(),
         )?;
 
-        let extension = config.options().extensions.get::<CustomExtension>();
-        assert!(extension.is_none());
+        let extension = config
+            .options()
+            .extensions
+            .get::<CustomExtension>()
+            .unwrap();
+        let default = CustomExtension::default();
+        assert_eq!(extension.foo, default.foo);
+        assert_eq!(extension.bar, default.bar);
+        assert_eq!(extension.baz, default.baz);
 
         Ok(())
     }
@@ -207,8 +244,15 @@ mod tests {
 
         set_distributed_option_extension_from_headers::<CustomExtension>(&mut config, &header_map)?;
 
-        let extension = config.options().extensions.get::<CustomExtension>();
-        assert!(extension.is_none());
+        let extension = config
+            .options()
+            .extensions
+            .get::<CustomExtension>()
+            .unwrap();
+        let default = CustomExtension::default();
+        assert_eq!(extension.foo, default.foo);
+        assert_eq!(extension.bar, default.bar);
+        assert_eq!(extension.baz, default.baz);
 
         Ok(())
     }

src/distributed_ext.rs

@@ -358,6 +358,13 @@ pub trait DistributedExt: Sized {
         &mut self,
         factor: f64,
     ) -> Result<(), DataFusionError>;
+
+    /// Enables metrics collection across network boundaries so that all the metrics gather in
+    /// each node are accessible from the head stage that started running the query.
+    fn with_distributed_metrics_collection(self, enabled: bool) -> Result<Self, DataFusionError>;
+
+    /// Same as [DistributedExt::with_distributed_metrics_collection] but with an in-place mutation.
+    fn set_distributed_metrics_collection(&mut self, enabled: bool) -> Result<(), DataFusionError>;
 }
 
 impl DistributedExt for SessionConfig {
@@ -372,7 +379,8 @@ impl DistributedExt for SessionConfig {
         &mut self,
         headers: &HeaderMap,
     ) -> Result<(), DataFusionError> {
-        set_distributed_option_extension_from_headers::<T>(self, headers)
+        set_distributed_option_extension_from_headers::<T>(self, headers)?;
+        Ok(())
     }
 
     fn set_distributed_user_codec<T: PhysicalExtensionCodec + 'static>(&mut self, codec: T) {
@@ -415,6 +423,12 @@ impl DistributedExt for SessionConfig {
         Ok(())
     }
 
+    fn set_distributed_metrics_collection(&mut self, enabled: bool) -> Result<(), DataFusionError> {
+        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        d_cfg.collect_metrics = enabled;
+        Ok(())
+    }
+
     delegate! {
         to self {
             #[call(set_distributed_option_extension)]
@@ -448,6 +462,10 @@ impl DistributedExt for SessionConfig {
             #[call(set_distributed_cardinality_effect_task_scale_factor)]
             #[expr($?;Ok(self))]
             fn with_distributed_cardinality_effect_task_scale_factor(mut self, factor: f64) -> Result<Self, DataFusionError>;
+
+            #[call(set_distributed_metrics_collection)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_metrics_collection(mut self, enabled: bool) -> Result<Self, DataFusionError>;
         }
     }
 }
@@ -494,6 +512,11 @@ impl DistributedExt for SessionStateBuilder {
             #[call(set_distributed_cardinality_effect_task_scale_factor)]
             #[expr($?;Ok(self))]
             fn with_distributed_cardinality_effect_task_scale_factor(mut self, factor: f64) -> Result<Self, DataFusionError>;
+
+            fn set_distributed_metrics_collection(&mut self, enabled: bool) -> Result<(), DataFusionError>;
+            #[call(set_distributed_metrics_collection)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_metrics_collection(mut self, enabled: bool) -> Result<Self, DataFusionError>;
         }
     }
 }
@@ -540,6 +563,11 @@ impl DistributedExt for SessionState {
             #[call(set_distributed_cardinality_effect_task_scale_factor)]
             #[expr($?;Ok(self))]
             fn with_distributed_cardinality_effect_task_scale_factor(mut self, factor: f64) -> Result<Self, DataFusionError>;
+
+            fn set_distributed_metrics_collection(&mut self, enabled: bool) -> Result<(), DataFusionError>;
+            #[call(set_distributed_metrics_collection)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_metrics_collection(mut self, enabled: bool) -> Result<Self, DataFusionError>;
         }
     }
 }
@@ -586,6 +614,11 @@ impl DistributedExt for SessionContext {
             #[call(set_distributed_cardinality_effect_task_scale_factor)]
             #[expr($?;Ok(self))]
             fn with_distributed_cardinality_effect_task_scale_factor(self, factor: f64) -> Result<Self, DataFusionError>;
+
+            fn set_distributed_metrics_collection(&mut self, enabled: bool) -> Result<(), DataFusionError>;
+            #[call(set_distributed_metrics_collection)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_metrics_collection(self, enabled: bool) -> Result<Self, DataFusionError>;
         }
     }
 }

src/distributed_planner/distributed_config.rs

@@ -24,6 +24,9 @@ extensions_options! {
         /// - If a node reduces the cardinality of the data, this factor will decrease.
         /// - In any other situation, this factor is left intact.
         pub cardinality_task_count_factor: f64, default = cardinality_task_count_factor_default()
+        /// Propagate collected metrics from all nodes in the plan across network boundaries
+        /// so that they can be reconstructed on the head node of the plan.
+        pub collect_metrics: bool, default = false
         /// Collection of [TaskEstimator]s that will be applied to leaf nodes in order to
         /// estimate how many tasks should be spawned for the [Stage] containing the leaf node.
         pub(crate) __private_task_estimator: CombinedTaskEstimator, default = CombinedTaskEstimator::default()

src/distributed_planner/task_estimator.rs

@@ -1,3 +1,4 @@
+use crate::config_extension_ext::set_distributed_option_extension;
 use crate::{ChannelResolver, DistributedConfig, PartitionIsolatorExec};
 use datafusion::catalog::memory::DataSourceExec;
 use datafusion::config::ConfigOptions;
@@ -103,10 +104,10 @@ pub(crate) fn set_distributed_task_estimator(
     } else {
         let mut estimators = CombinedTaskEstimator::default();
         estimators.user_provided.push(Arc::new(estimator));
-        opts.extensions.insert(DistributedConfig {
+        set_distributed_option_extension(cfg, DistributedConfig {
             __private_task_estimator: estimators,
             ..Default::default()
-        });
+        }).expect("Calling set_distributed_option_extension with a default DistributedConfig should never fail");
     }
 }
 

src/execution_plans/common.rs

@@ -1,6 +1,8 @@
+use crate::DistributedConfig;
 use datafusion::common::{DataFusionError, plan_err};
 use datafusion::physical_expr::Partitioning;
 use datafusion::physical_plan::{ExecutionPlan, PlanProperties};
+use http::HeaderMap;
 use std::borrow::Borrow;
 use std::sync::Arc;
 
@@ -40,3 +42,16 @@ pub(super) fn scale_partitioning(
         Partitioning::UnknownPartitioning(p) => Partitioning::UnknownPartitioning(f(*p)),
     }
 }
+
+/// Manual propagation of the [DistributedConfig] fields relevant for execution. Can be removed
+/// after <link to issue> is fixed, as this will become automatic.
+pub(super) fn manually_propagate_distributed_config(
+    mut headers: HeaderMap,
+    d_cfg: &DistributedConfig,
+) -> HeaderMap {
+    headers.insert(
+        "distributed.collect_metrics",
+        d_cfg.collect_metrics.to_string().parse().unwrap(),
+    );
+    headers
+}

src/execution_plans/network_coalesce.rs

@@ -1,13 +1,15 @@
 use crate::channel_resolver_ext::get_distributed_channel_resolver;
 use crate::config_extension_ext::ContextGrpcMetadata;
 use crate::distributed_planner::{InputStageInfo, NetworkBoundary, limit_tasks_err};
-use crate::execution_plans::common::{require_one_child, scale_partitioning_props};
+use crate::execution_plans::common::{
+    manually_propagate_distributed_config, require_one_child, scale_partitioning_props,
+};
 use crate::flight_service::DoGet;
 use crate::metrics::MetricsCollectingStream;
 use crate::metrics::proto::MetricsSetProto;
 use crate::protobuf::{StageKey, map_flight_to_datafusion_error, map_status_to_datafusion_error};
 use crate::stage::{MaybeEncodedPlan, Stage};
-use crate::{ChannelResolver, DistributedTaskContext};
+use crate::{ChannelResolver, DistributedConfig, DistributedTaskContext};
 use arrow_flight::Ticket;
 use arrow_flight::decode::FlightRecordBatchStream;
 use arrow_flight::error::FlightError;
@@ -18,7 +20,7 @@ use datafusion::error::DataFusionError;
 use datafusion::execution::{SendableRecordBatchStream, TaskContext};
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
-use futures::{TryFutureExt, TryStreamExt};
+use futures::{StreamExt, TryFutureExt, TryStreamExt};
 use http::Extensions;
 use prost::Message;
 use std::any::Any;
@@ -258,6 +260,9 @@ impl ExecutionPlan for NetworkCoalesceExec {
         // get the channel manager and current stage from our context
         let channel_resolver = get_distributed_channel_resolver(context.session_config())?;
 
+        let d_cfg = DistributedConfig::from_config_options(context.session_config().options())?;
+        let retrieve_metrics = d_cfg.collect_metrics;
+
         let input_stage = &self_ready.input_stage;
         let encoded_input_plan = input_stage.plan.encoded()?;
 
@@ -273,8 +278,10 @@ impl ExecutionPlan for NetworkCoalesceExec {
         let target_task = partition / partitions_per_task;
         let target_partition = partition % partitions_per_task;
 
+        // TODO: this propagation should be automatic <link to issue>
+        let context_headers = manually_propagate_distributed_config(context_headers, d_cfg);
         let ticket = Request::from_parts(
-            MetadataMap::from_headers(context_headers.clone()),
+            MetadataMap::from_headers(context_headers),
             Extensions::default(),
             Ticket {
                 ticket: DoGet {
@@ -311,13 +318,14 @@ impl ExecutionPlan for NetworkCoalesceExec {
                 .into_inner()
                 .map_err(|err| FlightError::Tonic(Box::new(err)));
 
-            let metrics_collecting_stream =
-                MetricsCollectingStream::new(stream, metrics_collection_capture);
+            let stream = if retrieve_metrics {
+                MetricsCollectingStream::new(stream, metrics_collection_capture).boxed()
+            } else {
+                stream.boxed()
+            };
 
-            Ok(
-                FlightRecordBatchStream::new_from_flight_data(metrics_collecting_stream)
-                    .map_err(map_flight_to_datafusion_error),
-            )
+            Ok(FlightRecordBatchStream::new_from_flight_data(stream)
+                .map_err(map_flight_to_datafusion_error))
         }
         .try_flatten_stream();