diff --git a/crates/core/src/delta_datafusion/mod.rs b/crates/core/src/delta_datafusion/mod.rs index c7f864c0c6..250e324125 100644 --- a/crates/core/src/delta_datafusion/mod.rs +++ b/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, 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::() .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), diff --git a/crates/core/src/delta_datafusion/planning_count_metrics_codec.rs b/crates/core/src/delta_datafusion/planning_count_metrics_codec.rs index c1f9163397..63aea569fb 100644 --- a/crates/core/src/delta_datafusion/planning_count_metrics_codec.rs +++ b/crates/core/src/delta_datafusion/planning_count_metrics_codec.rs @@ -1,7 +1,7 @@ //! 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. @@ -9,13 +9,32 @@ 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 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::() @@ -105,22 +141,19 @@ mod tests { ] .into_iter() .collect::(); - - 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() })); diff --git a/crates/core/src/delta_datafusion/table_provider.rs b/crates/core/src/delta_datafusion/table_provider.rs index 2dddfe8a01..5c4d1d131e 100644 --- a/crates/core/src/delta_datafusion/table_provider.rs +++ b/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) -> DeltaScan { + pub(super) fn into_delta_scan( + self, + parquet_scan: Arc, + 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 } } diff --git a/crates/core/src/delta_datafusion/table_provider/next/scan/codec.rs b/crates/core/src/delta_datafusion/table_provider/next/scan/codec.rs index 556eeecaeb..c1d3814fa0 100644 --- a/crates/core/src/delta_datafusion/table_provider/next/scan/codec.rs +++ b/crates/core/src/delta_datafusion/table_provider/next/scan/codec.rs @@ -12,7 +12,9 @@ use super::{DeltaScanExec, DeltaScanMetaExec, ProjectedScanContract, PublicFileI 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_codec::PlanningCountMetricsWire; +use crate::delta_datafusion::planning_count_metrics_codec::{ + DeltaPhysicalCodecConfig, PlanningCountMetricsWire, +}; use crate::kernel::Snapshot; use crate::kernel::size_limits::SnapshotLoadMetrics; use arrow::datatypes::SchemaRef; @@ -39,7 +41,16 @@ use serde::{Deserialize, Serialize}; /// to reconstruct the execution plan rather than the plan itself. This approach /// avoids the need for serde support in delta-kernel types. #[derive(Debug, Clone, Default)] -pub struct DeltaNextPhysicalCodec; +pub struct DeltaNextPhysicalCodec { + config: DeltaPhysicalCodecConfig, +} + +impl DeltaNextPhysicalCodec { + /// Creates a codec with the given serialization options. + pub fn new(config: DeltaPhysicalCodecConfig) -> Self { + Self { config } + } +} impl PhysicalExtensionCodec for DeltaNextPhysicalCodec { fn try_decode( @@ -53,7 +64,7 @@ impl PhysicalExtensionCodec for DeltaNextPhysicalCodec { })?; match wire { - ScanExecWire::Scan(wire) => wire.into_exec(inputs, ctx), + ScanExecWire::Scan(wire) => wire.into_exec(inputs, ctx, self.config), ScanExecWire::Meta(wire) => wire.into_exec(inputs, ctx), } } @@ -64,7 +75,7 @@ impl PhysicalExtensionCodec for DeltaNextPhysicalCodec { buf: &mut Vec, ) -> datafusion::common::Result<()> { let wire = if let Some(delta_scan) = node.downcast_ref::() { - ScanExecWire::Scan(DeltaScanExecWire::try_from(delta_scan)?) + ScanExecWire::Scan(DeltaScanExecWire::try_new(delta_scan, self.config)?) } else if let Some(meta_scan) = node.downcast_ref::() { ScanExecWire::Meta(DeltaScanMetaExecWire::try_from(meta_scan)?) } else { @@ -208,10 +219,11 @@ fn scan_plan_wire(scan_plan: &KernelScanPlan) -> Result for DeltaScanExecWire { - type Error = DataFusionError; - - fn try_from(exec: &DeltaScanExec) -> Result { +impl DeltaScanExecWire { + fn try_new( + exec: &DeltaScanExec, + codec_config: DeltaPhysicalCodecConfig, + ) -> Result { let scan_plan_wire = scan_plan_wire(&exec.scan_plan)?; let metrics = exec.metrics().unwrap_or_default(); @@ -237,7 +249,7 @@ impl TryFrom<&DeltaScanExec> for DeltaScanExecWire { 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), + planning_count_metrics: PlanningCountMetricsWire::encode(&metrics, codec_config), }) } } @@ -412,6 +424,7 @@ impl DeltaScanExecWire { self, inputs: &[Arc], task: &TaskContext, + codec_config: DeltaPhysicalCodecConfig, ) -> datafusion::common::Result> { if inputs.len() != 1 { return Err(DataFusionError::Internal(format!( @@ -492,7 +505,7 @@ impl DeltaScanExecWire { } } - let metrics = self.planning_count_metrics.into(); + let metrics = self.planning_count_metrics.decode(codec_config); let exec = DeltaScanExec::new( Arc::new(scan_plan), @@ -805,7 +818,7 @@ mod tests { ]; let session = create_session().into_inner(); let task = session.task_ctx(); - let codec = DeltaNextPhysicalCodec; + let codec = DeltaNextPhysicalCodec::default(); let transforms: HashMap<_, _> = [ ("a".to_string(), test_transform(1, false)), ("b".to_string(), test_transform(1, true)), @@ -817,7 +830,7 @@ mod tests { for plan in plans { let (mut wire, tag) = if let Some(scan) = extract_delta_scan_exec(&plan) { ( - ScanExecWire::Scan(DeltaScanExecWire::try_from(scan)?), + ScanExecWire::Scan(DeltaScanExecWire::try_new(scan, Default::default())?), "Scan", ) } else { @@ -950,7 +963,7 @@ mod tests { let delta_scan = extract_delta_scan_exec(&plan).expect("Expected DeltaScanExec"); - let codec = DeltaNextPhysicalCodec; + let codec = DeltaNextPhysicalCodec::default(); let mut buf = Vec::new(); codec.try_encode(plan.clone(), &mut buf)?; @@ -983,6 +996,75 @@ mod tests { Ok(()) } + #[rstest::rstest] + #[case::classic(false)] + #[case::next(true)] + #[tokio::test] + #[allow(deprecated)] + async fn test_codec_planning_metrics_follow_config( + #[case] next: bool, + #[values(true, false)] serialize_planning_metrics: bool, + ) -> TestResult { + use crate::delta_datafusion::{DeltaPhysicalCodec, table_provider::DeltaScan}; + use datafusion::physical_plan::{empty::EmptyExec, metrics::MetricBuilder}; + + let plan: Arc = if next { + create_delta_scan_exec(&[], None).await? + } else { + let schema = Arc::new(arrow::datatypes::Schema::empty()); + let scan = DeltaScan::new( + &url::Url::parse("s3://my_bucket/table")?, + Default::default(), + Arc::new(EmptyExec::new(schema.clone())), + schema, + ); + MetricBuilder::new(&scan.metrics) + .global_counter("files_scanned") + .add(4); + Arc::new(scan) + }; + let config = DeltaPhysicalCodecConfig { + serialize_planning_metrics, + }; + let codec: Box = if next { + Box::new(DeltaNextPhysicalCodec::new(config)) + } else { + Box::new(DeltaPhysicalCodec::new(config)) + }; + let metric_name = if next { + "count_files_scanned" + } else { + "files_scanned" + }; + let planned = plan + .metrics() + .unwrap() + .sum_by_name(metric_name) + .unwrap() + .as_usize(); + assert!(planned > 0); + let mut buf = Vec::new(); + codec.try_encode(plan.clone(), &mut buf)?; + let wire: serde_json::Value = serde_json::from_slice(&buf)?; + let wire = if next { &wire["Scan"] } else { &wire }; + assert_eq!( + wire.get("planning_count_metrics").is_some(), + serialize_planning_metrics + ); + + let inputs = plan.children().into_iter().cloned().collect::>(); + let decoded = codec.try_decode(&buf, &inputs, &create_session().into_inner().task_ctx())?; + assert_eq!( + decoded + .metrics() + .unwrap() + .sum_by_name(metric_name) + .map(|value| value.as_usize()), + serialize_planning_metrics.then_some(planned), + ); + Ok(()) + } + #[tokio::test] async fn test_codec_roundtrip_with_projection() -> TestResult { let projection = vec![0usize]; @@ -990,7 +1072,7 @@ mod tests { let delta_scan = extract_delta_scan_exec(&plan).expect("Expected DeltaScanExec"); - let codec = DeltaNextPhysicalCodec; + let codec = DeltaNextPhysicalCodec::default(); let mut buf = Vec::new(); codec.try_encode(plan.clone(), &mut buf)?; @@ -1025,7 +1107,7 @@ mod tests { let delta_scan = extract_delta_scan_exec(&plan).expect("Expected DeltaScanExec"); - let codec = DeltaNextPhysicalCodec; + let codec = DeltaNextPhysicalCodec::default(); let mut buf = Vec::new(); codec.try_encode(plan.clone(), &mut buf)?; @@ -1054,7 +1136,7 @@ mod tests { let delta_scan = extract_delta_scan_exec(&plan).expect("Expected DeltaScanExec"); - let wire = DeltaScanExecWire::try_from(delta_scan)?; + let wire = DeltaScanExecWire::try_new(delta_scan, Default::default())?; let json = serde_json::to_string(&wire)?; assert!(!json.is_empty(), "JSON should not be empty"); @@ -1081,7 +1163,7 @@ mod tests { async fn test_codec_decode_wrong_input_count() -> TestResult { let plan = create_delta_scan_exec(&[], None).await?; - let codec = DeltaNextPhysicalCodec; + let codec = DeltaNextPhysicalCodec::default(); let mut buf = Vec::new(); codec.try_encode(plan.clone(), &mut buf)?; @@ -1145,7 +1227,7 @@ mod tests { let plan = create_delta_scan_exec(&[], None).await?; let delta_scan = extract_delta_scan_exec(&plan).expect("Expected DeltaScanExec"); - let wire = DeltaScanExecWire::try_from(delta_scan)?; + let wire = DeltaScanExecWire::try_new(delta_scan, Default::default())?; assert!( wire.selection_vectors.is_empty() || !wire.selection_vectors.is_empty(), @@ -1169,7 +1251,7 @@ mod tests { let plan = create_delta_scan_exec(&[], None).await?; let delta_scan = extract_delta_scan_exec(&plan).expect("Expected DeltaScanExec"); - let codec = DeltaNextPhysicalCodec; + let codec = DeltaNextPhysicalCodec::default(); let mut buf = Vec::new(); codec.try_encode(plan.clone(), &mut buf)?; @@ -1207,7 +1289,7 @@ mod tests { "Table with deletion vectors should have non-empty selection_vectors" ); - let codec = DeltaNextPhysicalCodec; + let codec = DeltaNextPhysicalCodec::default(); let mut buf = Vec::new(); codec.try_encode(plan.clone(), &mut buf)?; @@ -1252,7 +1334,7 @@ mod tests { // Column mapping tables have transforms that inject partition values. // This test verifies full roundtrip serialization of Transform expressions. - let codec = DeltaNextPhysicalCodec; + let codec = DeltaNextPhysicalCodec::default(); let mut buf = Vec::new(); codec.try_encode(plan.clone(), &mut buf)?; @@ -1280,7 +1362,7 @@ mod tests { create_delta_scan_meta_exec_from_table(TestTables::WithColumnMapping, &[]).await?; let meta_scan = extract_delta_scan_meta_exec(&plan).expect("Expected DeltaScanMetaExec"); - let codec = DeltaNextPhysicalCodec; + let codec = DeltaNextPhysicalCodec::default(); let mut buf = Vec::new(); codec.try_encode(plan.clone(), &mut buf)?; @@ -1320,7 +1402,7 @@ mod tests { create_delta_scan_meta_exec_from_table(TestTables::WithColumnMapping, &filters).await?; let meta_scan = extract_delta_scan_meta_exec(&plan).expect("Expected DeltaScanMetaExec"); - let codec = DeltaNextPhysicalCodec; + let codec = DeltaNextPhysicalCodec::default(); let mut buf = Vec::new(); codec.try_encode(plan.clone(), &mut buf)?; @@ -1355,7 +1437,7 @@ mod tests { "Table with deletion vectors should have non-empty selection_vectors" ); - let codec = DeltaNextPhysicalCodec; + let codec = DeltaNextPhysicalCodec::default(); let mut buf = Vec::new(); codec.try_encode(plan.clone(), &mut buf)?; @@ -1399,7 +1481,7 @@ mod tests { create_delta_scan_meta_exec_from_table(TestTables::WithColumnMapping, &[]).await?; let meta_scan = extract_delta_scan_meta_exec(&plan).expect("Expected DeltaScanMetaExec"); - let codec = DeltaNextPhysicalCodec; + let codec = DeltaNextPhysicalCodec::default(); let mut buf = Vec::new(); codec.try_encode(plan.clone(), &mut buf)?; @@ -1431,7 +1513,7 @@ mod tests { Arc::new(CrossJoinExec::new(scan_plan.clone(), meta_plan.clone())) }; - let codec = DeltaNextPhysicalCodec; + let codec = DeltaNextPhysicalCodec::default(); let proto = protobuf::PhysicalPlanNode::try_from_physical_plan(join.clone(), &codec)?; let session = create_session().into_inner(); diff --git a/crates/core/tests/it_datafusion/integration_datafusion_deep.rs b/crates/core/tests/it_datafusion/integration_datafusion_deep.rs index c384684ae7..439533abb7 100644 --- a/crates/core/tests/it_datafusion/integration_datafusion_deep.rs +++ b/crates/core/tests/it_datafusion/integration_datafusion_deep.rs @@ -115,7 +115,7 @@ async fn test_hstack_deep_column_pruning() -> datafusion::common::Result<()> { let codec = ComposedPhysicalExtensionCodec::new( vec![ Arc::new(DefaultPhysicalExtensionCodec {}), - Arc::new(DeltaPhysicalCodec{}) + Arc::new(DeltaPhysicalCodec::default()) ] ); let proto = PhysicalPlanNode::try_from_physical_plan(physical_plan.clone(), &codec) @@ -235,7 +235,7 @@ async fn test_hstack_deep_column_pruning_next_codec() -> datafusion::common::Res let codec = ComposedPhysicalExtensionCodec::new( vec![ Arc::new(DefaultPhysicalExtensionCodec {}), - Arc::new(DeltaNextPhysicalCodec{}) + Arc::new(DeltaNextPhysicalCodec::default()) ] );