fix(cube): Make LogicalPlan protobuf deserialization roundtrip unnested unions

Author: srh

datafusion/proto/src/logical_plan/mod.rs

@@ -785,16 +785,12 @@ impl AsLogicalPlan for LogicalPlanNode {
                         "Protobuf deserialization error, Union was require at least two input.",
                     )));
                 }
-                let (first, rest) = union.inputs.split_first().unwrap();
-                let mut builder = LogicalPlanBuilder::from(
-                    first.try_into_logical_plan(ctx, extension_codec)?,
-                );
 
-                for i in rest {
-                    let plan = i.try_into_logical_plan(ctx, extension_codec)?;
-                    builder = builder.union(plan)?;
+                let mut input_plans = Vec::with_capacity(union.inputs.len());
+                for input in &union.inputs {
+                    input_plans.push(Arc::new(input.try_into_logical_plan(ctx, extension_codec)?));
                 }
-                builder.build()
+                Ok(LogicalPlan::Union(datafusion_expr::Union::try_new_with_loose_types(input_plans)?))
             }
             LogicalPlanType::CrossJoin(crossjoin) => {
                 let left = into_logical_plan!(crossjoin.left, ctx, extension_codec)?;

datafusion/proto/tests/cases/roundtrip_logical_plan.rs

@@ -25,8 +25,6 @@ use arrow::datatypes::{
 };
 use arrow::util::pretty::pretty_format_batches;
 use datafusion::datasource::file_format::json::JsonFormatFactory;
-use datafusion::optimizer::eliminate_nested_union::EliminateNestedUnion;
-use datafusion::optimizer::Optimizer;
 use datafusion_common::parsers::CompressionTypeVariant;
 use prost::Message;
 use std::any::Any;
@@ -2603,13 +2601,9 @@ async fn roundtrip_union_query() -> Result<()> {
     ctx.register_csv("t2", "tests/testdata/test.csv", CsvReadOptions::default())
         .await?;
     let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx)?;
-    // proto deserialization only supports 2-way union, hence this plan has nested unions
-    // apply the flatten unions optimizer rule to be able to compare
-    let optimizer = Optimizer::with_rules(vec![Arc::new(EliminateNestedUnion::new())]);
-    let unnested = optimizer.optimize(logical_round_trip, &(ctx.state()), |_x, _y| {})?;
     assert_eq!(
         format!("{}", plan.display_indent_schema()),
-        format!("{}", unnested.display_indent_schema()),
+        format!("{}", logical_round_trip.display_indent_schema()),
     );
     Ok(())
 }