Add --collect-metrics argument to benchmarks

Author: gabotechs

benchmarks/cdk/bin/datafusion-bench.ts

@@ -15,6 +15,7 @@ async function main () {
     .option('-i, --iterations <number>', 'Number of iterations', '3')
     .option('--files-per-task <number>', 'Files per task', '4')
     .option('--cardinality-task-sf <number>', 'Cardinality task scale factor', '2')
+    .option('--collect-metrics <boolean>', 'Propagates metric collection', 'true')
     .option('--query <number>', 'A specific query to run', undefined)
     .parse(process.argv);
 
@@ -24,6 +25,7 @@ async function main () {
   const iterations = parseInt(options.iterations);
   const filesPerTask = parseInt(options.filesPerTask);
   const cardinalityTaskSf = parseInt(options.cardinalityTaskSf);
+  const collectMetrics = options.collectMetrics === 'true' || options.collectMetrics === 1
 
   // Compare with previous results first
   const results: BenchmarkResults = { queries: [] };
@@ -33,7 +35,8 @@ async function main () {
   await query(createTablesSql(sf))
   await query(`
     SET distributed.files_per_task=${filesPerTask};
-    SET distributed.cardinality_task_count_factor=${cardinalityTaskSf}
+    SET distributed.cardinality_task_count_factor=${cardinalityTaskSf};
+    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));