Tpch bench to report metrics (#2453)

Author: onursatici

Cargo.lock

@@ -70,6 +70,7 @@ url = { workspace = true }
 uuid = { workspace = true, features = ["v4"] }
 vortex = { workspace = true, features = ["object_store", "parquet"] }
 vortex-datafusion = { workspace = true }
+vortex-metrics = { workspace = true }
 xshell = { workspace = true }
 
 [features]

bench-vortex/Cargo.toml

@@ -3,6 +3,7 @@ use std::time::{Duration, Instant};
 
 use bench_vortex::display::{print_measurements_json, render_table, DisplayFormat, RatioMode};
 use bench_vortex::measurements::QueryMeasurement;
+use bench_vortex::metrics::MetricsSetExt;
 use bench_vortex::tpch::dbgen::{DBGen, DBGenOptions};
 use bench_vortex::tpch::duckdb::{generate_tpch, DuckdbTpchOptions};
 use bench_vortex::tpch::{
@@ -11,6 +12,7 @@ use bench_vortex::tpch::{
 };
 use bench_vortex::{default_env_filter, feature_flagged_allocator, setup_logger, Format};
 use clap::{Parser, ValueEnum};
+use datafusion_physical_plan::metrics::{Label, MetricsSet};
 use indicatif::ProgressBar;
 use itertools::Itertools;
 use log::{info, warn};
@@ -48,6 +50,8 @@ struct Args {
     emulate_object_store: bool,
     #[arg(long, default_value_t, value_enum)]
     data_generator: DataGenerator,
+    #[arg(long)]
+    all_metrics: bool,
 }
 
 #[derive(ValueEnum, Default, Clone, Debug)]
@@ -132,6 +136,7 @@ fn main() -> ExitCode {
         args.emulate_object_store,
         args.scale_factor,
         url,
+        args.all_metrics,
     ))
 }
 
@@ -145,6 +150,7 @@ async fn bench_main(
     emulate_object_store: bool,
     scale_factor: u8,
     url: Url,
+    display_all_metrics: bool,
 ) -> ExitCode {
     let expected_row_counts = if scale_factor == 1 {
         EXPECTED_ROW_COUNTS_SF1
@@ -170,6 +176,8 @@ async fn bench_main(
     let mut row_counts = Vec::new();
     let mut measurements = Vec::new();
 
+    let mut metrics = MetricsSet::new();
+
     for format in formats.iter().copied() {
         // Load datasets
         let ctx = load_datasets(&url, format, emulate_object_store)
@@ -192,9 +200,18 @@ async fn bench_main(
             }
 
             for i in 0..2 {
-                let row_count = run_tpch_query(&ctx, &sql_queries, query_idx).await;
+                let (row_count, new_metrics) = run_tpch_query(&ctx, &sql_queries, query_idx).await;
                 if i == 0 {
                     row_counts.push((query_idx, format, row_count));
+                    for (idx, m) in new_metrics.into_iter().enumerate() {
+                        metrics.merge_all_with_label(
+                            m,
+                            &[
+                                Label::new("query_idx", query_idx.to_string()),
+                                Label::new("vortex_exec_idx", idx.to_string()),
+                            ],
+                        );
+                    }
                 }
             }
 
@@ -271,6 +288,12 @@ async fn bench_main(
 
     match display_format {
         DisplayFormat::Table => {
+            if !display_all_metrics {
+                metrics = metrics.aggregate();
+            }
+            for m in metrics.timestamps_removed().sorted_for_display().iter() {
+                println!("{}", m);
+            }
             render_table(measurements, &formats, RatioMode::Time).unwrap();
         }
         DisplayFormat::GhJson => {

bench-vortex/src/bin/tpch.rs

@@ -17,6 +17,7 @@ use datafusion::execution::cache::cache_unit::{DefaultFileStatisticsCache, Defau
 use datafusion::execution::object_store::DefaultObjectStoreRegistry;
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use datafusion::prelude::{SessionConfig, SessionContext};
+use datafusion_physical_plan::metrics::MetricsSet;
 use datafusion_physical_plan::{collect, ExecutionPlan};
 use rand::{Rng, SeedableRng as _};
 use tracing::level_filters::LevelFilter;
@@ -28,6 +29,7 @@ use vortex::error::VortexResult;
 use vortex::sampling_compressor::ALL_ENCODINGS_CONTEXT;
 use vortex::validity::Validity;
 use vortex::{ContextRef, IntoArray};
+use vortex_datafusion::persistent::metrics::VortexMetricsFinder;
 
 pub mod bench_run;
 pub mod blob;
@@ -36,6 +38,7 @@ pub mod compress;
 pub mod datasets;
 pub mod display;
 pub mod measurements;
+pub mod metrics;
 pub mod parquet_reader;
 pub mod random_access;
 pub mod tpch;
@@ -197,12 +200,18 @@ pub fn default_env_filter(is_verbose: bool) -> EnvFilter {
     }
 }
 
-pub async fn execute_query(ctx: &SessionContext, query: &str) -> VortexResult<Vec<RecordBatch>> {
+pub async fn execute_query(
+    ctx: &SessionContext,
+    query: &str,
+) -> VortexResult<(Vec<RecordBatch>, Vec<MetricsSet>)> {
     let plan = ctx.sql(query).await?;
     let (state, plan) = plan.into_parts();
     let physical_plan = state.create_physical_plan(&plan).await?;
     let result = collect(physical_plan.clone(), state.task_ctx()).await?;
-    Ok(result)
+    Ok((
+        result,
+        VortexMetricsFinder::find_all(physical_plan.as_ref()),
+    ))
 }
 
 pub async fn execute_physical_plan(

bench-vortex/src/lib.rs

@@ -0,0 +1,67 @@
+use std::sync::Arc;
+
+use datafusion_physical_plan::metrics::{Label, MetricValue, MetricsSet};
+use datafusion_physical_plan::Metric;
+use vortex::aliases::hash_map::HashMap;
+
+pub trait MetricsSetExt {
+    fn merge_all_with_label(&mut self, other: MetricsSet, labels: &[Label]);
+    fn aggregate(&self) -> Self;
+}
+
+impl MetricsSetExt for MetricsSet {
+    fn merge_all_with_label(&mut self, other: MetricsSet, labels: &[Label]) {
+        for m in other.iter() {
+            let mut new_metric =
+                Metric::new_with_labels(m.value().clone(), m.partition(), m.labels().into());
+            for label in labels.iter() {
+                new_metric = new_metric.with_label(label.clone());
+            }
+            self.push(Arc::new(new_metric));
+        }
+    }
+
+    fn aggregate(&self) -> Self {
+        let mut map = HashMap::new();
+        let filtered = self
+            .iter()
+            .filter(|m| !m.value().name().ends_with("p99")) // can't aggregate percentiles
+            .filter(|m| !m.value().name().ends_with("p95"));
+        for metric in filtered {
+            let key = metric.value().name();
+            map.entry(key)
+                .and_modify(|accum: &mut Metric| {
+                    aggregate_metric(accum.value_mut(), metric.value())
+                })
+                .or_insert_with(|| {
+                    let mut accum = Metric::new(metric.value().new_empty(), None);
+                    aggregate_metric(accum.value_mut(), metric.value());
+                    accum
+                });
+        }
+
+        let mut res = MetricsSet::new();
+        map.into_iter()
+            .map(|(_k, v)| Arc::new(v))
+            .for_each(|m| res.push(m));
+        res
+    }
+}
+
+fn aggregate_metric(metric: &mut MetricValue, to_aggregate: &MetricValue) {
+    match (metric, to_aggregate) {
+        (
+            MetricValue::Gauge { name, gauge },
+            MetricValue::Gauge {
+                gauge: other_gauge, ..
+            },
+        ) => match name {
+            _ if name.ends_with("max") => gauge.set_max(other_gauge.value()),
+            _ if name.ends_with("min") => {
+                gauge.set(gauge.value().min(other_gauge.value()));
+            }
+            _ => gauge.add(other_gauge.value()),
+        },
+        (metric, to_aggregate) => metric.aggregate(to_aggregate),
+    };
+}

bench-vortex/src/metrics.rs

@@ -1,32 +1,27 @@
 use datafusion::prelude::SessionContext;
+use datafusion_physical_plan::metrics::MetricsSet;
 
 use crate::execute_query;
 
-pub async fn run_tpch_query(ctx: &SessionContext, queries: &[String], idx: usize) -> usize {
+pub async fn run_tpch_query(
+    ctx: &SessionContext,
+    queries: &[String],
+    idx: usize,
+) -> (usize, Vec<MetricsSet>) {
     if idx == 15 {
         let mut result = None;
         for (i, q) in queries.iter().enumerate() {
             if i == 1 {
-                result = Some(
-                    execute_query(ctx, q)
-                        .await
-                        .unwrap()
-                        .iter()
-                        .map(|r| r.num_rows())
-                        .sum(),
-                );
+                let (record_batches, metrics) = execute_query(ctx, q).await.unwrap();
+                result = Some((record_batches.iter().map(|r| r.num_rows()).sum(), metrics));
             } else {
                 execute_query(ctx, q).await.unwrap();
             }
         }
         result.expect("Must have had a result in 2nd sql statement for query 15")
     } else {
         let q = &queries[0];
-        execute_query(ctx, q)
-            .await
-            .unwrap()
-            .iter()
-            .map(|r| r.num_rows())
-            .sum()
+        let (record_batches, metrics) = execute_query(ctx, q).await.unwrap();
+        (record_batches.iter().map(|r| r.num_rows()).sum(), metrics)
     }
 }

bench-vortex/src/tpch/execute.rs

@@ -1,17 +1,51 @@
+//! Vortex table provider metrics.
 use std::sync::Arc;
-use std::time::Duration;
 
 use datafusion_physical_plan::metrics::{
     Count, ExecutionPlanMetricsSet, Gauge, Label as DatafusionLabel,
-    MetricValue as DatafusionMetricValue, MetricsSet, Time,
+    MetricValue as DatafusionMetricValue, MetricsSet,
+};
+use datafusion_physical_plan::{
+    accept, ExecutionPlan, ExecutionPlanVisitor, Metric as DatafusionMetric,
 };
-use datafusion_physical_plan::Metric as DatafusionMetric;
 use vortex_metrics::{DefaultTags, Metric, MetricId, Tags, VortexMetrics};
 
-pub static PARTITION_LABEL: &str = "partition";
+use super::execution::VortexExec;
+
+pub(crate) static PARTITION_LABEL: &str = "partition";
+
+/// Extracts datafusion metrics from all VortexExec instances in
+/// a given physical plan.
+#[derive(Default)]
+pub struct VortexMetricsFinder(Vec<MetricsSet>);
+
+impl VortexMetricsFinder {
+    /// find all metrics for VortexExec nodes.
+    pub fn find_all(plan: &dyn ExecutionPlan) -> Vec<MetricsSet> {
+        let mut finder = Self::default();
+        match accept(plan, &mut finder) {
+            Ok(()) => finder.0,
+            Err(_) => Vec::new(),
+        }
+    }
+}
+
+impl ExecutionPlanVisitor for VortexMetricsFinder {
+    type Error = std::convert::Infallible;
+    fn pre_visit(&mut self, plan: &dyn ExecutionPlan) -> Result<bool, Self::Error> {
+        if let Some(metrics) = plan
+            .as_any()
+            .downcast_ref::<VortexExec>()
+            .and_then(|exec| exec.metrics())
+        {
+            self.0.push(metrics);
+        }
+        Ok(true)
+    }
+}
 
 #[derive(Clone, Debug, Default)]
-pub struct VortexExecMetrics {
+pub(crate) struct VortexExecMetrics {
     pub vortex: VortexMetrics,
     pub execution_plan: ExecutionPlanMetricsSet,
 }
@@ -91,16 +125,17 @@ fn metric_value_to_datafusion(name: &str, metric: &Metric) -> Vec<DatafusionMetr
             }
             let snapshot = timer.snapshot();
             if let Ok(max) = snapshot.max().try_into() {
-                res.push(df_time(format!("{name}_max"), max));
+                // NOTE(os): unlike Time metrics, gauges allow custom aggregation
+                res.push(df_gauge(format!("{name}_max"), max));
             }
             if let Ok(min) = snapshot.min().try_into() {
-                res.push(df_time(format!("{name}_min"), min));
+                res.push(df_gauge(format!("{name}_min"), min));
             }
             if let Some(p95) = f_to_u(snapshot.value(0.95)) {
-                res.push(df_time(format!("{name}_p95"), p95 as u64));
+                res.push(df_gauge(format!("{name}_p95"), p95));
             }
             if let Some(p99) = f_to_u(snapshot.value(0.95)) {
-                res.push(df_time(format!("{name}_p99"), p99 as u64));
+                res.push(df_gauge(format!("{name}_p99"), p99));
             }
             res
         }
@@ -127,15 +162,6 @@ fn df_gauge(name: String, value: usize) -> DatafusionMetricValue {
     }
 }
 
-fn df_time(name: String, nanos: u64) -> DatafusionMetricValue {
-    let time = Time::new();
-    time.add_duration(Duration::from_nanos(nanos));
-    DatafusionMetricValue::Time {
-        name: name.into(),
-        time,
-    }
-}
-
 fn f_to_u(f: f64) -> Option<usize> {
     (f.is_finite() && f >= usize::MIN as f64 && f <= usize::MAX as f64).then(|| f.trunc() as usize)
 }

vortex-datafusion/src/persistent/metrics.rs

@@ -3,7 +3,7 @@ mod cache;
 mod config;
 mod execution;
 mod format;
-mod metrics;
+pub mod metrics;
 mod opener;
 mod sink;
 

vortex-datafusion/src/persistent/mod.rs

@@ -10,7 +10,7 @@ use vortex_error::{vortex_err, vortex_panic, VortexExpect, VortexResult};
 use vortex_io::VortexReadAt;
 use vortex_layout::scan::ScanDriver;
 use vortex_layout::segments::{AsyncSegmentReader, SegmentId};
-use vortex_metrics::{Counter, MetricId, VortexMetrics};
+use vortex_metrics::{Counter, VortexMetrics};
 
 use crate::footer::{FileLayout, Segment};
 use crate::segments::channel::SegmentChannel;
@@ -364,14 +364,13 @@ struct CoalescingMetrics {
 
 impl From<VortexMetrics> for CoalescingMetrics {
     fn from(metrics: VortexMetrics) -> Self {
-        let byte_ranges = MetricId::new("vortex.scan.requests.bytes");
-        let requests = MetricId::new("vortex.scan.requests.count");
+        const BYTES: &str = "vortex.scan.requests.bytes";
+        const COUNT: &str = "vortex.scan.requests.count";
         Self {
-            bytes_uncoalesced: metrics.counter(byte_ranges.clone().with_tag("kind", "uncoalesced")),
-            bytes_coalesced: metrics.counter(byte_ranges.with_tag("kind", "coalesced")),
-            request_count_uncoalesced: metrics
-                .counter(requests.clone().with_tag("kind", "uncoalesced")),
-            request_count_coalesced: metrics.counter(requests.with_tag("kind", "coalesced")),
+            bytes_uncoalesced: metrics.counter(format!("{BYTES}.uncoalesced")),
+            bytes_coalesced: metrics.counter(format!("{BYTES}.coalesced")),
+            request_count_uncoalesced: metrics.counter(format!("{COUNT}.uncoalesced")),
+            request_count_coalesced: metrics.counter(format!("{COUNT}.coalesced")),
         }
     }
 }