datafusion_sql/
utils.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! SQL Utility Functions
19
20use std::vec;
21
22use arrow::datatypes::{
23    DECIMAL_DEFAULT_SCALE, DECIMAL128_MAX_PRECISION, DECIMAL256_MAX_PRECISION, DataType,
24};
25use datafusion_common::tree_node::{
26    Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter,
27};
28use datafusion_common::{
29    Column, DFSchemaRef, Diagnostic, HashMap, Result, ScalarValue,
30    assert_or_internal_err, exec_datafusion_err, exec_err, internal_err, plan_err,
31};
32use datafusion_expr::builder::get_struct_unnested_columns;
33use datafusion_expr::expr::{
34    Alias, GroupingSet, Unnest, WindowFunction, WindowFunctionParams,
35};
36use datafusion_expr::utils::{expr_as_column_expr, find_column_exprs};
37use datafusion_expr::{
38    ColumnUnnestList, Expr, ExprSchemable, LogicalPlan, SortExpr, col, expr_vec_fmt,
39};
40
41use indexmap::IndexMap;
42use sqlparser::ast::{Ident, Value};
43
44/// Make a best-effort attempt at resolving all columns in the expression tree
45pub(crate) fn resolve_columns(expr: &Expr, plan: &LogicalPlan) -> Result<Expr> {
46    expr.clone()
47        .transform_up(|nested_expr| {
48            match nested_expr {
49                Expr::Column(col) => {
50                    let (qualifier, field) =
51                        plan.schema().qualified_field_from_column(&col)?;
52                    Ok(Transformed::yes(Expr::Column(Column::from((
53                        qualifier, field,
54                    )))))
55                }
56                _ => {
57                    // keep recursing
58                    Ok(Transformed::no(nested_expr))
59                }
60            }
61        })
62        .data()
63}
64
65/// Rebuilds an `Expr` as a projection on top of a collection of `Expr`'s.
66///
67/// For example, the expression `a + b < 1` would require, as input, the 2
68/// individual columns, `a` and `b`. But, if the base expressions already
69/// contain the `a + b` result, then that may be used in lieu of the `a` and
70/// `b` columns.
71///
72/// This is useful in the context of a query like:
73///
74/// SELECT a + b < 1 ... GROUP BY a + b
75///
76/// where post-aggregation, `a + b` need not be a projection against the
77/// individual columns `a` and `b`, but rather it is a projection against the
78/// `a + b` found in the GROUP BY.
79pub(crate) fn rebase_expr(
80    expr: &Expr,
81    base_exprs: &[Expr],
82    plan: &LogicalPlan,
83) -> Result<Expr> {
84    expr.clone()
85        .transform_down(|nested_expr| {
86            if base_exprs.contains(&nested_expr) {
87                Ok(Transformed::yes(expr_as_column_expr(&nested_expr, plan)?))
88            } else {
89                Ok(Transformed::no(nested_expr))
90            }
91        })
92        .data()
93}
94
95#[derive(Debug, Clone, Copy, PartialEq, Eq)]
96pub(crate) enum CheckColumnsMustReferenceAggregatePurpose {
97    Projection,
98    Having,
99    Qualify,
100    OrderBy,
101    DistinctOn,
102}
103
104#[derive(Debug, Clone, Copy, PartialEq, Eq)]
105pub(crate) enum CheckColumnsSatisfyExprsPurpose {
106    Aggregate(CheckColumnsMustReferenceAggregatePurpose),
107}
108
109impl CheckColumnsSatisfyExprsPurpose {
110    fn message_prefix(&self) -> &'static str {
111        match self {
112            Self::Aggregate(CheckColumnsMustReferenceAggregatePurpose::Projection) => {
113                "Column in SELECT must be in GROUP BY or an aggregate function"
114            }
115            Self::Aggregate(CheckColumnsMustReferenceAggregatePurpose::Having) => {
116                "Column in HAVING must be in GROUP BY or an aggregate function"
117            }
118            Self::Aggregate(CheckColumnsMustReferenceAggregatePurpose::Qualify) => {
119                "Column in QUALIFY must be in GROUP BY or an aggregate function"
120            }
121            Self::Aggregate(CheckColumnsMustReferenceAggregatePurpose::OrderBy) => {
122                "Column in ORDER BY must be in GROUP BY or an aggregate function"
123            }
124            Self::Aggregate(CheckColumnsMustReferenceAggregatePurpose::DistinctOn) => {
125                "Column in DISTINCT ON must be in GROUP BY or an aggregate function"
126            }
127        }
128    }
129
130    fn diagnostic_message(&self, expr: &Expr) -> String {
131        format!(
132            "'{expr}' must appear in GROUP BY clause because it's not an aggregate expression"
133        )
134    }
135}
136
137/// Determines if the set of `Expr`'s are a valid projection on the input
138/// `Expr::Column`'s.
139pub(crate) fn check_columns_satisfy_exprs(
140    columns: &[Expr],
141    exprs: &[Expr],
142    purpose: CheckColumnsSatisfyExprsPurpose,
143) -> Result<()> {
144    columns.iter().try_for_each(|c| match c {
145        Expr::Column(_) => Ok(()),
146        _ => internal_err!("Expr::Column are required"),
147    })?;
148    let column_exprs = find_column_exprs(exprs);
149    for e in &column_exprs {
150        match e {
151            Expr::GroupingSet(GroupingSet::Rollup(exprs)) => {
152                for e in exprs {
153                    check_column_satisfies_expr(columns, e, purpose)?;
154                }
155            }
156            Expr::GroupingSet(GroupingSet::Cube(exprs)) => {
157                for e in exprs {
158                    check_column_satisfies_expr(columns, e, purpose)?;
159                }
160            }
161            Expr::GroupingSet(GroupingSet::GroupingSets(lists_of_exprs)) => {
162                for exprs in lists_of_exprs {
163                    for e in exprs {
164                        check_column_satisfies_expr(columns, e, purpose)?;
165                    }
166                }
167            }
168            _ => check_column_satisfies_expr(columns, e, purpose)?,
169        }
170    }
171    Ok(())
172}
173
174fn check_column_satisfies_expr(
175    columns: &[Expr],
176    expr: &Expr,
177    purpose: CheckColumnsSatisfyExprsPurpose,
178) -> Result<()> {
179    if !columns.contains(expr) {
180        let diagnostic = Diagnostic::new_error(
181            purpose.diagnostic_message(expr),
182            expr.spans().and_then(|spans| spans.first()),
183        )
184        .with_help(format!("Either add '{expr}' to GROUP BY clause, or use an aggregate function like ANY_VALUE({expr})"), None);
185
186        return plan_err!(
187            "{}: While expanding wildcard, column \"{}\" must appear in the GROUP BY clause or must be part of an aggregate function, currently only \"{}\" appears in the SELECT clause satisfies this requirement",
188            purpose.message_prefix(),
189            expr,
190            expr_vec_fmt!(columns);
191            diagnostic=diagnostic
192        );
193    }
194    Ok(())
195}
196
197/// Returns mapping of each alias (`String`) to the expression (`Expr`) it is
198/// aliasing.
199pub(crate) fn extract_aliases(exprs: &[Expr]) -> HashMap<String, Expr> {
200    exprs
201        .iter()
202        .filter_map(|expr| match expr {
203            Expr::Alias(Alias { expr, name, .. }) => Some((name.clone(), *expr.clone())),
204            _ => None,
205        })
206        .collect::<HashMap<String, Expr>>()
207}
208
209/// If `expr` is a bare unqualified `Column` whose name matches a SELECT
210/// alias, swap it for the alias's underlying expression. Nested occurrences
211/// are left alone: PostgreSQL only resolves a top-level identifier as an
212/// output alias in clauses like ORDER BY and DISTINCT ON.
213pub(crate) fn substitute_top_level_alias(
214    expr: Expr,
215    aliases: &HashMap<String, Expr>,
216) -> Expr {
217    if let Expr::Column(col) = &expr
218        && col.relation.is_none()
219        && let Some(underlying) = aliases.get(&col.name)
220    {
221        return underlying.clone();
222    }
223
224    expr
225}
226
227/// Applies [`substitute_top_level_alias`] to each sort expression.
228pub(crate) fn substitute_top_level_aliases_in_sorts(
229    sort_exprs: Vec<SortExpr>,
230    aliases: &HashMap<String, Expr>,
231) -> Vec<SortExpr> {
232    if aliases.is_empty() {
233        return sort_exprs;
234    }
235
236    sort_exprs
237        .into_iter()
238        .map(|sort_expr| {
239            sort_expr
240                .with_expr(substitute_top_level_alias(sort_expr.expr.clone(), aliases))
241        })
242        .collect()
243}
244
245/// Given an expression that's literal int encoding position, lookup the corresponding expression
246/// in the select_exprs list, if the index is within the bounds and it is indeed a position literal,
247/// otherwise, returns planning error.
248/// If input expression is not an int literal, returns expression as-is.
249pub(crate) fn resolve_positions_to_exprs(
250    expr: Expr,
251    select_exprs: &[Expr],
252) -> Result<Expr> {
253    match expr {
254        // sql_expr_to_logical_expr maps number to i64
255        // https://github.com/apache/datafusion/blob/8d175c759e17190980f270b5894348dc4cff9bbf/datafusion/src/sql/planner.rs#L882-L887
256        Expr::Literal(ScalarValue::Int64(Some(position)), _)
257            if position > 0_i64 && position <= select_exprs.len() as i64 =>
258        {
259            let index = (position - 1) as usize;
260            let select_expr = &select_exprs[index];
261            Ok(match select_expr {
262                Expr::Alias(Alias { expr, .. }) => *expr.clone(),
263                _ => select_expr.clone(),
264            })
265        }
266        Expr::Literal(ScalarValue::Int64(Some(position)), _) => plan_err!(
267            "Cannot find column with position {} in SELECT clause. Valid columns: 1 to {}",
268            position,
269            select_exprs.len()
270        ),
271        _ => Ok(expr),
272    }
273}
274
275/// Rebuilds an `Expr` with columns that refer to aliases replaced by the
276/// alias' underlying `Expr`.
277pub(crate) fn resolve_aliases_to_exprs(
278    expr: Expr,
279    aliases: &HashMap<String, Expr>,
280) -> Result<Expr> {
281    expr.transform_up(|nested_expr| match nested_expr {
282        Expr::Column(c) if c.relation.is_none() => {
283            if let Some(aliased_expr) = aliases.get(&c.name) {
284                Ok(Transformed::yes(aliased_expr.clone()))
285            } else {
286                Ok(Transformed::no(Expr::Column(c)))
287            }
288        }
289        _ => Ok(Transformed::no(nested_expr)),
290    })
291    .data()
292}
293
294/// Given a slice of window expressions sharing the same sort key, find their common partition
295/// keys.
296pub fn window_expr_common_partition_keys(window_exprs: &[Expr]) -> Result<&[Expr]> {
297    let all_partition_keys = window_exprs
298        .iter()
299        .map(|expr| match expr {
300            Expr::WindowFunction(window_fun) => {
301                let WindowFunction {
302                    params: WindowFunctionParams { partition_by, .. },
303                    ..
304                } = window_fun.as_ref();
305                Ok(partition_by)
306            }
307            Expr::Alias(Alias { expr, .. }) => match expr.as_ref() {
308                Expr::WindowFunction(window_fun) => {
309                    let WindowFunction {
310                        params: WindowFunctionParams { partition_by, .. },
311                        ..
312                    } = window_fun.as_ref();
313                    Ok(partition_by)
314                }
315                expr => exec_err!("Impossibly got non-window expr {expr:?}"),
316            },
317            expr => exec_err!("Impossibly got non-window expr {expr:?}"),
318        })
319        .collect::<Result<Vec<_>>>()?;
320    let result = all_partition_keys
321        .iter()
322        .min_by_key(|s| s.len())
323        .ok_or_else(|| exec_datafusion_err!("No window expressions found"))?;
324    Ok(result)
325}
326
327/// Returns a validated `DataType` for the specified precision and
328/// scale
329pub(crate) fn make_decimal_type(
330    precision: Option<u64>,
331    scale: Option<u64>,
332) -> Result<DataType> {
333    // postgres like behavior
334    let (precision, scale) = match (precision, scale) {
335        (Some(p), Some(s)) => (p as u8, s as i8),
336        (Some(p), None) => (p as u8, 0),
337        (None, Some(_)) => {
338            return plan_err!("Cannot specify only scale for decimal data type");
339        }
340        (None, None) => (DECIMAL128_MAX_PRECISION, DECIMAL_DEFAULT_SCALE),
341    };
342
343    if precision == 0
344        || precision > DECIMAL256_MAX_PRECISION
345        || scale.unsigned_abs() > precision
346    {
347        plan_err!(
348            "Decimal(precision = {precision}, scale = {scale}) should satisfy `0 < precision <= 76`, and `scale <= precision`."
349        )
350    } else if precision > DECIMAL128_MAX_PRECISION
351        && precision <= DECIMAL256_MAX_PRECISION
352    {
353        Ok(DataType::Decimal256(precision, scale))
354    } else {
355        Ok(DataType::Decimal128(precision, scale))
356    }
357}
358
359/// Normalize an owned identifier to a lowercase string, unless the identifier is quoted.
360pub(crate) fn normalize_ident(id: Ident) -> String {
361    match id.quote_style {
362        Some(_) => id.value,
363        None => id.value.to_ascii_lowercase(),
364    }
365}
366
367pub(crate) fn value_to_string(value: &Value) -> Option<String> {
368    match value {
369        Value::SingleQuotedString(s) => Some(s.to_string()),
370        Value::DollarQuotedString(s) => Some(s.to_string()),
371        Value::Number(_, _) | Value::Boolean(_) => Some(value.to_string()),
372        Value::UnicodeStringLiteral(s) => Some(s.to_string()),
373        Value::EscapedStringLiteral(s) => Some(s.to_string()),
374        Value::QuoteDelimitedStringLiteral(s)
375        | Value::NationalQuoteDelimitedStringLiteral(s) => Some(s.value.to_string()),
376        Value::DoubleQuotedString(_)
377        | Value::NationalStringLiteral(_)
378        | Value::SingleQuotedByteStringLiteral(_)
379        | Value::DoubleQuotedByteStringLiteral(_)
380        | Value::TripleSingleQuotedString(_)
381        | Value::TripleDoubleQuotedString(_)
382        | Value::TripleSingleQuotedByteStringLiteral(_)
383        | Value::TripleDoubleQuotedByteStringLiteral(_)
384        | Value::SingleQuotedRawStringLiteral(_)
385        | Value::DoubleQuotedRawStringLiteral(_)
386        | Value::TripleSingleQuotedRawStringLiteral(_)
387        | Value::TripleDoubleQuotedRawStringLiteral(_)
388        | Value::HexStringLiteral(_)
389        | Value::Null
390        | Value::Placeholder(_) => None,
391    }
392}
393
394pub(crate) fn rewrite_recursive_unnests_bottom_up(
395    input: &LogicalPlan,
396    unnest_placeholder_columns: &mut IndexMap<Column, Option<Vec<ColumnUnnestList>>>,
397    inner_projection_exprs: &mut Vec<Expr>,
398    original_exprs: &[Expr],
399) -> Result<Vec<Expr>> {
400    Ok(original_exprs
401        .iter()
402        .map(|expr| {
403            rewrite_recursive_unnest_bottom_up(
404                input,
405                unnest_placeholder_columns,
406                inner_projection_exprs,
407                expr,
408            )
409        })
410        .collect::<Result<Vec<_>>>()?
411        .into_iter()
412        .flatten()
413        .collect::<Vec<_>>())
414}
415
416pub const UNNEST_PLACEHOLDER: &str = "__unnest_placeholder";
417
418/*
419This is only useful when used with transform down up
420A full example of how the transformation works:
421 */
422struct RecursiveUnnestRewriter<'a> {
423    input_schema: &'a DFSchemaRef,
424    root_expr: &'a Expr,
425    // Useful to detect which child expr is a part of/ not a part of unnest operation
426    top_most_unnest: Option<Unnest>,
427    consecutive_unnest: Vec<Option<Unnest>>,
428    inner_projection_exprs: &'a mut Vec<Expr>,
429    columns_unnestings: &'a mut IndexMap<Column, Option<Vec<ColumnUnnestList>>>,
430    transformed_root_exprs: Option<Vec<Expr>>,
431}
432impl RecursiveUnnestRewriter<'_> {
433    /// This struct stores the history of expr
434    /// during its tree-traversal with a notation of
435    /// \[None,**Unnest(exprA)**,**Unnest(exprB)**,None,None\]
436    /// then this function will returns \[**Unnest(exprA)**,**Unnest(exprB)**\]
437    ///
438    /// The first item will be the inner most expr
439    fn get_latest_consecutive_unnest(&self) -> Vec<Unnest> {
440        self.consecutive_unnest
441            .iter()
442            .rev()
443            .skip_while(|item| item.is_none())
444            .take_while(|item| item.is_some())
445            .to_owned()
446            .cloned()
447            .map(|item| item.unwrap())
448            .collect()
449    }
450
451    /// Check if the current expression is at the root level for struct unnest purposes.
452    /// This is true if:
453    /// 1. The expression IS the root expression, OR
454    /// 2. The root expression is an Alias wrapping this expression
455    ///
456    /// This allows `unnest(struct_col) AS alias` to work, where the alias is simply
457    /// ignored for struct unnest (matching DuckDB behavior).
458    fn is_at_struct_allowed_root(&self, expr: &Expr) -> bool {
459        if expr == self.root_expr {
460            return true;
461        }
462        // Allow struct unnest when root is an alias wrapping the unnest
463        if let Expr::Alias(Alias { expr: inner, .. }) = self.root_expr {
464            return inner.as_ref() == expr;
465        }
466        false
467    }
468
469    fn transform(
470        &mut self,
471        level: usize,
472        alias_name: String,
473        expr_in_unnest: &Expr,
474        struct_allowed: bool,
475    ) -> Result<Vec<Expr>> {
476        let inner_expr_name = expr_in_unnest.schema_name().to_string();
477
478        // Full context, we are trying to plan the execution as InnerProjection->Unnest->OuterProjection
479        // inside unnest execution, each column inside the inner projection
480        // will be transformed into new columns. Thus we need to keep track of these placeholding column names
481        let placeholder_name = format!("{UNNEST_PLACEHOLDER}({inner_expr_name})");
482        let post_unnest_name =
483            format!("{UNNEST_PLACEHOLDER}({inner_expr_name},depth={level})");
484        // This is due to the fact that unnest transformation should keep the original
485        // column name as is, to comply with group by and order by
486        let placeholder_column = Column::from_name(placeholder_name.clone());
487        let field = expr_in_unnest.to_field(self.input_schema)?.1;
488        let data_type = field.data_type();
489
490        match data_type {
491            DataType::Struct(inner_fields) => {
492                assert_or_internal_err!(
493                    struct_allowed,
494                    "unnest on struct can only be applied at the root level of select expression"
495                );
496                push_projection_dedupl(
497                    self.inner_projection_exprs,
498                    expr_in_unnest.clone().alias(placeholder_name.clone()),
499                );
500                self.columns_unnestings
501                    .insert(Column::from_name(placeholder_name.clone()), None);
502                Ok(get_struct_unnested_columns(&placeholder_name, inner_fields)
503                    .into_iter()
504                    .map(Expr::Column)
505                    .collect())
506            }
507            DataType::List(_)
508            | DataType::FixedSizeList(_, _)
509            | DataType::LargeList(_)
510            | DataType::ListView(_)
511            | DataType::LargeListView(_) => {
512                push_projection_dedupl(
513                    self.inner_projection_exprs,
514                    expr_in_unnest.clone().alias(placeholder_name.clone()),
515                );
516
517                let post_unnest_expr = col(post_unnest_name.clone()).alias(alias_name);
518                let list_unnesting = self
519                    .columns_unnestings
520                    .entry(placeholder_column)
521                    .or_insert(Some(vec![]));
522                let unnesting = ColumnUnnestList {
523                    output_column: Column::from_name(post_unnest_name),
524                    depth: level,
525                };
526                let list_unnestings = list_unnesting.as_mut().unwrap();
527                if !list_unnestings.contains(&unnesting) {
528                    list_unnestings.push(unnesting);
529                }
530                Ok(vec![post_unnest_expr])
531            }
532            _ => {
533                internal_err!("unnest on non-list or struct type is not supported")
534            }
535        }
536    }
537}
538
539impl TreeNodeRewriter for RecursiveUnnestRewriter<'_> {
540    type Node = Expr;
541
542    /// This downward traversal needs to keep track of:
543    /// - Whether or not some unnest expr has been visited from the top until the current node
544    /// - If some unnest expr has been visited, maintain a stack of such information, this
545    ///   is used to detect if some recursive unnest expr exists (e.g **unnest(unnest(unnest(3d column))))**
546    fn f_down(&mut self, expr: Expr) -> Result<Transformed<Expr>> {
547        if let Expr::Unnest(ref unnest_expr) = expr {
548            let field = unnest_expr.expr.to_field(self.input_schema)?.1;
549            let data_type = field.data_type();
550            self.consecutive_unnest.push(Some(unnest_expr.clone()));
551            // if expr inside unnest is a struct, do not consider
552            // the next unnest as consecutive unnest (if any)
553            // meaning unnest(unnest(struct_arr_col)) can't
554            // be interpreted as unnest(struct_arr_col, depth:=2)
555            // but has to be split into multiple unnest logical plan instead
556            // a.k.a:
557            // - unnest(struct_col)
558            //      unnest(struct_arr_col) as struct_col
559
560            if let DataType::Struct(_) = data_type {
561                self.consecutive_unnest.push(None);
562            }
563            if self.top_most_unnest.is_none() {
564                self.top_most_unnest = Some(unnest_expr.clone());
565            }
566
567            Ok(Transformed::no(expr))
568        } else {
569            self.consecutive_unnest.push(None);
570            Ok(Transformed::no(expr))
571        }
572    }
573
574    /// The rewriting only happens when the traversal has reached the top-most unnest expr
575    /// within a sequence of consecutive unnest exprs node
576    ///
577    /// For example an expr of **unnest(unnest(column1)) + unnest(unnest(unnest(column2)))**
578    /// ```text
579    ///                         ┌──────────────────┐
580    ///                         │    binaryexpr    │
581    ///                         │                  │
582    ///                         └──────────────────┘
583    ///                f_down  / /            │ │
584    ///                       / / f_up        │ │
585    ///                      / /        f_down│ │f_up
586    ///                  unnest               │ │
587    ///                                       │ │
588    ///       f_down  / / f_up(rewriting)     │ │
589    ///              / /
590    ///             / /                      unnest
591    ///         unnest
592    ///                           f_down  / / f_up(rewriting)
593    /// f_down / /f_up                   / /
594    ///       / /                       / /
595    ///      / /                    unnest
596    ///   column1
597    ///                     f_down / /f_up
598    ///                           / /
599    ///                          / /
600    ///                       column2
601    /// ```
602    fn f_up(&mut self, expr: Expr) -> Result<Transformed<Expr>> {
603        if let Expr::Unnest(ref traversing_unnest) = expr {
604            if traversing_unnest == self.top_most_unnest.as_ref().unwrap() {
605                self.top_most_unnest = None;
606            }
607            // Find inside consecutive_unnest, the sequence of continuous unnest exprs
608
609            // Get the latest consecutive unnest exprs
610            // and check if current upward traversal is the returning to the root expr
611            // for example given a expr `unnest(unnest(col))` then the traversal happens like:
612            // down(unnest) -> down(unnest) -> down(col) -> up(col) -> up(unnest) -> up(unnest)
613            // the result of such traversal is unnest(col, depth:=2)
614            let unnest_stack = self.get_latest_consecutive_unnest();
615
616            // This traversal has reached the top most unnest again
617            // e.g Unnest(top) -> Unnest(2nd) -> Column(bottom)
618            // -> Unnest(2nd) -> Unnest(top) a.k.a here
619            // Thus
620            // Unnest(Unnest(some_col)) is rewritten into Unnest(some_col, depth:=2)
621            if traversing_unnest == unnest_stack.last().unwrap() {
622                let most_inner = unnest_stack.first().unwrap();
623                let inner_expr = most_inner.expr.as_ref();
624                // unnest(unnest(struct_arr_col)) is not allow to be done recursively
625                // it needs to be split into multiple unnest logical plan
626                // unnest(struct_arr)
627                //  unnest(struct_arr_col) as struct_arr
628                // instead of unnest(struct_arr_col, depth = 2)
629
630                let unnest_recursion = unnest_stack.len();
631                let struct_allowed =
632                    self.is_at_struct_allowed_root(&expr) && unnest_recursion == 1;
633
634                let mut transformed_exprs = self.transform(
635                    unnest_recursion,
636                    expr.schema_name().to_string(),
637                    inner_expr,
638                    struct_allowed,
639                )?;
640                // Only set transformed_root_exprs for struct unnest (which returns multiple expressions).
641                // For list unnest (single expression), we let the normal rewrite handle the alias.
642                if struct_allowed && transformed_exprs.len() > 1 {
643                    self.transformed_root_exprs = Some(transformed_exprs.clone());
644                }
645                return Ok(Transformed::new(
646                    transformed_exprs.swap_remove(0),
647                    true,
648                    TreeNodeRecursion::Continue,
649                ));
650            }
651        } else {
652            self.consecutive_unnest.push(None);
653        }
654
655        // For column exprs that are not descendants of any unnest node
656        // retain their projection
657        // e.g given expr tree unnest(col_a) + col_b, we have to retain projection of col_b
658        // this condition can be checked by maintaining an Option<top most unnest>
659        if matches!(&expr, Expr::Column(_)) && self.top_most_unnest.is_none() {
660            push_projection_dedupl(self.inner_projection_exprs, expr.clone());
661        }
662
663        Ok(Transformed::no(expr))
664    }
665}
666
667fn push_projection_dedupl(projection: &mut Vec<Expr>, expr: Expr) {
668    let schema_name = expr.schema_name().to_string();
669    if !projection
670        .iter()
671        .any(|e| e.schema_name().to_string() == schema_name)
672    {
673        projection.push(expr);
674    }
675}
676/// The context is we want to rewrite unnest() into InnerProjection->Unnest->OuterProjection
677/// Given an expression which contains unnest expr as one of its children,
678/// Try transform depends on unnest type
679/// - For list column: unnest(col) with type list -> unnest(col) with type list::item
680/// - For struct column: unnest(struct(field1, field2)) -> unnest(struct).field1, unnest(struct).field2
681///
682/// The transformed exprs will be used in the outer projection
683/// If along the path from root to bottom, there are multiple unnest expressions, the transformation
684/// is done only for the bottom expression
685pub(crate) fn rewrite_recursive_unnest_bottom_up(
686    input: &LogicalPlan,
687    unnest_placeholder_columns: &mut IndexMap<Column, Option<Vec<ColumnUnnestList>>>,
688    inner_projection_exprs: &mut Vec<Expr>,
689    original_expr: &Expr,
690) -> Result<Vec<Expr>> {
691    let mut rewriter = RecursiveUnnestRewriter {
692        input_schema: input.schema(),
693        root_expr: original_expr,
694        top_most_unnest: None,
695        consecutive_unnest: vec![],
696        inner_projection_exprs,
697        columns_unnestings: unnest_placeholder_columns,
698        transformed_root_exprs: None,
699    };
700
701    // This transformation is only done for list unnest
702    // struct unnest is done at the root level, and at the later stage
703    // because the syntax of TreeNode only support transform into 1 Expr, while
704    // Unnest struct will be transformed into multiple Exprs
705    // TODO: This can be resolved after this issue is resolved: https://github.com/apache/datafusion/issues/10102
706    //
707    // The transformation looks like:
708    // - unnest(array_col) will be transformed into Column("unnest_place_holder(array_col)")
709    // - unnest(array_col) + 1 will be transformed into Column("unnest_place_holder(array_col) + 1")
710    let Transformed {
711        data: transformed_expr,
712        transformed,
713        tnr: _,
714    } = original_expr.clone().rewrite(&mut rewriter)?;
715
716    if !transformed {
717        // TODO: remove the next line after `Expr::Wildcard` is removed
718        #[expect(deprecated)]
719        if matches!(&transformed_expr, Expr::Column(_))
720            || matches!(&transformed_expr, Expr::Wildcard { .. })
721        {
722            push_projection_dedupl(inner_projection_exprs, transformed_expr.clone());
723            Ok(vec![transformed_expr])
724        } else {
725            // We need to evaluate the expr in the inner projection,
726            // outer projection just select its name
727            let column_name = transformed_expr.schema_name().to_string();
728            push_projection_dedupl(inner_projection_exprs, transformed_expr);
729            Ok(vec![Expr::Column(Column::from_name(column_name))])
730        }
731    } else {
732        if let Some(transformed_root_exprs) = rewriter.transformed_root_exprs {
733            return Ok(transformed_root_exprs);
734        }
735        Ok(vec![transformed_expr])
736    }
737}
738
739#[cfg(test)]
740mod tests {
741    use std::{ops::Add, sync::Arc};
742
743    use arrow::datatypes::{DataType as ArrowDataType, Field, Fields, Schema};
744    use datafusion_common::{Column, DFSchema, Result};
745    use datafusion_expr::{
746        ColumnUnnestList, EmptyRelation, LogicalPlan, col, lit, unnest,
747    };
748    use datafusion_functions::core::expr_ext::FieldAccessor;
749    use datafusion_functions_aggregate::expr_fn::count;
750
751    use crate::utils::{resolve_positions_to_exprs, rewrite_recursive_unnest_bottom_up};
752    use indexmap::IndexMap;
753
754    fn column_unnests_eq(
755        l: Vec<&str>,
756        r: &IndexMap<Column, Option<Vec<ColumnUnnestList>>>,
757    ) {
758        let r_formatted: Vec<String> = r
759            .iter()
760            .map(|i| match i.1 {
761                None => format!("{}", i.0),
762                Some(vec) => format!(
763                    "{}=>[{}]",
764                    i.0,
765                    vec.iter()
766                        .map(|i| format!("{i}"))
767                        .collect::<Vec<String>>()
768                        .join(", ")
769                ),
770            })
771            .collect();
772        let l_formatted: Vec<String> = l.iter().map(|i| (*i).to_string()).collect();
773        assert_eq!(l_formatted, r_formatted);
774    }
775
776    #[test]
777    fn test_transform_bottom_unnest_recursive() -> Result<()> {
778        let schema = Schema::new(vec![
779            Field::new(
780                "3d_col",
781                ArrowDataType::List(Arc::new(Field::new(
782                    "2d_col",
783                    ArrowDataType::List(Arc::new(Field::new(
784                        "elements",
785                        ArrowDataType::Int64,
786                        true,
787                    ))),
788                    true,
789                ))),
790                true,
791            ),
792            Field::new("i64_col", ArrowDataType::Int64, true),
793        ]);
794
795        let dfschema = DFSchema::try_from(schema)?;
796
797        let input = LogicalPlan::EmptyRelation(EmptyRelation {
798            produce_one_row: false,
799            schema: Arc::new(dfschema),
800        });
801
802        let mut unnest_placeholder_columns = IndexMap::new();
803        let mut inner_projection_exprs = vec![];
804
805        // unnest(unnest(3d_col)) + unnest(unnest(3d_col))
806        let original_expr = unnest(unnest(col("3d_col")))
807            .add(unnest(unnest(col("3d_col"))))
808            .add(col("i64_col"));
809        let transformed_exprs = rewrite_recursive_unnest_bottom_up(
810            &input,
811            &mut unnest_placeholder_columns,
812            &mut inner_projection_exprs,
813            &original_expr,
814        )?;
815        // Only the bottom most unnest exprs are transformed
816        assert_eq!(
817            transformed_exprs,
818            vec![
819                col("__unnest_placeholder(3d_col,depth=2)")
820                    .alias("UNNEST(UNNEST(3d_col))")
821                    .add(
822                        col("__unnest_placeholder(3d_col,depth=2)")
823                            .alias("UNNEST(UNNEST(3d_col))")
824                    )
825                    .add(col("i64_col"))
826            ]
827        );
828        column_unnests_eq(
829            vec![
830                "__unnest_placeholder(3d_col)=>[__unnest_placeholder(3d_col,depth=2)|depth=2]",
831            ],
832            &unnest_placeholder_columns,
833        );
834
835        // Still reference struct_col in original schema but with alias,
836        // to avoid colliding with the projection on the column itself if any
837        assert_eq!(
838            inner_projection_exprs,
839            vec![
840                col("3d_col").alias("__unnest_placeholder(3d_col)"),
841                col("i64_col")
842            ]
843        );
844
845        // unnest(3d_col) as 2d_col
846        let original_expr_2 = unnest(col("3d_col")).alias("2d_col");
847        let transformed_exprs = rewrite_recursive_unnest_bottom_up(
848            &input,
849            &mut unnest_placeholder_columns,
850            &mut inner_projection_exprs,
851            &original_expr_2,
852        )?;
853
854        assert_eq!(
855            transformed_exprs,
856            vec![
857                (col("__unnest_placeholder(3d_col,depth=1)").alias("UNNEST(3d_col)"))
858                    .alias("2d_col")
859            ]
860        );
861        column_unnests_eq(
862            vec![
863                "__unnest_placeholder(3d_col)=>[__unnest_placeholder(3d_col,depth=2)|depth=2, __unnest_placeholder(3d_col,depth=1)|depth=1]",
864            ],
865            &unnest_placeholder_columns,
866        );
867        // Still reference struct_col in original schema but with alias,
868        // to avoid colliding with the projection on the column itself if any
869        assert_eq!(
870            inner_projection_exprs,
871            vec![
872                col("3d_col").alias("__unnest_placeholder(3d_col)"),
873                col("i64_col")
874            ]
875        );
876
877        Ok(())
878    }
879
880    #[test]
881    fn test_transform_bottom_unnest() -> Result<()> {
882        let schema = Schema::new(vec![
883            Field::new(
884                "struct_col",
885                ArrowDataType::Struct(Fields::from(vec![
886                    Field::new("field1", ArrowDataType::Int32, false),
887                    Field::new("field2", ArrowDataType::Int32, false),
888                ])),
889                false,
890            ),
891            Field::new(
892                "array_col",
893                ArrowDataType::List(Arc::new(Field::new_list_field(
894                    ArrowDataType::Int64,
895                    true,
896                ))),
897                true,
898            ),
899            Field::new("int_col", ArrowDataType::Int32, false),
900        ]);
901
902        let dfschema = DFSchema::try_from(schema)?;
903
904        let input = LogicalPlan::EmptyRelation(EmptyRelation {
905            produce_one_row: false,
906            schema: Arc::new(dfschema),
907        });
908
909        let mut unnest_placeholder_columns = IndexMap::new();
910        let mut inner_projection_exprs = vec![];
911
912        // unnest(struct_col)
913        let original_expr = unnest(col("struct_col"));
914        let transformed_exprs = rewrite_recursive_unnest_bottom_up(
915            &input,
916            &mut unnest_placeholder_columns,
917            &mut inner_projection_exprs,
918            &original_expr,
919        )?;
920        assert_eq!(
921            transformed_exprs,
922            vec![
923                col("__unnest_placeholder(struct_col).field1"),
924                col("__unnest_placeholder(struct_col).field2"),
925            ]
926        );
927        column_unnests_eq(
928            vec!["__unnest_placeholder(struct_col)"],
929            &unnest_placeholder_columns,
930        );
931        // Still reference struct_col in original schema but with alias,
932        // to avoid colliding with the projection on the column itself if any
933        assert_eq!(
934            inner_projection_exprs,
935            vec![col("struct_col").alias("__unnest_placeholder(struct_col)"),]
936        );
937
938        // unnest(array_col) + 1
939        let original_expr = unnest(col("array_col")).add(lit(1i64));
940        let transformed_exprs = rewrite_recursive_unnest_bottom_up(
941            &input,
942            &mut unnest_placeholder_columns,
943            &mut inner_projection_exprs,
944            &original_expr,
945        )?;
946        column_unnests_eq(
947            vec![
948                "__unnest_placeholder(struct_col)",
949                "__unnest_placeholder(array_col)=>[__unnest_placeholder(array_col,depth=1)|depth=1]",
950            ],
951            &unnest_placeholder_columns,
952        );
953        // Only transform the unnest children
954        assert_eq!(
955            transformed_exprs,
956            vec![
957                col("__unnest_placeholder(array_col,depth=1)")
958                    .alias("UNNEST(array_col)")
959                    .add(lit(1i64))
960            ]
961        );
962
963        // Keep appending to the current vector
964        // Still reference array_col in original schema but with alias,
965        // to avoid colliding with the projection on the column itself if any
966        assert_eq!(
967            inner_projection_exprs,
968            vec![
969                col("struct_col").alias("__unnest_placeholder(struct_col)"),
970                col("array_col").alias("__unnest_placeholder(array_col)")
971            ]
972        );
973
974        Ok(())
975    }
976
977    // Unnest -> field access -> unnest
978    #[test]
979    fn test_transform_non_consecutive_unnests() -> Result<()> {
980        // List of struct
981        // [struct{'subfield1':list(i64), 'subfield2':list(utf8)}]
982        let schema = Schema::new(vec![
983            Field::new(
984                "struct_list",
985                ArrowDataType::List(Arc::new(Field::new(
986                    "element",
987                    ArrowDataType::Struct(Fields::from(vec![
988                        Field::new(
989                            // list of i64
990                            "subfield1",
991                            ArrowDataType::List(Arc::new(Field::new(
992                                "i64_element",
993                                ArrowDataType::Int64,
994                                true,
995                            ))),
996                            true,
997                        ),
998                        Field::new(
999                            // list of utf8
1000                            "subfield2",
1001                            ArrowDataType::List(Arc::new(Field::new(
1002                                "utf8_element",
1003                                ArrowDataType::Utf8,
1004                                true,
1005                            ))),
1006                            true,
1007                        ),
1008                    ])),
1009                    true,
1010                ))),
1011                true,
1012            ),
1013            Field::new("int_col", ArrowDataType::Int32, false),
1014        ]);
1015
1016        let dfschema = DFSchema::try_from(schema)?;
1017
1018        let input = LogicalPlan::EmptyRelation(EmptyRelation {
1019            produce_one_row: false,
1020            schema: Arc::new(dfschema),
1021        });
1022
1023        let mut unnest_placeholder_columns = IndexMap::new();
1024        let mut inner_projection_exprs = vec![];
1025
1026        // An expr with multiple unnest
1027        let select_expr1 = unnest(unnest(col("struct_list")).field("subfield1"));
1028        let transformed_exprs = rewrite_recursive_unnest_bottom_up(
1029            &input,
1030            &mut unnest_placeholder_columns,
1031            &mut inner_projection_exprs,
1032            &select_expr1,
1033        )?;
1034        // Only the inner most/ bottom most unnest is transformed
1035        assert_eq!(
1036            transformed_exprs,
1037            vec![unnest(
1038                col("__unnest_placeholder(struct_list,depth=1)")
1039                    .alias("UNNEST(struct_list)")
1040                    .field("subfield1")
1041            )]
1042        );
1043
1044        column_unnests_eq(
1045            vec![
1046                "__unnest_placeholder(struct_list)=>[__unnest_placeholder(struct_list,depth=1)|depth=1]",
1047            ],
1048            &unnest_placeholder_columns,
1049        );
1050
1051        assert_eq!(
1052            inner_projection_exprs,
1053            vec![col("struct_list").alias("__unnest_placeholder(struct_list)")]
1054        );
1055
1056        // continue rewrite another expr in select
1057        let select_expr2 = unnest(unnest(col("struct_list")).field("subfield2"));
1058        let transformed_exprs = rewrite_recursive_unnest_bottom_up(
1059            &input,
1060            &mut unnest_placeholder_columns,
1061            &mut inner_projection_exprs,
1062            &select_expr2,
1063        )?;
1064        // Only the inner most/ bottom most unnest is transformed
1065        assert_eq!(
1066            transformed_exprs,
1067            vec![unnest(
1068                col("__unnest_placeholder(struct_list,depth=1)")
1069                    .alias("UNNEST(struct_list)")
1070                    .field("subfield2")
1071            )]
1072        );
1073
1074        // unnest place holder columns remain the same
1075        // because expr1 and expr2 derive from the same unnest result
1076        column_unnests_eq(
1077            vec![
1078                "__unnest_placeholder(struct_list)=>[__unnest_placeholder(struct_list,depth=1)|depth=1]",
1079            ],
1080            &unnest_placeholder_columns,
1081        );
1082
1083        assert_eq!(
1084            inner_projection_exprs,
1085            vec![col("struct_list").alias("__unnest_placeholder(struct_list)")]
1086        );
1087
1088        Ok(())
1089    }
1090
1091    #[test]
1092    fn test_resolve_positions_to_exprs() -> Result<()> {
1093        let select_exprs = vec![col("c1"), col("c2"), count(lit(1))];
1094
1095        // Assert 1 resolved as first column in select list
1096        let resolved = resolve_positions_to_exprs(lit(1i64), &select_exprs)?;
1097        assert_eq!(resolved, col("c1"));
1098
1099        // Assert error if index out of select clause bounds
1100        let resolved = resolve_positions_to_exprs(lit(-1i64), &select_exprs);
1101        assert!(resolved.is_err_and(|e| e.message().contains(
1102            "Cannot find column with position -1 in SELECT clause. Valid columns: 1 to 3"
1103        )));
1104
1105        let resolved = resolve_positions_to_exprs(lit(5i64), &select_exprs);
1106        assert!(resolved.is_err_and(|e| e.message().contains(
1107            "Cannot find column with position 5 in SELECT clause. Valid columns: 1 to 3"
1108        )));
1109
1110        // Assert expression returned as-is
1111        let resolved = resolve_positions_to_exprs(lit("text"), &select_exprs)?;
1112        assert_eq!(resolved, lit("text"));
1113
1114        let resolved = resolve_positions_to_exprs(col("fake"), &select_exprs)?;
1115        assert_eq!(resolved, col("fake"));
1116
1117        Ok(())
1118    }
1119}
datafusion_sql/utils.rs

datafusion_sql/
utils.rs
