You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by al...@apache.org on 2024/02/28 16:20:49 UTC
(arrow-datafusion) branch main updated: feat: support `unnest` in FROM clause (#9355)
This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new ae4b3a0f83 feat: support `unnest` in FROM clause (#9355)
ae4b3a0f83 is described below
commit ae4b3a0f8366ab18be5ef5cfa2b3cc3aca12baf1
Author: Jonah Gao <jo...@msn.com>
AuthorDate: Thu Feb 29 00:20:43 2024 +0800
feat: support `unnest` in FROM clause (#9355)
* feat: support `unnest` in FROM clause
* add test
* remove incorrect test
---
datafusion/sql/src/expr/function.rs | 59 ++++++++-------
datafusion/sql/src/relation/mod.rs | 27 ++++++-
datafusion/sql/src/select.rs | 2 +-
datafusion/sqllogictest/test_files/unnest.slt | 103 ++++++++++++++++++++++++--
4 files changed, 154 insertions(+), 37 deletions(-)
diff --git a/datafusion/sql/src/expr/function.rs b/datafusion/sql/src/expr/function.rs
index bc1d672522..bb73e69ba9 100644
--- 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
index 1d52899160..987fe85de3 100644
--- 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
index b5b6ada48f..59ef16d758 100644
--- 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
index 9990c00f75..ccdd078759 100644
--- a/datafusion/sqllogictest/test_files/unnest.slt
+++ b/datafusion/sqllogictest/test_files/unnest.slt
@@ -28,7 +28,7 @@ AS VALUES
([12], [null, 42, null], null)
;
-## Basic unnest expression
+## Basic unnest expression in select list
query I
select unnest([1,2,3]);
----
@@ -36,14 +36,35 @@ select unnest([1,2,3]);
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;