findepi commented on code in PR #13709:
URL: https://github.com/apache/datafusion/pull/13709#discussion_r1877914306
##########
datafusion/proto/tests/cases/roundtrip_logical_plan.rs:
##########
@@ -2555,3 +2557,47 @@ async fn roundtrip_recursive_query() {
format!("{}", pretty_format_batches(&output_round_trip).unwrap())
);
}
+
+#[tokio::test]
+async fn roundtrip_union_query() -> Result<()> {
+ let query = "SELECT a FROM t1
+ UNION (SELECT a from t1 UNION SELECT a from t2)";
+
+ let ctx = SessionContext::new();
+ ctx.register_csv("t1", "tests/testdata/test.csv",
CsvReadOptions::default())
+ .await?;
+ ctx.register_csv("t2", "tests/testdata/test.csv",
CsvReadOptions::default())
+ .await?;
+ let dataframe = ctx.sql(query).await?;
+ let plan = dataframe.into_optimized_plan()?;
+ // plan.display_indent_schema()
+ // Aggregate: groupBy=[[t1.a]], aggr=[[]] [a:Int64;N]
+ // Union [a:Int64;N]
+ // TableScan: t1 projection=[a] [a:Int64;N]
+ // TableScan: t1 projection=[a] [a:Int64;N]
+ // TableScan: t2 projection=[a] [a:Int64;N]
+
+ let bytes = logical_plan_to_bytes(&plan)?;
+
+ let ctx = SessionContext::new();
+ ctx.register_csv("t1", "tests/testdata/test.csv",
CsvReadOptions::default())
+ .await?;
+ ctx.register_csv("t2", "tests/testdata/test.csv",
CsvReadOptions::default())
+ .await?;
+ let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx)?;
+ // proto deserialisation only supports 2-way union, hence this plan has
nested unions
+ // logical_round_trip.display_indent_schema()
+ // Aggregate: groupBy=[[t1.a]], aggr=[[]] [a:Int64;N]
+ // Union [a:Int64;N]
+ // Union [a:Int64;N]
+ // TableScan: t1 projection=[a] [a:Int64;N]
+ // TableScan: t1 projection=[a] [a:Int64;N]
+ // TableScan: t2 projection=[a] [a:Int64;N]
Review Comment:
as a code comment (which can become outdated), is this useful?
##########
datafusion/proto/src/logical_plan/mod.rs:
##########
@@ -737,23 +737,18 @@ impl AsLogicalPlan for LogicalPlanNode {
builder.build()
}
LogicalPlanType::Union(union) => {
- let mut input_plans: Vec<LogicalPlan> = union
- .inputs
- .iter()
- .map(|i| i.try_into_logical_plan(ctx, extension_codec))
- .collect::<Result<_>>()?;
-
- if input_plans.len() < 2 {
+ if union.inputs.len() < 2 {
return Err( DataFusionError::Internal(String::from(
"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)?,
+ );
- let first = input_plans.pop().ok_or_else(||
DataFusionError::Internal(String::from(
Review Comment:
Would it be enough to change `pop` to `split_first` here?
--
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]
