diff --git a/datafusion/physical-plan/src/analyze.rs b/datafusion/physical-plan/src/analyze.rs index 92f747954719..fd8564e3d656 100644 --- a/datafusion/physical-plan/src/analyze.rs +++ b/datafusion/physical-plan/src/analyze.rs @@ -340,10 +340,7 @@ impl ExecutionPlan for AnalyzeExec { let Self { verbose, show_statistics, - // TODO: not on the wire. `AnalyzeExecBuilder` always resets this to - // `[Summary, Dev]`, so a non-default selection is lost on - // round-trip. Fixing it needs a new proto field. - metric_types: _, + metric_types, metric_categories, format, input, @@ -365,6 +362,13 @@ impl ExecutionPlan for AnalyzeExec { ExplainFormat::PostgresJSON => protobuf::ExplainFormat::Pgjson, ExplainFormat::Graphviz => protobuf::ExplainFormat::Graphviz, } as i32; + let metric_types = metric_types + .iter() + .map(|metric_type| match metric_type { + MetricType::Summary => protobuf::MetricType::Summary, + MetricType::Dev => protobuf::MetricType::Dev, + } as i32) + .collect(); Ok(Some(protobuf::PhysicalPlanNode { physical_plan_type: Some( protobuf::physical_plan_node::PhysicalPlanType::Analyze(Box::new( @@ -376,6 +380,8 @@ impl ExecutionPlan for AnalyzeExec { has_metric_categories, metric_categories, format, + has_metric_types: true, + metric_types, }, )), ), @@ -407,10 +413,33 @@ impl AnalyzeExec { has_metric_categories, metric_categories, format, + has_metric_types, + metric_types, } = analyze.as_ref(); let input = ctx.decode_required_child(input.as_deref(), "AnalyzeExec", "input")?; + let metric_types = if *has_metric_types { + Some( + metric_types + .iter() + .map(|metric_type| { + let metric_type = protobuf::MetricType::try_from(*metric_type) + .map_err(|_| { + datafusion_common::internal_datafusion_err!( + "Received an AnalyzeExecNode message with unknown MetricType {metric_type}" + ) + })?; + Ok(match metric_type { + protobuf::MetricType::Summary => MetricType::Summary, + protobuf::MetricType::Dev => MetricType::Dev, + }) + }) + .collect::>>()?, + ) + } else { + None + }; let metric_categories = if *has_metric_categories { Some( metric_categories @@ -437,16 +466,22 @@ impl AnalyzeExec { "AnalyzeExec is missing required field 'schema'" ) })?; + let mut builder = AnalyzeExec::builder( + *verbose, + *show_statistics, + input, + Arc::new(arrow::datatypes::Schema::try_from(schema)?), + ); + // Retain the builder default for plans written before `metric_types` + // and its presence bit were added. + if let Some(metric_types) = metric_types { + builder = builder.with_metric_types(metric_types); + } Ok(Arc::new( - AnalyzeExec::builder( - *verbose, - *show_statistics, - input, - Arc::new(arrow::datatypes::Schema::try_from(schema)?), - ) - .with_metric_categories(metric_categories) - .with_format(format) - .build(), + builder + .with_metric_categories(metric_categories) + .with_format(format) + .build(), )) } } diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 7b8f8e398539..aa2d0f716abc 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1397,6 +1397,10 @@ message AnalyzeExecNode { bool has_metric_categories = 5; repeated string metric_categories = 6; datafusion_common.ExplainFormat format = 7; + // Whether metric_types is present. False means use SUMMARY and DEV. + bool has_metric_types = 8; + // Types of metrics to display. + repeated datafusion_common.MetricType metric_types = 9; } message CrossJoinExecNode { diff --git a/datafusion/proto-models/src/generated/mod.rs b/datafusion/proto-models/src/generated/mod.rs index 4362b741d93a..7ed335e81351 100644 --- a/datafusion/proto-models/src/generated/mod.rs +++ b/datafusion/proto-models/src/generated/mod.rs @@ -18,6 +18,7 @@ // This code is generated so we don't want to fix any lint violations manually #[allow(clippy::allow_attributes)] #[allow(clippy::all)] +#[allow(clippy::uninlined_format_args)] #[allow(unused_qualifications)] #[rustfmt::skip] pub mod datafusion { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index bc88f5c9c2f4..5c50ff01be9b 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -1019,6 +1019,12 @@ impl serde::Serialize for AnalyzeExecNode { if self.format != 0 { len += 1; } + if self.has_metric_types { + len += 1; + } + if !self.metric_types.is_empty() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.AnalyzeExecNode", len)?; if self.verbose { struct_ser.serialize_field("verbose", &self.verbose)?; @@ -1043,6 +1049,16 @@ impl serde::Serialize for AnalyzeExecNode { .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.format)))?; struct_ser.serialize_field("format", &v)?; } + if self.has_metric_types { + struct_ser.serialize_field("hasMetricTypes", &self.has_metric_types)?; + } + if !self.metric_types.is_empty() { + let v = self.metric_types.iter().cloned().map(|v| { + super::datafusion_common::MetricType::try_from(v) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", v))) + }).collect::, _>>()?; + struct_ser.serialize_field("metricTypes", &v)?; + } struct_ser.end() } } @@ -1063,6 +1079,10 @@ impl<'de> serde::Deserialize<'de> for AnalyzeExecNode { "metric_categories", "metricCategories", "format", + "has_metric_types", + "hasMetricTypes", + "metric_types", + "metricTypes", ]; #[allow(clippy::enum_variant_names)] @@ -1074,6 +1094,8 @@ impl<'de> serde::Deserialize<'de> for AnalyzeExecNode { HasMetricCategories, MetricCategories, Format, + HasMetricTypes, + MetricTypes, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -1102,6 +1124,8 @@ impl<'de> serde::Deserialize<'de> for AnalyzeExecNode { "hasMetricCategories" | "has_metric_categories" => Ok(GeneratedField::HasMetricCategories), "metricCategories" | "metric_categories" => Ok(GeneratedField::MetricCategories), "format" => Ok(GeneratedField::Format), + "hasMetricTypes" | "has_metric_types" => Ok(GeneratedField::HasMetricTypes), + "metricTypes" | "metric_types" => Ok(GeneratedField::MetricTypes), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -1128,6 +1152,8 @@ impl<'de> serde::Deserialize<'de> for AnalyzeExecNode { let mut has_metric_categories__ = None; let mut metric_categories__ = None; let mut format__ = None; + let mut has_metric_types__ = None; + let mut metric_types__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Verbose => { @@ -1172,6 +1198,18 @@ impl<'de> serde::Deserialize<'de> for AnalyzeExecNode { } format__ = Some(map_.next_value::()? as i32); } + GeneratedField::HasMetricTypes => { + if has_metric_types__.is_some() { + return Err(serde::de::Error::duplicate_field("hasMetricTypes")); + } + has_metric_types__ = Some(map_.next_value()?); + } + GeneratedField::MetricTypes => { + if metric_types__.is_some() { + return Err(serde::de::Error::duplicate_field("metricTypes")); + } + metric_types__ = Some(map_.next_value::>()?.into_iter().map(|x| x as i32).collect()); + } } } Ok(AnalyzeExecNode { @@ -1182,6 +1220,8 @@ impl<'de> serde::Deserialize<'de> for AnalyzeExecNode { has_metric_categories: has_metric_categories__.unwrap_or_default(), metric_categories: metric_categories__.unwrap_or_default(), format: format__.unwrap_or_default(), + has_metric_types: has_metric_types__.unwrap_or_default(), + metric_types: metric_types__.unwrap_or_default(), }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 5bc46dfd7c2c..d6a56310b198 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2119,6 +2119,12 @@ pub struct AnalyzeExecNode { pub metric_categories: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, #[prost(enumeration = "super::datafusion_common::ExplainFormat", tag = "7")] pub format: i32, + /// Whether metric_types is present. False means use SUMMARY and DEV. + #[prost(bool, tag = "8")] + pub has_metric_types: bool, + /// Types of metrics to display. + #[prost(enumeration = "super::datafusion_common::MetricType", repeated, tag = "9")] + pub metric_types: ::prost::alloc::vec::Vec, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct CrossJoinExecNode { diff --git a/datafusion/proto/tests/cases/plans/misc.rs b/datafusion/proto/tests/cases/plans/misc.rs index 41bb051c2873..c990d977dd43 100644 --- a/datafusion/proto/tests/cases/plans/misc.rs +++ b/datafusion/proto/tests/cases/plans/misc.rs @@ -27,7 +27,7 @@ use datafusion::physical_plan::analyze::AnalyzeExec; use datafusion::physical_plan::empty::EmptyExec; use datafusion::physical_plan::explain::ExplainExec; use datafusion::physical_plan::expressions::{PhysicalSortExpr, col, lit}; -use datafusion::physical_plan::metrics::MetricCategory; +use datafusion::physical_plan::metrics::{MetricCategory, MetricType}; use datafusion::physical_plan::placeholder_row::PlaceholderRowExec; use datafusion::physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion::physical_plan::repartition::RepartitionExec; @@ -94,6 +94,180 @@ fn roundtrip_analyze() -> Result<()> { Ok(()) } +#[test] +fn roundtrip_analyze_metric_types() -> Result<()> { + use protobuf::MetricType as ProtoMetricType; + + let codec = DefaultPhysicalExtensionCodec {}; + let ctx = SessionContext::new(); + + for (metric_types, expected) in [ + ( + Some(vec![MetricType::Summary]), + vec![ProtoMetricType::Summary as i32], + ), + ( + Some(vec![MetricType::Dev]), + vec![ProtoMetricType::Dev as i32], + ), + (Some(vec![]), vec![]), + ( + None, + vec![ProtoMetricType::Summary as i32, ProtoMetricType::Dev as i32], + ), + ] { + let legacy = metric_types.is_none(); + let schema = Arc::new(Schema::new(vec![ + Field::new("plan_type", DataType::Utf8, false), + Field::new("plan", DataType::Utf8, false), + ])); + let input = Arc::new(PlaceholderRowExec::new(Arc::clone(&schema))); + let builder = AnalyzeExec::builder(false, false, input, schema); + let analyze = Arc::new(match metric_types { + Some(metric_types) => builder.with_metric_types(metric_types).build(), + None => builder.build(), + }); + let mut node = PhysicalPlanNode::try_from_physical_plan(analyze, &codec)?; + + let Some(protobuf::physical_plan_node::PhysicalPlanType::Analyze(analyze)) = + node.physical_plan_type.as_mut() + else { + unreachable!("expected AnalyzeExecNode") + }; + if legacy { + analyze.has_metric_types = false; + analyze.metric_types.clear(); + } + + let node = PhysicalPlanNode::decode(node.encode_to_vec().as_slice()) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + #[cfg(feature = "json")] + let node: PhysicalPlanNode = + serde_json::from_str(&serde_json::to_string(&node).unwrap()).unwrap(); + let roundtripped = node.try_into_physical_plan(&ctx.task_ctx(), &codec)?; + let mut node = PhysicalPlanNode::try_from_physical_plan(roundtripped, &codec)?; + { + let Some(protobuf::physical_plan_node::PhysicalPlanType::Analyze(analyze)) = + node.physical_plan_type.as_ref() + else { + unreachable!("expected AnalyzeExecNode") + }; + assert!(analyze.has_metric_types); + assert_eq!(analyze.metric_types, expected); + } + if expected == [ProtoMetricType::Summary as i32] { + let mut invalid_schema_node = node.clone(); + let Some(protobuf::physical_plan_node::PhysicalPlanType::Analyze(analyze)) = + invalid_schema_node.physical_plan_type.as_mut() + else { + unreachable!("expected AnalyzeExecNode") + }; + analyze + .schema + .as_mut() + .unwrap() + .columns + .push(protobuf::Field::default()); + let error = invalid_schema_node + .try_into_physical_plan(&ctx.task_ctx(), &codec) + .unwrap_err(); + assert!(error.strip_backtrace().contains("arrow_type")); + + let Some(protobuf::physical_plan_node::PhysicalPlanType::Analyze(analyze)) = + node.physical_plan_type.as_mut() + else { + unreachable!("expected AnalyzeExecNode") + }; + analyze.metric_types = vec![i32::MAX]; + let error = node + .try_into_physical_plan(&ctx.task_ctx(), &codec) + .unwrap_err(); + assert!(error.strip_backtrace().contains("unknown MetricType")); + } + } + Ok(()) +} + +#[cfg(feature = "json")] +#[test] +fn analyze_metric_types_json() { + use protobuf::{AnalyzeExecNode, MetricType as ProtoMetricType}; + + struct FailAfterFirstByte(bool); + + impl std::io::Write for FailAfterFirstByte { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + if buf.is_empty() { + return Ok(0); + } + if self.0 { + return Err(std::io::Error::other("expected write failure")); + } + self.0 = true; + Ok(1) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + let node = AnalyzeExecNode { + has_metric_types: true, + metric_types: vec![ProtoMetricType::Summary as i32, ProtoMetricType::Dev as i32], + ..Default::default() + }; + let json = serde_json::to_string(&node).unwrap(); + assert_eq!( + json, + r#"{"hasMetricTypes":true,"metricTypes":["METRIC_TYPE_SUMMARY","METRIC_TYPE_DEV"]}"# + ); + + let node: AnalyzeExecNode = serde_json::from_str( + r#"{"has_metric_types":true,"metric_types":["METRIC_TYPE_DEV"]}"#, + ) + .unwrap(); + assert!(node.has_metric_types); + assert_eq!(node.metric_types, [ProtoMetricType::Dev as i32]); + + for json in [ + r#"{"hasMetricTypes":true,"has_metric_types":false}"#, + r#"{"metricTypes":[],"metric_types":[]}"#, + ] { + let error = serde_json::from_str::(json).unwrap_err(); + assert!(error.to_string().contains("duplicate field")); + } + + for json in [ + r#"{"hasMetricTypes":"true"}"#, + r#"{"metricTypes":true}"#, + r#"{"metricTypes":["UNKNOWN"]}"#, + ] { + assert!(serde_json::from_str::(json).is_err()); + } + + let invalid = AnalyzeExecNode { + metric_types: vec![i32::MAX], + ..Default::default() + }; + let error = serde_json::to_string(&invalid).unwrap_err(); + assert!(error.to_string().contains("Invalid variant")); + + for node in [ + AnalyzeExecNode { + has_metric_types: true, + ..Default::default() + }, + AnalyzeExecNode { + metric_types: vec![ProtoMetricType::Summary as i32], + ..Default::default() + }, + ] { + let error = serde_json::to_writer(FailAfterFirstByte(false), &node).unwrap_err(); + assert!(error.to_string().contains("expected write failure")); + } +} + #[test] fn roundtrip_explain() -> Result<()> { let schema = Arc::new(Schema::new(vec![