feat: support unnest in FROM clause (#9355)

* feat: support `unnest` in FROM clause * add test * remove incorrect test
apache · Feb 28, 2024 · ae4b3a0 · ae4b3a0
1 parent 19d892a
commit ae4b3a0
Show file tree

Hide file tree

Showing 4 changed files with 154 additions and 37 deletions.
diff --git a/datafusion/sql/src/expr/function.rs b/datafusion/sql/src/expr/function.rs
@@ -80,36 +80,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         if name.eq("unnest") {
             let exprs =
                 self.function_args_to_expr(args.clone(), schema, planner_context)?;
-            // Currently only one argument is supported
-            let arg = match exprs.len() {
-                0 => {
-                    return plan_err!("unnest() requires at least one argument");
-                }
-                1 => &exprs[0],
-                _ => {
-                    return not_impl_err!(
-                        "unnest() does not support multiple arguments yet"
-                    );
-                }
-            };
-            // Check argument type, array types are supported
-            match arg.get_type(schema)? {
-                DataType::List(_)
-                | DataType::LargeList(_)
-                | DataType::FixedSizeList(_, _) => {}
-                DataType::Struct(_) => {
-                    return not_impl_err!("unnest() does not support struct yet");
-                }
-                DataType::Null => {
-                    return not_impl_err!("unnest() does not support null yet");
-                }
-                _ => {
-                    return plan_err!(
-                        "unnest() can only be applied to array, struct and null"
-                    );
-                }
-            }
-
+            Self::check_unnest_args(&exprs, schema)?;
             return Ok(Expr::Unnest(Unnest { exprs }));
         }
 
@@ -312,4 +283,32 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
             .map(|a| self.sql_fn_arg_to_logical_expr(a, schema, planner_context))
             .collect::<Result<Vec<Expr>>>()
     }
+
+    pub(crate) fn check_unnest_args(args: &[Expr], schema: &DFSchema) -> Result<()> {
+        // Currently only one argument is supported
+        let arg = match args.len() {
+            0 => {
+                return plan_err!("unnest() requires at least one argument");
+            }
+            1 => &args[0],
+            _ => {
+                return not_impl_err!("unnest() does not support multiple arguments yet");
+            }
+        };
+        // Check argument type, array types are supported
+        match arg.get_type(schema)? {
+            DataType::List(_)
+            | DataType::LargeList(_)
+            | DataType::FixedSizeList(_, _) => Ok(()),
+            DataType::Struct(_) => {
+                not_impl_err!("unnest() does not support struct yet")
+            }
+            DataType::Null => {
+                not_impl_err!("unnest() does not support null yet")
+            }
+            _ => {
+                plan_err!("unnest() can only be applied to array, struct and null")
+            }
+        }
+    }
 }
diff --git a/datafusion/sql/src/relation/mod.rs b/datafusion/sql/src/relation/mod.rs
@@ -17,7 +17,7 @@
 
 use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
 use datafusion_common::{not_impl_err, plan_err, DFSchema, Result, TableReference};
-use datafusion_expr::{LogicalPlan, LogicalPlanBuilder};
+use datafusion_expr::{expr::Unnest, Expr, LogicalPlan, LogicalPlanBuilder};
 use sqlparser::ast::{FunctionArg, FunctionArgExpr, TableFactor};
 
 mod join;
@@ -96,6 +96,31 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                 self.plan_table_with_joins(*table_with_joins, planner_context)?,
                 alias,
             ),
+            TableFactor::UNNEST {
+                alias,
+                array_exprs,
+                with_offset: false,
+                with_offset_alias: None,
+            } => {
+                // Unnest table factor has empty input
+                let schema = DFSchema::empty();
+                let input = LogicalPlanBuilder::empty(true).build()?;
+                let exprs = array_exprs
+                    .into_iter()
+                    .map(|expr| {
+                        self.sql_expr_to_logical_expr(expr, &schema, planner_context)
+                    })
+                    .collect::<Result<Vec<_>>>()?;
+                Self::check_unnest_args(&exprs, &schema)?;
+                let unnest_expr = Expr::Unnest(Unnest { exprs });
+                let logical_plan = self.try_process_unnest(input, vec![unnest_expr])?;
+                (logical_plan, alias)
+            }
+            TableFactor::UNNEST { .. } => {
+                return not_impl_err!(
+                    "UNNEST table factor with offset is not supported yet"
+                );
+            }
             // @todo Support TableFactory::TableFunction?
             _ => {
                 return not_impl_err!(

diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs
@@ -277,7 +277,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
     }
 
     // Try converting Expr::Unnest to LogicalPlan::Unnest if possible, otherwise do the final projection
-    fn try_process_unnest(
+    pub(super) fn try_process_unnest(
         &self,
         input: LogicalPlan,
         select_exprs: Vec<Expr>,

diff --git a/datafusion/sqllogictest/test_files/unnest.slt b/datafusion/sqllogictest/test_files/unnest.slt
@@ -28,22 +28,43 @@ AS VALUES
     ([12], [null, 42, null], null)
 ;
 
-## Basic unnest expression
+## Basic unnest expression in select list
 query I
 select unnest([1,2,3]);
 ----
 1
 2
 3
 
+## Basic unnest expression in from clause
+query I
+select * from unnest([1,2,3]);
+----
+1
+2
+3
+
+
+## Unnest null in select list
 query error DataFusion error: This feature is not implemented: unnest\(\) does not support null yet
 select unnest(null);
 
-## Unnest empty array
+## Unnest null in from clause
+query error DataFusion error: This feature is not implemented: unnest\(\) does not support null yet
+select * from unnest(null);
+
+
+## Unnest empty array in select list
 query ?
 select unnest([]);
 ----
 
+## Unnest empty array in from clause
+query ?
+select * from unnest([]);
+----
+
+
 ## Unnest column non-null
 query I
 select unnest(column1) from unnest_table;
@@ -78,19 +99,42 @@ select unnest(column3) from unnest_table;
 query error DataFusion error: This feature is not implemented: Only support single unnest expression for now
 select unnest(column1), unnest(column2) from unnest_table;
 
-## Unnest scalar
+
+## Unnest scalar in select list
 query error DataFusion error: Error during planning: unnest\(\) can only be applied to array, struct and null
 select unnest(1);
 
+## Unnest scalar in from clause
+query error DataFusion error: Error during planning: unnest\(\) can only be applied to array, struct and null
+select * from unnest(1);
+
 
-## Unnest empty expression
+## Unnest empty expression in select list
 query error DataFusion error: Error during planning: unnest\(\) requires at least one argument
 select unnest();
 
-## Unnest struct expression
+## Unnest empty expression in from clause
+query error DataFusion error: SQL error: ParserError\("Expected an expression:, found: \)"\)
+select * from unnest();
+
+
+## Unnest multiple expressions in select list
+query error DataFusion error: This feature is not implemented: unnest\(\) does not support multiple arguments yet
+select unnest([1,2], [2,3]);
+
+## Unnest multiple expressions in from clause
+query error DataFusion error: This feature is not implemented: unnest\(\) does not support multiple arguments yet
+select * from unnest([1,2], [2,3]);
+
+
+## Unnest struct expression in select list
 query error DataFusion error: This feature is not implemented: unnest\(\) does not support struct yet
 select unnest(struct(null));
 
+## Unnest struct expression in from clause
+query error DataFusion error: This feature is not implemented: unnest\(\) does not support struct yet
+select * from unnest(struct(null));
+
 
 ## Unnest array expression
 query I
@@ -99,18 +143,36 @@ select unnest(range(1, 3));
 1
 2
 
+query I
+select * from unnest(range(1, 3));
+----
+1
+2
+
 query I
 select unnest(arrow_cast(range(1, 3), 'LargeList(Int64)'));
 ----
 1
 2
 
+query I
+select * from unnest(arrow_cast(range(1, 3), 'LargeList(Int64)'));
+----
+1
+2
+
 query I
 select unnest(arrow_cast(range(1, 3), 'FixedSizeList(2, Int64)'));
 ----
 1
 2
 
+query I
+select * from unnest(arrow_cast(range(1, 3), 'FixedSizeList(2, Int64)'));
+----
+1
+2
+
 query I
 select unnest(array_remove(column1, 12)) from unnest_table;
 ----
@@ -121,5 +183,36 @@ select unnest(array_remove(column1, 12)) from unnest_table;
 5
 6
 
+
+## Unnest in from clause with alias
+query I
+select * from unnest([1,2]) as t;
+----
+1
+2
+
+query I
+select a from unnest([1,2]) as t(a);
+----
+1
+2
+
+
+## Unnest in from clause with offset is not supported
+query error DataFusion error: This feature is not implemented: UNNEST table factor with offset is not supported yet
+select * from unnest([1,2]) with offset;
+
+query error DataFusion error: This feature is not implemented: UNNEST table factor with offset is not supported yet
+select * from unnest([1,2]) with offset offset_alias;
+
+
+## More complex cases
+query I
+select * from unnest([1,2,(select sum(column3) from unnest_table)]);
+----
+1
+2
+6
+
 statement ok
 drop table unnest_table;