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