From e55a2fb3bb107d4be7eb8b10f7490302efc7469d Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Fri, 28 Jul 2023 12:19:20 +0800 Subject: [PATCH] cleanup Signed-off-by: jayzhan211 --- .../tests/sqllogictests/test_files/array.slt | 31 +++++++------------ datafusion/expr/src/logical_plan/builder.rs | 7 +++-- .../optimizer/src/unnest_expressions.rs | 30 +----------------- datafusion/sql/src/relation/unnest.rs | 9 +----- 4 files changed, 19 insertions(+), 58 deletions(-) diff --git a/datafusion/core/tests/sqllogictests/test_files/array.slt b/datafusion/core/tests/sqllogictests/test_files/array.slt index 3120c4119999b..24860accf5217 100644 --- a/datafusion/core/tests/sqllogictests/test_files/array.slt +++ b/datafusion/core/tests/sqllogictests/test_files/array.slt @@ -1717,18 +1717,6 @@ SELECT sum(data.a), sum(data.b) FROM unnest(make_array(1, 2, 3), make_array(4, 5 ---- 6 22 -# # TODO: sum for unnest as subquery not working yet -# # query error DataFusion error: Error during planning: The function Sum does not support inputs of type List\(Field \{ name: "item", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: \{\} \}\)\. -# # SELECT sum(a) FROM (select unnest(make_array(1, 2, 3)) as a); -# -# # TODO: select from part of the columns is not allowed. -# # query error DataFusion error: External error: Arrow error: Invalid argument error: must either specify a row count or at least one column -# # SELECT data.a FROM unnest(make_array(1, 2, 3), make_array(4, 5, 6, 7)) as data(a, b); -# -# # TODO: select from part of the columns is not allowed. -# # query error DataFusion error: External error: External error: Arrow error: Invalid argument error: must either specify a row count or at least one column -# # SELECT sum(data.b) FROM unnest(make_array(1, 2, 3), make_array(4, 5, 6, 7)) as data(a, b); - # Test unnest with multi-dimensional arrays query ? select unnest(make_array(make_array(10,20,30),make_array(1,NULL,10),make_array(4,5,6))) @@ -1743,6 +1731,18 @@ NULL 5 6 +# # TODO: sum for unnest as subquery not working yet +# # query error DataFusion error: Error during planning: The function Sum does not support inputs of type List\(Field \{ name: "item", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: \{\} \}\)\. +# # SELECT sum(a) FROM (select unnest(make_array(1, 2, 3)) as a); +# +# # TODO: select from part of the columns is not allowed. +# # query error DataFusion error: External error: Arrow error: Invalid argument error: must either specify a row count or at least one column +# # SELECT data.a FROM unnest(make_array(1, 2, 3), make_array(4, 5, 6, 7)) as data(a, b); +# +# # TODO: select from part of the columns is not allowed. +# # query error DataFusion error: External error: External error: Arrow error: Invalid argument error: must either specify a row count or at least one column +# # SELECT sum(data.b) FROM unnest(make_array(1, 2, 3), make_array(4, 5, 6, 7)) as data(a, b); + # # Postgres # select unnest(array[1,2,3]); # ---- @@ -1791,10 +1791,3 @@ NULL # 4 # 5 # (Must match dimensions) - -# query ? -# select unnest(make_array(1,2,3), make_array(4,5)); -# ---- -# 1 4 -# 2 5 -# 3 NULL diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 45f1296b91d69..c980f0bbb7938 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1141,11 +1141,14 @@ impl LogicalPlanBuilder { .build() } - pub fn unnest_arrays_v3( + pub fn unnest_arrays( self, array_exprs: Vec, - columns_name: Vec, + // columns_name: Vec, ) -> Result { + let columns_name = (0..array_exprs.len()) + .map(|idx| format!("col{idx}")) + .collect::>(); let (unnest_plans, columns_name) = build_unnest_plans(self.plan.clone(), array_exprs.clone(), columns_name)?; let plan = Self::join_unnest_plans(unnest_plans, columns_name)?; diff --git a/datafusion/optimizer/src/unnest_expressions.rs b/datafusion/optimizer/src/unnest_expressions.rs index 762ea8786606c..47a66d6d7bf75 100644 --- a/datafusion/optimizer/src/unnest_expressions.rs +++ b/datafusion/optimizer/src/unnest_expressions.rs @@ -56,59 +56,31 @@ impl UnnestExpressions { fn optimize_internal(plan: &LogicalPlan) -> Result { if let LogicalPlan::Projection(_) = plan { - let mut idx = 0; - let mut columns_name = vec![]; let mut array_exprs_to_unnest = vec![]; for expr in plan.expressions() { if let Expr::Unnest(Unnest { array_exprs }) = expr { assert_eq!(array_exprs.len(), 1); - columns_name.push(format!("col{idx}")); - idx += 1; - - // plans.push(unnest_plan); array_exprs_to_unnest.push(array_exprs[0].clone()); } else if let Expr::Alias(Alias { expr, .. }) = expr { if let Expr::Unnest(Unnest { array_exprs }) = expr.as_ref() { assert_eq!(array_exprs.len(), 1); - // let unnest_plan = LogicalPlanBuilder::from(plan.clone()) - // .unnest_arrays_v2( - // array_exprs.clone(), - // format!("col{idx}").as_str(), - // )? - // .build()?; - - columns_name.push(format!("col{idx}")); - idx += 1; - // plans.push(unnest_plan); array_exprs_to_unnest.push(array_exprs[0].clone()); } - } else { } } if array_exprs_to_unnest.len() > 0 { let unnest_plan = LogicalPlanBuilder::from(plan.clone()) - .unnest_arrays_v3( - array_exprs_to_unnest.clone(), - columns_name.clone(), - )? + .unnest_arrays(array_exprs_to_unnest.clone())? .build()?; Ok(unnest_plan) } else { Ok(plan.clone()) } - - // // call join_unnest_plans - // if plans.len() > 0 { - // let plan = Self::join_unnest_plans(plans, columns_name)?; - // Ok(plan) - // } else { - // Ok(plan.clone()) - // } } else { Ok(plan.clone()) } diff --git a/datafusion/sql/src/relation/unnest.rs b/datafusion/sql/src/relation/unnest.rs index 885a9e43f2a85..22dca8e6c7634 100644 --- a/datafusion/sql/src/relation/unnest.rs +++ b/datafusion/sql/src/relation/unnest.rs @@ -38,15 +38,8 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { let plan = LogicalPlanBuilder::empty(true).build()?; - let columns_name: Vec = - (0..exprs.len()).map(|idx| format!("col{idx}")).collect(); - LogicalPlanBuilder::from(plan) - .unnest_arrays_v3(exprs, columns_name)? + .unnest_arrays(exprs)? .build() - - // LogicalPlanBuilder::from(plan) - // .unnest_arrays(exprs, "col")? - // .build() } }