cleanup

Signed-off-by: jayzhan211 <[email protected]>

datafusion/common/src/scalar.rs

@@ -3183,7 +3183,6 @@ mod tests {
     use std::cmp::Ordering;
     use std::sync::Arc;
 
-    use crate::Result;
     use arrow::compute::kernels;
     use arrow::compute::{concat, is_null};
     use arrow::datatypes::ArrowPrimitiveType;

datafusion/core/src/physical_planner.rs

@@ -30,8 +30,8 @@ use crate::datasource::source_as_provider;
 use crate::execution::context::{ExecutionProps, SessionState};
 use crate::logical_expr::utils::generate_sort_key;
 use crate::logical_expr::{
-    Aggregate, EmptyRelation, Join, Projection, Sort, SubqueryAlias, TableScan, Unnest,
-    Window,
+    Aggregate, EmptyRelation, Join, Projection, Sort, SubqueryAlias, TableScan,
+    Unnest as UnnestPlan, Window,
 };
 use crate::logical_expr::{
     CrossJoin, Expr, LogicalPlan, Partitioning as LogicalPartitioning, PlanType,
@@ -79,8 +79,9 @@ use datafusion_common::{
 use datafusion_expr::expr::{
     self, AggregateFunction, AggregateUDF, Alias, Between, BinaryExpr, Cast,
     GetFieldAccess, GetIndexedField, GroupingSet, InList, Like, ScalarUDF, TryCast,
-    Unnest as UnnestExpr, WindowFunction,
+    Unnest, WindowFunction,
 };
+
 use datafusion_expr::expr_rewriter::{unalias, unnormalize_cols};
 use datafusion_expr::logical_plan::builder::wrap_projection_for_join_if_necessary;
 use datafusion_expr::{DmlStatement, StringifiedPlan, WriteOp};
@@ -212,7 +213,7 @@ fn create_physical_name(e: &Expr, is_first_expr: bool) -> Result<String> {
 
             Ok(name)
         }
-        Expr::Unnest(UnnestExpr { array_exprs, .. }) => {
+        Expr::Unnest(Unnest { array_exprs, .. }) => {
             create_function_physical_name("unnest", false, array_exprs)
         }
         Expr::ScalarFunction(func) => {
@@ -443,11 +444,9 @@ impl PhysicalPlanner for DefaultPhysicalPlanner {
         match self.handle_explain(logical_plan, session_state).await? {
             Some(plan) => Ok(plan),
             None => {
-                // println!("create init plan");
                 let plan = self
                     .create_initial_plan(logical_plan, session_state)
                     .await?;
-                // println!("create init plan done");
                 self.optimize_internal(plan, session_state, |_, _| {})
             }
         }
@@ -1217,7 +1216,7 @@ impl DefaultPhysicalPlanner {
 
                     Ok(Arc::new(GlobalLimitExec::new(input, *skip, *fetch)))
                 }
-                LogicalPlan::Unnest(Unnest { input, column, schema, options }) => {
+                LogicalPlan::Unnest(UnnestPlan { input, column, schema, options }) => {
                     let input = self.create_initial_plan(input, session_state).await?;
                     let column_exec = schema.index_of_column(column)
                         .map(|idx| Column::new(&column.name, idx))?;

datafusion/expr/src/expr.rs

@@ -1177,10 +1177,9 @@ impl fmt::Display for Expr {
                     write!(f, " NULLS LAST")
                 }
             }
-            Expr::Unnest(Unnest {
-                array_exprs,
-                ..
-            }) => fmt_function(f, "unnest", false, array_exprs, false),
+            Expr::Unnest(Unnest { array_exprs, .. }) => {
+                fmt_function(f, "unnest", false, array_exprs, false)
+            }
             Expr::ScalarFunction(func) => {
                 fmt_function(f, &func.fun.to_string(), false, &func.args, true)
             }
@@ -1514,10 +1513,9 @@ fn create_name(e: &Expr) -> Result<String> {
                 }
             }
         }
-        Expr::Unnest(Unnest {
-            array_exprs,
-            ..
-        }) => create_function_name("unnest", false, array_exprs),
+        Expr::Unnest(Unnest { array_exprs, .. }) => {
+            create_function_name("unnest", false, array_exprs)
+        }
         Expr::ScalarFunction(func) => {
             create_function_name(&func.fun.to_string(), false, &func.args)
         }

datafusion/expr/src/expr_schema.rs

@@ -88,10 +88,7 @@ impl ExprSchemable for Expr {
                     .collect::<Result<Vec<_>>>()?;
                 Ok((fun.return_type)(&data_types)?.as_ref().clone())
             }
-            Expr::Unnest(Unnest {
-                array_exprs,
-                ..
-            }) => {
+            Expr::Unnest(Unnest { array_exprs, .. }) => {
                 let data_types = array_exprs
                     .iter()
                     .map(|e| e.get_type(schema))

datafusion/expr/src/logical_plan/builder.rs

@@ -27,8 +27,8 @@ use crate::expr_rewriter::{
 use crate::type_coercion::binary::comparison_coercion;
 use crate::utils::{columnize_expr, compare_sort_expr};
 use crate::{
-    and, binary_expr, expr, BuiltInWindowFunction, DmlStatement, Operator,
-    WindowFrame, WriteOp,
+    and, binary_expr, expr, BuiltInWindowFunction, DmlStatement, Operator, WindowFrame,
+    WriteOp,
 };
 use crate::{col, ident, WindowFunction};
 use crate::{
@@ -1173,10 +1173,10 @@ impl LogicalPlanBuilder {
     pub fn unnest_arrays(
         self,
         array_exprs: Vec<Expr>,
-        options: UnnestOptions,
+        array_options: Vec<UnnestOptions>,
     ) -> Result<Self> {
         let (unnest_plans, columns_name) =
-            build_unnest_plans(self.plan.clone(), array_exprs.clone(), options.clone())?;
+            build_unnest_plans(self.plan.clone(), array_exprs.clone(), array_options)?;
 
         Self::join_unnest_plans(unnest_plans, columns_name)
     }
@@ -1626,23 +1626,23 @@ pub fn unnest_with_options(
 fn build_unnest_plans(
     input: LogicalPlan,
     array_exprs: Vec<Expr>,
-    options: UnnestOptions,
+    array_options: Vec<UnnestOptions>,
 ) -> Result<(Vec<LogicalPlan>, Vec<String>)> {
-    let array_exprs = align_arrays_with_nulls(&array_exprs)?;
+    let array_expression = align_arrays_with_nulls(&array_exprs)?;
 
     let project_plan_builder =
-        LogicalPlanBuilder::from(input.clone()).project(array_exprs.clone())?;
+        LogicalPlanBuilder::from(input.clone()).project(array_expression.clone())?;
 
     let mut unnest_plans = vec![];
     let mut columns_name = vec![];
 
-    for array_expr in array_exprs.iter() {
-        let column = array_expr.display_name()?;
+    for (expr, options) in array_expression.iter().zip(array_options.into_iter()) {
+        let column = expr.display_name()?;
         columns_name.push(column.clone());
 
         let unnest_plan = project_plan_builder
             .clone()
-            .unnest_column_with_options(column, options.clone())?
+            .unnest_column_with_options(column, options)?
             .build()?;
         unnest_plans.push(unnest_plan);
     }

datafusion/proto/src/logical_plan/to_proto.rs

@@ -34,7 +34,9 @@ use arrow::datatypes::{
     DataType, Field, IntervalMonthDayNanoType, IntervalUnit, Schema, SchemaRef, TimeUnit,
     UnionMode,
 };
-use datafusion_common::{Column, DFField, DFSchemaRef, OwnedTableReference, ScalarValue};
+use datafusion_common::{
+    not_impl_err, Column, DFField, DFSchemaRef, OwnedTableReference, ScalarValue,
+};
 use datafusion_expr::expr::{
     self, Alias, Between, BinaryExpr, Cast, GetFieldAccess, GetIndexedField, GroupingSet,
     InList, Like, Placeholder, ScalarFunction, ScalarUDF, Sort,
@@ -741,7 +743,9 @@ impl TryFrom<&Expr> for protobuf::LogicalExprNode {
                         .to_string(),
                 ))
             }
-            Expr::Unnest(..) => todo!("Unnest not supported"),
+            Expr::Unnest(..) => {
+                return not_impl_err!("try_from() for Unnest is not implemented")
+            }
             Expr::ScalarFunction(ScalarFunction { fun, args }) => {
                 let fun: protobuf::ScalarFunction = fun.try_into()?;
                 let args: Vec<Self> = args

datafusion/sql/src/query.rs

@@ -73,6 +73,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                 // Each `WITH` block can change the column names in the last
                 // projection (e.g. "WITH table(t1, t2) AS SELECT 1, 2").
                 let logical_plan = self.apply_table_alias(logical_plan, cte.alias)?;
+
                 planner_context.insert_cte(cte_name, logical_plan);
             }
         }

datafusion/sql/src/relation/mod.rs

@@ -62,13 +62,17 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                 self.plan_table_with_joins(*table_with_joins, planner_context)?,
                 alias,
             ),
+
+            // TODO: Support UnnestOptions
             TableFactor::UNNEST {
                 alias,
                 array_exprs,
                 with_offset: _,
                 with_offset_alias: _,
             } => {
-                let options: UnnestOptions = Default::default();
+                // TODO: Support UnnestOptions
+                let array_options: Vec<UnnestOptions> =
+                    vec![UnnestOptions::default(); array_exprs.len()];
 
                 // If column aliases are not supplied, then for a function returning a base data type,
                 // the column name is also the same as the function name.
@@ -77,12 +81,12 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                         alias.columns = vec![alias.name.clone()];
                     }
                     (
-                        self.plan_unnest(array_exprs, planner_context, options)?,
+                        self.plan_unnest(array_exprs, planner_context, array_options)?,
                         Some(alias),
                     )
                 } else {
                     (
-                        self.plan_unnest(array_exprs, planner_context, options)?,
+                        self.plan_unnest(array_exprs, planner_context, array_options)?,
                         None,
                     )
                 }

datafusion/sql/src/relation/unnest.rs

@@ -27,7 +27,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         &self,
         array_exprs: Vec<SQLExpr>,
         planner_context: &mut PlannerContext,
-        options: UnnestOptions,
+        array_options: Vec<UnnestOptions>,
     ) -> Result<LogicalPlan> {
         // No pre-defiend schema for Unnest
         let schema = DFSchema::empty();
@@ -40,7 +40,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         let plan = LogicalPlanBuilder::empty(true).build()?;
 
         LogicalPlanBuilder::from(plan)
-            .unnest_arrays(exprs, options)?
+            .unnest_arrays(exprs, array_options)?
             .build()
     }
 }

datafusion/sql/src/select.rs

@@ -17,6 +17,7 @@
 
 use std::collections::HashSet;
 use std::sync::Arc;
+use std::vec;
 
 use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
 use crate::utils::{
@@ -25,10 +26,9 @@ use crate::utils::{
 };
 
 use datafusion_common::{
-    get_target_functional_dependencies, not_impl_err, plan_err, DFSchemaRef,
+    get_target_functional_dependencies, not_impl_err, plan_err, Column, DFSchemaRef,
     DataFusionError, Result,
 };
-use datafusion_common::{Column, UnnestOptions};
 use datafusion_expr::expr::{Alias, Unnest};
 use datafusion_expr::expr_rewriter::{
     normalize_col, normalize_col_with_schemas_and_ambiguity_check,
@@ -318,18 +318,14 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         empty_from: bool,
         planner_context: &mut PlannerContext,
     ) -> Result<Vec<Expr>> {
-        let select_exprs = projection
+        projection
             .into_iter()
             .map(|expr| self.sql_select_to_rex(expr, plan, empty_from, planner_context))
             .flat_map(|result| match result {
                 Ok(vec) => vec.into_iter().map(Ok).collect(),
                 Err(err) => vec![Err(err)],
             })
-            .collect::<Result<Vec<Expr>>>();
-
-        // TODO: Align unnest expr here
-
-        select_exprs
+            .collect::<Result<Vec<Expr>>>()
     }
 
     /// Generate a relational expression from a select SQL expression
@@ -688,6 +684,7 @@ fn get_updated_group_by_exprs(
 
 // convert Expr::Unnest to LogicalPlan::Unnest
 fn process_unnest_expr(input: LogicalPlan, select_exprs: &[Expr]) -> Result<LogicalPlan> {
+    let mut array_options = vec![];
     let mut array_exprs_to_unnest = vec![];
     for expr in select_exprs.iter() {
         if let Expr::Unnest(Unnest {
@@ -696,22 +693,23 @@ fn process_unnest_expr(input: LogicalPlan, select_exprs: &[Expr]) -> Result<Logi
         }) = expr
         {
             array_exprs_to_unnest.push(array_exprs[0].clone());
+            array_options.push(options.clone());
         } else if let Expr::Alias(Alias { expr, .. }) = expr {
             if let Expr::Unnest(Unnest {
                 array_exprs,
                 options,
             }) = expr.as_ref()
             {
                 array_exprs_to_unnest.push(array_exprs[0].clone());
+                array_options.push(options.clone());
             }
         }
     }
     if array_exprs_to_unnest.is_empty() {
         Ok(input)
     } else {
-        let options: UnnestOptions = Default::default();
         LogicalPlanBuilder::from(input)
-            .unnest_arrays(array_exprs_to_unnest, options)?
+            .unnest_arrays(array_exprs_to_unnest, array_options)?
             .build()
     }
 }