feat!: track cumulative wall time in analyze plan

rust/lance-datafusion/src/exec.rs

@@ -10,6 +10,8 @@ use std::{
     time::Duration,
 };
 
+use chrono::{DateTime, Utc};
+
 use arrow_array::RecordBatch;
 use arrow_schema::Schema as ArrowSchema;
 use datafusion::{
@@ -26,6 +28,7 @@ use datafusion::{
         analyze::AnalyzeExec,
         display::DisplayableExecutionPlan,
         execution_plan::{Boundedness, CardinalityEffect, EmissionType},
+        metrics::MetricValue,
         stream::RecordBatchStreamAdapter,
         streaming::PartitionStream,
         DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream,
@@ -567,22 +570,87 @@ pub fn format_plan(plan: Arc<dyn ExecutionPlan>) -> String {
     struct CalculateVisitor {
         highest_index: usize,
         index_to_cumulative_cpu: HashMap<usize, usize>,
+        index_to_cumulative_wall: HashMap<usize, Duration>,
+    }
+
+    /// Result of calculating metrics for a subtree
+    struct SubtreeMetrics {
+        cumulative_cpu: usize,
+        min_start: Option<DateTime<Utc>>,
+        max_end: Option<DateTime<Utc>>,
     }
+
     impl CalculateVisitor {
-        fn calculate_cumulative_cpu(&mut self, plan: &Arc<dyn ExecutionPlan>) -> usize {
+        fn calculate_metrics(&mut self, plan: &Arc<dyn ExecutionPlan>) -> SubtreeMetrics {
             self.highest_index += 1;
             let plan_index = self.highest_index;
+
+            // Get CPU time for this node
             let elapsed_cpu: usize = match plan.metrics() {
                 Some(metrics) => metrics.elapsed_compute().unwrap_or_default(),
                 None => 0,
             };
+
+            // Get timestamps for this node
+            let (mut min_start, mut max_end) = Self::node_timerange(plan);
+
+            // Accumulate from children
             let mut cumulative_cpu = elapsed_cpu;
             for child in plan.children() {
-                cumulative_cpu += self.calculate_cumulative_cpu(child);
+                let child_metrics = self.calculate_metrics(child);
+                cumulative_cpu += child_metrics.cumulative_cpu;
+                min_start = Self::min_option(min_start, child_metrics.min_start);
+                max_end = Self::max_option(max_end, child_metrics.max_end);
             }
+
+            // Calculate wall clock duration for this subtree
+            let wall_duration = match (min_start, max_end) {
+                (Some(start), Some(end)) => (end - start).to_std().unwrap_or_default(),
+                _ => Duration::ZERO,
+            };
+
             self.index_to_cumulative_cpu
                 .insert(plan_index, cumulative_cpu);
-            cumulative_cpu
+            self.index_to_cumulative_wall
+                .insert(plan_index, wall_duration);
+
+            SubtreeMetrics {
+                cumulative_cpu,
+                min_start,
+                max_end,
+            }
+        }
+
+        fn node_timerange(
+            plan: &Arc<dyn ExecutionPlan>,
+        ) -> (Option<DateTime<Utc>>, Option<DateTime<Utc>>) {
+            let metrics = match plan.metrics() {
+                Some(m) => m,
+                None => return (None, None),
+            };
+            let min_start = metrics
+                .iter()
+                .filter_map(|m| match m.value() {
+                    MetricValue::StartTimestamp(ts) => ts.value(),
+                    _ => None,
+                })
+                .min();
+            let max_end = metrics
+                .iter()
+                .filter_map(|m| match m.value() {
+                    MetricValue::EndTimestamp(ts) => ts.value(),
+                    _ => None,
+                })
+                .max();
+            (min_start, max_end)
+        }
+
+        fn min_option(a: Option<DateTime<Utc>>, b: Option<DateTime<Utc>>) -> Option<DateTime<Utc>> {
+            [a, b].into_iter().flatten().min()
+        }
+
+        fn max_option(a: Option<DateTime<Utc>>, b: Option<DateTime<Utc>>) -> Option<DateTime<Utc>> {
+            [a, b].into_iter().flatten().max()
         }
     }
 
@@ -617,6 +685,11 @@ pub fn format_plan(plan: Arc<dyn ExecutionPlan>) -> String {
                 .unwrap();
             let cumulative_cpu_duration = Duration::from_nanos((*cumulative_cpu) as u64);
             write!(f, ", cumulative_cpu={cumulative_cpu_duration:?}")?;
+            let cumulative_wall = calcs
+                .index_to_cumulative_wall
+                .get(&self.highest_index)
+                .unwrap();
+            write!(f, ", cumulative_wall={cumulative_wall:?}")?;
             writeln!(f)?;
             self.indent += 1;
             for child in plan.children() {
@@ -635,8 +708,9 @@ pub fn format_plan(plan: Arc<dyn ExecutionPlan>) -> String {
             let mut calcs = CalculateVisitor {
                 highest_index: 0,
                 index_to_cumulative_cpu: HashMap::new(),
+                index_to_cumulative_wall: HashMap::new(),
             };
-            calcs.calculate_cumulative_cpu(&self.plan);
+            calcs.calculate_metrics(&self.plan);
             let mut prints = PrintVisitor {
                 highest_index: 0,
                 indent: 0,

rust/lance/src/dataset/write/merge_insert.rs

@@ -4312,9 +4312,9 @@ MergeInsert: on=[id], when_matched=UpdateAll, when_not_matched=InsertAll, when_n
         );
 
         // Also validate the full string structure with pattern matching
-        let expected_pattern = "[...MergeInsert: on=[id], when_matched=UpdateAll, when_not_matched=InsertAll, when_not_matched_by_source=Keep, metrics=...bytes_written=...num_deleted_rows=0, num_files_written=...num_inserted_rows=1, num_updated_rows=1], cumulative_cpu=...
+        let expected_pattern = "[...MergeInsert: on=[id], when_matched=UpdateAll, when_not_matched=InsertAll, when_not_matched_by_source=Keep, metrics=...bytes_written=...num_deleted_rows=0, num_files_written=...num_inserted_rows=1, num_updated_rows=1], cumulative_cpu=..., cumulative_wall=...
     ...
-    StreamingTableExec: partition_sizes=1, projection=[id, name], metrics=[], cumulative_cpu=...]";
+    StreamingTableExec: partition_sizes=1, projection=[id, name], metrics=[], cumulative_cpu=..., cumulative_wall=...]";
         assert_string_matches(&analysis, expected_pattern).unwrap();
         assert!(analysis.contains("bytes_written"));
         assert!(analysis.contains("num_files_written"));