more err handling

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

datafusion/sql/src/expr/function.rs

@@ -17,15 +17,15 @@
 
 use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
 use datafusion_common::{
-    not_impl_err, plan_datafusion_err, plan_err, DFSchema, DataFusionError, Dependency,
-    Result,
+    exec_err, not_impl_err, plan_datafusion_err, plan_err, DFSchema, DataFusionError,
+    Dependency, Result,
 };
 use datafusion_expr::expr::{ScalarFunction, Unnest};
 use datafusion_expr::function::suggest_valid_function;
 use datafusion_expr::window_frame::{check_window_frame, regularize_window_order_by};
 use datafusion_expr::{
-    expr, AggregateFunction, BuiltinScalarFunction, Expr, WindowFrame,
-    WindowFunctionDefinition,
+    expr, AggregateFunction, BuiltinScalarFunction, Expr, ScalarFunctionDefinition,
+    WindowFrame, WindowFunctionDefinition,
 };
 use sqlparser::ast::{
     Expr as SQLExpr, Function as SQLFunction, FunctionArg, FunctionArgExpr, WindowType,
@@ -74,6 +74,43 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         if name.eq("unnest") {
             let exprs =
                 self.function_args_to_expr(args.clone(), schema, planner_context)?;
+
+            match exprs.len() {
+                0 => {
+                    return exec_err!("unnest() requires at least one argument");
+                }
+                1 => {
+                    if let Expr::ScalarFunction(ScalarFunction {
+                        func_def:
+                            ScalarFunctionDefinition::BuiltIn(
+                                BuiltinScalarFunction::MakeArray,
+                            ),
+                        ..
+                    }) = exprs[0]
+                    {
+                        // valid
+                    } else if let Expr::Column(_) = exprs[0] {
+                        // valid
+                    } else if let Expr::ScalarFunction(ScalarFunction {
+                        func_def:
+                            ScalarFunctionDefinition::BuiltIn(BuiltinScalarFunction::Struct),
+                        ..
+                    }) = exprs[0]
+                    {
+                        return not_impl_err!("unnest() does not support struct yet");
+                    } else {
+                        return plan_err!(
+                            "unnest() can only be applied to array and structs and null"
+                        );
+                    }
+                }
+                _ => {
+                    return not_impl_err!(
+                        "unnest() does not support multiple arguments yet"
+                    );
+                }
+            }
+
             return Ok(Expr::Unnest(Unnest { exprs }));
         }
 

datafusion/sql/src/select.rs

@@ -280,24 +280,25 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         input: LogicalPlan,
         select_exprs: Vec<Expr>,
     ) -> Result<LogicalPlan> {
-        let mut array_exprs_to_unnest = vec![];
+        let mut exprs_to_unnest = vec![];
 
         for expr in select_exprs.iter() {
-            if let Expr::Unnest(Unnest { exprs: array_expr }) = expr {
-                array_exprs_to_unnest.push(array_expr[0].clone());
+            if let Expr::Unnest(Unnest { exprs }) = expr {
+                exprs_to_unnest.push(exprs[0].clone());
             }
         }
 
         // Do the final projection
-        if array_exprs_to_unnest.is_empty() {
+        if exprs_to_unnest.is_empty() {
             LogicalPlanBuilder::from(input)
                 .project(select_exprs)?
                 .build()
         } else {
-            // Only support single unnest expression for now
-            assert_eq!(array_exprs_to_unnest.len(), 1);
+            if exprs_to_unnest.len() > 1 {
+                return not_impl_err!("Only support single unnest expression for now");
+            }
 
-            let expr = array_exprs_to_unnest[0].clone();
+            let expr = exprs_to_unnest[0].clone();
             let column = expr.display_name()?;
 
             LogicalPlanBuilder::from(input)

datafusion/sqllogictest/test_files/unnest.slt

@@ -22,9 +22,9 @@
 statement ok
 CREATE TABLE unnest_table
 AS VALUES
-    ([1,2,3]),
-    ([4,5]),
-    ([6])
+    ([1,2,3], [7], 1),
+    ([4,5], [8,9,10], 2),
+    ([6], [11,12], 3)
 ;
 
 query I
@@ -55,5 +55,18 @@ select unnest(column1) from unnest_table;
 5
 6
 
+query error DataFusion error: This feature is not implemented: Only support single unnest expression for now
+select unnest(column1), unnest(column2) from unnest_table;
+
+query error DataFusion error: Error during planning: unnest\(\) can only be applied to array and structs and null
+select unnest(1);
+
+# TODO: This should be an error, but unnest is able to process scalar values now.
+# query error
+# select unnest(column3) from unnest_table;
+
+query error DataFusion error: Execution error: unnest\(\) requires at least one argument
+select unnest();
+
 statement ok
 drop table unnest_table;