vortex-metrics v2 (#6187)

Closes #6078.

This PR reworks our metrics crate to simplify it and makes it much
harder to make it use unbounded amounts of memory. The implementation is
inspired by our current use of `witchcraft-metrics`, the `metrics`
crate, DataFusion's metrics and some reading of the OTel spec.

It includes several changes:
1. `VortexMetrics` is replaced by the new `MetricsRegistry` trait, which
allows user to integrate metrics to whatever metrics backend they have.
1. Removes metrics from the session, making it the responsibility of the
user to wire it correctly between various components.
1. No more hierarchical metrics instances, just wire labels if they
exist.
1. Removes the `witchcraft-metrics` dependency. It does many things we
don't care about or that are opinionated in a way that I don't think is
intuitive or useful for most users. It also makes our external API less
dependent on other crates we have no control over.
1. Histograms are now f64 and don't use exponential decay, Timers only
speak in `Duration`, and counters are u64s. It also introduces a `Gauge`
type which I don't think we actually use and I might just delete.

I've tested it by running our DF benchmarks with `--show-metrics` and
that seems to not change, if anyone else has made use of the metrics, I
would love to help test and debug them in another context.

---------

Signed-off-by: Adam Gutglick <adam@spiraldb.com>

Author: AdamGS

Cargo.lock

@@ -234,7 +234,6 @@ url = "2.5.7"
 uuid = { version = "1.19", features = ["js"] }
 walkdir = "2.5.0"
 wasm-bindgen-futures = "0.4.39"
-witchcraft-metrics = "1.0.1"
 xshell = "0.2.6"
 zigzag = "0.1.0"
 zip = "6.0.0"

Cargo.toml

@@ -330,9 +330,9 @@ fn print_metrics(plans: &[(usize, Format, Arc<dyn ExecutionPlan>)]) {
 
         eprintln!("metrics for query={query_idx}, {format}:");
         for (scan_idx, metrics_set) in metric_sets.iter().enumerate() {
-            eprintln!("  scan[{scan_idx}]:");
-            for metric in metrics_set.clone().aggregate().sorted_for_display().iter() {
-                eprintln!("    {metric}");
+            eprintln!("\tscan[{scan_idx}]:");
+            for metric in metrics_set.aggregate().sorted_for_display().iter() {
+                eprintln!("\t\t{metric}");
             }
         }
     }

benchmarks/datafusion-bench/src/main.rs

@@ -32,7 +32,6 @@ Each Vortex crate defines a session variable that holds a registry for its exten
 | `ExprSession`     | `vortex-array`   | Scalar expression vtables                    |
 | `LayoutSession`   | `vortex-layout`  | Layout encoding vtables (Flat, Chunked, ...) |
 | `RuntimeSession`  | `vortex-io`      | Async runtime handle                         |
-| `VortexMetrics`   | `vortex-metrics` | Counters, histograms, and timers             |
 | `CudaSession`     | `vortex-cuda`    | CUDA context, kernels, and stream pool       |
 
 Session variables are created lazily on first access with their `Default` implementation, which

docs/developer-guide/internals/session.md

@@ -3,6 +3,7 @@
 
 //! Vortex table provider metrics.
 use std::sync::Arc;
+use std::time::Duration;
 
 use datafusion_datasource::file_scan_config::FileScanConfig;
 use datafusion_datasource::source::DataSourceExec;
@@ -15,13 +16,16 @@ use datafusion_physical_plan::metrics::Gauge;
 use datafusion_physical_plan::metrics::Label as DatafusionLabel;
 use datafusion_physical_plan::metrics::MetricValue as DatafusionMetricValue;
 use datafusion_physical_plan::metrics::MetricsSet;
+use datafusion_physical_plan::metrics::Time;
+use vortex::error::VortexExpect;
+use vortex::metrics::Label;
 use vortex::metrics::Metric;
-use vortex::metrics::MetricId;
-use vortex::metrics::Tags;
+use vortex::metrics::MetricValue;
 
 use crate::persistent::source::VortexSource;
 
 pub(crate) static PARTITION_LABEL: &str = "partition";
+pub(crate) static PATH_LABEL: &str = "file_path";
 
 /// Extracts datafusion metrics from all VortexExec instances in
 /// a given physical plan.
@@ -43,9 +47,8 @@ impl ExecutionPlanVisitor for VortexMetricsFinder {
     type Error = std::convert::Infallible;
     fn pre_visit(&mut self, plan: &dyn ExecutionPlan) -> Result<bool, Self::Error> {
         if let Some(exec) = plan.as_any().downcast_ref::<DataSourceExec>() {
-            if let Some(metrics) = exec.metrics() {
-                self.0.push(metrics);
-            }
+            // Start with exec metrics or create a new set
+            let mut set = exec.metrics().unwrap_or_default();
 
             // Include our own metrics from VortexSource
             if let Some(file_scan) = exec.data_source().as_any().downcast_ref::<FileScanConfig>()
@@ -54,92 +57,99 @@ impl ExecutionPlanVisitor for VortexMetricsFinder {
                     .as_any()
                     .downcast_ref::<VortexSource>()
             {
-                let mut set = MetricsSet::new();
                 for metric in scan
-                    .vx_metrics()
+                    .metrics_registry()
                     .snapshot()
                     .iter()
-                    .flat_map(|(id, metric)| metric_to_datafusion(id, metric))
+                    .flat_map(metric_to_datafusion)
                 {
                     set.push(Arc::new(metric));
                 }
-
-                self.0.push(set);
             }
-        }
 
-        Ok(true)
+            self.0.push(set);
+
+            Ok(false)
+        } else {
+            Ok(true)
+        }
     }
 }
 
-fn metric_to_datafusion(id: MetricId, metric: &Metric) -> impl Iterator<Item = DatafusionMetric> {
-    let (partition, labels) = tags_to_datafusion(id.tags());
-    metric_value_to_datafusion(id.name(), metric)
+fn metric_to_datafusion(metric: &Metric) -> impl Iterator<Item = DatafusionMetric> {
+    let (partition, labels) = labels_to_datafusion(metric.labels());
+    metric_value_to_datafusion(metric.name(), metric.value())
         .into_iter()
         .map(move |metric_value| {
             DatafusionMetric::new_with_labels(metric_value, partition, labels.clone())
         })
 }
 
-fn tags_to_datafusion(tags: &Tags) -> (Option<usize>, Vec<DatafusionLabel>) {
+fn labels_to_datafusion(tags: &[Label]) -> (Option<usize>, Vec<DatafusionLabel>) {
     tags.iter()
-        .fold((None, Vec::new()), |(mut partition, mut labels), (k, v)| {
-            if k == PARTITION_LABEL {
-                partition = v.parse().ok();
+        .fold((None, Vec::new()), |(mut partition, mut labels), metric| {
+            if metric.key() == PARTITION_LABEL {
+                partition = metric.value().parse().ok();
             } else {
-                labels.push(DatafusionLabel::new(k.to_string(), v.to_string()));
+                labels.push(DatafusionLabel::new(
+                    metric.key().to_string(),
+                    metric.value().to_string(),
+                ));
             }
             (partition, labels)
         })
 }
 
-fn metric_value_to_datafusion(name: &str, metric: &Metric) -> Vec<DatafusionMetricValue> {
+fn metric_value_to_datafusion(name: &str, metric: &MetricValue) -> Vec<DatafusionMetricValue> {
     match metric {
-        Metric::Counter(counter) => counter
-            .count()
+        MetricValue::Counter(counter) => counter
+            .value()
             .try_into()
             .into_iter()
             .map(|count| df_counter(name.to_string(), count))
             .collect(),
-        Metric::Histogram(hist) => {
+        MetricValue::Histogram(hist) => {
             let mut res = Vec::new();
-            if let Ok(count) = hist.count().try_into() {
-                res.push(df_counter(format!("{name}_count"), count));
-            }
-            let snapshot = hist.snapshot();
-            if let Ok(max) = snapshot.max().try_into() {
-                res.push(df_gauge(format!("{name}_max"), max));
-            }
-            if let Ok(min) = snapshot.min().try_into() {
-                res.push(df_gauge(format!("{name}_min"), min));
-            }
-            if let Some(p90) = f_to_u(snapshot.value(0.90)) {
-                res.push(df_gauge(format!("{name}_p95"), p90));
-            }
-            if let Some(p99) = f_to_u(snapshot.value(0.99)) {
-                res.push(df_gauge(format!("{name}_p99"), p99));
+
+            res.push(df_counter(format!("{name}_count"), hist.count()));
+
+            if !hist.is_empty() {
+                if let Some(max) = f_to_u(hist.quantile(1.0).vortex_expect("must not be empty")) {
+                    res.push(df_gauge(format!("{name}_max"), max));
+                }
+
+                if let Some(min) = f_to_u(hist.quantile(0.0).vortex_expect("must not be empty")) {
+                    res.push(df_gauge(format!("{name}_min"), min));
+                }
+
+                if let Some(p95) = f_to_u(hist.quantile(0.95).vortex_expect("must not be empty")) {
+                    res.push(df_gauge(format!("{name}_p95"), p95));
+                }
+                if let Some(p99) = f_to_u(hist.quantile(0.99).vortex_expect("must not be empty")) {
+                    res.push(df_gauge(format!("{name}_p99"), p99));
+                }
             }
+
             res
         }
-        Metric::Timer(timer) => {
+        MetricValue::Timer(timer) => {
             let mut res = Vec::new();
-            if let Ok(count) = timer.count().try_into() {
-                res.push(df_counter(format!("{name}_count"), count));
-            }
-            let snapshot = timer.snapshot();
-            if let Ok(max) = snapshot.max().try_into() {
-                // 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_gauge(format!("{name}_min"), min));
-            }
-            if let Some(p95) = f_to_u(snapshot.value(0.95)) {
-                res.push(df_gauge(format!("{name}_p95"), p95));
-            }
-            if let Some(p99) = f_to_u(snapshot.value(0.95)) {
-                res.push(df_gauge(format!("{name}_p99"), p99));
+            res.push(df_counter(format!("{name}_count"), timer.count()));
+
+            if !timer.is_empty() {
+                let max = timer.quantile(1.0).vortex_expect("must not be empty");
+                res.push(df_timer(format!("{name}_max"), max));
+
+                let min = timer.quantile(0.0).vortex_expect("must not be empty");
+                res.push(df_timer(format!("{name}_min"), min));
+
+                let p95 = timer.quantile(0.95).vortex_expect("must not be empty");
+                res.push(df_timer(format!("{name}_p95"), p95));
+
+                let p99 = timer.quantile(0.99).vortex_expect("must not be empty");
+                res.push(df_timer(format!("{name}_p99"), p99));
             }
+
             res
         }
         // TODO(os): add more metric types when added to VortexMetrics
@@ -165,6 +175,15 @@ fn df_gauge(name: String, value: usize) -> DatafusionMetricValue {
     }
 }
 
+fn df_timer(name: String, value: Duration) -> DatafusionMetricValue {
+    let time = Time::new();
+    time.add_duration(value);
+    DatafusionMetricValue::Time {
+        name: name.into(),
+        time,
+    }
+}
+
 #[expect(
     clippy::cast_possible_truncation,
     reason = "truncation is checked before cast"
@@ -174,3 +193,74 @@ fn f_to_u(f: f64) -> Option<usize> {
         // After the range check, truncation is guaranteed to keep the value in usize bounds.
         f.trunc() as usize)
 }
+
+#[cfg(test)]
+mod tests {
+
+    use datafusion_datasource::source::DataSourceExec;
+    use datafusion_physical_plan::ExecutionPlanVisitor;
+    use datafusion_physical_plan::accept;
+
+    use super::VortexMetricsFinder;
+    use crate::common_tests::TestSessionContext;
+
+    /// Counts the number of DataSourceExec nodes in a plan.
+    struct DataSourceExecCounter(usize);
+
+    impl ExecutionPlanVisitor for DataSourceExecCounter {
+        type Error = std::convert::Infallible;
+        fn pre_visit(
+            &mut self,
+            plan: &dyn datafusion_physical_plan::ExecutionPlan,
+        ) -> Result<bool, Self::Error> {
+            if plan.as_any().downcast_ref::<DataSourceExec>().is_some() {
+                self.0 += 1;
+                Ok(false)
+            } else {
+                Ok(true)
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn metrics_finder_returns_one_set_per_data_source_exec() -> anyhow::Result<()> {
+        let ctx = TestSessionContext::default();
+
+        ctx.session
+            .sql(
+                "CREATE EXTERNAL TABLE my_tbl \
+                (c1 VARCHAR NOT NULL, c2 INT NOT NULL) \
+                STORED AS vortex \
+                LOCATION 'files/'",
+            )
+            .await?;
+
+        ctx.session
+            .sql("INSERT INTO my_tbl VALUES ('a', 1), ('b', 2)")
+            .await?
+            .collect()
+            .await?;
+
+        let df = ctx.session.sql("SELECT * FROM my_tbl").await?;
+        let (state, plan) = df.into_parts();
+        let physical_plan = state.create_physical_plan(&plan).await?;
+
+        // Count DataSourceExec nodes
+        let mut counter = DataSourceExecCounter(0);
+        accept(physical_plan.as_ref(), &mut counter)?;
+
+        // Get metrics sets
+        let metrics_sets = VortexMetricsFinder::find_all(physical_plan.as_ref());
+
+        assert!(!metrics_sets.is_empty());
+        assert_eq!(
+            metrics_sets.len(),
+            counter.0,
+            "Expected one MetricsSet per DataSourceExec, got {} sets for {} DataSourceExec nodes",
+            metrics_sets.len(),
+            counter.0
+        );
+
+        Ok(())
+    }
+}