From 506326a74edc2b8a357d0f5a9a1ea430df40747c Mon Sep 17 00:00:00 2001 From: Huaijin Date: Tue, 25 Aug 2026 21:59:57 +0800 Subject: [PATCH 1/3] fix(proto): preserve AnalyzeExec metric types across serialization --- datafusion/physical-plan/src/analyze.rs | 61 +++++++-- .../proto-models/proto/datafusion.proto | 4 + datafusion/proto-models/src/generated/mod.rs | 1 + .../proto-models/src/generated/pbjson.rs | 40 ++++++ .../proto-models/src/generated/prost.rs | 6 + datafusion/proto/tests/cases/plans/misc.rs | 118 ++++++++++++++++-- 6 files changed, 208 insertions(+), 22 deletions(-) diff --git a/datafusion/physical-plan/src/analyze.rs b/datafusion/physical-plan/src/analyze.rs index 92f7479547197..fd8564e3d6564 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 7b8f8e398539c..aa2d0f716abc5 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 4362b741d93a9..7ed335e813513 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 bc88f5c9c2f45..5c50ff01be9b1 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 5bc46dfd7c2c8..d6a56310b198b 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 41bb051c28730..66c6978b7762e 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; @@ -63,34 +63,134 @@ fn roundtrip_analyze() -> Result<()> { let metric_categories = vec![MetricCategory::Rows, MetricCategory::Timing]; let analyze = Arc::new( AnalyzeExec::builder(true, true, input, Arc::clone(&schema)) + // `Summary` is non-default and was previously reset to + // `[Summary, Dev]` by the protobuf round-trip. + .with_metric_types(vec![MetricType::Summary]) .with_metric_categories(Some(metric_categories.clone())) .with_format(ExplainFormat::Tree) .build(), ); let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; let roundtripped = roundtrip_test_and_return( analyze, &ctx, - &DefaultPhysicalExtensionCodec {}, + &codec, &DefaultPhysicalProtoConverter {}, )?; - let roundtripped = roundtripped.downcast_ref::().unwrap(); + let roundtripped_analyze = roundtripped.downcast_ref::().unwrap(); - assert_eq!(roundtripped.schema(), schema); - assert!(roundtripped.verbose()); - assert!(roundtripped.show_statistics()); + assert_eq!(roundtripped_analyze.schema(), schema); + assert!(roundtripped_analyze.verbose()); + assert!(roundtripped_analyze.show_statistics()); assert_eq!( - roundtripped.metric_categories(), + roundtripped_analyze.metric_categories(), Some(metric_categories.as_slice()) ); - assert_eq!(roundtripped.format(), &ExplainFormat::Tree); + assert_eq!(roundtripped_analyze.format(), &ExplainFormat::Tree); assert!( - roundtripped + roundtripped_analyze .input() .downcast_ref::() .is_some() ); + + let node = PhysicalPlanNode::try_from_physical_plan(roundtripped, &codec)?; + let Some(protobuf::physical_plan_node::PhysicalPlanType::Analyze(analyze)) = + node.physical_plan_type + else { + unreachable!("expected AnalyzeExecNode") + }; + assert!(analyze.has_metric_types); + assert_eq!( + analyze.metric_types, + vec![protobuf::MetricType::Summary as i32] + ); + Ok(()) +} + +#[test] +fn roundtrip_analyze_dev_and_empty_metric_types() -> Result<()> { + let codec = DefaultPhysicalExtensionCodec {}; + for (metric_types, expected) in [ + ( + vec![MetricType::Dev], + vec![protobuf::MetricType::Dev as i32], + ), + (vec![], vec![]), + ] { + 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 analyze = Arc::new( + AnalyzeExec::builder(false, false, input, schema) + .with_metric_types(metric_types) + .build(), + ); + + let roundtripped = roundtrip_test_and_return( + analyze, + &SessionContext::new(), + &codec, + &DefaultPhysicalProtoConverter {}, + )?; + let node = PhysicalPlanNode::try_from_physical_plan(roundtripped, &codec)?; + let Some(protobuf::physical_plan_node::PhysicalPlanType::Analyze(analyze)) = + node.physical_plan_type + else { + unreachable!("expected AnalyzeExecNode") + }; + assert!(analyze.has_metric_types); + assert_eq!(analyze.metric_types, expected); + } + Ok(()) +} + +#[test] +fn decode_analyze_without_metric_types_uses_defaults() -> Result<()> { + 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 analyze = Arc::new(AnalyzeExec::builder(false, false, input, schema).build()); + let codec = DefaultPhysicalExtensionCodec {}; + 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") + }; + analyze.has_metric_types = false; + analyze.metric_types.clear(); + + // Round-trip through bytes to simulate a plan written before the field was + // added to AnalyzeExecNode. + let node = PhysicalPlanNode::decode(node.encode_to_vec().as_slice()) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + let ctx = SessionContext::new(); + let roundtripped = node.try_into_physical_plan(&ctx.task_ctx(), &codec)?; + + // Re-encode the plan to inspect the restored private state without adding + // a public accessor solely for this test. + let node = PhysicalPlanNode::try_from_physical_plan(roundtripped, &codec)?; + let Some(protobuf::physical_plan_node::PhysicalPlanType::Analyze(analyze)) = + node.physical_plan_type + else { + unreachable!("expected AnalyzeExecNode") + }; + assert!(analyze.has_metric_types); + assert_eq!( + analyze.metric_types, + vec![ + protobuf::MetricType::Summary as i32, + protobuf::MetricType::Dev as i32, + ] + ); Ok(()) } From 858933f08fa60ce9d3b296631bc07a41475bc1e0 Mon Sep 17 00:00:00 2001 From: Huaijin Date: Tue, 25 Aug 2026 22:55:54 +0800 Subject: [PATCH 2/3] update --- datafusion/proto/tests/cases/plans/misc.rs | 144 ++++++++------------- 1 file changed, 57 insertions(+), 87 deletions(-) diff --git a/datafusion/proto/tests/cases/plans/misc.rs b/datafusion/proto/tests/cases/plans/misc.rs index 66c6978b7762e..1ecec9b1e0d84 100644 --- a/datafusion/proto/tests/cases/plans/misc.rs +++ b/datafusion/proto/tests/cases/plans/misc.rs @@ -63,137 +63,107 @@ fn roundtrip_analyze() -> Result<()> { let metric_categories = vec![MetricCategory::Rows, MetricCategory::Timing]; let analyze = Arc::new( AnalyzeExec::builder(true, true, input, Arc::clone(&schema)) - // `Summary` is non-default and was previously reset to - // `[Summary, Dev]` by the protobuf round-trip. - .with_metric_types(vec![MetricType::Summary]) .with_metric_categories(Some(metric_categories.clone())) .with_format(ExplainFormat::Tree) .build(), ); let ctx = SessionContext::new(); - let codec = DefaultPhysicalExtensionCodec {}; let roundtripped = roundtrip_test_and_return( analyze, &ctx, - &codec, + &DefaultPhysicalExtensionCodec {}, &DefaultPhysicalProtoConverter {}, )?; - let roundtripped_analyze = roundtripped.downcast_ref::().unwrap(); + let roundtripped = roundtripped.downcast_ref::().unwrap(); - assert_eq!(roundtripped_analyze.schema(), schema); - assert!(roundtripped_analyze.verbose()); - assert!(roundtripped_analyze.show_statistics()); + assert_eq!(roundtripped.schema(), schema); + assert!(roundtripped.verbose()); + assert!(roundtripped.show_statistics()); assert_eq!( - roundtripped_analyze.metric_categories(), + roundtripped.metric_categories(), Some(metric_categories.as_slice()) ); - assert_eq!(roundtripped_analyze.format(), &ExplainFormat::Tree); + assert_eq!(roundtripped.format(), &ExplainFormat::Tree); assert!( - roundtripped_analyze + roundtripped .input() .downcast_ref::() .is_some() ); - - let node = PhysicalPlanNode::try_from_physical_plan(roundtripped, &codec)?; - let Some(protobuf::physical_plan_node::PhysicalPlanType::Analyze(analyze)) = - node.physical_plan_type - else { - unreachable!("expected AnalyzeExecNode") - }; - assert!(analyze.has_metric_types); - assert_eq!( - analyze.metric_types, - vec![protobuf::MetricType::Summary as i32] - ); Ok(()) } #[test] -fn roundtrip_analyze_dev_and_empty_metric_types() -> Result<()> { +fn roundtrip_analyze_metric_types() -> Result<()> { + use protobuf::MetricType as ProtoMetricType; + let codec = DefaultPhysicalExtensionCodec {}; + let ctx = SessionContext::new(); + for (metric_types, expected) in [ ( - vec![MetricType::Dev], - vec![protobuf::MetricType::Dev as i32], + 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], ), - (vec![], vec![]), ] { + 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 analyze = Arc::new( - AnalyzeExec::builder(false, false, input, schema) - .with_metric_types(metric_types) - .build(), - ); - - let roundtripped = roundtrip_test_and_return( - analyze, - &SessionContext::new(), - &codec, - &DefaultPhysicalProtoConverter {}, - )?; - let node = PhysicalPlanNode::try_from_physical_plan(roundtripped, &codec)?; + 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 + 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_mut() else { unreachable!("expected AnalyzeExecNode") }; assert!(analyze.has_metric_types); assert_eq!(analyze.metric_types, expected); + if expected == [ProtoMetricType::Summary as i32] { + 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(()) } -#[test] -fn decode_analyze_without_metric_types_uses_defaults() -> Result<()> { - 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 analyze = Arc::new(AnalyzeExec::builder(false, false, input, schema).build()); - let codec = DefaultPhysicalExtensionCodec {}; - 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") - }; - analyze.has_metric_types = false; - analyze.metric_types.clear(); - - // Round-trip through bytes to simulate a plan written before the field was - // added to AnalyzeExecNode. - let node = PhysicalPlanNode::decode(node.encode_to_vec().as_slice()) - .map_err(|e| DataFusionError::External(Box::new(e)))?; - let ctx = SessionContext::new(); - let roundtripped = node.try_into_physical_plan(&ctx.task_ctx(), &codec)?; - - // Re-encode the plan to inspect the restored private state without adding - // a public accessor solely for this test. - let node = PhysicalPlanNode::try_from_physical_plan(roundtripped, &codec)?; - let Some(protobuf::physical_plan_node::PhysicalPlanType::Analyze(analyze)) = - node.physical_plan_type - else { - unreachable!("expected AnalyzeExecNode") - }; - assert!(analyze.has_metric_types); - assert_eq!( - analyze.metric_types, - vec![ - protobuf::MetricType::Summary as i32, - protobuf::MetricType::Dev as i32, - ] - ); - Ok(()) -} - #[test] fn roundtrip_explain() -> Result<()> { let schema = Arc::new(Schema::new(vec![ From 36447201b8ba1e78ab66e8d6e33b863fcd2859d9 Mon Sep 17 00:00:00 2001 From: Huaijin Date: Wed, 26 Aug 2026 10:04:31 +0800 Subject: [PATCH 3/3] fix coverage --- datafusion/proto/tests/cases/plans/misc.rs | 118 +++++++++++++++++++-- 1 file changed, 111 insertions(+), 7 deletions(-) diff --git a/datafusion/proto/tests/cases/plans/misc.rs b/datafusion/proto/tests/cases/plans/misc.rs index 1ecec9b1e0d84..c990d977dd432 100644 --- a/datafusion/proto/tests/cases/plans/misc.rs +++ b/datafusion/proto/tests/cases/plans/misc.rs @@ -146,14 +146,38 @@ fn roundtrip_analyze_metric_types() -> Result<()> { 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_mut() - else { - unreachable!("expected AnalyzeExecNode") - }; - assert!(analyze.has_metric_types); - assert_eq!(analyze.metric_types, expected); + { + 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) @@ -164,6 +188,86 @@ fn roundtrip_analyze_metric_types() -> Result<()> { 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![