// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements. See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership. The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License. You may obtain a copy of the License at

//

// http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied. See the License for the

// specific language governing permissions and limitations

// under the License.

use std::collections::HashSet;

use std::ops::ControlFlow;

use std::sync::Arc;

use crate::expr::{QUALIFY_HELP, reject_window_functions};

use crate::planner::{ContextProvider, PlannerContext, SqlToRel};

use crate::query::to_order_by_exprs_with_select;

use crate::utils::{

CheckColumnsMustReferenceAggregatePurpose, CheckColumnsSatisfyExprsPurpose,

DedupedProjection, check_columns_satisfy_exprs, extract_aliases, rebase_expr,

resolve_aliases_to_exprs, resolve_columns, resolve_positions_to_exprs,

rewrite_recursive_unnest_bottom_up, rewrite_recursive_unnests_bottom_up,

substitute_top_level_alias, substitute_top_level_aliases_in_sorts,

};

use arrow::datatypes::DataType;

use datafusion_common::error::DataFusionErrorBuilder;

use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};

use datafusion_common::{

Column, DFSchema, DFSchemaRef, HashMap, Result, not_impl_err, plan_err,

};

use datafusion_common::{NullHandling, RecursionUnnestOption, UnnestOptions};

use datafusion_expr::ExprSchemable;

use datafusion_expr::builder::get_struct_unnested_columns;

use datafusion_expr::expr::Unnest as UnnestExpr;

use datafusion_expr::expr::{PlannedReplaceSelectItem, WildcardOptions};

use datafusion_expr::expr_rewriter::{

normalize_col, normalize_col_with_schemas_and_ambiguity_check, normalize_sorts,

};

use datafusion_expr::select_expr::SelectExpr;

use datafusion_expr::utils::{

expr_as_column_expr, expr_to_columns, find_aggregate_exprs, find_window_exprs,

};

use datafusion_expr::{

Aggregate, Expr, Filter, GroupingSet, LogicalPlan, LogicalPlanBuilder,

LogicalPlanBuilderOptions, Partitioning, SortExpr,

};

use indexmap::IndexMap;

use sqlparser::ast::{

Distinct, Expr as SQLExpr, GroupByExpr, NamedWindowExpr, OrderBy,

SelectItemQualifiedWildcardKind, WildcardAdditionalOptions, WindowType,

visit_expressions_mut,

};

use sqlparser::ast::{NamedWindowDefinition, Select, SelectItem, TableWithJoins};

/// Result of the `aggregate` function, containing the aggregate plan and

/// rewritten expressions that reference the aggregate output columns.

struct AggregatePlanResult {

/// The aggregate logical plan

plan: LogicalPlan,

/// SELECT expressions rewritten to reference aggregate output columns

select_exprs: Vec<Expr>,

/// HAVING expression rewritten to reference aggregate output columns

having_expr: Option<Expr>,

/// QUALIFY expression rewritten to reference aggregate output columns

qualify_expr: Option<Expr>,

/// ORDER BY expressions rewritten to reference aggregate output columns

order_by_exprs: Vec<SortExpr>,

/// DISTINCT ON expressions rewritten to reference aggregate output columns

on_exprs: Vec<Expr>,

}

struct DistinctOnUnnestPlanResult {

plan: LogicalPlan,

select_exprs: Vec<Expr>,

on_exprs: Vec<Expr>,

order_by_exprs: Vec<SortExpr>,

}

struct RewrittenUnnestExprGroups {

plan: LogicalPlan,

expr_groups: Vec<Vec<Expr>>,

}

fn flatten_expr_groups(expr_groups: Vec<Vec<Expr>>) -> Vec<Expr> {

expr_groups.into_iter().flatten().collect()

}

impl<S: ContextProvider> SqlToRel<'_, S> {

/// Generate a logic plan from an SQL select

pub(super) fn select_to_plan(

&self,

mut select: Select,

query_order_by: Option<OrderBy>,

planner_context: &mut PlannerContext,

) -> Result<LogicalPlan> {

// Check for unsupported syntax first

if !select.cluster_by.is_empty() {

return not_impl_err!("CLUSTER BY");

}

if !select.lateral_views.is_empty() {

return not_impl_err!("LATERAL VIEWS");

}

if select.top.is_some() {

return not_impl_err!("TOP");

}

if !select.sort_by.is_empty() {

return not_impl_err!("SORT BY");

}

// Capture and clear set expression schema so it doesn't leak

// into subqueries planned during FROM clause handling.

let set_expr_left_schema = planner_context.set_set_expr_left_schema(None);

// Process `from` clause

let plan = self.plan_from_tables(select.from, planner_context)?;

let empty_from = matches!(plan, LogicalPlan::EmptyRelation(_));

// Process `where` clause

let base_plan = self.plan_selection(select.selection, plan, planner_context)?;

// Handle named windows before processing the projection expression

check_conflicting_windows(&select.named_window)?;

self.match_window_definitions(&mut select.projection, &select.named_window)?;

// Process the SELECT expressions

let select_exprs = self.prepare_select_exprs(

&base_plan,

select.projection,

empty_from,

planner_context,

)?;

// Having and group by clause may reference aliases defined in select projection

let projected_plan =

self.project(base_plan.clone(), select_exprs, set_expr_left_schema)?;

let select_exprs = projected_plan.expressions();

let order_by =

to_order_by_exprs_with_select(query_order_by, Some(&select_exprs))?;

// Place the fields of the base plan at the front so that when there are references

// with the same name, the fields of the base plan will be searched first.

// See https://github.com/apache/datafusion/issues/9162

let mut combined_schema = base_plan.schema().as_ref().clone();

combined_schema.merge(projected_plan.schema());

// Order-by expressions prioritize referencing columns from the select list,

// then from the FROM clause.

let order_by_rex = self.order_by_to_sort_expr(

order_by,

projected_plan.schema().as_ref(),

planner_context,

true,

Some(base_plan.schema().as_ref()),

)?;

let order_by_rex = normalize_sorts(order_by_rex, &projected_plan)?;

// This alias map is resolved and looked up in both having exprs and group by exprs

let alias_map = extract_aliases(&select_exprs);

// DISTINCT ON expressions are parsed alongside HAVING / QUALIFY so

// they participate in aggregate / window discovery and get rebased

// through the same pipeline. The SQL nodes are taken out of

// `select.distinct` so the later match on `Distinct::On` still fires

// but does not move the original Vec.

//

// Resolution precedence matches PostgreSQL and ORDER BY: SELECT

// aliases win over input columns. For example,

// SELECT DISTINCT ON (b) a AS b ... GROUP BY a

// resolves `b` to the alias for `a`, not to a same-named input

// column.

let on_exprs_sql: Vec<SQLExpr> = match &mut select.distinct {

Some(Distinct::On(exprs)) => std::mem::take(exprs),

_ => Vec::new(),

};

let mut on_expr_schema = projected_plan.schema().as_ref().clone();

on_expr_schema.merge(base_plan.schema());

let on_exprs_pre_aggr: Vec<Expr> = on_exprs_sql

.into_iter()

.map(|e| {

let expr =

self.sql_expr_to_logical_expr(e, &on_expr_schema, planner_context)?;

// PostgreSQL only substitutes an output alias when the whole

// ON expression is a bare identifier. `b` resolves to the

// alias; `b + 0` keeps `b` as the input column.

let expr = substitute_top_level_alias(expr, &alias_map);

let expr = normalize_col(expr, &projected_plan)?;

let (expr, _) = expr.infer_placeholder_types(&on_expr_schema)?;

Ok(expr)

})

.collect::<Result<Vec<Expr>>>()?;

// Optionally the HAVING expression.

let having_expr_opt = select

.having

.map::<Result<Expr>, _>(|having_expr| {

self.warn_on_null_equality_predicate(&having_expr);

let window_span = self.window_function_span(&having_expr, &alias_map);

let having_expr = self.sql_expr_to_logical_expr(

having_expr,

&combined_schema,

planner_context,

)?;

// This step "dereferences" any aliases in the HAVING clause.

//

// This is how we support queries with HAVING expressions that

// refer to aliased columns.

//

// For example:

//

// SELECT c1, MAX(c2) AS m FROM t GROUP BY c1 HAVING m > 10;

//

// are rewritten as, respectively:

//

// SELECT c1, MAX(c2) AS m FROM t GROUP BY c1 HAVING MAX(c2) > 10;

//

let having_expr = resolve_aliases_to_exprs(having_expr, &alias_map)?;

// HAVING is evaluated before window functions are computed, so

// they may not appear there (checked after alias resolution so

// that an alias of a window function is rejected too)

reject_window_functions(

&having_expr,

"HAVING",

QUALIFY_HELP,

window_span,

)?;

let having_expr = normalize_col(having_expr, &projected_plan)?;

let (having_expr, _) =

having_expr.infer_placeholder_types(&combined_schema)?;

Ok(having_expr)

})

.transpose()?;

// All of the group by expressions

let group_by_exprs = if let GroupByExpr::Expressions(exprs, _) = select.group_by {

exprs

.into_iter()

.map(|e| {

// Aliases from the projection can conflict with same-named expressions in the input

let mut alias_map = alias_map.clone();

for f in base_plan.schema().fields() {

alias_map.remove(f.name());

}

let window_span = self.window_function_span(&e, &alias_map);

let group_by_expr = self.sql_expr_to_logical_expr(

e,

&combined_schema,

planner_context,

)?;

let group_by_expr =

resolve_aliases_to_exprs(group_by_expr, &alias_map)?;

let group_by_expr =

resolve_positions_to_exprs(group_by_expr, &select_exprs)?;

// Window functions are computed after grouping, so they may

// not be grouped on (checked after aliases and positions are

// resolved so `GROUP BY rn` / `GROUP BY 1` are rejected too)

reject_window_functions(

&group_by_expr,

"GROUP BY",

"Compute the window function in a subquery and group by its result",

window_span,

)?;

let group_by_expr = normalize_col(group_by_expr, &projected_plan)?;

self.validate_schema_satisfies_exprs(

base_plan.schema(),

std::slice::from_ref(&group_by_expr),

)?;

let (group_by_expr, _) =

group_by_expr.infer_placeholder_types(&combined_schema)?;

Ok(group_by_expr)

})

.collect::<Result<Vec<Expr>>>()?

} else {

// 'group by all' groups wrt. all select expressions except 'AggregateFunction's.

// Filter and collect non-aggregate select expressions.

select_exprs

.iter()

.filter(|select_expr| {

find_aggregate_exprs(std::iter::once(*select_expr)).is_empty()

})

.cloned()

.collect()

};

// Optionally the QUALIFY expression.

let qualify_expr_opt = select

.qualify

.map::<Result<Expr>, _>(|qualify_expr| {

let qualify_expr = self.sql_expr_to_logical_expr(

qualify_expr,

&combined_schema,

planner_context,

)?;

// This step "dereferences" any aliases in the QUALIFY clause.

//

// This is how we support queries with QUALIFY expressions that

// refer to aliased columns.

//

// For example:

//

// select row_number() over (PARTITION BY id) as rk from users qualify rk > 1;

//

// are rewritten as, respectively:

//

// select row_number() over (PARTITION BY id) as rk from users qualify row_number() over (PARTITION BY id) > 1;

//

let qualify_expr = resolve_aliases_to_exprs(qualify_expr, &alias_map)?;

let qualify_expr = normalize_col(qualify_expr, &projected_plan)?;

let (qualify_expr, _) =

qualify_expr.infer_placeholder_types(&combined_schema)?;

Ok(qualify_expr)

})

.transpose()?;

// The outer expressions we will search through for aggregates.

// First, find aggregates in SELECT, HAVING, and QUALIFY

let select_having_qualify_aggrs = find_aggregate_exprs(

select_exprs

.iter()

.chain(having_expr_opt.iter())

.chain(qualify_expr_opt.iter()),

);

// Find aggregates in ORDER BY

let order_by_aggrs = find_aggregate_exprs(order_by_rex.iter().map(|s| &s.expr));

// Find aggregates in DISTINCT ON

let on_aggrs = find_aggregate_exprs(on_exprs_pre_aggr.iter());

// Combine: all aggregates from SELECT/HAVING/QUALIFY, plus ORDER BY

// and DISTINCT ON aggregates that aren't already covered.

let mut aggr_exprs = select_having_qualify_aggrs;

for extra_aggr in order_by_aggrs.into_iter().chain(on_aggrs) {

if !aggr_exprs.iter().any(|e| e == &extra_aggr) {

aggr_exprs.push(extra_aggr);

}

}

// Process group by, aggregation or having

let AggregatePlanResult {

plan,

select_exprs: mut select_exprs_post_aggr,

having_expr: having_expr_post_aggr,

qualify_expr: qualify_expr_post_aggr,

order_by_exprs: mut order_by_rex,

on_exprs: mut on_exprs_post_aggr,

} = if !group_by_exprs.is_empty() || !aggr_exprs.is_empty() {

self.aggregate(

&base_plan,

&select_exprs,

having_expr_opt.as_ref(),

qualify_expr_opt.as_ref(),

&order_by_rex,

&on_exprs_pre_aggr,

&group_by_exprs,

&aggr_exprs,

)?

} else {

match having_expr_opt {

Some(having_expr) => {

return plan_err!(

"HAVING clause references: {having_expr} must appear in the GROUP BY clause or be used in an aggregate function"

);

}

None => AggregatePlanResult {

plan: base_plan.clone(),

select_exprs: select_exprs.clone(),

having_expr: having_expr_opt,

qualify_expr: qualify_expr_opt,

order_by_exprs: order_by_rex,

on_exprs: on_exprs_pre_aggr,

},

}

};

let plan = if let Some(having_expr_post_aggr) = having_expr_post_aggr {

LogicalPlanBuilder::from(plan)

.having(having_expr_post_aggr)?

.build()?

} else {

plan

};

// All of the window expressions (deduplicated and rewritten to reference aggregates as

// columns from input). Window functions may be sourced from the SELECT list, QUALIFY

// expression, ORDER BY, or DISTINCT ON.

let window_func_exprs = find_window_exprs(

select_exprs_post_aggr

.iter()

.chain(qualify_expr_post_aggr.iter())

.chain(order_by_rex.iter().map(|s| &s.expr))

.chain(on_exprs_post_aggr.iter()),

);

// Process window functions after aggregation as they can reference

// aggregate functions in their body

let plan = if window_func_exprs.is_empty() {

plan

} else {

let plan = LogicalPlanBuilder::window_plan(plan, window_func_exprs.clone())?;

// Re-write the projection

select_exprs_post_aggr = select_exprs_post_aggr

.iter()

.map(|expr| rebase_expr(expr, &window_func_exprs, &plan))

.collect::<Result<Vec<Expr>>>()?;

order_by_rex = order_by_rex

.into_iter()

.map(|sort_expr| {

Ok(sort_expr.with_expr(rebase_expr(

&sort_expr.expr,

&window_func_exprs,

&plan,

)?))

})

.collect::<Result<Vec<_>>>()?;

on_exprs_post_aggr = on_exprs_post_aggr

.iter()

.map(|expr| rebase_expr(expr, &window_func_exprs, &plan))

.collect::<Result<Vec<Expr>>>()?;

plan

};

// Process QUALIFY clause after window functions

// QUALIFY filters the results of window functions, similar to how HAVING filters aggregates

let plan = if let Some(qualify_expr) = qualify_expr_post_aggr {

// Validate that QUALIFY is used with window functions in SELECT or QUALIFY

let qualify_window_func_exprs = find_window_exprs(

select_exprs_post_aggr

.iter()

.chain(std::iter::once(&qualify_expr)),

);

if qualify_window_func_exprs.is_empty() {

return plan_err!(

"QUALIFY clause requires window functions in the SELECT list or QUALIFY clause"

);

}

// now attempt to resolve columns and replace with fully-qualified columns

let windows_projection_exprs = window_func_exprs

.iter()

.map(|expr| resolve_columns(expr, &plan))

.collect::<Result<Vec<Expr>>>()?;

// Rewrite the qualify expression to reference columns from the window plan

let qualify_expr_post_window =

rebase_expr(&qualify_expr, &windows_projection_exprs, &plan)?;

// Validate that the qualify expression can be resolved from the window plan schema

self.validate_schema_satisfies_exprs(

plan.schema(),

std::slice::from_ref(&qualify_expr_post_window),

)?;

LogicalPlanBuilder::from(plan)

.filter(qualify_expr_post_window)?

.build()?

} else {

plan

};

// Process distinct clause. For `DISTINCT ON` combined with

// aggregation, GROUP BY, or window functions we apply DistinctOn

// *before* the final projection so grouping columns and ORDER BY

// tie-breakers that aren't in the user SELECT stay in scope.

// DistinctOn provides the projection in that case (its select_expr

// list is wrapped in FIRST_VALUE during lowering).

let plan = match select.distinct {

None | Some(Distinct::All) => {

self.try_process_unnest(plan, select_exprs_post_aggr)?

}

Some(Distinct::Distinct) => {

let plan = self.try_process_unnest(plan, select_exprs_post_aggr)?;

LogicalPlanBuilder::from(plan).distinct()?.build()?

}

Some(Distinct::On(_)) => {

if aggr_exprs.is_empty()

&& group_by_exprs.is_empty()

&& window_func_exprs.is_empty()

{

// Fast path: no aggregation context. Fuse projection

// and deduplication into a single DistinctOn over

// `base_plan`. The sort attached to DistinctOn via

// `with_sort_expr` later normalizes against base_plan,

// so a bare ORDER BY alias (e.g. `ORDER BY x` where

// SELECT has `a AS x`) must be swapped back to the

// underlying input expression first.

order_by_rex =

substitute_top_level_aliases_in_sorts(order_by_rex, &alias_map);

LogicalPlanBuilder::from(base_plan)

.distinct_on(on_exprs_post_aggr, select_exprs, None)?

.build()?

} else {

// General path: DistinctOn layered over the post-

// aggregate / post-window plan (no extra Projection

// node — DistinctOn's lowering wraps each select_expr

// in FIRST_VALUE, which acts as the projection).

//

// The DistinctOn input has the post-aggregate raw

// column names (e.g. `max(t.c4)`), not the user-facing

// SELECT aliases (`agg2`). ORDER BY may reference

// those aliases — substitute them back to the

// underlying post-aggregate expression so they

// resolve against the DistinctOn input.

let select_alias_map = extract_aliases(&select_exprs_post_aggr);

order_by_rex = substitute_top_level_aliases_in_sorts(

order_by_rex,

&select_alias_map,

);

let DistinctOnUnnestPlanResult {

plan,

select_exprs: select_exprs_post_aggr,

on_exprs: on_exprs_post_aggr,

order_by_exprs: rewritten_order_by_rex,

} = self.try_process_distinct_on_unnest(

plan,

select_exprs_post_aggr,

on_exprs_post_aggr,

order_by_rex,

)?;

order_by_rex = rewritten_order_by_rex;

LogicalPlanBuilder::from(plan)

.distinct_on(on_exprs_post_aggr, select_exprs_post_aggr, None)?

.build()?

}

}

};

// DISTRIBUTE BY

let plan = if !select.distribute_by.is_empty() {

let x = select

.distribute_by

.iter()

.map(|e| {

self.sql_expr_to_logical_expr(

e.clone(),

&combined_schema,

planner_context,

)

})

.collect::<Result<Vec<_>>>()?;

LogicalPlanBuilder::from(plan)

.repartition(Partitioning::DistributeBy(x))?

.build()?

} else {

plan

};

let plan = self.order_by(plan, order_by_rex)?;

Ok(plan)

}

/// Try converting Expr(Unnest(Expr)) to Projection/Unnest/Projection

pub(super) fn try_process_unnest(

&self,

input: LogicalPlan,

select_exprs: Vec<Expr>,

) -> Result<LogicalPlan> {

let RewrittenUnnestExprGroups { plan, expr_groups } = self

.rewrite_unnest_expr_groups(

input,

select_exprs.into_iter().map(|expr| vec![expr]).collect(),

)?;

LogicalPlanBuilder::from(plan)

.project(flatten_expr_groups(expr_groups))?

.build()

}

/// Rewrites SELECT-list UNNESTs while keeping hidden DISTINCT ON / ORDER

/// BY inputs available to the DistinctOn node.

fn try_process_distinct_on_unnest(

&self,

input: LogicalPlan,

select_exprs: Vec<Expr>,

on_exprs: Vec<Expr>,

order_by_exprs: Vec<SortExpr>,

) -> Result<DistinctOnUnnestPlanResult> {

let select_len = select_exprs.len();

let on_len = on_exprs.len();

let mut expr_groups = select_exprs

.into_iter()

.map(|expr| vec![expr])

.collect::<Vec<_>>();

expr_groups.extend(on_exprs.into_iter().map(|expr| vec![expr]));

expr_groups.extend(

order_by_exprs

.iter()

.map(|sort_expr| vec![sort_expr.expr.clone()]),

);

let RewrittenUnnestExprGroups {

plan,

mut expr_groups,

} = self.rewrite_unnest_expr_groups(input, expr_groups)?;

let rewritten_select_exprs =

flatten_expr_groups(expr_groups.drain(..select_len).collect());

let rewritten_on_exprs = expr_groups

.drain(..on_len)

.map(|exprs| self.expect_single_distinct_on_expr(exprs, "DISTINCT ON"))

.collect::<Result<Vec<_>>>()?;

let rewritten_order_by_exprs = order_by_exprs

.into_iter()

.zip(expr_groups)

.map(|(sort_expr, exprs)| {

Ok(sort_expr

.with_expr(self.expect_single_distinct_on_expr(exprs, "ORDER BY")?))

})

.collect::<Result<Vec<_>>>()?;

Ok(DistinctOnUnnestPlanResult {

plan,

select_exprs: rewritten_select_exprs,

on_exprs: rewritten_on_exprs,

order_by_exprs: rewritten_order_by_exprs,

})

}

fn expect_single_distinct_on_expr(

&self,

exprs: Vec<Expr>,

clause: &str,

) -> Result<Expr> {

if exprs.len() == 1 {

return Ok(exprs.into_iter().next().expect("len checked above"));

}

not_impl_err!(

"{clause} expressions that expand to multiple columns are not supported with DISTINCT ON"

)

}

fn rewrite_unnest_expr_groups(

&self,

input: LogicalPlan,

expr_groups: Vec<Vec<Expr>>,

) -> Result<RewrittenUnnestExprGroups> {

// Try process group by unnest

let input = self.try_process_aggregate_unnest(input)?;

let mut intermediate_plan = input;

let mut intermediate_expr_groups = expr_groups;

loop {

let mut unnest_columns = IndexMap::new();

// from which columns used for projection, before the unnest happen

// including non unnest columns and unnest columns

let mut inner_projection_exprs = DedupedProjection::default();

let mut outer_expr_groups =

Vec::with_capacity(intermediate_expr_groups.len());

for expr_group in &intermediate_expr_groups {

let mut outer_expr_group = vec![];

for expr in expr_group {

let mut rewritten_exprs = rewrite_recursive_unnest_bottom_up(

&intermediate_plan,

&mut unnest_columns,

&mut inner_projection_exprs,

expr,

)?;

if let Some(columns) =

self.get_struct_unnest_columns(&intermediate_plan, expr)?

{

rewritten_exprs = rewritten_exprs

.into_iter()

.zip(columns)

.map(|(expr, column)| expr.alias(column.flat_name()))

.collect();

}

outer_expr_group.extend(rewritten_exprs);

}

outer_expr_groups.push(outer_expr_group);

}

// No more unnest is possible

if unnest_columns.is_empty() {

return Ok(RewrittenUnnestExprGroups {

plan: intermediate_plan,

expr_groups: intermediate_expr_groups,

});

}

// The default SQL `UNNEST` matches DuckDB/PostgreSQL: drop both

// NULL and empty input lists. Outer-unnest (modelled as

// `Unnest { outer: true }`) overrides that and selects

// `NullHandling::PreserveAndExpandEmpty`. Mixing the two in a

// single SELECT is a planning error because `UnnestOptions` is

// per-`UnnestExec`, not per-column.

let null_handling = collect_unnest_null_handling(&intermediate_expr_groups)?;

let mut unnest_options =

UnnestOptions::new().with_null_handling(null_handling);

let mut unnest_col_vec = Vec::with_capacity(unnest_columns.len());

for (col, maybe_list_unnest) in unnest_columns.into_iter() {

if let Some(list_unnest) = maybe_list_unnest {

unnest_options = list_unnest.into_iter().fold(

unnest_options,

|options, unnest_list| {

options.with_recursions(RecursionUnnestOption {

input_column: col.clone(),

output_column: unnest_list.output_column,

depth: unnest_list.depth,

})

},

);

}

unnest_col_vec.push(col);

}

intermediate_plan = LogicalPlanBuilder::from(intermediate_plan)

.project(inner_projection_exprs.into_exprs())?

.unnest_columns_with_options(unnest_col_vec, unnest_options)?

.build()?;

intermediate_expr_groups = outer_expr_groups;

}

}

fn get_struct_unnest_columns(

&self,

input: &LogicalPlan,

expr: &Expr,

) -> Result<Option<Vec<Column>>> {

let unnest_expr = match expr {

Expr::Unnest(unnest_expr) => Some(unnest_expr),

Expr::Alias(alias) => match alias.expr.as_ref() {

Expr::Unnest(unnest_expr) => Some(unnest_expr),

_ => None,

},

_ => None,

};

let Some(unnest_expr) = unnest_expr else {

return Ok(None);

};

let field = unnest_expr.expr.to_field(input.schema())?.1;

let DataType::Struct(inner_fields) = field.data_type() else {

return Ok(None);

};

Ok(Some(get_struct_unnested_columns(

&unnest_expr.expr.schema_name().to_string(),

inner_fields,

)))

}

fn try_process_aggregate_unnest(&self, input: LogicalPlan) -> Result<LogicalPlan> {

match input {

// Fast path if there are no unnest in group by

LogicalPlan::Aggregate(ref agg)

if !&agg.group_expr.iter().any(has_unnest_expr_recursively) =>

{

Ok(input)

}

LogicalPlan::Aggregate(agg) => {

let agg_expr = agg.aggr_expr.clone();

let (new_input, new_group_by_exprs) =

self.try_process_group_by_unnest(agg)?;

let options = LogicalPlanBuilderOptions::new()

.with_add_implicit_group_by_exprs(true);

LogicalPlanBuilder::from(new_input)

.with_options(options)

.aggregate(new_group_by_exprs, agg_expr)?

.build()

}

LogicalPlan::Filter(mut filter) => {

filter.input =

Arc::new(self.try_process_aggregate_unnest(Arc::unwrap_or_clone(

filter.input,

))?);

Ok(LogicalPlan::Filter(filter))

}

_ => Ok(input),

}

}

/// Try converting Unnest(Expr) of group by to Unnest/Projection.

/// Return the new input and group_by_exprs of Aggregate.

/// Select exprs can be different from agg exprs, for example:

fn try_process_group_by_unnest(

&self,

agg: Aggregate,

) -> Result<(LogicalPlan, Vec<Expr>)> {

let mut aggr_expr_using_columns: Option<HashSet<Expr>> = None;

let Aggregate {

input,

group_expr,

aggr_expr,

..

} = agg;

// Process unnest of group_by_exprs, and input of agg will be rewritten

// for example:

//

// ```

// Aggregate: groupBy=[[UNNEST(Column(Column { relation: Some(Bare { table: "tab" }), name: "array_col" }))]], aggr=[[]]

// TableScan: tab

// ```

//

// will be transformed into

//

// ```

// Aggregate: groupBy=[[unnest(tab.array_col)]], aggr=[[]]

// Unnest: lists[unnest(tab.array_col)] structs[]

// Projection: tab.array_col AS unnest(tab.array_col)

// TableScan: tab

// ```

let mut intermediate_plan = Arc::unwrap_or_clone(input);

let mut intermediate_select_exprs = group_expr;

loop {

let mut unnest_columns = IndexMap::new();

let mut inner_projection_exprs = DedupedProjection::default();

let outer_projection_exprs = rewrite_recursive_unnests_bottom_up(

&intermediate_plan,

&mut unnest_columns,

&mut inner_projection_exprs,

&intermediate_select_exprs,

)?;

if unnest_columns.is_empty() {

break;

} else {

let mut unnest_options = UnnestOptions::new().with_preserve_nulls(false);

#[allow(clippy::allow_attributes, clippy::mutable_key_type)]

// Expr contains Arc with interior mutability but is intentionally used as hash key

let mut projection_exprs = match &aggr_expr_using_columns {

Some(exprs) => (*exprs).clone(),

None => {

#[allow(clippy::allow_attributes, clippy::mutable_key_type)]

let mut columns = HashSet::new();

for expr in &aggr_expr {

expr.apply(|expr| {

if let Expr::Column(c) = expr {

columns.insert(Expr::Column(c.clone()));

}

Ok(TreeNodeRecursion::Continue)

})

// As the closure always returns Ok, this "can't" error

.expect("Unexpected error");

}

aggr_expr_using_columns = Some(columns.clone());

columns

}

};

projection_exprs.extend(inner_projection_exprs.into_exprs());

let mut unnest_col_vec = Vec::with_capacity(unnest_columns.len());

for (col, maybe_list_unnest) in unnest_columns.into_iter() {

if let Some(list_unnest) = maybe_list_unnest {

unnest_options = list_unnest.into_iter().fold(

unnest_options,

|options, unnest_list| {

options.with_recursions(RecursionUnnestOption {

input_column: col.clone(),

output_column: unnest_list.output_column,

depth: unnest_list.depth,

})

},

);

}

unnest_col_vec.push(col);

}

intermediate_plan = LogicalPlanBuilder::from(intermediate_plan)

.project(projection_exprs)?

.unnest_columns_with_options(unnest_col_vec, unnest_options)?

.build()?;

intermediate_select_exprs = outer_projection_exprs;

}

}

Ok((intermediate_plan, intermediate_select_exprs))

}

pub(crate) fn plan_selection(

&self,

selection: Option<SQLExpr>,

plan: LogicalPlan,

planner_context: &mut PlannerContext,

) -> Result<LogicalPlan> {

match selection {

Some(predicate_expr) => {

let fallback_schemas = plan.fallback_normalize_schemas();

self.warn_on_null_equality_predicate(&predicate_expr);

let window_span =

self.window_function_span(&predicate_expr, &HashMap::new());

let filter_expr =

self.sql_to_expr(predicate_expr, plan.schema(), planner_context)?;

// Check for aggregation functions

let aggregate_exprs =

find_aggregate_exprs(std::slice::from_ref(&filter_expr));

if !aggregate_exprs.is_empty() {

return plan_err!(

"Aggregate functions are not allowed in the WHERE clause. Consider using HAVING instead"

);

}

// WHERE is evaluated before window functions are computed, so

// they may not appear there either

reject_window_functions(

&filter_expr,

"WHERE",

QUALIFY_HELP,

window_span,

)?;

let mut using_columns = HashSet::new();

expr_to_columns(&filter_expr, &mut using_columns)?;

let mut schema_stack: Vec<Vec<&DFSchema>> =

vec![vec![plan.schema()], fallback_schemas];

for sc in planner_context.outer_schemas_iter() {

schema_stack.push(vec![sc.as_ref()]);

}

let filter_expr = normalize_col_with_schemas_and_ambiguity_check(

filter_expr,

schema_stack

.iter()

.map(|sc| sc.as_slice())

.collect::<Vec<&[&DFSchema]>>()

.as_slice(),

&[using_columns],

)?;

Ok(LogicalPlan::Filter(Filter::try_new(

filter_expr,

Arc::new(plan),

)?))

}

None => Ok(plan),

}

}

pub(crate) fn plan_from_tables(

&self,

mut from: Vec<TableWithJoins>,

planner_context: &mut PlannerContext,

) -> Result<LogicalPlan> {

match from.len() {

0 => Ok(LogicalPlanBuilder::empty(true).build()?),

1 => {

let input = from.remove(0);

self.plan_table_with_joins(input, planner_context)

}

_ => {

let mut from = from.into_iter();

let mut left = LogicalPlanBuilder::from({

let input = from.next().unwrap();

self.plan_table_with_joins(input, planner_context)?

});

let old_outer_from_schema = {

let left_schema = Some(Arc::clone(left.schema()));

planner_context.set_outer_from_schema(left_schema)

};

for input in from {

// Join `input` with the current result (`left`).

let right = self.plan_table_with_joins(input, planner_context)?;

left = left.cross_join(right)?;

// Update the outer FROM schema.

let left_schema = Some(Arc::clone(left.schema()));

planner_context.set_outer_from_schema(left_schema);

}

planner_context.set_outer_from_schema(old_outer_from_schema);

left.build()

}

}

}

/// Returns the `Expr`'s corresponding to a SQL query's SELECT expressions.

pub(crate) fn prepare_select_exprs(

&self,

plan: &LogicalPlan,

projection: Vec<SelectItem>,