fix(datafusion): preserve Delta planning count metrics

crates/core/src/delta_datafusion/mod.rs

@@ -105,6 +105,7 @@ pub mod expr;
 mod file_id;
 mod find_files;
 pub mod logical;
+mod planning_count_metrics;
 pub mod physical;
 pub mod planner;
 mod session;

crates/core/src/delta_datafusion/planning_count_metrics.rs

@@ -0,0 +1,126 @@
+//! Serializable planning count metrics attached to Delta scan plans.
+
+use std::borrow::Cow;
+use std::sync::Arc;
+
+use datafusion::physical_plan::metrics::{
+    Count, ExecutionPlanMetricsSet, Label, Metric, MetricValue, MetricsSet,
+};
+use serde::{Deserialize, Serialize};
+
+// Consumed by datafusion-distributed's MetricsWrapperExec before task aggregation.
+const DFD_AGGREGATION_SCOPE_LABEL: &str = "dfd_aggregation_scope";
+const COORDINATOR_AGGREGATION_SCOPE: &str = "coordinator";
+
+#[derive(Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(transparent)]
+pub(super) struct PlanningCountMetricsWire(Vec<PlanningCountMetricWire>);
+
+impl PlanningCountMetricsWire {
+    pub(super) fn is_empty(&self) -> bool {
+        self.0.is_empty()
+    }
+}
+
+impl From<&MetricsSet> for PlanningCountMetricsWire {
+    fn from(metrics: &MetricsSet) -> Self {
+        Self(
+            metrics
+                .iter()
+                .filter_map(|metric| match metric.value() {
+                    MetricValue::Count { name, count } => Some(PlanningCountMetricWire {
+                        name: name.to_string(),
+                        value: count.value(),
+                    }),
+                    _ => None,
+                })
+                .collect(),
+        )
+    }
+}
+
+impl From<PlanningCountMetricsWire> for ExecutionPlanMetricsSet {
+    fn from(wire: PlanningCountMetricsWire) -> Self {
+        wire.0
+            .into_iter()
+            .map(|metric| {
+                let count = Count::new();
+                count.add(metric.value);
+                Arc::new(Metric::new_with_labels(
+                    MetricValue::Count {
+                        name: Cow::Owned(metric.name),
+                        count,
+                    },
+                    None,
+                    vec![Label::new(
+                        DFD_AGGREGATION_SCOPE_LABEL,
+                        COORDINATOR_AGGREGATION_SCOPE,
+                    )],
+                ))
+            })
+            .collect::<MetricsSet>()
+            .into()
+    }
+}
+
+#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
+struct PlanningCountMetricWire {
+    name: String,
+    value: usize,
+}
+
+#[cfg(test)]
+mod tests {
+    use datafusion::physical_plan::metrics::{Label, MetricType};
+
+    use super::*;
+
+    fn count(value: usize) -> Count {
+        let count = Count::new();
+        count.add(value);
+        count
+    }
+
+    #[test]
+    fn planning_count_metrics_wire_roundtrips_named_counts_only() {
+        let metrics = vec![
+            Arc::new(Metric::new_with_labels(
+                MetricValue::Count {
+                    name: Cow::Borrowed("files_scanned"),
+                    count: count(4),
+                },
+                Some(7),
+                vec![Label::new("source", "planning")],
+            )),
+            Arc::new(Metric::new(
+                MetricValue::Count {
+                    name: Cow::Borrowed("files_pruned"),
+                    count: count(2),
+                },
+                None,
+            )),
+            Arc::new(Metric::new(MetricValue::OutputRows(count(80)), None)),
+        ]
+        .into_iter()
+        .collect::<MetricsSet>();
+
+        let wire = PlanningCountMetricsWire::from(&metrics);
+        let encoded = serde_json::to_vec(&wire).unwrap();
+        let decoded: PlanningCountMetricsWire = serde_json::from_slice(&encoded).unwrap();
+        assert_eq!(decoded, wire);
+
+        let decoded = ExecutionPlanMetricsSet::from(decoded).clone_inner();
+        assert_eq!(decoded.iter().count(), 2);
+        assert_eq!(decoded.sum_by_name("files_scanned").unwrap().as_usize(), 4);
+        assert_eq!(decoded.sum_by_name("files_pruned").unwrap().as_usize(), 2);
+        assert!(decoded.iter().all(|metric| {
+            let labels = metric.labels();
+            metric.partition().is_none()
+                && labels.len() == 1
+                && labels[0].name() == DFD_AGGREGATION_SCOPE_LABEL
+                && labels[0].value() == COORDINATOR_AGGREGATION_SCOPE
+                && metric.metric_type() == MetricType::Dev
+                && metric.metric_category().is_none()
+        }));
+    }
+}

crates/core/src/delta_datafusion/table_provider.rs

@@ -21,6 +21,7 @@ use serde::{Deserialize, Serialize};
 use url::Url;
 use uuid::Uuid;
 
+use crate::delta_datafusion::planning_count_metrics::PlanningCountMetricsWire;
 use crate::delta_datafusion::table_provider::next::SnapshotWrapper;
 use crate::delta_datafusion::{DataFusionMixins as _, FindFilesExprProperties};
 use crate::kernel::{Add, EagerSnapshot, Snapshot, Version};
@@ -560,6 +561,8 @@ pub(super) struct DeltaScanWire {
     pub(crate) table_url: Url,
     pub(crate) config: DeltaScanConfig,
     pub(crate) logical_schema: Arc<Schema>,
+    #[serde(default, skip_serializing_if = "PlanningCountMetricsWire::is_empty")]
+    pub(crate) planning_count_metrics: PlanningCountMetricsWire,
 }
 
 impl From<&DeltaScan> for DeltaScanWire {
@@ -568,18 +571,21 @@ impl From<&DeltaScan> for DeltaScanWire {
             table_url: scan.table_url.clone(),
             config: scan.config.clone(),
             logical_schema: scan.logical_schema.clone(),
+            planning_count_metrics: PlanningCountMetricsWire::from(&scan.metrics.clone_inner()),
         }
     }
 }
 
 impl DeltaScanWire {
     pub(super) fn into_delta_scan(self, parquet_scan: Arc<dyn ExecutionPlan>) -> DeltaScan {
-        DeltaScan::new(
+        let mut scan = DeltaScan::new(
             &self.table_url,
             self.config,
             parquet_scan,
             self.logical_schema,
-        )
+        );
+        scan.metrics = self.planning_count_metrics.into();
+        scan
     }
 }
 

crates/core/src/delta_datafusion/table_provider/next/scan/codec.rs

@@ -12,6 +12,7 @@
 use crate::DeltaTableConfig;
 use crate::delta_datafusion::DeltaScanConfig;
 use crate::delta_datafusion::engine::{to_datafusion_expr, to_delta_expression};
+use crate::delta_datafusion::planning_count_metrics::PlanningCountMetricsWire;
 use crate::kernel::Snapshot;
 use crate::kernel::size_limits::SnapshotLoadMetrics;
 use arrow::datatypes::SchemaRef;
@@ -24,7 +25,7 @@
 use datafusion::prelude::Expr;
 use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder};
 use datafusion_datasource::source::DataSourceExec;
 use datafusion_proto::bytes::Serializeable;
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use delta_kernel::expressions::{
     ColumnName, Expression, ExpressionFieldPatch, ExpressionStructPatch,
@@ -123,6 +124,8 @@
     input_file_id_column: String,
     file_id_column: Option<String>,
     public_file_ids: std::collections::HashMap<String, String>,
+    #[serde(default, skip_serializing_if = "PlanningCountMetricsWire::is_empty")]
+    planning_count_metrics: PlanningCountMetricsWire,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -210,6 +213,9 @@
 
     fn try_from(exec: &DeltaScanExec) -> Result<Self, Self::Error> {
         let scan_plan_wire = scan_plan_wire(&exec.scan_plan)?;
+        let metrics = exec.metrics().ok_or_else(|| {
+            DataFusionError::Internal("DeltaScanExec did not expose its metrics".to_string())
+        })?;
 
         let transforms = serialize_transforms(&exec.transforms)?;
 
@@ -233,6 +239,7 @@
             input_file_id_column: exec.input_file_id_column.clone(),
             file_id_column: exec.file_id_column.clone(),
             public_file_ids,
+            planning_count_metrics: PlanningCountMetricsWire::from(&metrics),
         })
     }
 }
@@ -487,14 +494,16 @@
             }
         }
 
+        let metrics = self.planning_count_metrics.into();
+
         let exec = DeltaScanExec::new(
             Arc::new(scan_plan),
             execution_plan,
             Arc::new(transforms),
             Arc::new(selection_vectors),
             Arc::new(public_file_ids),
             Default::default(),
-            Default::default(),
+            metrics,
         );
 
         Ok(Arc::new(exec))