cleanup

Signed-off-by: jayzhan211 <[email protected]>
apache · Jul 28, 2023 · e55a2fb · e55a2fb
1 parent a107905
commit e55a2fb
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 58 deletions.
diff --git 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
@@ -1141,11 +1141,14 @@ impl LogicalPlanBuilder {
             .build()
     }
 
-    pub fn unnest_arrays_v3(
+    pub fn unnest_arrays(
         self,
         array_exprs: Vec<Expr>,
-        columns_name: Vec<String>,
+        // columns_name: Vec<String>,
     ) -> Result<Self> {
+        let columns_name = (0..array_exprs.len())
+            .map(|idx| format!("col{idx}"))
+            .collect::<Vec<String>>();
         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
@@ -56,59 +56,31 @@ impl UnnestExpressions {
 
     fn optimize_internal(plan: &LogicalPlan) -> Result<LogicalPlan> {
         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
@@ -38,15 +38,8 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
 
         let plan = LogicalPlanBuilder::empty(true).build()?;
 
-        let columns_name: Vec<String> =
-            (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()
     }
 }