Make metrics collection optional

Author: gabotechs

benchmarks/cdk/bin/datafusion-bench.ts

@@ -16,6 +16,7 @@ async function main () {
     .option('--files-per-task <number>', 'Files per task', '4')
     .option('--cardinality-task-sf <number>', 'Cardinality task scale factor', '2')
     .option('--shuffle-batch-size <number>', 'Shuffle batch coalescing size (number of rows)', '8192')
+    .option('--collect-metrics <boolean>', 'Propagates metric collection', 'true')
     .option('--query <number>', 'A specific query to run', undefined)
     .parse(process.argv);
 
@@ -26,6 +27,7 @@ async function main () {
   const filesPerTask = parseInt(options.filesPerTask);
   const cardinalityTaskSf = parseInt(options.cardinalityTaskSf);
   const shuffleBatchSize = parseInt(options.shuffleBatchSize);
+  const collectMetrics = options.collectMetrics === 'true' || options.collectMetrics === 1
 
   // Compare with previous results first
   const results: BenchmarkResults = { queries: [] };
@@ -36,7 +38,8 @@ async function main () {
   await query(`
     SET distributed.files_per_task=${filesPerTask};
     SET distributed.cardinality_task_count_factor=${cardinalityTaskSf};
-    SET distributed.shuffle_batch_size=${shuffleBatchSize}
+    SET distributed.shuffle_batch_size=${shuffleBatchSize};
+    SET distributed.collect_metrics=${collectMetrics}
   `)
 
   for (let id of IDS) {

benchmarks/src/tpch/run.rs

@@ -120,6 +120,10 @@ pub struct RunOpt {
     /// Task count scale factor for when nodes in stages change the cardinality of the data
     #[structopt(long)]
     cardinality_task_sf: Option<f64>,
+
+    /// Collects metrics across network boundaries
+    #[structopt(long)]
+    collect_metrics: bool,
 }
 
 #[async_trait]
@@ -147,7 +151,8 @@ impl DistributedSessionBuilder for RunOpt {
             )?
             .with_distributed_cardinality_effect_task_scale_factor(
                 self.cardinality_task_sf.unwrap_or(1.0),
-            )?;
+            )?
+            .with_distributed_metrics_collection(self.collect_metrics)?;
 
         if self.mem_table {
             builder = builder.with_physical_optimizer_rule(Arc::new(InMemoryDataSourceRule));

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!(

src/distributed_ext.rs

@@ -360,6 +360,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 {
@@ -374,7 +381,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) {
@@ -417,6 +425,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)]
@@ -450,6 +464,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>;
         }
     }
 }
@@ -496,6 +514,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>;
         }
     }
 }
@@ -542,6 +565,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>;
         }
     }
 }
@@ -588,6 +616,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

@@ -31,6 +31,9 @@ extensions_options! {
         /// batches over the wire.
         /// If set to 0, batch coalescing is disabled on network shuffle operations.
         pub shuffle_batch_size: usize, default = 8192
+        /// 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 https://github.com/datafusion-contrib/datafusion-distributed/issues/247 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).left_stream()
+            } else {
+                stream.right_stream()
+            };
 
-            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();
 

src/execution_plans/network_shuffle.rs

@@ -1,13 +1,17 @@
 use crate::channel_resolver_ext::get_distributed_channel_resolver;
 use crate::config_extension_ext::ContextGrpcMetadata;
-use crate::execution_plans::common::{require_one_child, scale_partitioning};
+use crate::execution_plans::common::{
+    manually_propagate_distributed_config, require_one_child, scale_partitioning,
+};
 use crate::flight_service::DoGet;
 use crate::metrics::MetricsCollectingStream;
 use crate::metrics::proto::MetricsSetProto;
 use crate::protobuf::StageKey;
 use crate::protobuf::{map_flight_to_datafusion_error, map_status_to_datafusion_error};
 use crate::stage::{MaybeEncodedPlan, Stage};
-use crate::{ChannelResolver, DistributedTaskContext, InputStageInfo, NetworkBoundary};
+use crate::{
+    ChannelResolver, DistributedConfig, DistributedTaskContext, InputStageInfo, NetworkBoundary,
+};
 use arrow_flight::Ticket;
 use arrow_flight::decode::FlightRecordBatchStream;
 use arrow_flight::error::FlightError;
@@ -325,6 +329,9 @@ impl ExecutionPlan for NetworkShuffleExec {
         // 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()?;
 
@@ -337,6 +344,8 @@ impl ExecutionPlan for NetworkShuffleExec {
         let task_context = DistributedTaskContext::from_ctx(&context);
         let off = self_ready.properties.partitioning.partition_count() * task_context.task_index;
 
+        // TODO: this propagation should be automatic <link to issue>
+        let context_headers = manually_propagate_distributed_config(context_headers, d_cfg);
         let stream = input_stage_tasks.into_iter().enumerate().map(|(i, task)| {
             let channel_resolver = Arc::clone(&channel_resolver);
 
@@ -370,13 +379,14 @@ impl ExecutionPlan for NetworkShuffleExec {
                     .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).left_stream()
+                } else {
+                    stream.right_stream()
+                };
 
-                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()
             .boxed()