One MetricsSet per VortexSource node

Signed-off-by: Adam Gutglick <[email protected]>

Author: AdamGS

benchmarks/datafusion-bench/src/main.rs

@@ -321,7 +321,7 @@ 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!("\tscan[{scan_idx}]:");
-            for metric in metrics_set.clone().aggregate().sorted_for_display().iter() {
+            for metric in metrics_set.aggregate().sorted_for_display().iter() {
                 eprintln!("\t\t{metric}");
             }
         }

vortex-datafusion/src/persistent/metrics.rs

@@ -47,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>()
@@ -58,7 +57,6 @@ impl ExecutionPlanVisitor for VortexMetricsFinder {
                     .as_any()
                     .downcast_ref::<VortexSource>()
             {
-                let mut set = MetricsSet::new();
                 for metric in scan
                     .metrics_registry()
                     .snapshot()
@@ -67,12 +65,14 @@ impl ExecutionPlanVisitor for VortexMetricsFinder {
                 {
                     set.push(Arc::new(metric));
                 }
-
-                self.0.push(set);
             }
-        }
 
-        Ok(true)
+            self.0.push(set);
+
+            Ok(false)
+        } else {
+            Ok(true)
+        }
     }
 }
 
@@ -193,3 +193,83 @@ 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::execution::SessionStateBuilder;
+    use datafusion::prelude::SessionContext;
+    use datafusion_datasource::source::DataSourceExec;
+    use datafusion_physical_plan::ExecutionPlanVisitor;
+    use datafusion_physical_plan::accept;
+    use tempfile::TempDir;
+
+    use super::VortexMetricsFinder;
+    use crate::VortexFormatFactory;
+    use crate::persistent::register_vortex_format_factory;
+
+    /// 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 dir = TempDir::new()?;
+
+        let factory = VortexFormatFactory::new();
+        let mut session_state_builder = SessionStateBuilder::new().with_default_features();
+        register_vortex_format_factory(factory, &mut session_state_builder);
+        let session = SessionContext::new_with_state(session_state_builder.build());
+
+        session
+            .sql(&format!(
+                "CREATE EXTERNAL TABLE my_tbl \
+                (c1 VARCHAR NOT NULL, c2 INT NOT NULL) \
+                STORED AS vortex \
+                LOCATION '{}'",
+                dir.path().to_str().unwrap()
+            ))
+            .await?;
+
+        session
+            .sql("INSERT INTO my_tbl VALUES ('a', 1), ('b', 2)")
+            .await?
+            .collect()
+            .await?;
+
+        let df = 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.len() > 0);
+        assert_eq!(
+            metrics_sets.len(),
+            counter.0,
+            "Expected one MetricsSet per DataSourceExec, got {} sets for {} DataSourceExec nodes",
+            metrics_sets.len(),
+            counter.0
+        );
+
+        Ok(())
+    }
+}