Skip to main content

datafusion_expr/logical_plan/
builder.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//! This module provides a builder for creating LogicalPlans
19
20use std::borrow::Cow;
21use std::cmp::Ordering;
22use std::collections::{HashMap, HashSet};
23use std::iter::once;
24use std::sync::Arc;
25
26use crate::dml::CopyTo;
27use crate::expr::{Alias, PlannedReplaceSelectItem, Sort as SortExpr};
28use crate::expr_rewriter::{
29    coerce_plan_expr_for_schema, normalize_col,
30    normalize_col_with_schemas_and_ambiguity_check, normalize_cols, normalize_sorts,
31    rewrite_sort_cols_by_aggs,
32};
33use crate::logical_plan::{
34    Aggregate, Analyze, Distinct, DistinctOn, EmptyRelation, Explain, Filter, Join,
35    JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Prepare,
36    Projection, Repartition, Sort, SubqueryAlias, TableScanBuilder, Union, Unnest,
37    Values, Window,
38};
39use crate::select_expr::SelectExpr;
40use crate::utils::{
41    can_hash, columnize_expr, compare_sort_expr, expand_qualified_wildcard,
42    expand_wildcard, expr_to_columns, find_valid_equijoin_key_pair,
43    group_window_expr_by_sort_keys,
44};
45use crate::{
46    DmlStatement, ExplainOption, Expr, ExprSchemable, Operator, RecursiveQuery,
47    Statement, TableProviderFilterPushDown, TableSource, WriteOp, and, binary_expr, lit,
48};
49
50use super::dml::InsertOp;
51use arrow::compute::can_cast_types;
52use arrow::datatypes::{DataType, Field, FieldRef, Fields, Schema, SchemaRef};
53use datafusion_common::display::ToStringifiedPlan;
54use datafusion_common::file_options::file_type::FileType;
55use datafusion_common::metadata::FieldMetadata;
56use datafusion_common::{
57    Column, Constraints, DFSchema, DFSchemaRef, NullEquality, Result, ScalarValue,
58    TableReference, ToDFSchema, UnnestOptions, exec_err,
59    get_target_functional_dependencies, internal_datafusion_err, plan_datafusion_err,
60    plan_err,
61};
62use datafusion_expr_common::type_coercion::binary::type_union_resolution;
63
64use indexmap::IndexSet;
65
66/// Default table name for unnamed table
67pub const UNNAMED_TABLE: &str = "?table?";
68
69/// Options for [`LogicalPlanBuilder`]
70#[derive(Default, Debug, Clone)]
71pub struct LogicalPlanBuilderOptions {
72    /// Flag indicating whether the plan builder should add
73    /// functionally dependent expressions as additional aggregation groupings.
74    add_implicit_group_by_exprs: bool,
75}
76
77impl LogicalPlanBuilderOptions {
78    pub fn new() -> Self {
79        Default::default()
80    }
81
82    /// Should the builder add functionally dependent expressions as additional aggregation groupings.
83    pub fn with_add_implicit_group_by_exprs(mut self, add: bool) -> Self {
84        self.add_implicit_group_by_exprs = add;
85        self
86    }
87}
88
89/// Builder for logical plans
90///
91/// # Example building a simple plan
92/// ```
93/// # use datafusion_expr::{lit, col, LogicalPlanBuilder, logical_plan::table_scan};
94/// # use datafusion_common::Result;
95/// # use arrow::datatypes::{Schema, DataType, Field};
96/// #
97/// # fn main() -> Result<()> {
98/// #
99/// # fn employee_schema() -> Schema {
100/// #    Schema::new(vec![
101/// #           Field::new("id", DataType::Int32, false),
102/// #           Field::new("first_name", DataType::Utf8, false),
103/// #           Field::new("last_name", DataType::Utf8, false),
104/// #           Field::new("state", DataType::Utf8, false),
105/// #           Field::new("salary", DataType::Int32, false),
106/// #       ])
107/// #   }
108/// #
109/// // Create a plan similar to
110/// // SELECT last_name
111/// // FROM employees
112/// // WHERE salary < 1000
113/// let plan = table_scan(Some("employee"), &employee_schema(), None)?
114///  // Keep only rows where salary < 1000
115///  .filter(col("salary").lt(lit(1000)))?
116///  // only show "last_name" in the final results
117///  .project(vec![col("last_name")])?
118///  .build()?;
119///
120/// // Convert from plan back to builder
121/// let builder = LogicalPlanBuilder::from(plan);
122///
123/// # Ok(())
124/// # }
125/// ```
126#[derive(Debug, Clone)]
127pub struct LogicalPlanBuilder {
128    plan: Arc<LogicalPlan>,
129    options: LogicalPlanBuilderOptions,
130}
131
132impl LogicalPlanBuilder {
133    /// Create a builder from an existing plan
134    pub fn new(plan: LogicalPlan) -> Self {
135        Self {
136            plan: Arc::new(plan),
137            options: LogicalPlanBuilderOptions::default(),
138        }
139    }
140
141    /// Create a builder from an existing plan
142    pub fn new_from_arc(plan: Arc<LogicalPlan>) -> Self {
143        Self {
144            plan,
145            options: LogicalPlanBuilderOptions::default(),
146        }
147    }
148
149    pub fn with_options(mut self, options: LogicalPlanBuilderOptions) -> Self {
150        self.options = options;
151        self
152    }
153
154    /// Return the output schema of the plan build so far
155    pub fn schema(&self) -> &DFSchemaRef {
156        self.plan.schema()
157    }
158
159    /// Return the LogicalPlan of the plan build so far
160    pub fn plan(&self) -> &LogicalPlan {
161        &self.plan
162    }
163
164    /// Create an empty relation.
165    ///
166    /// `produce_one_row` set to true means this empty node needs to produce a placeholder row.
167    pub fn empty(produce_one_row: bool) -> Self {
168        Self::new(LogicalPlan::EmptyRelation(EmptyRelation {
169            produce_one_row,
170            schema: DFSchemaRef::new(DFSchema::empty()),
171        }))
172    }
173
174    /// Convert a regular plan into a recursive query.
175    /// `is_distinct` indicates whether the recursive term should be de-duplicated (`UNION`) after each iteration or not (`UNION ALL`).
176    pub fn to_recursive_query(
177        self,
178        name: String,
179        recursive_term: LogicalPlan,
180        is_distinct: bool,
181    ) -> Result<Self> {
182        // Ensure that the static term and the recursive term have the same number of fields
183        let static_fields_len = self.plan.schema().fields().len();
184        let recursive_fields_len = recursive_term.schema().fields().len();
185        if static_fields_len != recursive_fields_len {
186            return plan_err!(
187                "Non-recursive term and recursive term must have the same number of columns ({} != {})",
188                static_fields_len,
189                recursive_fields_len
190            );
191        }
192        // Ensure that the recursive term has the same field types as the static term
193        let coerced_recursive_term =
194            coerce_plan_expr_for_schema(recursive_term, self.plan.schema())?;
195        let recursive_query = RecursiveQuery::try_new(
196            name,
197            self.plan,
198            Arc::new(coerced_recursive_term),
199            is_distinct,
200        )?;
201        Ok(Self::from(LogicalPlan::RecursiveQuery(recursive_query)))
202    }
203
204    /// Create a values list based relation, and the schema is inferred from data, consuming
205    /// `value`. See the [Postgres VALUES](https://www.postgresql.org/docs/current/queries-values.html)
206    /// documentation for more details.
207    ///
208    /// so it's usually better to override the default names with a table alias list.
209    ///
210    /// If the values include params/binders such as $1, $2, $3, etc, then the `param_data_types` should be provided.
211    pub fn values(values: Vec<Vec<Expr>>) -> Result<Self> {
212        if values.is_empty() {
213            return plan_err!("Values list cannot be empty");
214        }
215        let n_cols = values[0].len();
216        if n_cols == 0 {
217            return plan_err!("Values list cannot be zero length");
218        }
219        for (i, row) in values.iter().enumerate() {
220            if row.len() != n_cols {
221                return plan_err!(
222                    "Inconsistent data length across values list: got {} values in row {} but expected {}",
223                    row.len(),
224                    i,
225                    n_cols
226                );
227            }
228        }
229
230        // Infer from data itself
231        Self::infer_data(values)
232    }
233
234    /// Create a values list based relation, and the schema is inferred from data itself or table schema if provided, consuming
235    /// `value`. See the [Postgres VALUES](https://www.postgresql.org/docs/current/queries-values.html)
236    /// documentation for more details.
237    ///
238    /// By default, it assigns the names column1, column2, etc. to the columns of a VALUES table.
239    /// The column names are not specified by the SQL standard and different database systems do it differently,
240    /// so it's usually better to override the default names with a table alias list.
241    ///
242    /// If the values include params/binders such as $1, $2, $3, etc, then the `param_data_types` should be provided.
243    pub fn values_with_schema(
244        values: Vec<Vec<Expr>>,
245        schema: &DFSchemaRef,
246    ) -> Result<Self> {
247        if values.is_empty() {
248            return plan_err!("Values list cannot be empty");
249        }
250        let n_cols = schema.fields().len();
251        if n_cols == 0 {
252            return plan_err!("Values list cannot be zero length");
253        }
254        for (i, row) in values.iter().enumerate() {
255            if row.len() != n_cols {
256                return plan_err!(
257                    "Inconsistent data length across values list: got {} values in row {} but expected {}",
258                    row.len(),
259                    i,
260                    n_cols
261                );
262            }
263        }
264
265        // Check the type of value against the schema
266        Self::infer_values_from_schema(values, schema)
267    }
268
269    fn infer_values_from_schema(
270        values: Vec<Vec<Expr>>,
271        schema: &DFSchema,
272    ) -> Result<Self> {
273        let n_cols = values[0].len();
274        let mut fields = ValuesFields::new();
275        for j in 0..n_cols {
276            let field_type = schema.field(j).data_type();
277            let field_nullable = schema.field(j).is_nullable();
278            for row in values.iter() {
279                let value = &row[j];
280                let data_type = value.get_type(schema)?;
281
282                if !data_type.equals_datatype(field_type)
283                    && !can_cast_types(&data_type, field_type)
284                {
285                    return exec_err!(
286                        "Types don't match and no valid cast exists, received data of type {} for field of type {}",
287                        data_type,
288                        field_type
289                    );
290                }
291            }
292            fields.push(field_type.to_owned(), field_nullable);
293        }
294
295        Self::infer_inner(values, fields, schema)
296    }
297
298    fn infer_data(values: Vec<Vec<Expr>>) -> Result<Self> {
299        let n_cols = values[0].len();
300        let schema = DFSchema::empty();
301        let mut fields = ValuesFields::new();
302
303        for j in 0..n_cols {
304            let mut common_type: Option<DataType> = None;
305            let mut common_metadata: Option<FieldMetadata> = None;
306            let mut nullable = false;
307            for (i, row) in values.iter().enumerate() {
308                let value = &row[j];
309                let metadata = value.metadata(&schema)?;
310                if let Some(ref cm) = common_metadata {
311                    if &metadata != cm {
312                        return plan_err!(
313                            "Inconsistent metadata across values list at row {i} column {j}. Was {:?} but found {:?}",
314                            cm,
315                            metadata
316                        );
317                    }
318                } else {
319                    common_metadata = Some(metadata.clone());
320                }
321                if !nullable && value.nullable(&schema)? {
322                    nullable = true;
323                }
324                let data_type = value.get_type(&schema)?;
325                if data_type == DataType::Null {
326                    continue;
327                }
328
329                if let Some(prev_type) = common_type {
330                    // Widen the running type so that it can hold both the
331                    // previously seen rows and this row's value.
332                    let data_types = vec![prev_type.clone(), data_type.clone()];
333                    let Some(new_type) = type_union_resolution(&data_types) else {
334                        return plan_err!(
335                            "Inconsistent data type across values list at row {i} column {j}. Was {prev_type} but found {data_type}"
336                        );
337                    };
338                    common_type = Some(new_type);
339                } else {
340                    common_type = Some(data_type);
341                }
342            }
343            // If common_type is not set, every value in this column had type
344            // NULL. A DataType::Null field is always nullable.
345            let (data_type, nullable) = match common_type {
346                Some(t) => (t, nullable),
347                None => (DataType::Null, true),
348            };
349            fields.push_with_metadata(data_type, nullable, common_metadata);
350        }
351
352        Self::infer_inner(values, fields, &schema)
353    }
354
355    fn infer_inner(
356        mut values: Vec<Vec<Expr>>,
357        fields: ValuesFields,
358        schema: &DFSchema,
359    ) -> Result<Self> {
360        let fields = fields.into_fields();
361        // wrap cast if data type is not same as common type.
362        for row in &mut values {
363            for (j, field_type) in fields.iter().map(|f| f.data_type()).enumerate() {
364                if let Expr::Literal(ScalarValue::Null, metadata) = &row[j] {
365                    row[j] = Expr::Literal(
366                        ScalarValue::try_from(field_type)?,
367                        metadata.clone(),
368                    );
369                } else {
370                    row[j] = std::mem::take(&mut row[j]).cast_to(field_type, schema)?;
371                }
372            }
373        }
374
375        let dfschema = DFSchema::from_unqualified_fields(fields, HashMap::new())?;
376        let schema = DFSchemaRef::new(dfschema);
377
378        Ok(Self::new(LogicalPlan::Values(Values { schema, values })))
379    }
380
381    /// Convert a table provider into a builder with a TableScan
382    ///
383    /// Note that if you pass a string as `table_name`, it is treated
384    /// as a SQL identifier, as described on [`TableReference`] and
385    /// thus is normalized
386    ///
387    /// # Example:
388    /// ```
389    /// # use datafusion_expr::{lit, col, LogicalPlanBuilder,
390    /// #  logical_plan::builder::LogicalTableSource, logical_plan::table_scan
391    /// # };
392    /// # use std::sync::Arc;
393    /// # use arrow::datatypes::{Schema, DataType, Field};
394    /// # use datafusion_common::TableReference;
395    /// #
396    /// # let employee_schema = Arc::new(Schema::new(vec![
397    /// #           Field::new("id", DataType::Int32, false),
398    /// # ])) as _;
399    /// # let table_source = Arc::new(LogicalTableSource::new(employee_schema));
400    /// // Scan table_source with the name "mytable" (after normalization)
401    /// # let table = table_source.clone();
402    /// let scan = LogicalPlanBuilder::scan("MyTable", table, None);
403    ///
404    /// // Scan table_source with the name "MyTable" by enclosing in quotes
405    /// # let table = table_source.clone();
406    /// let scan = LogicalPlanBuilder::scan(r#""MyTable""#, table, None);
407    ///
408    /// // Scan table_source with the name "MyTable" by forming the table reference
409    /// # let table = table_source.clone();
410    /// let table_reference = TableReference::bare("MyTable");
411    /// let scan = LogicalPlanBuilder::scan(table_reference, table, None);
412    /// ```
413    pub fn scan(
414        table_name: impl Into<TableReference>,
415        table_source: Arc<dyn TableSource>,
416        projection: Option<Vec<usize>>,
417    ) -> Result<Self> {
418        Self::scan_with_filters(table_name, table_source, projection, vec![])
419    }
420
421    /// Create a [CopyTo] for copying the contents of this builder to the specified file(s)
422    pub fn copy_to(
423        input: LogicalPlan,
424        output_url: String,
425        file_type: Arc<dyn FileType>,
426        options: HashMap<String, String>,
427        partition_by: Vec<String>,
428    ) -> Result<Self> {
429        Ok(Self::new(LogicalPlan::Copy(CopyTo::new(
430            Arc::new(input),
431            output_url,
432            partition_by,
433            file_type,
434            options,
435        ))))
436    }
437
438    /// Create a [`DmlStatement`] for inserting the contents of this builder into the named table.
439    ///
440    /// Note,  use a [`DefaultTableSource`] to insert into a [`TableProvider`]
441    ///
442    /// [`DefaultTableSource`]: https://docs.rs/datafusion/latest/datafusion/datasource/default_table_source/struct.DefaultTableSource.html
443    /// [`TableProvider`]: https://docs.rs/datafusion/latest/datafusion/catalog/trait.TableProvider.html
444    ///
445    /// # Example:
446    /// ```
447    /// # use datafusion_expr::{lit, LogicalPlanBuilder,
448    /// #  logical_plan::builder::LogicalTableSource,
449    /// # };
450    /// # use std::sync::Arc;
451    /// # use arrow::datatypes::{Schema, DataType, Field};
452    /// # use datafusion_expr::dml::InsertOp;
453    /// #
454    /// # fn test() -> datafusion_common::Result<()> {
455    /// # let employee_schema = Arc::new(Schema::new(vec![
456    /// #     Field::new("id", DataType::Int32, false),
457    /// # ])) as _;
458    /// # let table_source = Arc::new(LogicalTableSource::new(employee_schema));
459    /// // VALUES (1), (2)
460    /// let input = LogicalPlanBuilder::values(vec![vec![lit(1)], vec![lit(2)]])?.build()?;
461    /// // INSERT INTO MyTable VALUES (1), (2)
462    /// let insert_plan = LogicalPlanBuilder::insert_into(
463    ///     input,
464    ///     "MyTable",
465    ///     table_source,
466    ///     InsertOp::Append,
467    /// )?;
468    /// # Ok(())
469    /// # }
470    /// ```
471    pub fn insert_into(
472        input: LogicalPlan,
473        table_name: impl Into<TableReference>,
474        target: Arc<dyn TableSource>,
475        insert_op: InsertOp,
476    ) -> Result<Self> {
477        Ok(Self::new(LogicalPlan::Dml(DmlStatement::new(
478            table_name.into(),
479            target,
480            WriteOp::Insert(insert_op),
481            Arc::new(input),
482        ))))
483    }
484
485    /// Convert a table provider into a builder with a TableScan
486    pub fn scan_with_filters(
487        table_name: impl Into<TableReference>,
488        table_source: Arc<dyn TableSource>,
489        projection: Option<Vec<usize>>,
490        filters: Vec<Expr>,
491    ) -> Result<Self> {
492        Self::scan_with_filters_inner(table_name, table_source, projection, filters, None)
493    }
494
495    /// Convert a table provider into a builder with a TableScan with filter and fetch
496    pub fn scan_with_filters_fetch(
497        table_name: impl Into<TableReference>,
498        table_source: Arc<dyn TableSource>,
499        projection: Option<Vec<usize>>,
500        filters: Vec<Expr>,
501        fetch: Option<usize>,
502    ) -> Result<Self> {
503        Self::scan_with_filters_inner(
504            table_name,
505            table_source,
506            projection,
507            filters,
508            fetch,
509        )
510    }
511
512    fn scan_with_filters_inner(
513        table_name: impl Into<TableReference>,
514        table_source: Arc<dyn TableSource>,
515        projection: Option<Vec<usize>>,
516        filters: Vec<Expr>,
517        fetch: Option<usize>,
518    ) -> Result<Self> {
519        let table_scan = TableScanBuilder::new(table_name, table_source)
520            .with_projection(projection)
521            .with_filters(filters)
522            .with_fetch(fetch)
523            .build()?;
524
525        // Inline TableScan
526        if table_scan.filters.is_empty()
527            && let Some(p) = table_scan.source.get_logical_plan()
528        {
529            let sub_plan = p.into_owned();
530
531            if let Some(proj) = table_scan.projection {
532                let projection_exprs = proj
533                    .into_iter()
534                    .map(|i| {
535                        Expr::Column(Column::from(sub_plan.schema().qualified_field(i)))
536                    })
537                    .collect::<Vec<_>>();
538                return Self::new(sub_plan)
539                    .project(projection_exprs)?
540                    .alias(table_scan.table_name);
541            }
542
543            // Ensures that the reference to the inlined table remains the
544            // same, meaning we don't have to change any of the parent nodes
545            // that reference this table.
546            return Self::new(sub_plan).alias(table_scan.table_name);
547        }
548
549        Ok(Self::new(LogicalPlan::TableScan(table_scan)))
550    }
551
552    /// Wrap a plan in a window
553    pub fn window_plan(
554        input: LogicalPlan,
555        window_exprs: impl IntoIterator<Item = Expr>,
556    ) -> Result<LogicalPlan> {
557        let mut plan = input;
558        let mut groups = group_window_expr_by_sort_keys(window_exprs)?;
559        // To align with the behavior of PostgreSQL, we want the sort_keys sorted as same rule as PostgreSQL that first
560        // we compare the sort key themselves and if one window's sort keys are a prefix of another
561        // put the window with more sort keys first. so more deeply sorted plans gets nested further down as children.
562        // The sort_by() implementation here is a stable sort.
563        // Note that by this rule if there's an empty over, it'll be at the top level
564        groups.sort_by(|(key_a, _), (key_b, _)| {
565            for ((first, _), (second, _)) in key_a.iter().zip(key_b.iter()) {
566                let key_ordering = compare_sort_expr(first, second, plan.schema());
567                match key_ordering {
568                    Ordering::Less => {
569                        return Ordering::Less;
570                    }
571                    Ordering::Greater => {
572                        return Ordering::Greater;
573                    }
574                    Ordering::Equal => {}
575                }
576            }
577            key_b.len().cmp(&key_a.len())
578        });
579        for (_, exprs) in groups {
580            let window_exprs = exprs.into_iter().collect::<Vec<_>>();
581            // Partition and sorting is done at physical level, see the EnforceDistribution
582            // and EnforceSorting rules.
583            plan = LogicalPlanBuilder::from(plan)
584                .window(window_exprs)?
585                .build()?;
586        }
587        Ok(plan)
588    }
589
590    /// Apply a projection without alias.
591    pub fn project(
592        self,
593        expr: impl IntoIterator<Item = impl Into<SelectExpr>>,
594    ) -> Result<Self> {
595        project(Arc::unwrap_or_clone(self.plan), expr).map(Self::new)
596    }
597
598    /// Apply a projection without alias with optional validation
599    /// (true to validate, false to not validate)
600    pub fn project_with_validation(
601        self,
602        expr: Vec<(impl Into<SelectExpr>, bool)>,
603    ) -> Result<Self> {
604        project_with_validation(Arc::unwrap_or_clone(self.plan), expr, None)
605            .map(Self::new)
606    }
607
608    /// Apply a projection, aliasing non-Column/non-Alias expressions to
609    /// match the field names from the provided schema.
610    pub fn project_with_validation_and_schema(
611        self,
612        expr: impl IntoIterator<Item = impl Into<SelectExpr>>,
613        schema: &DFSchemaRef,
614    ) -> Result<Self> {
615        project_with_validation(
616            Arc::unwrap_or_clone(self.plan),
617            expr.into_iter().map(|e| (e, true)),
618            Some(schema),
619        )
620        .map(Self::new)
621    }
622
623    /// Select the given column indices
624    pub fn select(self, indices: impl IntoIterator<Item = usize>) -> Result<Self> {
625        let exprs: Vec<_> = indices
626            .into_iter()
627            .map(|x| Expr::Column(Column::from(self.plan.schema().qualified_field(x))))
628            .collect();
629        self.project(exprs)
630    }
631
632    /// Apply a filter
633    pub fn filter(self, expr: impl Into<Expr>) -> Result<Self> {
634        let expr = normalize_col(expr.into(), &self.plan)?;
635        Filter::try_new(expr, self.plan)
636            .map(LogicalPlan::Filter)
637            .map(Self::new)
638    }
639
640    /// Apply a filter which is used for a having clause
641    pub fn having(self, expr: impl Into<Expr>) -> Result<Self> {
642        let expr = normalize_col(expr.into(), &self.plan)?;
643        Filter::try_new(expr, self.plan)
644            .map(LogicalPlan::Filter)
645            .map(Self::from)
646    }
647
648    /// Make a builder for a prepare logical plan from the builder's plan
649    pub fn prepare(self, name: String, fields: Vec<FieldRef>) -> Result<Self> {
650        Ok(Self::new(LogicalPlan::Statement(Statement::Prepare(
651            Prepare {
652                name,
653                fields,
654                input: self.plan,
655            },
656        ))))
657    }
658
659    /// Limit the number of rows returned
660    ///
661    /// `skip` - Number of rows to skip before fetch any row.
662    ///
663    /// `fetch` - Maximum number of rows to fetch, after skipping `skip` rows,
664    ///          if specified.
665    pub fn limit(self, skip: usize, fetch: Option<usize>) -> Result<Self> {
666        let skip_expr = if skip == 0 {
667            None
668        } else {
669            Some(lit(skip as i64))
670        };
671        let fetch_expr = fetch.map(|f| lit(f as i64));
672        self.limit_by_expr(skip_expr, fetch_expr)
673    }
674
675    /// Limit the number of rows returned
676    ///
677    /// Similar to `limit` but uses expressions for `skip` and `fetch`
678    pub fn limit_by_expr(self, skip: Option<Expr>, fetch: Option<Expr>) -> Result<Self> {
679        Ok(Self::new(LogicalPlan::Limit(Limit {
680            skip: skip.map(Box::new),
681            fetch: fetch.map(Box::new),
682            input: self.plan,
683        })))
684    }
685
686    /// Apply an alias
687    pub fn alias(self, alias: impl Into<TableReference>) -> Result<Self> {
688        subquery_alias(Arc::unwrap_or_clone(self.plan), alias).map(Self::new)
689    }
690
691    /// Add missing sort columns to all downstream projection
692    ///
693    /// Thus, if you have a LogicalPlan that selects A and B and have
694    /// not requested a sort by C, this code will add C recursively to
695    /// all input projections.
696    ///
697    /// Adding a new column is not correct if there is a `Distinct`
698    /// node, which produces only distinct values of its
699    /// inputs. Adding a new column to its input will result in
700    /// potentially different results than with the original column.
701    ///
702    /// For example, if the input is like:
703    ///
704    /// Distinct(A, B)
705    ///
706    /// If the input looks like
707    ///
708    /// a | b | c
709    /// --+---+---
710    /// 1 | 2 | 3
711    /// 1 | 2 | 4
712    ///
713    /// Distinct (A, B) --> (1,2)
714    ///
715    /// But Distinct (A, B, C) --> (1, 2, 3), (1, 2, 4)
716    ///  (which will appear as a (1, 2), (1, 2) if a and b are projected
717    ///
718    /// See <https://github.com/apache/datafusion/issues/5065> for more details
719    fn add_missing_columns(
720        curr_plan: LogicalPlan,
721        missing_cols: &IndexSet<Column>,
722        is_distinct: bool,
723    ) -> Result<LogicalPlan> {
724        match curr_plan {
725            LogicalPlan::Projection(Projection {
726                input,
727                mut expr,
728                schema: _,
729            }) if missing_cols.iter().all(|c| input.schema().has_column(c)) => {
730                let mut missing_exprs = missing_cols
731                    .iter()
732                    .map(|c| normalize_col(Expr::Column(c.clone()), &input))
733                    .collect::<Result<Vec<_>>>()?;
734
735                // Do not let duplicate columns to be added, some of the
736                // missing_cols may be already present but without the new
737                // projected alias.
738                missing_exprs.retain(|e| !expr.contains(e));
739                if is_distinct {
740                    Self::ambiguous_distinct_check(&missing_exprs, missing_cols, &expr)?;
741                }
742                expr.extend(missing_exprs);
743                project(Arc::unwrap_or_clone(input), expr)
744            }
745            _ => {
746                let is_distinct =
747                    is_distinct || matches!(curr_plan, LogicalPlan::Distinct(_));
748                let new_inputs = curr_plan
749                    .inputs()
750                    .into_iter()
751                    .map(|input_plan| {
752                        Self::add_missing_columns(
753                            (*input_plan).clone(),
754                            missing_cols,
755                            is_distinct,
756                        )
757                    })
758                    .collect::<Result<Vec<_>>>()?;
759                curr_plan.with_new_exprs(curr_plan.expressions(), new_inputs)
760            }
761        }
762    }
763
764    fn ambiguous_distinct_check(
765        missing_exprs: &[Expr],
766        missing_cols: &IndexSet<Column>,
767        projection_exprs: &[Expr],
768    ) -> Result<()> {
769        if missing_exprs.is_empty() {
770            return Ok(());
771        }
772
773        // if the missing columns are all only aliases for things in
774        // the existing select list, it is ok
775        //
776        // This handles the special case for
777        // SELECT col as <alias> ORDER BY <alias>
778        //
779        // As described in https://github.com/apache/datafusion/issues/5293
780        let all_aliases = missing_exprs.iter().all(|e| {
781            projection_exprs.iter().any(|proj_expr| {
782                if let Expr::Alias(Alias { expr, .. }) = proj_expr {
783                    e == expr.as_ref()
784                } else {
785                    false
786                }
787            })
788        });
789        if all_aliases {
790            return Ok(());
791        }
792
793        let missing_col_names = missing_cols
794            .iter()
795            .map(|col| col.flat_name())
796            .collect::<String>();
797
798        plan_err!(
799            "For SELECT DISTINCT, ORDER BY expressions {missing_col_names} must appear in select list"
800        )
801    }
802
803    /// Apply a sort by provided expressions with default direction
804    pub fn sort_by(
805        self,
806        expr: impl IntoIterator<Item = impl Into<Expr>> + Clone,
807    ) -> Result<Self> {
808        self.sort(
809            expr.into_iter()
810                .map(|e| e.into().sort(true, false))
811                .collect::<Vec<SortExpr>>(),
812        )
813    }
814
815    pub fn sort(
816        self,
817        sorts: impl IntoIterator<Item = impl Into<SortExpr>> + Clone,
818    ) -> Result<Self> {
819        self.sort_with_limit(sorts, None)
820    }
821
822    /// Apply a sort
823    pub fn sort_with_limit(
824        self,
825        sorts: impl IntoIterator<Item = impl Into<SortExpr>> + Clone,
826        fetch: Option<usize>,
827    ) -> Result<Self> {
828        let sorts = rewrite_sort_cols_by_aggs(sorts, &self.plan)?;
829
830        let schema = self.plan.schema();
831
832        // Collect sort columns that are missing in the input plan's schema
833        let mut missing_cols: IndexSet<Column> = IndexSet::new();
834        sorts.iter().try_for_each::<_, Result<()>>(|sort| {
835            let columns = sort.expr.column_refs();
836
837            missing_cols.extend(
838                columns
839                    .into_iter()
840                    .filter(|c| !schema.has_column(c))
841                    .cloned(),
842            );
843
844            Ok(())
845        })?;
846
847        if missing_cols.is_empty() {
848            return Ok(Self::new(LogicalPlan::Sort(Sort {
849                expr: normalize_sorts(sorts, &self.plan)?,
850                input: self.plan,
851                fetch,
852            })));
853        }
854
855        // remove pushed down sort columns
856        let new_expr = schema.columns().into_iter().map(Expr::Column).collect();
857
858        let is_distinct = false;
859        let plan = Self::add_missing_columns(
860            Arc::unwrap_or_clone(self.plan),
861            &missing_cols,
862            is_distinct,
863        )?;
864
865        let sort_plan = LogicalPlan::Sort(Sort {
866            expr: normalize_sorts(sorts, &plan)?,
867            input: Arc::new(plan),
868            fetch,
869        });
870
871        Projection::try_new(new_expr, Arc::new(sort_plan))
872            .map(LogicalPlan::Projection)
873            .map(Self::new)
874    }
875
876    /// Apply a union, preserving duplicate rows
877    pub fn union(self, plan: LogicalPlan) -> Result<Self> {
878        union(Arc::unwrap_or_clone(self.plan), plan).map(Self::new)
879    }
880
881    /// Apply a union by name, preserving duplicate rows
882    pub fn union_by_name(self, plan: LogicalPlan) -> Result<Self> {
883        union_by_name(Arc::unwrap_or_clone(self.plan), plan).map(Self::new)
884    }
885
886    /// Apply a union by name, removing duplicate rows
887    pub fn union_by_name_distinct(self, plan: LogicalPlan) -> Result<Self> {
888        let left_plan: LogicalPlan = Arc::unwrap_or_clone(self.plan);
889        let right_plan: LogicalPlan = plan;
890
891        Ok(Self::new(LogicalPlan::Distinct(Distinct::All(Arc::new(
892            union_by_name(left_plan, right_plan)?,
893        )))))
894    }
895
896    /// Apply a union, removing duplicate rows
897    pub fn union_distinct(self, plan: LogicalPlan) -> Result<Self> {
898        let left_plan: LogicalPlan = Arc::unwrap_or_clone(self.plan);
899        let right_plan: LogicalPlan = plan;
900
901        Ok(Self::new(LogicalPlan::Distinct(Distinct::All(Arc::new(
902            union(left_plan, right_plan)?,
903        )))))
904    }
905
906    /// Apply deduplication: Only distinct (different) values are returned)
907    pub fn distinct(self) -> Result<Self> {
908        Ok(Self::new(LogicalPlan::Distinct(Distinct::All(self.plan))))
909    }
910
911    /// Project first values of the specified expression list according to the provided
912    /// sorting expressions grouped by the `DISTINCT ON` clause expressions.
913    pub fn distinct_on(
914        self,
915        on_expr: Vec<Expr>,
916        select_expr: Vec<Expr>,
917        sort_expr: Option<Vec<SortExpr>>,
918    ) -> Result<Self> {
919        Ok(Self::new(LogicalPlan::Distinct(Distinct::On(
920            DistinctOn::try_new(on_expr, select_expr, sort_expr, self.plan)?,
921        ))))
922    }
923
924    /// Apply a join to `right` using explicitly specified columns and an
925    /// optional filter expression.
926    ///
927    /// See [`join_on`](Self::join_on) for a more concise way to specify the
928    /// join condition. Since DataFusion will automatically identify and
929    /// optimize equality predicates there is no performance difference between
930    /// this function and `join_on`
931    ///
932    /// `left_cols` and `right_cols` are used to form "equijoin" predicates (see
933    /// example below), which are then combined with the optional `filter`
934    /// expression.
935    ///
936    /// Note that in case of outer join, the `filter` is applied to only matched rows.
937    pub fn join(
938        self,
939        right: LogicalPlan,
940        join_type: JoinType,
941        join_keys: (Vec<impl Into<Column>>, Vec<impl Into<Column>>),
942        filter: Option<Expr>,
943    ) -> Result<Self> {
944        self.join_detailed(
945            right,
946            join_type,
947            join_keys,
948            filter,
949            NullEquality::NullEqualsNothing,
950        )
951    }
952
953    /// Apply a join using the specified expressions.
954    ///
955    /// Note that DataFusion automatically optimizes joins, including
956    /// identifying and optimizing equality predicates.
957    ///
958    /// # Example
959    ///
960    /// ```
961    /// # use datafusion_expr::{Expr, col, LogicalPlanBuilder,
962    /// #  logical_plan::builder::LogicalTableSource, logical_plan::JoinType,};
963    /// # use std::sync::Arc;
964    /// # use arrow::datatypes::{Schema, DataType, Field};
965    /// # use datafusion_common::Result;
966    /// # fn main() -> Result<()> {
967    /// let example_schema = Arc::new(Schema::new(vec![
968    ///     Field::new("a", DataType::Int32, false),
969    ///     Field::new("b", DataType::Int32, false),
970    ///     Field::new("c", DataType::Int32, false),
971    /// ]));
972    /// let table_source = Arc::new(LogicalTableSource::new(example_schema));
973    /// let left_table = table_source.clone();
974    /// let right_table = table_source.clone();
975    ///
976    /// let right_plan = LogicalPlanBuilder::scan("right", right_table, None)?.build()?;
977    ///
978    /// // Form the expression `(left.a != right.a)` AND `(left.b != right.b)`
979    /// let exprs = vec![
980    ///     col("left.a").eq(col("right.a")),
981    ///     col("left.b").not_eq(col("right.b")),
982    /// ];
983    ///
984    /// // Perform the equivalent of `left INNER JOIN right ON (a != a2 AND b != b2)`
985    /// // finding all pairs of rows from `left` and `right` where
986    /// // where `a = a2` and `b != b2`.
987    /// let plan = LogicalPlanBuilder::scan("left", left_table, None)?
988    ///     .join_on(right_plan, JoinType::Inner, exprs)?
989    ///     .build()?;
990    /// # Ok(())
991    /// # }
992    /// ```
993    pub fn join_on(
994        self,
995        right: LogicalPlan,
996        join_type: JoinType,
997        on_exprs: impl IntoIterator<Item = Expr>,
998    ) -> Result<Self> {
999        let filter = on_exprs.into_iter().reduce(Expr::and);
1000
1001        self.join_detailed(
1002            right,
1003            join_type,
1004            (Vec::<Column>::new(), Vec::<Column>::new()),
1005            filter,
1006            NullEquality::NullEqualsNothing,
1007        )
1008    }
1009
1010    pub(crate) fn normalize(plan: &LogicalPlan, column: Column) -> Result<Column> {
1011        if column.relation.is_some() {
1012            // column is already normalized
1013            return Ok(column);
1014        }
1015
1016        let schema = plan.schema();
1017        let fallback_schemas = plan.fallback_normalize_schemas();
1018        let using_columns = plan.using_columns()?;
1019        column.normalize_with_schemas_and_ambiguity_check(
1020            &[&[schema], &fallback_schemas],
1021            &using_columns,
1022        )
1023    }
1024
1025    /// Apply a join with on constraint and specified null equality.
1026    ///
1027    /// The behavior is the same as [`join`](Self::join) except that it allows
1028    /// specifying the null equality behavior.
1029    ///
1030    /// The `null_equality` dictates how `null` values are joined.
1031    pub fn join_detailed(
1032        self,
1033        right: LogicalPlan,
1034        join_type: JoinType,
1035        join_keys: (Vec<impl Into<Column>>, Vec<impl Into<Column>>),
1036        filter: Option<Expr>,
1037        null_equality: NullEquality,
1038    ) -> Result<Self> {
1039        self.join_detailed_with_options(
1040            right,
1041            join_type,
1042            join_keys,
1043            filter,
1044            null_equality,
1045            false,
1046        )
1047    }
1048
1049    pub fn join_detailed_with_options(
1050        self,
1051        right: LogicalPlan,
1052        join_type: JoinType,
1053        join_keys: (Vec<impl Into<Column>>, Vec<impl Into<Column>>),
1054        filter: Option<Expr>,
1055        null_equality: NullEquality,
1056        null_aware: bool,
1057    ) -> Result<Self> {
1058        if join_keys.0.len() != join_keys.1.len() {
1059            return plan_err!("left_keys and right_keys were not the same length");
1060        }
1061
1062        let filter = if let Some(expr) = filter {
1063            let filter = normalize_col_with_schemas_and_ambiguity_check(
1064                expr,
1065                &[&[self.schema(), right.schema()]],
1066                &[],
1067            )?;
1068            Some(filter)
1069        } else {
1070            None
1071        };
1072
1073        let (left_keys, right_keys): (Vec<Result<Column>>, Vec<Result<Column>>) =
1074            join_keys
1075                .0
1076                .into_iter()
1077                .zip(join_keys.1)
1078                .map(|(l, r)| {
1079                    let l = l.into();
1080                    let r = r.into();
1081
1082                    match (&l.relation, &r.relation) {
1083                        (Some(lr), Some(rr)) => {
1084                            let l_is_left =
1085                                self.plan.schema().field_with_qualified_name(lr, &l.name);
1086                            let l_is_right =
1087                                right.schema().field_with_qualified_name(lr, &l.name);
1088                            let r_is_left =
1089                                self.plan.schema().field_with_qualified_name(rr, &r.name);
1090                            let r_is_right =
1091                                right.schema().field_with_qualified_name(rr, &r.name);
1092
1093                            match (l_is_left, l_is_right, r_is_left, r_is_right) {
1094                                (_, Ok(_), Ok(_), _) => (Ok(r), Ok(l)),
1095                                (Ok(_), _, _, Ok(_)) => (Ok(l), Ok(r)),
1096                                _ => (
1097                                    Self::normalize(&self.plan, l),
1098                                    Self::normalize(&right, r),
1099                                ),
1100                            }
1101                        }
1102                        (Some(lr), None) => {
1103                            let l_is_left =
1104                                self.plan.schema().field_with_qualified_name(lr, &l.name);
1105                            let l_is_right =
1106                                right.schema().field_with_qualified_name(lr, &l.name);
1107
1108                            match (l_is_left, l_is_right) {
1109                                (Ok(_), _) => (Ok(l), Self::normalize(&right, r)),
1110                                (_, Ok(_)) => (Self::normalize(&self.plan, r), Ok(l)),
1111                                _ => (
1112                                    Self::normalize(&self.plan, l),
1113                                    Self::normalize(&right, r),
1114                                ),
1115                            }
1116                        }
1117                        (None, Some(rr)) => {
1118                            let r_is_left =
1119                                self.plan.schema().field_with_qualified_name(rr, &r.name);
1120                            let r_is_right =
1121                                right.schema().field_with_qualified_name(rr, &r.name);
1122
1123                            match (r_is_left, r_is_right) {
1124                                (Ok(_), _) => (Ok(r), Self::normalize(&right, l)),
1125                                (_, Ok(_)) => (Self::normalize(&self.plan, l), Ok(r)),
1126                                _ => (
1127                                    Self::normalize(&self.plan, l),
1128                                    Self::normalize(&right, r),
1129                                ),
1130                            }
1131                        }
1132                        (None, None) => {
1133                            let mut swap = false;
1134                            let left_key = Self::normalize(&self.plan, l.clone())
1135                                .or_else(|_| {
1136                                    swap = true;
1137                                    Self::normalize(&right, l)
1138                                });
1139                            if swap {
1140                                (Self::normalize(&self.plan, r), left_key)
1141                            } else {
1142                                (left_key, Self::normalize(&right, r))
1143                            }
1144                        }
1145                    }
1146                })
1147                .unzip();
1148
1149        let left_keys = left_keys.into_iter().collect::<Result<Vec<Column>>>()?;
1150        let right_keys = right_keys.into_iter().collect::<Result<Vec<Column>>>()?;
1151
1152        let on: Vec<_> = left_keys
1153            .into_iter()
1154            .zip(right_keys)
1155            .map(|(l, r)| (Expr::Column(l), Expr::Column(r)))
1156            .collect();
1157        let join_schema =
1158            build_join_schema(self.plan.schema(), right.schema(), &join_type)?;
1159
1160        // Inner type without join condition is cross join
1161        if join_type != JoinType::Inner && on.is_empty() && filter.is_none() {
1162            return plan_err!("join condition should not be empty");
1163        }
1164
1165        Ok(Self::new(LogicalPlan::Join(Join {
1166            left: self.plan,
1167            right: Arc::new(right),
1168            on,
1169            filter,
1170            join_type,
1171            join_constraint: JoinConstraint::On,
1172            schema: DFSchemaRef::new(join_schema),
1173            null_equality,
1174            null_aware,
1175        })))
1176    }
1177
1178    /// Apply a join with using constraint, which duplicates all join columns in output schema.
1179    pub fn join_using(
1180        self,
1181        right: LogicalPlan,
1182        join_type: JoinType,
1183        using_keys: Vec<Column>,
1184    ) -> Result<Self> {
1185        let left_keys: Vec<Column> = using_keys
1186            .clone()
1187            .into_iter()
1188            .map(|c| Self::normalize(&self.plan, c))
1189            .collect::<Result<_>>()?;
1190        let right_keys: Vec<Column> = using_keys
1191            .into_iter()
1192            .map(|c| Self::normalize(&right, c))
1193            .collect::<Result<_>>()?;
1194
1195        let on: Vec<(_, _)> = left_keys.into_iter().zip(right_keys).collect();
1196        let mut join_on: Vec<(Expr, Expr)> = vec![];
1197        let mut filters: Option<Expr> = None;
1198        for (l, r) in &on {
1199            if self.plan.schema().has_column(l)
1200                && right.schema().has_column(r)
1201                && can_hash(
1202                    datafusion_common::ExprSchema::field_from_column(
1203                        self.plan.schema(),
1204                        l,
1205                    )?
1206                    .data_type(),
1207                )
1208            {
1209                join_on.push((Expr::Column(l.clone()), Expr::Column(r.clone())));
1210            } else if self.plan.schema().has_column(l)
1211                && right.schema().has_column(r)
1212                && can_hash(
1213                    datafusion_common::ExprSchema::field_from_column(
1214                        self.plan.schema(),
1215                        r,
1216                    )?
1217                    .data_type(),
1218                )
1219            {
1220                join_on.push((Expr::Column(r.clone()), Expr::Column(l.clone())));
1221            } else {
1222                let expr = binary_expr(
1223                    Expr::Column(l.clone()),
1224                    Operator::Eq,
1225                    Expr::Column(r.clone()),
1226                );
1227                match filters {
1228                    None => filters = Some(expr),
1229                    Some(filter_expr) => filters = Some(and(expr, filter_expr)),
1230                }
1231            }
1232        }
1233
1234        if join_on.is_empty() {
1235            let join = Self::from(self.plan).cross_join(right)?;
1236            join.filter(filters.ok_or_else(|| {
1237                internal_datafusion_err!("filters should not be None here")
1238            })?)
1239        } else {
1240            let join = Join::try_new(
1241                self.plan,
1242                Arc::new(right),
1243                join_on,
1244                filters,
1245                join_type,
1246                JoinConstraint::Using,
1247                NullEquality::NullEqualsNothing,
1248                false, // null_aware
1249            )?;
1250
1251            Ok(Self::new(LogicalPlan::Join(join)))
1252        }
1253    }
1254
1255    /// Apply a cross join
1256    pub fn cross_join(self, right: LogicalPlan) -> Result<Self> {
1257        let join = Join::try_new(
1258            self.plan,
1259            Arc::new(right),
1260            vec![],
1261            None,
1262            JoinType::Inner,
1263            JoinConstraint::On,
1264            NullEquality::NullEqualsNothing,
1265            false, // null_aware
1266        )?;
1267
1268        Ok(Self::new(LogicalPlan::Join(join)))
1269    }
1270
1271    /// Repartition
1272    pub fn repartition(self, partitioning_scheme: Partitioning) -> Result<Self> {
1273        Ok(Self::new(LogicalPlan::Repartition(Repartition {
1274            input: self.plan,
1275            partitioning_scheme,
1276        })))
1277    }
1278
1279    /// Apply a window functions to extend the schema
1280    pub fn window(
1281        self,
1282        window_expr: impl IntoIterator<Item = impl Into<Expr>>,
1283    ) -> Result<Self> {
1284        let window_expr = normalize_cols(window_expr, &self.plan)?;
1285        validate_unique_names("Windows", &window_expr)?;
1286        Ok(Self::new(LogicalPlan::Window(Window::try_new(
1287            window_expr,
1288            self.plan,
1289        )?)))
1290    }
1291
1292    /// Apply an aggregate: grouping on the `group_expr` expressions
1293    /// and calculating `aggr_expr` aggregates for each distinct
1294    /// value of the `group_expr`;
1295    pub fn aggregate(
1296        self,
1297        group_expr: impl IntoIterator<Item = impl Into<Expr>>,
1298        aggr_expr: impl IntoIterator<Item = impl Into<Expr>>,
1299    ) -> Result<Self> {
1300        let group_expr = normalize_cols(group_expr, &self.plan)?;
1301        let aggr_expr = normalize_cols(aggr_expr, &self.plan)?;
1302
1303        let group_expr = if self.options.add_implicit_group_by_exprs {
1304            add_group_by_exprs_from_dependencies(group_expr, self.plan.schema())?
1305        } else {
1306            group_expr
1307        };
1308
1309        Aggregate::try_new(self.plan, group_expr, aggr_expr)
1310            .map(LogicalPlan::Aggregate)
1311            .map(Self::new)
1312    }
1313
1314    /// Create an expression to represent the explanation of the plan
1315    ///
1316    /// if `analyze` is true, runs the actual plan and produces
1317    /// information about metrics during run.
1318    ///
1319    /// if `verbose` is true, prints out additional details.
1320    pub fn explain(self, verbose: bool, analyze: bool) -> Result<Self> {
1321        // Keep the format default to Indent
1322        self.explain_option_format(
1323            ExplainOption::default()
1324                .with_verbose(verbose)
1325                .with_analyze(analyze),
1326        )
1327    }
1328
1329    /// Create an expression to represent the explanation of the plan
1330    /// The`explain_option` is used to specify the format and verbosity of the explanation.
1331    /// Details see [`ExplainOption`].
1332    pub fn explain_option_format(self, explain_option: ExplainOption) -> Result<Self> {
1333        let schema = LogicalPlan::explain_schema();
1334        let schema = schema.to_dfschema_ref()?;
1335
1336        if explain_option.analyze {
1337            Ok(Self::new(LogicalPlan::Analyze(Analyze {
1338                verbose: explain_option.verbose,
1339                format: explain_option.format,
1340                input: self.plan,
1341                schema,
1342                analyze_level: explain_option.analyze_level,
1343                analyze_categories: explain_option.analyze_categories,
1344            })))
1345        } else {
1346            let stringified_plans =
1347                vec![self.plan.to_stringified(PlanType::InitialLogicalPlan)];
1348
1349            Ok(Self::new(LogicalPlan::Explain(Explain {
1350                verbose: explain_option.verbose,
1351                plan: self.plan,
1352                explain_format: explain_option.format,
1353                stringified_plans,
1354                schema,
1355                logical_optimization_succeeded: false,
1356                show_statistics: explain_option.show_statistics,
1357            })))
1358        }
1359    }
1360
1361    /// Process intersect set operator
1362    pub fn intersect(
1363        left_plan: LogicalPlan,
1364        right_plan: LogicalPlan,
1365        is_all: bool,
1366    ) -> Result<LogicalPlan> {
1367        LogicalPlanBuilder::intersect_or_except(
1368            left_plan,
1369            right_plan,
1370            JoinType::LeftSemi,
1371            is_all,
1372        )
1373    }
1374
1375    /// Process except set operator
1376    pub fn except(
1377        left_plan: LogicalPlan,
1378        right_plan: LogicalPlan,
1379        is_all: bool,
1380    ) -> Result<LogicalPlan> {
1381        LogicalPlanBuilder::intersect_or_except(
1382            left_plan,
1383            right_plan,
1384            JoinType::LeftAnti,
1385            is_all,
1386        )
1387    }
1388
1389    /// Process intersect or except
1390    fn intersect_or_except(
1391        left_plan: LogicalPlan,
1392        right_plan: LogicalPlan,
1393        join_type: JoinType,
1394        is_all: bool,
1395    ) -> Result<LogicalPlan> {
1396        let left_len = left_plan.schema().fields().len();
1397        let right_len = right_plan.schema().fields().len();
1398
1399        if left_len != right_len {
1400            return plan_err!(
1401                "INTERSECT/EXCEPT query must have the same number of columns. Left is {left_len} and right is {right_len}."
1402            );
1403        }
1404
1405        // Requalify sides if needed to avoid duplicate qualified field names
1406        // (e.g., when both sides reference the same table)
1407        let left_builder = LogicalPlanBuilder::from(left_plan);
1408        let right_builder = LogicalPlanBuilder::from(right_plan);
1409        let (left_builder, right_builder, _requalified) =
1410            requalify_sides_if_needed(left_builder, right_builder)?;
1411        let left_plan = left_builder.build()?;
1412        let right_plan = right_builder.build()?;
1413
1414        let join_keys = left_plan
1415            .schema()
1416            .fields()
1417            .iter()
1418            .zip(right_plan.schema().fields().iter())
1419            .map(|(left_field, right_field)| {
1420                (
1421                    (Column::from_name(left_field.name())),
1422                    (Column::from_name(right_field.name())),
1423                )
1424            })
1425            .unzip();
1426        if is_all {
1427            LogicalPlanBuilder::from(left_plan)
1428                .join_detailed(
1429                    right_plan,
1430                    join_type,
1431                    join_keys,
1432                    None,
1433                    NullEquality::NullEqualsNull,
1434                )?
1435                .build()
1436        } else {
1437            LogicalPlanBuilder::from(left_plan)
1438                .distinct()?
1439                .join_detailed(
1440                    right_plan,
1441                    join_type,
1442                    join_keys,
1443                    None,
1444                    NullEquality::NullEqualsNull,
1445                )?
1446                .build()
1447        }
1448    }
1449
1450    /// Build the plan
1451    pub fn build(self) -> Result<LogicalPlan> {
1452        Ok(Arc::unwrap_or_clone(self.plan))
1453    }
1454
1455    /// Apply a join with both explicit equijoin and non equijoin predicates.
1456    ///
1457    /// Note this is a low level API that requires identifying specific
1458    /// predicate types. Most users should use  [`join_on`](Self::join_on) that
1459    /// automatically identifies predicates appropriately.
1460    ///
1461    /// `equi_exprs` defines equijoin predicates, of the form `l = r)` for each
1462    /// `(l, r)` tuple. `l`, the first element of the tuple, must only refer
1463    /// to columns from the existing input. `r`, the second element of the tuple,
1464    /// must only refer to columns from the right input.
1465    ///
1466    /// `filter` contains any other filter expression to apply during the
1467    /// join. Note that `equi_exprs` predicates are evaluated more efficiently
1468    /// than the filter expressions, so they are preferred.
1469    pub fn join_with_expr_keys(
1470        self,
1471        right: LogicalPlan,
1472        join_type: JoinType,
1473        equi_exprs: (Vec<impl Into<Expr>>, Vec<impl Into<Expr>>),
1474        filter: Option<Expr>,
1475    ) -> Result<Self> {
1476        if equi_exprs.0.len() != equi_exprs.1.len() {
1477            return plan_err!("left_keys and right_keys were not the same length");
1478        }
1479
1480        let join_key_pairs = equi_exprs
1481            .0
1482            .into_iter()
1483            .zip(equi_exprs.1)
1484            .map(|(l, r)| {
1485                let left_key = l.into();
1486                let right_key = r.into();
1487                let mut left_using_columns  = HashSet::new();
1488                expr_to_columns(&left_key, &mut left_using_columns)?;
1489                let normalized_left_key = normalize_col_with_schemas_and_ambiguity_check(
1490                    left_key,
1491                    &[&[self.plan.schema()]],
1492                    &[],
1493                )?;
1494
1495                let mut right_using_columns = HashSet::new();
1496                expr_to_columns(&right_key, &mut right_using_columns)?;
1497                let normalized_right_key = normalize_col_with_schemas_and_ambiguity_check(
1498                    right_key,
1499                    &[&[right.schema()]],
1500                    &[],
1501                )?;
1502
1503                // find valid equijoin
1504                find_valid_equijoin_key_pair(
1505                        &normalized_left_key,
1506                        &normalized_right_key,
1507                        self.plan.schema(),
1508                        right.schema(),
1509                    )?.ok_or_else(||
1510                        plan_datafusion_err!(
1511                            "can't create join plan, join key should belong to one input, error key: ({normalized_left_key},{normalized_right_key})"
1512                        ))
1513            })
1514            .collect::<Result<Vec<_>>>()?;
1515
1516        let join = Join::try_new(
1517            self.plan,
1518            Arc::new(right),
1519            join_key_pairs,
1520            filter,
1521            join_type,
1522            JoinConstraint::On,
1523            NullEquality::NullEqualsNothing,
1524            false, // null_aware
1525        )?;
1526
1527        Ok(Self::new(LogicalPlan::Join(join)))
1528    }
1529
1530    /// Unnest the given column.
1531    pub fn unnest_column(self, column: impl Into<Column>) -> Result<Self> {
1532        unnest(Arc::unwrap_or_clone(self.plan), vec![column.into()]).map(Self::new)
1533    }
1534
1535    /// Unnest the given column given [`UnnestOptions`]
1536    pub fn unnest_column_with_options(
1537        self,
1538        column: impl Into<Column>,
1539        options: UnnestOptions,
1540    ) -> Result<Self> {
1541        unnest_with_options(
1542            Arc::unwrap_or_clone(self.plan),
1543            vec![column.into()],
1544            options,
1545        )
1546        .map(Self::new)
1547    }
1548
1549    /// Unnest the given columns with the given [`UnnestOptions`]
1550    pub fn unnest_columns_with_options(
1551        self,
1552        columns: Vec<Column>,
1553        options: UnnestOptions,
1554    ) -> Result<Self> {
1555        unnest_with_options(Arc::unwrap_or_clone(self.plan), columns, options)
1556            .map(Self::new)
1557    }
1558}
1559
1560impl From<LogicalPlan> for LogicalPlanBuilder {
1561    fn from(plan: LogicalPlan) -> Self {
1562        LogicalPlanBuilder::new(plan)
1563    }
1564}
1565
1566impl From<Arc<LogicalPlan>> for LogicalPlanBuilder {
1567    fn from(plan: Arc<LogicalPlan>) -> Self {
1568        LogicalPlanBuilder::new_from_arc(plan)
1569    }
1570}
1571
1572/// Container used when building fields for a `VALUES` node.
1573#[derive(Default)]
1574struct ValuesFields {
1575    inner: Vec<Field>,
1576}
1577
1578impl ValuesFields {
1579    pub fn new() -> Self {
1580        Self::default()
1581    }
1582
1583    pub fn push(&mut self, data_type: DataType, nullable: bool) {
1584        self.push_with_metadata(data_type, nullable, None);
1585    }
1586
1587    pub fn push_with_metadata(
1588        &mut self,
1589        data_type: DataType,
1590        nullable: bool,
1591        metadata: Option<FieldMetadata>,
1592    ) {
1593        // Naming follows the convention described here:
1594        // https://www.postgresql.org/docs/current/queries-values.html
1595        let name = format!("column{}", self.inner.len() + 1);
1596        let mut field = Field::new(name, data_type, nullable);
1597        if let Some(metadata) = metadata {
1598            field.set_metadata(metadata.to_hashmap());
1599        }
1600        self.inner.push(field);
1601    }
1602
1603    pub fn into_fields(self) -> Fields {
1604        self.inner.into()
1605    }
1606}
1607
1608/// Returns aliases to make field names unique.
1609///
1610/// Returns a vector of optional aliases, one per input field. `None` means keep the original name,
1611/// `Some(alias)` means rename to the alias to ensure uniqueness.
1612///
1613/// Used when creating [`SubqueryAlias`] or similar operations that strip table qualifiers but need
1614/// to maintain unique column names.
1615///
1616/// # Example
1617/// Input fields: `[a, a, b, b, a, a:1]` ([`DFSchema`] valid when duplicate fields have different qualifiers)
1618/// Returns: `[None, Some("a:1"), None, Some("b:1"), Some("a:2"), Some("a:1:1")]`
1619pub fn unique_field_aliases(fields: &Fields) -> Vec<Option<String>> {
1620    // Some field names might already come to this function with the count (number of times it appeared)
1621    // as a suffix e.g. id:1, so there's still a chance of name collisions, for example,
1622    // if these three fields passed to this function: "col:1", "col" and "col", the function
1623    // would rename them to -> col:1, col, col:1 causing a posterior error when building the DFSchema.
1624    // That's why we need the `seen` set, so the fields are always unique.
1625
1626    // Tracks a mapping between a field name and the number of appearances of that field.
1627    let mut name_map = HashMap::<&str, usize>::new();
1628    // Tracks all the fields and aliases that were previously seen.
1629    let mut seen = HashSet::<Cow<String>>::new();
1630
1631    fields
1632        .iter()
1633        .map(|field| {
1634            let original_name = field.name();
1635            let mut name = Cow::Borrowed(original_name);
1636
1637            let count = name_map.entry(original_name).or_insert(0);
1638
1639            // Loop until we find a name that hasn't been used.
1640            while seen.contains(&name) {
1641                *count += 1;
1642                name = Cow::Owned(format!("{original_name}:{count}"));
1643            }
1644
1645            seen.insert(name.clone());
1646
1647            match name {
1648                Cow::Borrowed(_) => None,
1649                Cow::Owned(alias) => Some(alias),
1650            }
1651        })
1652        .collect()
1653}
1654
1655fn mark_field(schema: &DFSchema) -> (Option<TableReference>, Arc<Field>) {
1656    let mut table_references = schema
1657        .iter()
1658        .filter_map(|(qualifier, _)| qualifier)
1659        .collect::<Vec<_>>();
1660    table_references.dedup();
1661    let table_reference = if table_references.len() == 1 {
1662        table_references.pop().cloned()
1663    } else {
1664        None
1665    };
1666
1667    (
1668        table_reference,
1669        Arc::new(Field::new("mark", DataType::Boolean, false)),
1670    )
1671}
1672
1673/// Creates a schema for a join operation.
1674/// The fields from the left side are first
1675pub fn build_join_schema(
1676    left: &DFSchema,
1677    right: &DFSchema,
1678    join_type: &JoinType,
1679) -> Result<DFSchema> {
1680    fn nullify_fields<'a>(
1681        fields: impl Iterator<Item = (Option<&'a TableReference>, &'a Arc<Field>)>,
1682    ) -> Vec<(Option<TableReference>, Arc<Field>)> {
1683        fields
1684            .map(|(q, f)| {
1685                // TODO: find a good way to do that
1686                let field = f.as_ref().clone().with_nullable(true);
1687                (q.cloned(), Arc::new(field))
1688            })
1689            .collect()
1690    }
1691
1692    let right_fields = right.iter();
1693    let left_fields = left.iter();
1694
1695    let qualified_fields: Vec<(Option<TableReference>, Arc<Field>)> = match join_type {
1696        JoinType::Inner => {
1697            // left then right
1698            let left_fields = left_fields
1699                .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1700                .collect::<Vec<_>>();
1701            let right_fields = right_fields
1702                .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1703                .collect::<Vec<_>>();
1704            left_fields.into_iter().chain(right_fields).collect()
1705        }
1706        JoinType::Left => {
1707            // left then right, right set to nullable in case of not matched scenario
1708            let left_fields = left_fields
1709                .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1710                .collect::<Vec<_>>();
1711            left_fields
1712                .into_iter()
1713                .chain(nullify_fields(right_fields))
1714                .collect()
1715        }
1716        JoinType::Right => {
1717            // left then right, left set to nullable in case of not matched scenario
1718            let right_fields = right_fields
1719                .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1720                .collect::<Vec<_>>();
1721            nullify_fields(left_fields)
1722                .into_iter()
1723                .chain(right_fields)
1724                .collect()
1725        }
1726        JoinType::Full => {
1727            // left then right, all set to nullable in case of not matched scenario
1728            nullify_fields(left_fields)
1729                .into_iter()
1730                .chain(nullify_fields(right_fields))
1731                .collect()
1732        }
1733        JoinType::LeftSemi | JoinType::LeftAnti => {
1734            // Only use the left side for the schema
1735            left_fields
1736                .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1737                .collect()
1738        }
1739        JoinType::LeftMark => left_fields
1740            .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1741            .chain(once(mark_field(right)))
1742            .collect(),
1743        JoinType::RightSemi | JoinType::RightAnti => {
1744            // Only use the right side for the schema
1745            right_fields
1746                .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1747                .collect()
1748        }
1749        JoinType::RightMark => right_fields
1750            .map(|(q, f)| (q.cloned(), Arc::clone(f)))
1751            .chain(once(mark_field(left)))
1752            .collect(),
1753    };
1754    let func_dependencies = left.functional_dependencies().join(
1755        right.functional_dependencies(),
1756        join_type,
1757        left.fields().len(),
1758    );
1759
1760    let (schema1, schema2) = match join_type {
1761        JoinType::Right
1762        | JoinType::RightSemi
1763        | JoinType::RightAnti
1764        | JoinType::RightMark => (left, right),
1765        _ => (right, left),
1766    };
1767
1768    let metadata = schema1
1769        .metadata()
1770        .clone()
1771        .into_iter()
1772        .chain(schema2.metadata().clone())
1773        .collect();
1774
1775    let dfschema = DFSchema::new_with_metadata(qualified_fields, metadata)?;
1776    dfschema.with_functional_dependencies(func_dependencies)
1777}
1778
1779/// (Re)qualify the sides of a join if needed, i.e. if the columns from one side would otherwise
1780/// conflict with the columns from the other.
1781/// This is especially useful for queries that come as Substrait, since Substrait doesn't currently allow specifying
1782/// aliases, neither for columns nor for tables.  DataFusion requires columns to be uniquely identifiable, in some
1783/// places (see e.g. DFSchema::check_names).
1784/// The function returns:
1785/// - The requalified or original left logical plan
1786/// - The requalified or original right logical plan
1787/// - If a requalification was needed or not
1788pub fn requalify_sides_if_needed(
1789    left: LogicalPlanBuilder,
1790    right: LogicalPlanBuilder,
1791) -> Result<(LogicalPlanBuilder, LogicalPlanBuilder, bool)> {
1792    let left_cols = left.schema().columns();
1793    let right_cols = right.schema().columns();
1794
1795    // Requalify if merging the schemas would cause an error during join.
1796    // This can happen in several cases:
1797    // 1. Duplicate qualified fields: both sides have same relation.name
1798    // 2. Duplicate unqualified fields: both sides have same unqualified name
1799    // 3. Ambiguous reference: one side qualified, other unqualified, same name
1800    //
1801    // Implementation note: This uses a simple O(n*m) nested loop rather than
1802    // a HashMap-based O(n+m) approach. The nested loop is preferred because:
1803    // - Schemas are typically small (in TPCH benchmark, max is 16 columns),
1804    //   so n*m is negligible
1805    // - Early return on first conflict makes common case very fast
1806    // - Code is simpler and easier to reason about
1807    // - Called only during plan construction, not in execution hot path
1808    for l in &left_cols {
1809        for r in &right_cols {
1810            if l.name != r.name {
1811                continue;
1812            }
1813
1814            // Same name - check if this would cause a conflict
1815            match (&l.relation, &r.relation) {
1816                // Both qualified with same relation - duplicate qualified field
1817                (Some(l_rel), Some(r_rel)) if l_rel == r_rel => {
1818                    return Ok((
1819                        left.alias(TableReference::bare("left"))?,
1820                        right.alias(TableReference::bare("right"))?,
1821                        true,
1822                    ));
1823                }
1824                // Both unqualified - duplicate unqualified field
1825                (None, None) => {
1826                    return Ok((
1827                        left.alias(TableReference::bare("left"))?,
1828                        right.alias(TableReference::bare("right"))?,
1829                        true,
1830                    ));
1831                }
1832                // One qualified, one not - ambiguous reference
1833                (Some(_), None) | (None, Some(_)) => {
1834                    return Ok((
1835                        left.alias(TableReference::bare("left"))?,
1836                        right.alias(TableReference::bare("right"))?,
1837                        true,
1838                    ));
1839                }
1840                // Different qualifiers - OK, no conflict
1841                _ => {}
1842            }
1843        }
1844    }
1845
1846    // No conflicts found
1847    Ok((left, right, false))
1848}
1849/// Add additional "synthetic" group by expressions based on functional
1850/// dependencies.
1851///
1852/// For example, if we are grouping on `[c1]`, and we know from
1853/// functional dependencies that column `c1` determines `c2`, this function
1854/// adds `c2` to the group by list.
1855///
1856/// This allows MySQL style selects like
1857/// `SELECT col FROM t WHERE pk = 5` if col is unique
1858pub fn add_group_by_exprs_from_dependencies(
1859    mut group_expr: Vec<Expr>,
1860    schema: &DFSchemaRef,
1861) -> Result<Vec<Expr>> {
1862    // Names of the fields produced by the GROUP BY exprs for example, `GROUP BY
1863    // c1 + 1` produces an output field named `"c1 + 1"`
1864    let mut group_by_field_names = group_expr
1865        .iter()
1866        .map(|e| e.schema_name().to_string())
1867        .collect::<Vec<_>>();
1868
1869    if let Some(target_indices) =
1870        get_target_functional_dependencies(schema, &group_by_field_names)
1871    {
1872        for idx in target_indices {
1873            let expr = Expr::Column(Column::from(schema.qualified_field(idx)));
1874            let expr_name = expr.schema_name().to_string();
1875            if !group_by_field_names.contains(&expr_name) {
1876                group_by_field_names.push(expr_name);
1877                group_expr.push(expr);
1878            }
1879        }
1880    }
1881    Ok(group_expr)
1882}
1883
1884/// Errors if one or more expressions have equal names.
1885pub fn validate_unique_names<'a>(
1886    node_name: &str,
1887    expressions: impl IntoIterator<Item = &'a Expr>,
1888) -> Result<()> {
1889    let mut unique_names = HashMap::new();
1890
1891    expressions.into_iter().enumerate().try_for_each(|(position, expr)| {
1892        let name = expr.schema_name().to_string();
1893        match unique_names.get(&name) {
1894            None => {
1895                unique_names.insert(name, (position, expr));
1896                Ok(())
1897            },
1898            Some((existing_position, existing_expr)) => {
1899                plan_err!("{node_name} require unique expression names \
1900                             but the expression \"{existing_expr}\" at position {existing_position} and \"{expr}\" \
1901                             at position {position} have the same name. Consider aliasing (\"AS\") one of them."
1902                            )
1903            }
1904        }
1905    })
1906}
1907
1908/// Union two [`LogicalPlan`]s.
1909///
1910/// Constructs the UNION plan, but does not perform type-coercion. Therefore the
1911/// subtree expressions will not be properly typed until the optimizer pass.
1912///
1913/// If a properly typed UNION plan is needed, refer to [`TypeCoercionRewriter::coerce_union`]
1914/// or alternatively, merge the union input schema using [`coerce_union_schema`] and
1915/// apply the expression rewrite with [`coerce_plan_expr_for_schema`].
1916///
1917/// [`TypeCoercionRewriter::coerce_union`]: https://docs.rs/datafusion-optimizer/latest/datafusion_optimizer/analyzer/type_coercion/struct.TypeCoercionRewriter.html#method.coerce_union
1918/// [`coerce_union_schema`]: https://docs.rs/datafusion-optimizer/latest/datafusion_optimizer/analyzer/type_coercion/fn.coerce_union_schema.html
1919pub fn union(left_plan: LogicalPlan, right_plan: LogicalPlan) -> Result<LogicalPlan> {
1920    Ok(LogicalPlan::Union(Union::try_new_with_loose_types(vec![
1921        Arc::new(left_plan),
1922        Arc::new(right_plan),
1923    ])?))
1924}
1925
1926/// Like [`union`], but combine rows from different tables by name, rather than
1927/// by position.
1928pub fn union_by_name(
1929    left_plan: LogicalPlan,
1930    right_plan: LogicalPlan,
1931) -> Result<LogicalPlan> {
1932    Ok(LogicalPlan::Union(Union::try_new_by_name(vec![
1933        Arc::new(left_plan),
1934        Arc::new(right_plan),
1935    ])?))
1936}
1937
1938/// Create Projection
1939/// # Errors
1940/// This function errors under any of the following conditions:
1941/// * Two or more expressions have the same name
1942/// * An invalid expression is used (e.g. a `sort` expression)
1943pub fn project(
1944    plan: LogicalPlan,
1945    expr: impl IntoIterator<Item = impl Into<SelectExpr>>,
1946) -> Result<LogicalPlan> {
1947    project_with_validation(plan, expr.into_iter().map(|e| (e, true)), None)
1948}
1949
1950/// Create Projection. Similar to project except that the expressions
1951/// passed in have a flag to indicate if that expression requires
1952/// validation (normalize & columnize) (true) or not (false)
1953/// # Errors
1954/// This function errors under any of the following conditions:
1955/// * Two or more expressions have the same name
1956/// * An invalid expression is used (e.g. a `sort` expression)
1957fn project_with_validation(
1958    plan: LogicalPlan,
1959    expr: impl IntoIterator<Item = (impl Into<SelectExpr>, bool)>,
1960    schema: Option<&DFSchemaRef>,
1961) -> Result<LogicalPlan> {
1962    let mut projected_expr = vec![];
1963    let mut has_wildcard = false;
1964    for (e, validate) in expr {
1965        let e = e.into();
1966        match e {
1967            SelectExpr::Wildcard(opt) => {
1968                has_wildcard = true;
1969                let expanded = expand_wildcard(plan.schema(), &plan, Some(&opt))?;
1970
1971                // If there is a REPLACE statement, replace that column with the given
1972                // replace expression. Column name remains the same.
1973                let expanded = if let Some(replace) = opt.replace {
1974                    replace_columns(expanded, &replace)?
1975                } else {
1976                    expanded
1977                };
1978
1979                for e in expanded {
1980                    if validate {
1981                        projected_expr
1982                            .push(columnize_expr(normalize_col(e, &plan)?, &plan)?)
1983                    } else {
1984                        projected_expr.push(e)
1985                    }
1986                }
1987            }
1988            SelectExpr::QualifiedWildcard(table_ref, opt) => {
1989                has_wildcard = true;
1990                let expanded =
1991                    expand_qualified_wildcard(&table_ref, plan.schema(), Some(&opt))?;
1992
1993                // If there is a REPLACE statement, replace that column with the given
1994                // replace expression. Column name remains the same.
1995                let expanded = if let Some(replace) = opt.replace {
1996                    replace_columns(expanded, &replace)?
1997                } else {
1998                    expanded
1999                };
2000
2001                for e in expanded {
2002                    if validate {
2003                        projected_expr
2004                            .push(columnize_expr(normalize_col(e, &plan)?, &plan)?)
2005                    } else {
2006                        projected_expr.push(e)
2007                    }
2008                }
2009            }
2010            SelectExpr::Expression(e) => {
2011                if validate {
2012                    projected_expr.push(columnize_expr(normalize_col(e, &plan)?, &plan)?)
2013                } else {
2014                    projected_expr.push(e)
2015                }
2016            }
2017        }
2018    }
2019
2020    if has_wildcard && projected_expr.is_empty() && !plan.schema().fields().is_empty() {
2021        return plan_err!(
2022            "SELECT list is empty after resolving * expressions, \
2023             the wildcard expanded to zero columns"
2024        );
2025    }
2026
2027    // When inside a set expression, alias non-Column/non-Alias expressions
2028    // to match the left side's field names, avoiding duplicate name errors.
2029    if let Some(schema) = &schema {
2030        for (expr, field) in projected_expr.iter_mut().zip(schema.fields()) {
2031            if !matches!(expr, Expr::Column(_) | Expr::Alias(_)) {
2032                *expr = std::mem::take(expr).alias(field.name());
2033            }
2034        }
2035    }
2036
2037    validate_unique_names("Projections", projected_expr.iter())?;
2038
2039    Projection::try_new(projected_expr, Arc::new(plan)).map(LogicalPlan::Projection)
2040}
2041
2042/// If there is a REPLACE statement in the projected expression in the form of
2043/// "REPLACE (some_column_within_an_expr AS some_column)", this function replaces
2044/// that column with the given replace expression. Column name remains the same.
2045/// Multiple REPLACEs are also possible with comma separations.
2046fn replace_columns(
2047    mut exprs: Vec<Expr>,
2048    replace: &PlannedReplaceSelectItem,
2049) -> Result<Vec<Expr>> {
2050    for expr in exprs.iter_mut() {
2051        if let Expr::Column(Column { name, .. }) = expr
2052            && let Some((_, new_expr)) = replace
2053                .items()
2054                .iter()
2055                .zip(replace.expressions().iter())
2056                .find(|(item, _)| item.column_name.value == *name)
2057        {
2058            *expr = new_expr.clone().alias(name.clone())
2059        }
2060    }
2061    Ok(exprs)
2062}
2063
2064/// Create a SubqueryAlias to wrap a LogicalPlan.
2065pub fn subquery_alias(
2066    plan: LogicalPlan,
2067    alias: impl Into<TableReference>,
2068) -> Result<LogicalPlan> {
2069    SubqueryAlias::try_new(Arc::new(plan), alias).map(LogicalPlan::SubqueryAlias)
2070}
2071
2072/// Create a LogicalPlanBuilder representing a scan of a table with the provided name and schema.
2073/// This is mostly used for testing and documentation.
2074pub fn table_scan(
2075    name: Option<impl Into<TableReference>>,
2076    table_schema: &Schema,
2077    projection: Option<Vec<usize>>,
2078) -> Result<LogicalPlanBuilder> {
2079    table_scan_with_filters(name, table_schema, projection, vec![])
2080}
2081
2082/// Create a LogicalPlanBuilder representing a scan of a table with the provided name and schema,
2083/// and inlined filters.
2084/// This is mostly used for testing and documentation.
2085pub fn table_scan_with_filters(
2086    name: Option<impl Into<TableReference>>,
2087    table_schema: &Schema,
2088    projection: Option<Vec<usize>>,
2089    filters: Vec<Expr>,
2090) -> Result<LogicalPlanBuilder> {
2091    let table_source = table_source(table_schema);
2092    let name = name
2093        .map(|n| n.into())
2094        .unwrap_or_else(|| TableReference::bare(UNNAMED_TABLE));
2095    LogicalPlanBuilder::scan_with_filters(name, table_source, projection, filters)
2096}
2097
2098/// Create a LogicalPlanBuilder representing a scan of a table with the provided name and schema,
2099/// filters, and inlined fetch.
2100/// This is mostly used for testing and documentation.
2101pub fn table_scan_with_filter_and_fetch(
2102    name: Option<impl Into<TableReference>>,
2103    table_schema: &Schema,
2104    projection: Option<Vec<usize>>,
2105    filters: Vec<Expr>,
2106    fetch: Option<usize>,
2107) -> Result<LogicalPlanBuilder> {
2108    let table_source = table_source(table_schema);
2109    let name = name
2110        .map(|n| n.into())
2111        .unwrap_or_else(|| TableReference::bare(UNNAMED_TABLE));
2112    LogicalPlanBuilder::scan_with_filters_fetch(
2113        name,
2114        table_source,
2115        projection,
2116        filters,
2117        fetch,
2118    )
2119}
2120
2121pub fn table_source(table_schema: &Schema) -> Arc<dyn TableSource> {
2122    // TODO should we take SchemaRef and avoid cloning?
2123    let table_schema = Arc::new(table_schema.clone());
2124    Arc::new(LogicalTableSource {
2125        table_schema,
2126        constraints: Default::default(),
2127    })
2128}
2129
2130pub fn table_source_with_constraints(
2131    table_schema: &Schema,
2132    constraints: Constraints,
2133) -> Arc<dyn TableSource> {
2134    // TODO should we take SchemaRef and avoid cloning?
2135    let table_schema = Arc::new(table_schema.clone());
2136    Arc::new(LogicalTableSource {
2137        table_schema,
2138        constraints,
2139    })
2140}
2141
2142/// Wrap projection for a plan, if the join keys contains normal expression.
2143pub fn wrap_projection_for_join_if_necessary(
2144    join_keys: &[Expr],
2145    input: LogicalPlan,
2146) -> Result<(LogicalPlan, Vec<Column>, bool)> {
2147    let input_schema = input.schema();
2148    let alias_join_keys: Vec<Expr> = join_keys
2149        .iter()
2150        .map(|key| {
2151            // The display_name() of cast expression will ignore the cast info, and show the inner expression name.
2152            // If we do not add alias, it will throw same field name error in the schema when adding projection.
2153            // For example:
2154            //    input scan : [a, b, c],
2155            //    join keys: [cast(a as int)]
2156            //
2157            //  then a and cast(a as int) will use the same field name - `a` in projection schema.
2158            //  https://github.com/apache/datafusion/issues/4478
2159            if matches!(key, Expr::Cast(_)) || matches!(key, Expr::TryCast(_)) {
2160                let alias = format!("{key}");
2161                key.clone().alias(alias)
2162            } else {
2163                key.clone()
2164            }
2165        })
2166        .collect::<Vec<_>>();
2167
2168    let need_project = join_keys.iter().any(|key| !matches!(key, Expr::Column(_)));
2169    let plan = if need_project {
2170        // Include all columns from the input and extend them with the join keys
2171        let mut projection = input_schema
2172            .columns()
2173            .into_iter()
2174            .map(Expr::Column)
2175            .collect::<Vec<_>>();
2176        #[allow(clippy::allow_attributes, clippy::mutable_key_type)]
2177        // Expr contains Arc with interior mutability but is intentionally used as hash key
2178        let join_key_items = alias_join_keys
2179            .iter()
2180            .flat_map(|expr| expr.try_as_col().is_none().then_some(expr))
2181            .cloned()
2182            .collect::<HashSet<Expr>>();
2183        projection.extend(join_key_items);
2184
2185        LogicalPlanBuilder::from(input)
2186            .project(projection.into_iter().map(SelectExpr::from))?
2187            .build()?
2188    } else {
2189        input
2190    };
2191
2192    let join_on = alias_join_keys
2193        .into_iter()
2194        .map(|key| {
2195            if let Some(col) = key.try_as_col() {
2196                Ok(col.clone())
2197            } else {
2198                let name = key.schema_name().to_string();
2199                Ok(Column::from_name(name))
2200            }
2201        })
2202        .collect::<Result<Vec<_>>>()?;
2203
2204    Ok((plan, join_on, need_project))
2205}
2206
2207/// Basic TableSource implementation intended for use in tests and documentation. It is expected
2208/// that users will provide their own TableSource implementations or use DataFusion's
2209/// DefaultTableSource.
2210pub struct LogicalTableSource {
2211    table_schema: SchemaRef,
2212    constraints: Constraints,
2213}
2214
2215impl LogicalTableSource {
2216    /// Create a new LogicalTableSource
2217    pub fn new(table_schema: SchemaRef) -> Self {
2218        Self {
2219            table_schema,
2220            constraints: Constraints::default(),
2221        }
2222    }
2223
2224    pub fn with_constraints(mut self, constraints: Constraints) -> Self {
2225        self.constraints = constraints;
2226        self
2227    }
2228}
2229
2230impl TableSource for LogicalTableSource {
2231    fn schema(&self) -> SchemaRef {
2232        Arc::clone(&self.table_schema)
2233    }
2234
2235    fn constraints(&self) -> Option<&Constraints> {
2236        Some(&self.constraints)
2237    }
2238
2239    fn supports_filters_pushdown(
2240        &self,
2241        filters: &[&Expr],
2242    ) -> Result<Vec<TableProviderFilterPushDown>> {
2243        Ok(vec![TableProviderFilterPushDown::Exact; filters.len()])
2244    }
2245}
2246
2247/// Create a [`LogicalPlan::Unnest`] plan
2248pub fn unnest(input: LogicalPlan, columns: Vec<Column>) -> Result<LogicalPlan> {
2249    unnest_with_options(input, columns, UnnestOptions::default())
2250}
2251
2252pub fn get_struct_unnested_columns(
2253    col_name: &String,
2254    inner_fields: &Fields,
2255) -> Vec<Column> {
2256    inner_fields
2257        .iter()
2258        .map(|f| Column::from_name(format!("{}.{}", col_name, f.name())))
2259        .collect()
2260}
2261
2262/// Create a [`LogicalPlan::Unnest`] plan with options
2263/// This function receive a list of columns to be unnested
2264/// because multiple unnest can be performed on the same column (e.g unnest with different depth)
2265/// The new schema will contains post-unnest fields replacing the original field
2266///
2267/// For example:
2268/// Input schema as
2269/// ```text
2270/// +---------------------+-------------------+
2271/// | col1                | col2              |
2272/// +---------------------+-------------------+
2273/// | Struct(INT64,INT32) | List(List(Int64)) |
2274/// +---------------------+-------------------+
2275/// ```
2276///
2277///
2278///
2279/// Then unnesting columns with:
2280/// - (col1,Struct)
2281/// - (col2,List(\[depth=1,depth=2\]))
2282///
2283/// will generate a new schema as
2284/// ```text
2285/// +---------+---------+---------------------+---------------------+
2286/// | col1.c0 | col1.c1 | unnest_col2_depth_1 | unnest_col2_depth_2 |
2287/// +---------+---------+---------------------+---------------------+
2288/// | Int64   | Int32   | List(Int64)         |  Int64              |
2289/// +---------+---------+---------------------+---------------------+
2290/// ```
2291pub fn unnest_with_options(
2292    input: LogicalPlan,
2293    columns_to_unnest: Vec<Column>,
2294    options: UnnestOptions,
2295) -> Result<LogicalPlan> {
2296    Ok(LogicalPlan::Unnest(Unnest::try_new(
2297        Arc::new(input),
2298        columns_to_unnest,
2299        options,
2300    )?))
2301}
2302
2303#[cfg(test)]
2304mod tests {
2305    use std::vec;
2306
2307    use super::*;
2308    use crate::lit_with_metadata;
2309    use crate::logical_plan::StringifiedPlan;
2310    use crate::{col, expr, expr_fn::exists, in_subquery, scalar_subquery};
2311
2312    use crate::test::function_stub::sum;
2313    use datafusion_common::{
2314        Constraint, DataFusionError, RecursionUnnestOption, SchemaError,
2315    };
2316    use insta::assert_snapshot;
2317
2318    #[test]
2319    fn plan_builder_simple() -> Result<()> {
2320        let plan =
2321            table_scan(Some("employee_csv"), &employee_schema(), Some(vec![0, 3]))?
2322                .filter(col("state").eq(lit("CO")))?
2323                .project(vec![col("id")])?
2324                .build()?;
2325
2326        assert_snapshot!(plan, @r#"
2327        Projection: employee_csv.id
2328          Filter: employee_csv.state = Utf8("CO")
2329            TableScan: employee_csv projection=[id, state]
2330        "#);
2331
2332        Ok(())
2333    }
2334
2335    #[test]
2336    fn plan_builder_schema() {
2337        let schema = employee_schema();
2338        let projection = None;
2339        let plan =
2340            LogicalPlanBuilder::scan("employee_csv", table_source(&schema), projection)
2341                .unwrap();
2342        assert_snapshot!(plan.schema().as_ref(), @"fields:[employee_csv.id, employee_csv.first_name, employee_csv.last_name, employee_csv.state, employee_csv.salary], metadata:{}");
2343
2344        // Note scan of "EMPLOYEE_CSV" is treated as a SQL identifier
2345        // (and thus normalized to "employee"csv") as well
2346        let projection = None;
2347        let plan =
2348            LogicalPlanBuilder::scan("EMPLOYEE_CSV", table_source(&schema), projection)
2349                .unwrap();
2350        assert_snapshot!(plan.schema().as_ref(), @"fields:[employee_csv.id, employee_csv.first_name, employee_csv.last_name, employee_csv.state, employee_csv.salary], metadata:{}");
2351    }
2352
2353    #[test]
2354    fn plan_builder_empty_name() {
2355        let schema = employee_schema();
2356        let projection = None;
2357        let err =
2358            LogicalPlanBuilder::scan("", table_source(&schema), projection).unwrap_err();
2359        assert_snapshot!(
2360            err.strip_backtrace(),
2361            @"Error during planning: table_name cannot be empty"
2362        );
2363    }
2364
2365    #[test]
2366    fn plan_builder_sort() -> Result<()> {
2367        let plan =
2368            table_scan(Some("employee_csv"), &employee_schema(), Some(vec![3, 4]))?
2369                .sort(vec![
2370                    expr::Sort::new(col("state"), true, true),
2371                    expr::Sort::new(col("salary"), false, false),
2372                ])?
2373                .build()?;
2374
2375        assert_snapshot!(plan, @r"
2376        Sort: employee_csv.state ASC NULLS FIRST, employee_csv.salary DESC NULLS LAST
2377          TableScan: employee_csv projection=[state, salary]
2378        ");
2379
2380        Ok(())
2381    }
2382
2383    #[test]
2384    fn plan_builder_union() -> Result<()> {
2385        let plan =
2386            table_scan(Some("employee_csv"), &employee_schema(), Some(vec![3, 4]))?;
2387
2388        let plan = plan
2389            .clone()
2390            .union(plan.clone().build()?)?
2391            .union(plan.clone().build()?)?
2392            .union(plan.build()?)?
2393            .build()?;
2394
2395        assert_snapshot!(plan, @r"
2396        Union
2397          Union
2398            Union
2399              TableScan: employee_csv projection=[state, salary]
2400              TableScan: employee_csv projection=[state, salary]
2401            TableScan: employee_csv projection=[state, salary]
2402          TableScan: employee_csv projection=[state, salary]
2403        ");
2404
2405        Ok(())
2406    }
2407
2408    #[test]
2409    fn plan_builder_union_distinct() -> Result<()> {
2410        let plan =
2411            table_scan(Some("employee_csv"), &employee_schema(), Some(vec![3, 4]))?;
2412
2413        let plan = plan
2414            .clone()
2415            .union_distinct(plan.clone().build()?)?
2416            .union_distinct(plan.clone().build()?)?
2417            .union_distinct(plan.build()?)?
2418            .build()?;
2419
2420        assert_snapshot!(plan, @r"
2421        Distinct:
2422          Union
2423            Distinct:
2424              Union
2425                Distinct:
2426                  Union
2427                    TableScan: employee_csv projection=[state, salary]
2428                    TableScan: employee_csv projection=[state, salary]
2429                TableScan: employee_csv projection=[state, salary]
2430            TableScan: employee_csv projection=[state, salary]
2431        ");
2432
2433        Ok(())
2434    }
2435
2436    #[test]
2437    fn plan_builder_simple_distinct() -> Result<()> {
2438        let plan =
2439            table_scan(Some("employee_csv"), &employee_schema(), Some(vec![0, 3]))?
2440                .filter(col("state").eq(lit("CO")))?
2441                .project(vec![col("id")])?
2442                .distinct()?
2443                .build()?;
2444
2445        assert_snapshot!(plan, @r#"
2446        Distinct:
2447          Projection: employee_csv.id
2448            Filter: employee_csv.state = Utf8("CO")
2449              TableScan: employee_csv projection=[id, state]
2450        "#);
2451
2452        Ok(())
2453    }
2454
2455    #[test]
2456    fn exists_subquery() -> Result<()> {
2457        let foo = test_table_scan_with_name("foo")?;
2458        let bar = test_table_scan_with_name("bar")?;
2459
2460        let subquery = LogicalPlanBuilder::from(foo)
2461            .project(vec![col("a")])?
2462            .filter(col("a").eq(col("bar.a")))?
2463            .build()?;
2464
2465        let outer_query = LogicalPlanBuilder::from(bar)
2466            .project(vec![col("a")])?
2467            .filter(exists(Arc::new(subquery)))?
2468            .build()?;
2469
2470        assert_snapshot!(outer_query, @r"
2471        Filter: EXISTS (<subquery>)
2472          Subquery:
2473            Filter: foo.a = bar.a
2474              Projection: foo.a
2475                TableScan: foo
2476          Projection: bar.a
2477            TableScan: bar
2478        ");
2479
2480        Ok(())
2481    }
2482
2483    #[test]
2484    fn filter_in_subquery() -> Result<()> {
2485        let foo = test_table_scan_with_name("foo")?;
2486        let bar = test_table_scan_with_name("bar")?;
2487
2488        let subquery = LogicalPlanBuilder::from(foo)
2489            .project(vec![col("a")])?
2490            .filter(col("a").eq(col("bar.a")))?
2491            .build()?;
2492
2493        // SELECT a FROM bar WHERE a IN (SELECT a FROM foo WHERE a = bar.a)
2494        let outer_query = LogicalPlanBuilder::from(bar)
2495            .project(vec![col("a")])?
2496            .filter(in_subquery(col("a"), Arc::new(subquery)))?
2497            .build()?;
2498
2499        assert_snapshot!(outer_query, @r"
2500        Filter: bar.a IN (<subquery>)
2501          Subquery:
2502            Filter: foo.a = bar.a
2503              Projection: foo.a
2504                TableScan: foo
2505          Projection: bar.a
2506            TableScan: bar
2507        ");
2508
2509        Ok(())
2510    }
2511
2512    #[test]
2513    fn select_scalar_subquery() -> Result<()> {
2514        let foo = test_table_scan_with_name("foo")?;
2515        let bar = test_table_scan_with_name("bar")?;
2516
2517        let subquery = LogicalPlanBuilder::from(foo)
2518            .project(vec![col("b")])?
2519            .filter(col("a").eq(col("bar.a")))?
2520            .build()?;
2521
2522        // SELECT (SELECT a FROM foo WHERE a = bar.a) FROM bar
2523        let outer_query = LogicalPlanBuilder::from(bar)
2524            .project(vec![scalar_subquery(Arc::new(subquery))])?
2525            .build()?;
2526
2527        assert_snapshot!(outer_query, @r"
2528        Projection: (<subquery>)
2529          Subquery:
2530            Filter: foo.a = bar.a
2531              Projection: foo.b
2532                TableScan: foo
2533          TableScan: bar
2534        ");
2535
2536        Ok(())
2537    }
2538
2539    #[test]
2540    fn projection_non_unique_names() -> Result<()> {
2541        let plan = table_scan(
2542            Some("employee_csv"),
2543            &employee_schema(),
2544            // project id and first_name by column index
2545            Some(vec![0, 1]),
2546        )?
2547        // two columns with the same name => error
2548        .project(vec![col("id"), col("first_name").alias("id")]);
2549
2550        match plan {
2551            Err(DataFusionError::SchemaError(err, _)) => {
2552                if let SchemaError::AmbiguousReference { field } = *err {
2553                    let Column {
2554                        relation,
2555                        name,
2556                        spans: _,
2557                    } = *field;
2558                    let Some(TableReference::Bare { table }) = relation else {
2559                        return plan_err!(
2560                            "wrong relation: {relation:?}, expected table name"
2561                        );
2562                    };
2563                    assert_eq!(*"employee_csv", *table);
2564                    assert_eq!("id", &name);
2565                    Ok(())
2566                } else {
2567                    plan_err!("Plan should have returned an DataFusionError::SchemaError")
2568                }
2569            }
2570            _ => plan_err!("Plan should have returned an DataFusionError::SchemaError"),
2571        }
2572    }
2573
2574    fn employee_schema() -> Schema {
2575        Schema::new(vec![
2576            Field::new("id", DataType::Int32, false),
2577            Field::new("first_name", DataType::Utf8, false),
2578            Field::new("last_name", DataType::Utf8, false),
2579            Field::new("state", DataType::Utf8, false),
2580            Field::new("salary", DataType::Int32, false),
2581        ])
2582    }
2583
2584    #[test]
2585    fn stringified_plan() {
2586        let stringified_plan =
2587            StringifiedPlan::new(PlanType::InitialLogicalPlan, "...the plan...");
2588        assert!(stringified_plan.should_display(true));
2589        assert!(!stringified_plan.should_display(false)); // not in non verbose mode
2590
2591        let stringified_plan =
2592            StringifiedPlan::new(PlanType::FinalLogicalPlan, "...the plan...");
2593        assert!(stringified_plan.should_display(true));
2594        assert!(stringified_plan.should_display(false)); // display in non verbose mode too
2595
2596        let stringified_plan =
2597            StringifiedPlan::new(PlanType::InitialPhysicalPlan, "...the plan...");
2598        assert!(stringified_plan.should_display(true));
2599        assert!(!stringified_plan.should_display(false)); // not in non verbose mode
2600
2601        let stringified_plan =
2602            StringifiedPlan::new(PlanType::FinalPhysicalPlan, "...the plan...");
2603        assert!(stringified_plan.should_display(true));
2604        assert!(stringified_plan.should_display(false)); // display in non verbose mode
2605
2606        let stringified_plan = StringifiedPlan::new(
2607            PlanType::OptimizedLogicalPlan {
2608                optimizer_name: "random opt pass".into(),
2609            },
2610            "...the plan...",
2611        );
2612        assert!(stringified_plan.should_display(true));
2613        assert!(!stringified_plan.should_display(false));
2614    }
2615
2616    fn test_table_scan_with_name(name: &str) -> Result<LogicalPlan> {
2617        let schema = Schema::new(vec![
2618            Field::new("a", DataType::UInt32, false),
2619            Field::new("b", DataType::UInt32, false),
2620            Field::new("c", DataType::UInt32, false),
2621        ]);
2622        table_scan(Some(name), &schema, None)?.build()
2623    }
2624
2625    #[test]
2626    fn plan_builder_intersect_different_num_columns_error() -> Result<()> {
2627        let plan1 =
2628            table_scan(TableReference::none(), &employee_schema(), Some(vec![3]))?;
2629        let plan2 =
2630            table_scan(TableReference::none(), &employee_schema(), Some(vec![3, 4]))?;
2631
2632        let err_msg1 =
2633            LogicalPlanBuilder::intersect(plan1.build()?, plan2.build()?, true)
2634                .unwrap_err();
2635
2636        assert_snapshot!(err_msg1.strip_backtrace(), @"Error during planning: INTERSECT/EXCEPT query must have the same number of columns. Left is 1 and right is 2.");
2637
2638        Ok(())
2639    }
2640
2641    #[test]
2642    fn plan_builder_unnest() -> Result<()> {
2643        // Cannot unnest on a scalar column
2644        let err = nested_table_scan("test_table")?
2645            .unnest_column("scalar")
2646            .unwrap_err();
2647
2648        let DataFusionError::Internal(desc) = err else {
2649            return plan_err!("Plan should have returned an DataFusionError::Internal");
2650        };
2651
2652        let desc = (*desc
2653            .split(DataFusionError::BACK_TRACE_SEP)
2654            .collect::<Vec<&str>>()
2655            .first()
2656            .unwrap_or(&""))
2657        .to_string();
2658
2659        assert_snapshot!(desc, @"trying to unnest on invalid data type UInt32");
2660
2661        // Unnesting the strings list.
2662        let plan = nested_table_scan("test_table")?
2663            .unnest_column("strings")?
2664            .build()?;
2665
2666        assert_snapshot!(plan, @r"
2667        Unnest: lists[test_table.strings|depth=1] structs[]
2668          TableScan: test_table
2669        ");
2670
2671        // Check unnested field is a scalar
2672        let field = plan.schema().field_with_name(None, "strings").unwrap();
2673        assert_eq!(&DataType::Utf8, field.data_type());
2674
2675        // Unnesting the singular struct column result into 2 new columns for each subfield
2676        let plan = nested_table_scan("test_table")?
2677            .unnest_column("struct_singular")?
2678            .build()?;
2679
2680        assert_snapshot!(plan, @r"
2681        Unnest: lists[] structs[test_table.struct_singular]
2682          TableScan: test_table
2683        ");
2684
2685        for field_name in &["a", "b"] {
2686            // Check unnested struct field is a scalar
2687            let field = plan
2688                .schema()
2689                .field_with_name(None, &format!("struct_singular.{field_name}"))
2690                .unwrap();
2691            assert_eq!(&DataType::UInt32, field.data_type());
2692        }
2693
2694        // Unnesting multiple fields in separate plans
2695        let plan = nested_table_scan("test_table")?
2696            .unnest_column("strings")?
2697            .unnest_column("structs")?
2698            .unnest_column("struct_singular")?
2699            .build()?;
2700
2701        assert_snapshot!(plan, @r"
2702        Unnest: lists[] structs[test_table.struct_singular]
2703          Unnest: lists[test_table.structs|depth=1] structs[]
2704            Unnest: lists[test_table.strings|depth=1] structs[]
2705              TableScan: test_table
2706        ");
2707
2708        // Check unnested struct list field should be a struct.
2709        let field = plan.schema().field_with_name(None, "structs").unwrap();
2710        assert!(matches!(field.data_type(), DataType::Struct(_)));
2711
2712        // Unnesting multiple fields at the same time, using infer syntax
2713        let cols = vec!["strings", "structs", "struct_singular"]
2714            .into_iter()
2715            .map(|c| c.into())
2716            .collect();
2717
2718        let plan = nested_table_scan("test_table")?
2719            .unnest_columns_with_options(cols, UnnestOptions::default())?
2720            .build()?;
2721
2722        assert_snapshot!(plan, @r"
2723        Unnest: lists[test_table.strings|depth=1, test_table.structs|depth=1] structs[test_table.struct_singular]
2724          TableScan: test_table
2725        ");
2726
2727        // Unnesting missing column should fail.
2728        let plan = nested_table_scan("test_table")?.unnest_column("missing");
2729        assert!(plan.is_err());
2730
2731        // Simultaneously unnesting a list (with different depth) and a struct column
2732        let plan = nested_table_scan("test_table")?
2733            .unnest_columns_with_options(
2734                vec!["stringss".into(), "struct_singular".into()],
2735                UnnestOptions::default()
2736                    .with_recursions(RecursionUnnestOption {
2737                        input_column: "stringss".into(),
2738                        output_column: "stringss_depth_1".into(),
2739                        depth: 1,
2740                    })
2741                    .with_recursions(RecursionUnnestOption {
2742                        input_column: "stringss".into(),
2743                        output_column: "stringss_depth_2".into(),
2744                        depth: 2,
2745                    }),
2746            )?
2747            .build()?;
2748
2749        assert_snapshot!(plan, @r"
2750        Unnest: lists[test_table.stringss|depth=1, test_table.stringss|depth=2] structs[test_table.struct_singular]
2751          TableScan: test_table
2752        ");
2753
2754        // Check output columns has correct type
2755        let field = plan
2756            .schema()
2757            .field_with_name(None, "stringss_depth_1")
2758            .unwrap();
2759        assert_eq!(
2760            &DataType::new_list(DataType::Utf8, false),
2761            field.data_type()
2762        );
2763        let field = plan
2764            .schema()
2765            .field_with_name(None, "stringss_depth_2")
2766            .unwrap();
2767        assert_eq!(&DataType::Utf8, field.data_type());
2768        // unnesting struct is still correct
2769        for field_name in &["a", "b"] {
2770            let field = plan
2771                .schema()
2772                .field_with_name(None, &format!("struct_singular.{field_name}"))
2773                .unwrap();
2774            assert_eq!(&DataType::UInt32, field.data_type());
2775        }
2776
2777        Ok(())
2778    }
2779
2780    fn nested_table_scan(table_name: &str) -> Result<LogicalPlanBuilder> {
2781        // Create a schema with a scalar field, a list of strings, a list of structs
2782        // and a singular struct
2783        let struct_field_in_list = Field::new_struct(
2784            "item",
2785            vec![
2786                Field::new("a", DataType::UInt32, false),
2787                Field::new("b", DataType::UInt32, false),
2788            ],
2789            false,
2790        );
2791        let string_field = Field::new_list_field(DataType::Utf8, false);
2792        let strings_field = Field::new_list("item", string_field.clone(), false);
2793        let schema = Schema::new(vec![
2794            Field::new("scalar", DataType::UInt32, false),
2795            Field::new_list("strings", string_field, false),
2796            Field::new_list("structs", struct_field_in_list, false),
2797            Field::new(
2798                "struct_singular",
2799                DataType::Struct(Fields::from(vec![
2800                    Field::new("a", DataType::UInt32, false),
2801                    Field::new("b", DataType::UInt32, false),
2802                ])),
2803                false,
2804            ),
2805            Field::new_list("stringss", strings_field, false),
2806        ]);
2807
2808        table_scan(Some(table_name), &schema, None)
2809    }
2810
2811    #[test]
2812    fn test_union_after_join() -> Result<()> {
2813        let values = vec![vec![lit(1)]];
2814
2815        let left = LogicalPlanBuilder::values(values.clone())?
2816            .alias("left")?
2817            .build()?;
2818        let right = LogicalPlanBuilder::values(values)?
2819            .alias("right")?
2820            .build()?;
2821
2822        let join = LogicalPlanBuilder::from(left).cross_join(right)?.build()?;
2823
2824        let plan = LogicalPlanBuilder::from(join.clone())
2825            .union(join)?
2826            .build()?;
2827
2828        assert_snapshot!(plan, @r"
2829        Union
2830          Cross Join:
2831            SubqueryAlias: left
2832              Values: (Int32(1))
2833            SubqueryAlias: right
2834              Values: (Int32(1))
2835          Cross Join:
2836            SubqueryAlias: left
2837              Values: (Int32(1))
2838            SubqueryAlias: right
2839              Values: (Int32(1))
2840        ");
2841
2842        Ok(())
2843    }
2844
2845    #[test]
2846    fn plan_builder_from_logical_plan() -> Result<()> {
2847        let plan =
2848            table_scan(Some("employee_csv"), &employee_schema(), Some(vec![3, 4]))?
2849                .sort(vec![
2850                    expr::Sort::new(col("state"), true, true),
2851                    expr::Sort::new(col("salary"), false, false),
2852                ])?
2853                .build()?;
2854
2855        let plan_expected = format!("{plan}");
2856        let plan_builder: LogicalPlanBuilder = Arc::new(plan).into();
2857        assert_eq!(plan_expected, format!("{}", plan_builder.plan));
2858
2859        Ok(())
2860    }
2861
2862    #[test]
2863    fn plan_builder_aggregate_without_implicit_group_by_exprs() -> Result<()> {
2864        let constraints =
2865            Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0])]);
2866        let table_source = table_source_with_constraints(&employee_schema(), constraints);
2867
2868        let plan =
2869            LogicalPlanBuilder::scan("employee_csv", table_source, Some(vec![0, 3, 4]))?
2870                .aggregate(vec![col("id")], vec![sum(col("salary"))])?
2871                .build()?;
2872
2873        assert_snapshot!(plan, @r"
2874        Aggregate: groupBy=[[employee_csv.id]], aggr=[[sum(employee_csv.salary)]]
2875          TableScan: employee_csv projection=[id, state, salary]
2876        ");
2877
2878        Ok(())
2879    }
2880
2881    #[test]
2882    fn plan_builder_aggregate_with_implicit_group_by_exprs() -> Result<()> {
2883        let constraints =
2884            Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0])]);
2885        let table_source = table_source_with_constraints(&employee_schema(), constraints);
2886
2887        let options =
2888            LogicalPlanBuilderOptions::new().with_add_implicit_group_by_exprs(true);
2889        let plan =
2890            LogicalPlanBuilder::scan("employee_csv", table_source, Some(vec![0, 3, 4]))?
2891                .with_options(options)
2892                .aggregate(vec![col("id")], vec![sum(col("salary"))])?
2893                .build()?;
2894
2895        assert_snapshot!(plan, @r"
2896        Aggregate: groupBy=[[employee_csv.id, employee_csv.state, employee_csv.salary]], aggr=[[sum(employee_csv.salary)]]
2897          TableScan: employee_csv projection=[id, state, salary]
2898        ");
2899
2900        Ok(())
2901    }
2902
2903    #[test]
2904    fn plan_builder_aggregate_rejects_nested_aggregates() -> Result<()> {
2905        // https://github.com/apache/datafusion/issues/23812
2906        let err = table_scan(
2907            Some("employee_csv"),
2908            &employee_schema(),
2909            Some(vec![0, 3, 4]),
2910        )?
2911        .aggregate(vec![col("id")], vec![sum(sum(col("salary")))])
2912        .expect_err("nested aggregates should be rejected");
2913
2914        assert_snapshot!(
2915            err.strip_backtrace(),
2916            @"Error during planning: Aggregate function calls cannot be nested: 'sum(employee_csv.salary)' is nested inside 'sum(sum(employee_csv.salary))'"
2917        );
2918
2919        Ok(())
2920    }
2921
2922    #[test]
2923    fn plan_builder_window_rejects_nested_window_functions() -> Result<()> {
2924        // https://github.com/apache/datafusion/issues/23812
2925        let sum_over = |arg| {
2926            Expr::from(expr::WindowFunction::new(
2927                crate::WindowFunctionDefinition::AggregateUDF(
2928                    crate::test::function_stub::sum_udaf(),
2929                ),
2930                vec![arg],
2931            ))
2932        };
2933        let err = table_scan(Some("employee_csv"), &employee_schema(), Some(vec![4]))?
2934            .window(vec![sum_over(sum_over(col("salary")))])
2935            .expect_err("nested window functions should be rejected");
2936
2937        assert_snapshot!(
2938            err.strip_backtrace(),
2939            @"Error during planning: Window function calls cannot be nested: 'sum(employee_csv.salary) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING' is nested inside 'sum(sum(employee_csv.salary) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING'"
2940        );
2941
2942        Ok(())
2943    }
2944
2945    #[test]
2946    fn test_join_metadata() -> Result<()> {
2947        let left_schema = DFSchema::new_with_metadata(
2948            vec![(None, Arc::new(Field::new("a", DataType::Int32, false)))],
2949            HashMap::from([("key".to_string(), "left".to_string())]),
2950        )?;
2951        let right_schema = DFSchema::new_with_metadata(
2952            vec![(None, Arc::new(Field::new("b", DataType::Int32, false)))],
2953            HashMap::from([("key".to_string(), "right".to_string())]),
2954        )?;
2955
2956        let join_schema =
2957            build_join_schema(&left_schema, &right_schema, &JoinType::Left)?;
2958        assert_eq!(
2959            join_schema.metadata(),
2960            &HashMap::from([("key".to_string(), "left".to_string())])
2961        );
2962        let join_schema =
2963            build_join_schema(&left_schema, &right_schema, &JoinType::Right)?;
2964        assert_eq!(
2965            join_schema.metadata(),
2966            &HashMap::from([("key".to_string(), "right".to_string())])
2967        );
2968
2969        Ok(())
2970    }
2971
2972    #[test]
2973    fn test_values_metadata() -> Result<()> {
2974        let metadata: HashMap<String, String> =
2975            [("ARROW:extension:metadata".to_string(), "test".to_string())]
2976                .into_iter()
2977                .collect();
2978        let metadata = FieldMetadata::from(metadata);
2979        let values = LogicalPlanBuilder::values(vec![
2980            vec![lit_with_metadata(1, Some(metadata.clone()))],
2981            vec![lit_with_metadata(2, Some(metadata.clone()))],
2982        ])?
2983        .build()?;
2984        assert_eq!(*values.schema().field(0).metadata(), metadata.to_hashmap());
2985
2986        // Do not allow VALUES with different metadata mixed together
2987        let metadata2: HashMap<String, String> =
2988            [("ARROW:extension:metadata".to_string(), "test2".to_string())]
2989                .into_iter()
2990                .collect();
2991        let metadata2 = FieldMetadata::from(metadata2);
2992        assert!(
2993            LogicalPlanBuilder::values(vec![
2994                vec![lit_with_metadata(1, Some(metadata.clone()))],
2995                vec![lit_with_metadata(2, Some(metadata2.clone()))],
2996            ])
2997            .is_err()
2998        );
2999
3000        Ok(())
3001    }
3002
3003    #[test]
3004    fn test_unique_field_aliases() {
3005        let t1_field_1 = Field::new("a", DataType::Int32, false);
3006        let t2_field_1 = Field::new("a", DataType::Int32, false);
3007        let t2_field_3 = Field::new("a", DataType::Int32, false);
3008        let t2_field_4 = Field::new("a:1", DataType::Int32, false);
3009        let t1_field_2 = Field::new("b", DataType::Int32, false);
3010        let t2_field_2 = Field::new("b", DataType::Int32, false);
3011
3012        let fields = vec![
3013            t1_field_1, t2_field_1, t1_field_2, t2_field_2, t2_field_3, t2_field_4,
3014        ];
3015        let fields = Fields::from(fields);
3016
3017        let remove_redundant = unique_field_aliases(&fields);
3018
3019        // Input [a, a, b, b, a, a:1] becomes [None, a:1, None, b:1, a:2, a:1:1]
3020        // First occurrence of each field name keeps original name (None), duplicates get
3021        // incremental suffixes (:1, :2, etc.).
3022        // Crucially in this case the 2nd occurrence of `a` gets rewritten to `a:1` which later
3023        // conflicts with the last column which is _actually_ called `a:1` so we need to rename it
3024        // as well to `a:1:1`.
3025        assert_eq!(
3026            remove_redundant,
3027            vec![
3028                None,
3029                Some("a:1".to_string()),
3030                None,
3031                Some("b:1".to_string()),
3032                Some("a:2".to_string()),
3033                Some("a:1:1".to_string()),
3034            ]
3035        );
3036    }
3037
3038    #[test]
3039    fn test_values_with_schema_type_mismatch_error_message() {
3040        // Date32 field, but the value is a Boolean, which cannot be cast to Date32.
3041        let schema = Arc::new(
3042            DFSchema::from_unqualified_fields(
3043                vec![Field::new("a", DataType::Date32, false)].into(),
3044                HashMap::new(),
3045            )
3046            .unwrap(),
3047        );
3048
3049        let err = LogicalPlanBuilder::values_with_schema(vec![vec![lit(true)]], &schema)
3050            .unwrap_err();
3051
3052        assert_eq!(
3053            err.strip_backtrace(),
3054            "Execution error: Types don't match and no valid cast exists, \
3055         received data of type Boolean for field of type Date32"
3056        );
3057    }
3058}