kosiew commented on code in PR #24669:
URL: https://github.com/apache/datafusion/pull/24669#discussion_r3886127505


##########
datafusion/proto/tests/cases/plans/misc.rs:
##########
@@ -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] {

Review Comment:
   Could we move the invalid-schema and unknown-enum assertions into a 
dedicated malformed `AnalyzeExecNode` decode test? Right now they only run when 
the successful round-trip case happens to select `Summary`, which makes the 
independent decode-error contract a little harder to see. This is just a test 
organization suggestion, and the existing coverage looks equivalent.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to