diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index 4de3952dc7ea1..95a44dace31a8 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -23,7 +23,7 @@ use crate::planner::{ }; use crate::utils::{ check_columns_satisfy_exprs, extract_aliases, rebase_expr, resolve_aliases_to_exprs, - resolve_columns, resolve_positions_to_exprs, transform_bottom_unnest, + resolve_columns, resolve_positions_to_exprs, transform_bottom_unnests, }; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; @@ -318,20 +318,12 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { // - unnest(struct_col) will be transformed into unnest(struct_col).field1, unnest(struct_col).field2 // - unnest(array_col) will be transformed into unnest(array_col).element // - unnest(array_col) + 1 will be transformed into unnest(array_col).element +1 - let outer_projection_exprs: Vec = intermediate_select_exprs - .iter() - .map(|expr| { - transform_bottom_unnest( - &intermediate_plan, - &mut unnest_columns, - &mut inner_projection_exprs, - expr, - ) - }) - .collect::>>()? - .into_iter() - .flatten() - .collect(); + let outer_projection_exprs = transform_bottom_unnests( + &intermediate_plan, + &mut unnest_columns, + &mut inner_projection_exprs, + &intermediate_select_exprs, + )?; // No more unnest is possible if unnest_columns.is_empty() { @@ -417,20 +409,12 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { let mut unnest_columns = vec![]; let mut inner_projection_exprs = vec![]; - let outer_projection_exprs: Vec = intermediate_select_exprs - .iter() - .map(|expr| { - transform_bottom_unnest( - &intermediate_plan, - &mut unnest_columns, - &mut inner_projection_exprs, - expr, - ) - }) - .collect::>>()? - .into_iter() - .flatten() - .collect(); + let outer_projection_exprs = transform_bottom_unnests( + &intermediate_plan, + &mut unnest_columns, + &mut inner_projection_exprs, + &intermediate_select_exprs, + )?; if unnest_columns.is_empty() { break; diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index a9a782902ac95..3b044646e6cb2 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -287,6 +287,28 @@ pub(crate) fn value_to_string(value: &Value) -> Option { } } +pub(crate) fn transform_bottom_unnests( + input: &LogicalPlan, + unnest_placeholder_columns: &mut Vec, + inner_projection_exprs: &mut Vec, + original_exprs: &[Expr], +) -> Result> { + Ok(original_exprs + .iter() + .map(|expr| { + transform_bottom_unnest( + input, + unnest_placeholder_columns, + inner_projection_exprs, + expr, + ) + }) + .collect::>>()? + .into_iter() + .flatten() + .collect::>()) +} + /// The context is we want to rewrite unnest() into InnerProjection->Unnest->OuterProjection /// Given an expression which contains unnest expr as one of its children, /// Try transform depends on unnest type