[HSTACK] feat(datafusion): configurable planning-metrics serialization

crates/core/src/delta_datafusion/mod.rs

@@ -75,6 +75,7 @@ pub(crate) use data_validation::{
 };
 pub use engine::DataFusionEngine;
 pub(crate) use find_files::*;
+pub use planning_count_metrics_codec::DeltaPhysicalCodecConfig;
 pub(crate) use table_provider::next::normalize_path_as_file_id;
 pub use table_provider::{
     DeltaScanConfig, DeltaScanConfigBuilder, TableProviderBuilder, next::DeltaScanExec,
@@ -470,8 +471,18 @@ pub(crate) fn to_correct_scalar_value(
             Use DeltaLogicalCodec for table-provider serialization until a DeltaScanExec \
             physical codec is available."
 )]
-#[derive(Debug)]
-pub struct DeltaPhysicalCodec {}
+#[derive(Debug, Default)]
+pub struct DeltaPhysicalCodec {
+    config: DeltaPhysicalCodecConfig,
+}
+
+#[allow(deprecated)]
+impl DeltaPhysicalCodec {
+    /// Creates a codec with the given serialization options.
+    pub fn new(config: DeltaPhysicalCodecConfig) -> Self {
+        Self { config }
+    }
+}
 
 #[allow(deprecated)]
 impl PhysicalExtensionCodec for DeltaPhysicalCodec {
@@ -483,7 +494,7 @@ impl PhysicalExtensionCodec for DeltaPhysicalCodec {
     ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
         let wire: DeltaScanWire = serde_json::from_reader(buf)
             .map_err(|_| DataFusionError::Internal("Unable to decode DeltaScan".to_string()))?;
-        let delta_scan = wire.into_delta_scan((*inputs)[0].clone());
+        let delta_scan = wire.into_delta_scan((*inputs)[0].clone(), self.config);
         Ok(Arc::new(delta_scan))
     }
 
@@ -496,7 +507,7 @@ impl PhysicalExtensionCodec for DeltaPhysicalCodec {
             .downcast_ref::<DeltaScan>()
             .ok_or_else(|| DataFusionError::Internal("Not a legacy delta scan!".to_string()))?;
 
-        let wire = DeltaScanWire::from(delta_scan);
+        let wire = DeltaScanWire::new(delta_scan, self.config);
         serde_json::to_writer(buf, &wire).map_err(|_| {
             DataFusionError::Internal("Unable to encode legacy delta scan!".to_string())
         })?;
@@ -736,7 +747,7 @@ mod tests {
     #[allow(deprecated)]
     fn roundtrip_test_delta_exec_plan() {
         let ctx = SessionContext::new();
-        let codec = DeltaPhysicalCodec {};
+        let codec = DeltaPhysicalCodec::default();
 
         let schema = Arc::new(Schema::new(vec![
             Field::new("a", ArrowDataType::Utf8, false),

crates/core/src/delta_datafusion/planning_count_metrics_codec.rs

@@ -1,21 +1,40 @@
 //! Codec for planning-time count metrics attached to Delta scan plans.
 //!
 //! Only [`MetricValue::Count`] metrics are serialized. Runtime metrics are collected by the
-//! execution plan on the worker and are not copied from the coordinator.
+//! execution plan wherever it runs and are never copied from the encoded plan.
 // TODO: Move this codec into `table_provider::next::scan::codec` after removing the legacy table
 // provider.
 
 use std::borrow::Cow;
 use std::sync::Arc;
 
 use datafusion::physical_plan::metrics::{
-    Count, ExecutionPlanMetricsSet, Label, Metric, MetricValue, MetricsSet,
+    Count, ExecutionPlanMetricsSet, Metric, MetricValue, MetricsSet,
 };
 use serde::{Deserialize, Serialize};
 
-// Consumed by datafusion-distributed's MetricsWrapperExec before task aggregation.
-const CREATED_AT_LABEL: &str = "created_at";
-const CREATED_AT_PLANNING: &str = "planning";
+/// Serialization options shared by the Delta physical codecs
+/// ([`DeltaNextPhysicalCodec`](super::DeltaNextPhysicalCodec) and the legacy
+/// `DeltaPhysicalCodec`).
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct DeltaPhysicalCodecConfig {
+    /// Write planning-time scan count metrics (for example the number of scanned or pruned
+    /// files) with an encoded scan, and restore them when decoding. Enabled by default.
+    ///
+    /// Disable it for plans shipped to execution-only workers, such as distributed tasks whose
+    /// runtime metrics are merged back into the coordinator's plan: the coordinator already owns
+    /// the planning metrics, so a copy in every task would be counted once per task. When
+    /// disabled, planning metrics are neither written nor restored from a received plan.
+    pub serialize_planning_metrics: bool,
+}
+
+impl Default for DeltaPhysicalCodecConfig {
+    fn default() -> Self {
+        Self {
+            serialize_planning_metrics: true,
+        }
+    }
+}
 
 #[derive(Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(transparent)]
@@ -25,6 +44,24 @@ impl PlanningCountMetricsWire {
     pub(super) fn is_empty(&self) -> bool {
         self.0.is_empty()
     }
+
+    /// Wire form of a scan's planning metrics; empty when `config` disables them.
+    pub(super) fn encode(metrics: &MetricsSet, config: DeltaPhysicalCodecConfig) -> Self {
+        if config.serialize_planning_metrics {
+            Self::from(metrics)
+        } else {
+            Self::default()
+        }
+    }
+
+    /// Metrics for a decoded scan; empty when `config` disables planning metrics.
+    pub(super) fn decode(self, config: DeltaPhysicalCodecConfig) -> ExecutionPlanMetricsSet {
+        if config.serialize_planning_metrics {
+            self.into()
+        } else {
+            ExecutionPlanMetricsSet::new()
+        }
+    }
 }
 
 impl From<&MetricsSet> for PlanningCountMetricsWire {
@@ -51,13 +88,12 @@ impl From<PlanningCountMetricsWire> for ExecutionPlanMetricsSet {
             .map(|metric| {
                 let count = Count::new();
                 count.add(metric.value);
-                Arc::new(Metric::new_with_labels(
+                Arc::new(Metric::new(
                     MetricValue::Count {
                         name: Cow::Owned(metric.name),
                         count,
                     },
                     None,
-                    vec![Label::new(CREATED_AT_LABEL, CREATED_AT_PLANNING)],
                 ))
             })
             .collect::<MetricsSet>()
@@ -105,22 +141,19 @@ mod tests {
         ]
         .into_iter()
         .collect::<MetricsSet>();
-
-        let wire = PlanningCountMetricsWire::from(&metrics);
+        let config = DeltaPhysicalCodecConfig::default();
+        let wire = PlanningCountMetricsWire::encode(&metrics, config);
         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();
+        let decoded = decoded.decode(config).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() == CREATED_AT_LABEL
-                && labels[0].value() == CREATED_AT_PLANNING
+                && metric.labels().is_empty()
                 && metric.metric_type() == MetricType::Dev
                 && metric.metric_category().is_none()
         }));

crates/core/src/delta_datafusion/table_provider.rs

@@ -21,7 +21,9 @@ use serde::{Deserialize, Serialize};
 use url::Url;
 use uuid::Uuid;
 
-use crate::delta_datafusion::planning_count_metrics_codec::PlanningCountMetricsWire;
+use crate::delta_datafusion::planning_count_metrics_codec::{
+    DeltaPhysicalCodecConfig, PlanningCountMetricsWire,
+};
 use crate::delta_datafusion::table_provider::next::SnapshotWrapper;
 use crate::delta_datafusion::{DataFusionMixins as _, FindFilesExprProperties};
 use crate::kernel::{Add, EagerSnapshot, Snapshot, Version};
@@ -565,26 +567,31 @@ pub(super) struct DeltaScanWire {
     pub(crate) planning_count_metrics: PlanningCountMetricsWire,
 }
 
-impl From<&DeltaScan> for DeltaScanWire {
-    fn from(scan: &DeltaScan) -> Self {
+impl DeltaScanWire {
+    pub(super) fn new(scan: &DeltaScan, codec_config: DeltaPhysicalCodecConfig) -> Self {
         Self {
             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()),
+            planning_count_metrics: PlanningCountMetricsWire::encode(
+                &scan.metrics.clone_inner(),
+                codec_config,
+            ),
         }
     }
-}
 
-impl DeltaScanWire {
-    pub(super) fn into_delta_scan(self, parquet_scan: Arc<dyn ExecutionPlan>) -> DeltaScan {
+    pub(super) fn into_delta_scan(
+        self,
+        parquet_scan: Arc<dyn ExecutionPlan>,
+        codec_config: DeltaPhysicalCodecConfig,
+    ) -> DeltaScan {
         let mut scan = DeltaScan::new(
             &self.table_url,
             self.config,
             parquet_scan,
             self.logical_schema,
         );
-        scan.metrics = self.planning_count_metrics.into();
+        scan.metrics = self.planning_count_metrics.decode(codec_config);
         scan
     }
 }