Skip to main content

datafusion_optimizer/optimize_projections/
mod.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//! [`OptimizeProjections`] identifies and eliminates unused columns
19
20mod required_indices;
21
22use crate::optimizer::ApplyOrder;
23use crate::{OptimizerConfig, OptimizerRule};
24use std::sync::Arc;
25
26use datafusion_common::{
27    Column, DFSchema, HashMap, JoinType, Result, assert_eq_or_internal_err,
28    get_required_group_by_exprs_indices, internal_datafusion_err, internal_err,
29};
30use datafusion_expr::expr::Alias;
31use datafusion_expr::{
32    Aggregate, Distinct, EmptyRelation, Expr, Projection, TableScanBuilder, Unnest,
33    Window, logical_plan::LogicalPlan,
34};
35
36use crate::optimize_projections::required_indices::RequiredIndices;
37use crate::utils::NamePreserver;
38use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeContainer};
39
40/// Optimizer rule to prune unnecessary columns from intermediate schemas
41/// inside the [`LogicalPlan`]. This rule:
42/// - Removes unnecessary columns that do not appear at the output and/or are
43///   not used during any computation step.
44/// - Adds projections to decrease table column size before operators that
45///   benefit from a smaller memory footprint at its input.
46/// - Removes unnecessary [`LogicalPlan::Projection`]s from the [`LogicalPlan`].
47///
48/// `OptimizeProjections` is an optimizer rule that identifies and eliminates
49/// columns from a logical plan that are not used by downstream operations.
50/// This can improve query performance and reduce unnecessary data processing.
51///
52/// The rule analyzes the input logical plan, determines the necessary column
53/// indices, and then removes any unnecessary columns. It also removes any
54/// unnecessary projections from the plan tree.
55///
56/// ## Schema, Field Properties, and Metadata Handling
57///
58/// The `OptimizeProjections` rule preserves schema and field metadata in most optimization scenarios:
59///
60/// **Schema-level metadata preservation by plan type**:
61/// - **Window and Aggregate plans**: Schema metadata is preserved
62/// - **Projection plans**: Schema metadata is preserved per [`projection_schema`](datafusion_expr::logical_plan::projection_schema).
63/// - **Other logical plans**: Schema metadata is preserved unless [`LogicalPlan::recompute_schema`]
64///   is called on plan types that drop metadata
65///
66/// **Field-level properties and metadata**: Individual field properties are preserved when fields
67/// are retained in the optimized plan, determined by [`exprlist_to_fields`](datafusion_expr::utils::exprlist_to_fields)
68/// and [`ExprSchemable::to_field`](datafusion_expr::expr_schema::ExprSchemable::to_field).
69///
70/// **Field precedence**: When the same field appears multiple times, the optimizer
71/// maintains one occurrence and removes duplicates (refer to `RequiredIndices::compact()`),
72/// preserving the properties and metadata of that occurrence.
73#[derive(Default, Debug)]
74pub struct OptimizeProjections {}
75
76impl OptimizeProjections {
77    #[expect(missing_docs)]
78    pub fn new() -> Self {
79        Self {}
80    }
81}
82
83impl OptimizerRule for OptimizeProjections {
84    fn name(&self) -> &str {
85        "optimize_projections"
86    }
87
88    fn apply_order(&self) -> Option<ApplyOrder> {
89        None
90    }
91
92    fn supports_rewrite(&self) -> bool {
93        true
94    }
95
96    fn rewrite(
97        &self,
98        plan: LogicalPlan,
99        config: &dyn OptimizerConfig,
100    ) -> Result<Transformed<LogicalPlan>> {
101        // All output fields are necessary:
102        let indices = RequiredIndices::new_for_all_exprs(&plan);
103        optimize_projections(plan, config, indices)
104    }
105}
106
107/// Removes unnecessary columns (e.g. columns that do not appear in the output
108/// schema and/or are not used during any computation step such as expression
109/// evaluation) from the logical plan and its inputs.
110///
111/// # Parameters
112///
113/// - `plan`: A reference to the input `LogicalPlan` to optimize.
114/// - `config`: A reference to the optimizer configuration.
115/// - `indices`: A slice of column indices that represent the necessary column
116///   indices for downstream (parent) plan nodes.
117///
118/// # Returns
119///
120/// A `Result` object with the following semantics:
121///
122/// - `Ok(Some(LogicalPlan))`: An optimized `LogicalPlan` without unnecessary
123///   columns.
124/// - `Ok(None)`: Signal that the given logical plan did not require any change.
125/// - `Err(error)`: An error occurred during the optimization process.
126#[cfg_attr(feature = "recursive_protection", recursive::recursive)]
127fn optimize_projections(
128    plan: LogicalPlan,
129    config: &dyn OptimizerConfig,
130    indices: RequiredIndices,
131) -> Result<Transformed<LogicalPlan>> {
132    // Recursively rewrite any nodes that may be able to avoid computation given
133    // their parents' required indices.
134    match plan {
135        LogicalPlan::Projection(proj) => {
136            return merge_consecutive_projections(proj)?
137                .transform_data(|proj| {
138                    rewrite_projection_given_requirements(proj, config, &indices)
139                })?
140                .transform_data(|plan| optimize_subqueries(plan, config));
141        }
142        LogicalPlan::Aggregate(aggregate) => {
143            // Split parent requirements to GROUP BY and aggregate sections:
144            let n_group_exprs = aggregate.group_expr_len()?;
145            // Offset aggregate indices so that they point to valid indices at
146            // `aggregate.aggr_expr`:
147            let (group_by_reqs, aggregate_reqs) = indices.split_off(n_group_exprs);
148
149            // Get absolutely necessary GROUP BY fields.
150            //
151            // When the input has no functional dependencies, we can
152            // short-circuit this analysis.
153            let new_group_bys = if aggregate
154                .input
155                .schema()
156                .functional_dependencies()
157                .is_empty()
158            {
159                aggregate.group_expr
160            } else {
161                let group_by_expr_existing = aggregate
162                    .group_expr
163                    .iter()
164                    .map(|group_by_expr| group_by_expr.schema_name().to_string())
165                    .collect::<Vec<_>>();
166
167                if let Some(simplest_groupby_indices) =
168                    get_required_group_by_exprs_indices(
169                        aggregate.input.schema(),
170                        &group_by_expr_existing,
171                    )
172                {
173                    // Some of the fields in the GROUP BY may be required by
174                    // the parent even if these fields are unnecessary in
175                    // terms of functional dependency.
176                    group_by_reqs
177                        .append(&simplest_groupby_indices)
178                        .get_at_indices(&aggregate.group_expr)
179                } else {
180                    aggregate.group_expr
181                }
182            };
183
184            // Only use the absolutely necessary aggregate expressions required
185            // by the parent:
186            let new_aggr_expr = aggregate_reqs.get_at_indices(&aggregate.aggr_expr);
187
188            if new_group_bys.is_empty() && new_aggr_expr.is_empty() {
189                // Global aggregation with no aggregate functions always produces 1 row and no columns.
190                return Ok(Transformed::yes(LogicalPlan::EmptyRelation(
191                    EmptyRelation {
192                        produce_one_row: true,
193                        schema: Arc::new(DFSchema::empty()),
194                    },
195                )));
196            }
197
198            let all_exprs_iter = new_group_bys.iter().chain(new_aggr_expr.iter());
199            let schema = aggregate.input.schema();
200            let necessary_indices =
201                RequiredIndices::new().with_exprs(schema, all_exprs_iter);
202            let necessary_exprs = necessary_indices.get_required_exprs(schema);
203
204            return optimize_projections(
205                Arc::unwrap_or_clone(aggregate.input),
206                config,
207                necessary_indices,
208            )?
209            .transform_data(|aggregate_input| {
210                // Simplify the input of the aggregation by adding a projection so
211                // that its input only contains absolutely necessary columns for
212                // the aggregate expressions. Note that necessary_indices refer to
213                // fields in `aggregate.input.schema()`.
214                add_projection_on_top_if_helpful(aggregate_input, necessary_exprs)
215            })?
216            .map_data(|aggregate_input| {
217                // Create a new aggregate plan with the updated input and only the
218                // absolutely necessary fields:
219                Aggregate::try_new(
220                    Arc::new(aggregate_input),
221                    new_group_bys,
222                    new_aggr_expr,
223                )
224                .map(LogicalPlan::Aggregate)
225            })?
226            .transform_data(|plan| optimize_subqueries(plan, config));
227        }
228        LogicalPlan::Window(window) => {
229            let input_schema = Arc::clone(window.input.schema());
230            // Split parent requirements to child and window expression sections:
231            let n_input_fields = input_schema.fields().len();
232            // Offset window expression indices so that they point to valid
233            // indices at `window.window_expr`:
234            let (child_reqs, window_reqs) = indices.split_off(n_input_fields);
235
236            // Only use window expressions that are absolutely necessary according
237            // to parent requirements:
238            let new_window_expr = window_reqs.get_at_indices(&window.window_expr);
239
240            // Get all the required column indices at the input, either by the
241            // parent or window expression requirements.
242            let required_indices = child_reqs.with_exprs(&input_schema, &new_window_expr);
243
244            return optimize_projections(
245                Arc::unwrap_or_clone(window.input),
246                config,
247                required_indices.clone(),
248            )?
249            .transform_data(|window_child| {
250                if new_window_expr.is_empty() {
251                    // When no window expression is necessary, use the input directly:
252                    Ok(Transformed::no(window_child))
253                } else {
254                    // Calculate required expressions at the input of the window.
255                    // Please note that we use `input_schema`, because `required_indices`
256                    // refers to that schema
257                    let required_exprs =
258                        required_indices.get_required_exprs(&input_schema);
259                    let window_child =
260                        add_projection_on_top_if_helpful(window_child, required_exprs)?
261                            .data;
262                    Window::try_new(new_window_expr, Arc::new(window_child))
263                        .map(LogicalPlan::Window)
264                        .map(Transformed::yes)
265                }
266            })?
267            .transform_data(|plan| optimize_subqueries(plan, config));
268        }
269        LogicalPlan::TableScan(table_scan) => {
270            // Get indices referred to in the original (schema with all fields)
271            // given projected indices.
272            let projection = match &table_scan.projection {
273                Some(projection) => indices.into_mapped_indices(|idx| projection[idx]),
274                None => indices.into_inner(),
275            };
276            let new_scan = TableScanBuilder::from(table_scan)
277                .with_projection(Some(projection))
278                .build()?;
279
280            return Transformed::yes(LogicalPlan::TableScan(new_scan))
281                .transform_data(|plan| optimize_subqueries(plan, config));
282        }
283        // Other node types are handled below
284        _ => {}
285    };
286
287    // For other plan node types, calculate indices for columns they use and
288    // try to rewrite their children
289    let mut child_required_indices: Vec<RequiredIndices> = match &plan {
290        LogicalPlan::Sort(_)
291        | LogicalPlan::Filter(_)
292        | LogicalPlan::Repartition(_)
293        | LogicalPlan::Union(_)
294        | LogicalPlan::SubqueryAlias(_)
295        | LogicalPlan::Distinct(Distinct::On(_)) => {
296            // Pass index requirements from the parent as well as column indices
297            // that appear in this plan's expressions to its child. All these
298            // operators benefit from "small" inputs, so the projection_beneficial
299            // flag is `true`.
300            plan.inputs()
301                .into_iter()
302                .map(|input| {
303                    indices
304                        .clone()
305                        .with_projection_beneficial()
306                        .with_plan_exprs(&plan, input.schema())
307                })
308                .collect::<Result<_>>()?
309        }
310        LogicalPlan::Limit(_) => {
311            // Pass index requirements from the parent as well as column indices
312            // that appear in this plan's expressions to its child. These operators
313            // do not benefit from "small" inputs, so the projection_beneficial
314            // flag is `false`.
315            plan.inputs()
316                .into_iter()
317                .map(|input| indices.clone().with_plan_exprs(&plan, input.schema()))
318                .collect::<Result<_>>()?
319        }
320        LogicalPlan::Copy(_)
321        | LogicalPlan::Ddl(_)
322        | LogicalPlan::Dml(_)
323        | LogicalPlan::Explain(_)
324        | LogicalPlan::Analyze(_)
325        | LogicalPlan::Subquery(_)
326        | LogicalPlan::Statement(_)
327        | LogicalPlan::Distinct(Distinct::All(_)) => {
328            // These plans require all their fields, and their children should
329            // be treated as final plans -- otherwise, we may have schema a
330            // mismatch.
331            // TODO: For some subquery variants (e.g. a subquery arising from an
332            //       EXISTS expression), we may not need to require all indices.
333            plan.inputs()
334                .into_iter()
335                .map(RequiredIndices::new_for_all_exprs)
336                .collect()
337        }
338        LogicalPlan::Extension(extension) => {
339            if let Some(necessary_children_indices) =
340                extension.node.necessary_children_exprs(indices.indices())
341            {
342                let children = extension.node.inputs();
343                assert_eq_or_internal_err!(
344                    children.len(),
345                    necessary_children_indices.len(),
346                    "Inconsistent length between children and necessary children indices. \
347                Make sure `.necessary_children_exprs` implementation of the \
348                `UserDefinedLogicalNode` is consistent with actual children length \
349                for the node."
350                );
351                children
352                    .into_iter()
353                    .zip(necessary_children_indices)
354                    .map(|(child, necessary_indices)| {
355                        RequiredIndices::new_from_indices(necessary_indices)
356                            .with_plan_exprs(&plan, child.schema())
357                    })
358                    .collect::<Result<Vec<_>>>()?
359            } else {
360                // Requirements from parent cannot be routed down to user defined logical plan safely
361                // Assume it requires all input exprs here
362                plan.inputs()
363                    .into_iter()
364                    .map(RequiredIndices::new_for_all_exprs)
365                    .collect()
366            }
367        }
368        LogicalPlan::EmptyRelation(_)
369        | LogicalPlan::Values(_)
370        | LogicalPlan::DescribeTable(_) => {
371            // These operators have no inputs, so stop the optimization process.
372            return Ok(Transformed::no(plan));
373        }
374        LogicalPlan::RecursiveQuery(_) => {
375            // optimize the static and recursive terms: treat each recursive CTE term like a
376            // standalone subquery: optimize its internals, but do not push parent required indices
377            // through the RecursiveQuery boundary, as this can otherwise lead to bugs
378            // (see: https://github.com/apache/datafusion/issues/22249)
379            return plan.map_children(|c| {
380                let indices = RequiredIndices::new_for_all_exprs(&c);
381                optimize_projections(c, config, indices)
382            });
383        }
384        LogicalPlan::Join(join) => {
385            let left_len = join.left.schema().fields().len();
386            let right_len = join.right.schema().fields().len();
387            let (left_req_indices, right_req_indices) =
388                split_join_requirements(left_len, right_len, indices, &join.join_type);
389            let mut left_indices =
390                left_req_indices.with_plan_exprs(&plan, join.left.schema())?;
391            let mut right_indices =
392                right_req_indices.with_plan_exprs(&plan, join.right.schema())?;
393            // Ensure an empty mark join still has a column to qualify mark
394            match join.join_type {
395                JoinType::LeftMark if right_indices.indices().is_empty() => {
396                    right_indices = right_indices.append(&[0]);
397                }
398                JoinType::RightMark if left_indices.indices().is_empty() => {
399                    left_indices = left_indices.append(&[0]);
400                }
401                _ => {}
402            }
403            // Joins benefit from "small" input tables (lower memory usage).
404            // Therefore, each child benefits from projection:
405            vec![
406                left_indices.with_projection_beneficial(),
407                right_indices.with_projection_beneficial(),
408            ]
409        }
410        // these nodes are explicitly rewritten in the match statement above
411        LogicalPlan::Projection(_)
412        | LogicalPlan::Aggregate(_)
413        | LogicalPlan::Window(_)
414        | LogicalPlan::TableScan(_) => {
415            return internal_err!(
416                "OptimizeProjection: should have handled in the match statement above"
417            );
418        }
419        LogicalPlan::Unnest(Unnest {
420            input,
421            dependency_indices,
422            ..
423        }) => {
424            // at least provide the indices for the exec-columns as a starting point
425            let required_indices =
426                RequiredIndices::new().with_plan_exprs(&plan, input.schema())?;
427
428            // Add additional required indices from the parent
429            let mut additional_necessary_child_indices = Vec::new();
430            indices.indices().iter().for_each(|idx| {
431                if let Some(index) = dependency_indices.get(*idx) {
432                    additional_necessary_child_indices.push(*index);
433                }
434            });
435            vec![required_indices.append(&additional_necessary_child_indices)]
436        }
437    };
438
439    // Required indices are currently ordered (child0, child1, ...)
440    // but the loop pops off the last element, so we need to reverse the order
441    child_required_indices.reverse();
442    assert_eq_or_internal_err!(
443        child_required_indices.len(),
444        plan.inputs().len(),
445        "OptimizeProjection: child_required_indices length mismatch with plan inputs"
446    );
447
448    // Rewrite children of the plan
449    let transformed_plan = plan.map_children(|child| {
450        let required_indices = child_required_indices.pop().ok_or_else(|| {
451            internal_datafusion_err!(
452                "Unexpected number of required_indices in OptimizeProjections rule"
453            )
454        })?;
455
456        let projection_beneficial = required_indices.projection_beneficial();
457        let project_exprs = required_indices.get_required_exprs(child.schema());
458
459        optimize_projections(child, config, required_indices)?.transform_data(
460            |new_input| {
461                if projection_beneficial {
462                    add_projection_on_top_if_helpful(new_input, project_exprs)
463                } else {
464                    Ok(Transformed::no(new_input))
465                }
466            },
467        )
468    })?;
469
470    let transformed_plan =
471        transformed_plan.transform_data(|plan| optimize_subqueries(plan, config))?;
472
473    // If any of the children are transformed, we need to potentially update the plan's schema
474    if transformed_plan.transformed {
475        transformed_plan.map_data(|plan| plan.recompute_schema())
476    } else {
477        Ok(transformed_plan)
478    }
479}
480
481/// Optimizes uncorrelated subquery plans embedded in expressions of the given
482/// plan node (e.g., `Expr::ScalarSubquery`). `map_children` only visits direct
483/// plan inputs, so subqueries must be handled separately.
484fn optimize_subqueries(
485    plan: LogicalPlan,
486    config: &dyn OptimizerConfig,
487) -> Result<Transformed<LogicalPlan>> {
488    plan.map_uncorrelated_subqueries(|subquery_plan| {
489        let indices = RequiredIndices::new_for_all_exprs(&subquery_plan);
490        optimize_projections(subquery_plan, config, indices)
491    })
492}
493
494/// Given a projection `proj`, this function attempts to merge it with a previous
495/// projection if it exists and if merging is beneficial. Merging is considered
496/// beneficial when expressions in the current projection are non-trivial and
497/// appear more than once in its input fields. This can act as a caching mechanism
498/// for non-trivial computations.
499///
500/// ## Metadata Handling During Projection Merging
501///
502/// **Alias metadata preservation**: When merging projections, alias metadata from both
503/// the current and previous projections is carefully preserved. The presence of metadata
504/// precludes alias trimming.
505///
506/// **Schema, Fields, and metadata**: If a projection is rewritten, the schema and metadata
507/// are preserved. Individual field properties and metadata flows through expression rewriting
508/// and are preserved when fields are referenced in the merged projection.
509/// Refer to [`projection_schema`](datafusion_expr::logical_plan::projection_schema)
510/// for more details.
511///
512/// # Parameters
513///
514/// * `proj` - A reference to the `Projection` to be merged.
515///
516/// # Returns
517///
518/// A `Result` object with the following semantics:
519///
520/// - `Ok(Some(Projection))`: Merge was beneficial and successful. Contains the
521///   merged projection.
522/// - `Ok(None)`: Signals that merge is not beneficial (and has not taken place).
523/// - `Err(error)`: An error occurred during the function call.
524fn merge_consecutive_projections(proj: Projection) -> Result<Transformed<Projection>> {
525    // Collapse the whole chain in one pass; otherwise an N-deep chain needs
526    // N outer optimizer passes to fully fold.
527    let mut current = proj;
528    let mut transformed_any = false;
529    loop {
530        let Transformed {
531            data, transformed, ..
532        } = merge_consecutive_projections_one_level(current)?;
533        current = data;
534        if !transformed {
535            break;
536        }
537        transformed_any = true;
538    }
539    Ok(if transformed_any {
540        Transformed::yes(current)
541    } else {
542        Transformed::no(current)
543    })
544}
545
546fn merge_consecutive_projections_one_level(
547    proj: Projection,
548) -> Result<Transformed<Projection>> {
549    let Projection {
550        expr,
551        input,
552        schema,
553        ..
554    } = proj;
555    let LogicalPlan::Projection(prev_projection) = input.as_ref() else {
556        return Projection::try_new_with_schema(expr, input, schema).map(Transformed::no);
557    };
558
559    // A fast path: if the previous projection is same as the current projection
560    // we can directly remove the current projection and return child projection.
561    if prev_projection.expr == expr {
562        return Projection::try_new_with_schema(
563            expr,
564            Arc::clone(&prev_projection.input),
565            schema,
566        )
567        .map(Transformed::yes);
568    }
569
570    // Count usages (referrals) of each projection expression in its input fields:
571    let mut column_referral_map = HashMap::<&Column, usize>::new();
572    expr.iter()
573        .for_each(|expr| expr.add_column_ref_counts(&mut column_referral_map));
574
575    // If an expression is non-trivial (KeepInPlace) and appears more than once, do not merge
576    // them as consecutive projections will benefit from a compute-once approach.
577    // For details, see: https://github.com/apache/datafusion/issues/8296
578    if column_referral_map.into_iter().any(|(col, usage)| {
579        usage > 1
580            && !prev_projection.expr[prev_projection.schema.index_of_column(col).unwrap()]
581                .placement()
582                .should_push_to_leaves()
583    }) {
584        // no change
585        return Projection::try_new_with_schema(expr, input, schema).map(Transformed::no);
586    }
587
588    let LogicalPlan::Projection(prev_projection) = Arc::unwrap_or_clone(input) else {
589        // We know it is a `LogicalPlan::Projection` from check above
590        unreachable!();
591    };
592
593    // Try to rewrite the expressions in the current projection using the
594    // previous projection as input:
595    let name_preserver = NamePreserver::new_for_projection();
596    let mut original_names = vec![];
597    let new_exprs = expr.map_elements(|expr| {
598        original_names.push(name_preserver.save(&expr));
599
600        // do not rewrite top level Aliases (rewriter will remove all aliases within exprs)
601        match expr {
602            Expr::Alias(Alias {
603                expr,
604                relation,
605                name,
606                metadata,
607            }) => rewrite_expr(*expr, &prev_projection).map(|result| {
608                result.update_data(|expr| {
609                    // After substitution, the inner expression may now have the
610                    // same schema_name as the alias (e.g. when an extraction
611                    // alias like `__extracted_1 AS f(x)` is resolved back to
612                    // `f(x)`). Wrapping in a redundant self-alias causes a
613                    // cosmetic `f(x) AS f(x)` due to Display vs schema_name
614                    // formatting differences. Drop the alias when it matches.
615                    if metadata.is_none() && expr.schema_name().to_string() == name {
616                        expr
617                    } else {
618                        Expr::Alias(Alias {
619                            expr: Box::new(expr),
620                            relation,
621                            name,
622                            metadata,
623                        })
624                    }
625                })
626            }),
627            e => rewrite_expr(e, &prev_projection),
628        }
629    })?;
630
631    // if the expressions could be rewritten, create a new projection with the
632    // new expressions
633    if new_exprs.transformed {
634        // Add any needed aliases back to the expressions
635        let new_exprs = new_exprs
636            .data
637            .into_iter()
638            .zip(original_names)
639            .map(|(expr, original_name)| original_name.restore(expr))
640            .collect::<Vec<_>>();
641        Projection::try_new(new_exprs, prev_projection.input).map(Transformed::yes)
642    } else {
643        // not rewritten, so put the projection back together
644        let input = Arc::new(LogicalPlan::Projection(prev_projection));
645        Projection::try_new_with_schema(new_exprs.data, input, schema)
646            .map(Transformed::no)
647    }
648}
649
650/// Rewrites a projection expression using the projection before it (i.e. its input)
651/// This is a subroutine to the `merge_consecutive_projections` function.
652///
653/// # Parameters
654///
655/// * `expr` - A reference to the expression to rewrite.
656/// * `input` - A reference to the input of the projection expression (itself
657///   a projection).
658///
659/// # Returns
660///
661/// A `Result` object with the following semantics:
662///
663/// - `Ok(Some(Expr))`: Rewrite was successful. Contains the rewritten result.
664/// - `Ok(None)`: Signals that `expr` can not be rewritten.
665/// - `Err(error)`: An error occurred during the function call.
666///
667/// # Notes
668/// This rewrite also removes any unnecessary layers of aliasing. "Unnecessary" is
669/// defined as not contributing new information, such as metadata.
670///
671/// Without trimming, we can end up with unnecessary indirections inside expressions
672/// during projection merges.
673///
674/// Consider:
675///
676/// ```text
677/// Projection(a1 + b1 as sum1)
678/// --Projection(a as a1, b as b1)
679/// ----Source(a, b)
680/// ```
681///
682/// After merge, we want to produce:
683///
684/// ```text
685/// Projection(a + b as sum1)
686/// --Source(a, b)
687/// ```
688///
689/// Without trimming, we would end up with:
690///
691/// ```text
692/// Projection((a as a1 + b as b1) as sum1)
693/// --Source(a, b)
694/// ```
695fn rewrite_expr(expr: Expr, input: &Projection) -> Result<Transformed<Expr>> {
696    expr.transform_up(|expr| {
697        match expr {
698            //  remove any intermediate aliases if they do not carry metadata
699            Expr::Alias(alias) => {
700                match alias
701                    .metadata
702                    .as_ref()
703                    .map(|h| h.is_empty())
704                    .unwrap_or(true)
705                {
706                    true => Ok(Transformed::yes(*alias.expr)),
707                    false => Ok(Transformed::no(Expr::Alias(alias))),
708                }
709            }
710            Expr::Column(col) => {
711                // Find index of column:
712                let idx = input.schema.index_of_column(&col)?;
713                // get the corresponding unaliased input expression
714                //
715                // For example:
716                // * the input projection is [`a + b` as c, `d + e` as f]
717                // * the current column is an expression "f"
718                //
719                // return the expression `d + e` (not `d + e` as f)
720                let input_expr = input.expr[idx].clone().unalias_nested().data;
721                Ok(Transformed::yes(input_expr))
722            }
723            // Unsupported type for consecutive projection merge analysis.
724            _ => Ok(Transformed::no(expr)),
725        }
726    })
727}
728
729/// Splits requirement indices for a join into left and right children based on
730/// the join type.
731///
732/// This function takes the length of the left child, a slice of requirement
733/// indices, and the type of join (e.g. `INNER`, `LEFT`, `RIGHT`) as arguments.
734/// Depending on the join type, it divides the requirement indices into those
735/// that apply to the left child and those that apply to the right child.
736///
737/// - For `INNER`, `LEFT`, `RIGHT`, `FULL`, `LEFTMARK`, and `RIGHTMARK` joins,
738///   the requirements are split between left and right children. The right
739///   child indices are adjusted to point to valid positions within the right
740///   child by subtracting the length of the left child.
741///
742/// - For `LEFT ANTI`, `LEFT SEMI`, `RIGHT SEMI` and `RIGHT ANTI` joins, all
743///   requirements are re-routed to either the left child or the right child
744///   directly, depending on the join type.
745///
746/// # Parameters
747///
748/// * `left_len` - The length of the left child.
749/// * `right_len` - The length of the right child.
750/// * `indices` - A slice of requirement indices.
751/// * `join_type` - The type of join (e.g. `INNER`, `LEFT`, `RIGHT`).
752///
753/// # Returns
754///
755/// A tuple containing two vectors of `usize` indices: The first vector represents
756/// the requirements for the left child, and the second vector represents the
757/// requirements for the right child. The indices are appropriately split and
758/// adjusted based on the join type.
759fn split_join_requirements(
760    left_len: usize,
761    right_len: usize,
762    indices: RequiredIndices,
763    join_type: &JoinType,
764) -> (RequiredIndices, RequiredIndices) {
765    match join_type {
766        // In these cases requirements are split between left/right children:
767        JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => {
768            // Decrease right side indices by `left_len` so that they point to valid
769            // positions within the right child:
770            indices.split_off(left_len)
771        }
772        JoinType::LeftMark => {
773            // LeftMark output: [left_cols(0..left_len), mark]
774            // The mark column is synthetic (produced by the join itself),
775            // so discard it and route only to the left child.
776            let (left_indices, _mark) = indices.split_off(left_len);
777            (left_indices, RequiredIndices::new())
778        }
779        JoinType::RightMark => {
780            // Same as LeftMark, but for the right child.
781            let (right_indices, _mark) = indices.split_off(right_len);
782            (RequiredIndices::new(), right_indices)
783        }
784        // All requirements can be re-routed to left child directly.
785        JoinType::LeftAnti | JoinType::LeftSemi => (indices, RequiredIndices::new()),
786        // All requirements can be re-routed to right side directly.
787        // No need to change index, join schema is right child schema.
788        JoinType::RightSemi | JoinType::RightAnti => (RequiredIndices::new(), indices),
789    }
790}
791
792/// Adds a projection on top of a logical plan if doing so reduces the number
793/// of columns for the parent operator.
794///
795/// This function takes a `LogicalPlan` and a list of projection expressions.
796/// If the projection is beneficial (it reduces the number of columns in the
797/// plan) a new `LogicalPlan` with the projection is created and returned, along
798/// with a `true` flag. If the projection doesn't reduce the number of columns,
799/// the original plan is returned with a `false` flag.
800///
801/// # Parameters
802///
803/// * `plan` - The input `LogicalPlan` to potentially add a projection to.
804/// * `project_exprs` - A list of expressions for the projection.
805///
806/// # Returns
807///
808/// A `Transformed` indicating if a projection was added
809fn add_projection_on_top_if_helpful(
810    plan: LogicalPlan,
811    project_exprs: Vec<Expr>,
812) -> Result<Transformed<LogicalPlan>> {
813    // Make sure projection decreases the number of columns, otherwise it is unnecessary.
814    if project_exprs.len() >= plan.schema().fields().len() {
815        Ok(Transformed::no(plan))
816    } else {
817        Projection::try_new(project_exprs, Arc::new(plan))
818            .map(LogicalPlan::Projection)
819            .map(Transformed::yes)
820    }
821}
822
823/// Rewrite the given projection according to the fields required by its
824/// ancestors.
825///
826/// # Parameters
827///
828/// * `proj` - A reference to the original projection to rewrite.
829/// * `config` - A reference to the optimizer configuration.
830/// * `indices` - A slice of indices representing the columns required by the
831///   ancestors of the given projection.
832///
833/// # Returns
834///
835/// A `Result` object with the following semantics:
836///
837/// - `Ok(Some(LogicalPlan))`: Contains the rewritten projection
838/// - `Ok(None)`: No rewrite necessary.
839/// - `Err(error)`: An error occurred during the function call.
840fn rewrite_projection_given_requirements(
841    proj: Projection,
842    config: &dyn OptimizerConfig,
843    indices: &RequiredIndices,
844) -> Result<Transformed<LogicalPlan>> {
845    let Projection { expr, input, .. } = proj;
846
847    let exprs_used = indices.get_at_indices(&expr);
848
849    let required_indices =
850        RequiredIndices::new().with_exprs(input.schema(), exprs_used.iter());
851
852    // rewrite the children projection, and if they are changed rewrite the
853    // projection down
854    optimize_projections(Arc::unwrap_or_clone(input), config, required_indices)?
855        .transform_data(|input| {
856            if is_projection_unnecessary(&input, &exprs_used)? {
857                Ok(Transformed::yes(input))
858            } else {
859                Projection::try_new(exprs_used, Arc::new(input))
860                    .map(LogicalPlan::Projection)
861                    .map(Transformed::yes)
862            }
863        })
864}
865
866/// Projection is unnecessary, when
867/// - input schema of the projection, output schema of the projection are same, and
868/// - all projection expressions are either Column or Literal
869pub fn is_projection_unnecessary(
870    input: &LogicalPlan,
871    proj_exprs: &[Expr],
872) -> Result<bool> {
873    // First check if the number of expressions is equal to the number of fields in the input schema.
874    if proj_exprs.len() != input.schema().fields().len() {
875        return Ok(false);
876    }
877    Ok(input.schema().iter().zip(proj_exprs.iter()).all(
878        |((field_relation, field_name), expr)| {
879            // Check if the expression is a column and if it matches the field name
880            if let Expr::Column(col) = expr {
881                col.relation.as_ref() == field_relation && col.name.eq(field_name.name())
882            } else {
883                false
884            }
885        },
886    ))
887}
888
889#[cfg(test)]
890mod tests {
891    use std::cmp::Ordering;
892    use std::collections::HashMap;
893    use std::fmt::Formatter;
894    use std::ops::Add;
895    use std::sync::Arc;
896    use std::vec;
897
898    use crate::optimize_projections::OptimizeProjections;
899    use crate::optimizer::Optimizer;
900    use crate::test::{
901        assert_fields_eq, scan_empty, test_table_scan, test_table_scan_fields,
902        test_table_scan_with_name,
903    };
904    use crate::{OptimizerContext, OptimizerRule};
905    use arrow::datatypes::{DataType, Field, Schema};
906    use datafusion_common::{
907        Column, DFSchema, DFSchemaRef, JoinType, Result, TableReference,
908    };
909    use datafusion_expr::ExprFunctionExt;
910    use datafusion_expr::{
911        BinaryExpr, Expr, Extension, Like, LogicalPlan, Operator, Projection,
912        UserDefinedLogicalNodeCore, WindowFunctionDefinition, binary_expr,
913        build_join_schema,
914        builder::table_scan_with_filters,
915        col,
916        expr::{self, Cast},
917        lit,
918        logical_plan::{builder::LogicalPlanBuilder, table_scan},
919        not, try_cast, when,
920    };
921    use insta::assert_snapshot;
922
923    use crate::assert_optimized_plan_eq_snapshot;
924    use datafusion_functions_aggregate::count::count_udaf;
925    use datafusion_functions_aggregate::expr_fn::{count, max, min};
926    use datafusion_functions_aggregate::min_max::max_udaf;
927
928    macro_rules! assert_optimized_plan_equal {
929        (
930            $plan:expr,
931            @ $expected:literal $(,)?
932        ) => {{
933            let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
934            let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![Arc::new(OptimizeProjections::new())];
935            assert_optimized_plan_eq_snapshot!(
936                optimizer_ctx,
937                rules,
938                $plan,
939                @ $expected,
940            )
941        }};
942    }
943
944    #[derive(Debug, Hash, PartialEq, Eq)]
945    struct NoOpUserDefined {
946        exprs: Vec<Expr>,
947        schema: DFSchemaRef,
948        input: Arc<LogicalPlan>,
949    }
950
951    impl NoOpUserDefined {
952        fn new(schema: DFSchemaRef, input: Arc<LogicalPlan>) -> Self {
953            Self {
954                exprs: vec![],
955                schema,
956                input,
957            }
958        }
959
960        fn with_exprs(mut self, exprs: Vec<Expr>) -> Self {
961            self.exprs = exprs;
962            self
963        }
964    }
965
966    // Manual implementation needed because of `schema` field. Comparison excludes this field.
967    impl PartialOrd for NoOpUserDefined {
968        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
969            match self.exprs.partial_cmp(&other.exprs) {
970                Some(Ordering::Equal) => self.input.partial_cmp(&other.input),
971                cmp => cmp,
972            }
973            // TODO (https://github.com/apache/datafusion/issues/17477) avoid recomparing all fields
974            .filter(|cmp| *cmp != Ordering::Equal || self == other)
975        }
976    }
977
978    impl UserDefinedLogicalNodeCore for NoOpUserDefined {
979        fn name(&self) -> &str {
980            "NoOpUserDefined"
981        }
982
983        fn inputs(&self) -> Vec<&LogicalPlan> {
984            vec![&self.input]
985        }
986
987        fn schema(&self) -> &DFSchemaRef {
988            &self.schema
989        }
990
991        fn expressions(&self) -> Vec<Expr> {
992            self.exprs.clone()
993        }
994
995        fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result {
996            write!(f, "NoOpUserDefined")
997        }
998
999        fn with_exprs_and_inputs(
1000            &self,
1001            exprs: Vec<Expr>,
1002            mut inputs: Vec<LogicalPlan>,
1003        ) -> Result<Self> {
1004            Ok(Self {
1005                exprs,
1006                input: Arc::new(inputs.swap_remove(0)),
1007                schema: Arc::clone(&self.schema),
1008            })
1009        }
1010
1011        fn necessary_children_exprs(
1012            &self,
1013            output_columns: &[usize],
1014        ) -> Option<Vec<Vec<usize>>> {
1015            // Since schema is same. Output columns requires their corresponding version in the input columns.
1016            Some(vec![output_columns.to_vec()])
1017        }
1018
1019        fn supports_limit_pushdown(&self) -> bool {
1020            false // Disallow limit push-down by default
1021        }
1022    }
1023
1024    #[derive(Debug, Hash, PartialEq, Eq)]
1025    struct UserDefinedCrossJoin {
1026        exprs: Vec<Expr>,
1027        schema: DFSchemaRef,
1028        left_child: Arc<LogicalPlan>,
1029        right_child: Arc<LogicalPlan>,
1030    }
1031
1032    impl UserDefinedCrossJoin {
1033        fn new(left_child: Arc<LogicalPlan>, right_child: Arc<LogicalPlan>) -> Self {
1034            let left_schema = left_child.schema();
1035            let right_schema = right_child.schema();
1036            let schema = Arc::new(
1037                build_join_schema(left_schema, right_schema, &JoinType::Inner).unwrap(),
1038            );
1039            Self {
1040                exprs: vec![],
1041                schema,
1042                left_child,
1043                right_child,
1044            }
1045        }
1046    }
1047
1048    // Manual implementation needed because of `schema` field. Comparison excludes this field.
1049    impl PartialOrd for UserDefinedCrossJoin {
1050        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
1051            match self.exprs.partial_cmp(&other.exprs) {
1052                Some(Ordering::Equal) => {
1053                    match self.left_child.partial_cmp(&other.left_child) {
1054                        Some(Ordering::Equal) => {
1055                            self.right_child.partial_cmp(&other.right_child)
1056                        }
1057                        cmp => cmp,
1058                    }
1059                }
1060                cmp => cmp,
1061            }
1062            // TODO (https://github.com/apache/datafusion/issues/17477) avoid recomparing all fields
1063            .filter(|cmp| *cmp != Ordering::Equal || self == other)
1064        }
1065    }
1066
1067    impl UserDefinedLogicalNodeCore for UserDefinedCrossJoin {
1068        fn name(&self) -> &str {
1069            "UserDefinedCrossJoin"
1070        }
1071
1072        fn inputs(&self) -> Vec<&LogicalPlan> {
1073            vec![&self.left_child, &self.right_child]
1074        }
1075
1076        fn schema(&self) -> &DFSchemaRef {
1077            &self.schema
1078        }
1079
1080        fn expressions(&self) -> Vec<Expr> {
1081            self.exprs.clone()
1082        }
1083
1084        fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result {
1085            write!(f, "UserDefinedCrossJoin")
1086        }
1087
1088        fn with_exprs_and_inputs(
1089            &self,
1090            exprs: Vec<Expr>,
1091            mut inputs: Vec<LogicalPlan>,
1092        ) -> Result<Self> {
1093            assert_eq!(inputs.len(), 2);
1094            Ok(Self {
1095                exprs,
1096                left_child: Arc::new(inputs.remove(0)),
1097                right_child: Arc::new(inputs.remove(0)),
1098                schema: Arc::clone(&self.schema),
1099            })
1100        }
1101
1102        fn necessary_children_exprs(
1103            &self,
1104            output_columns: &[usize],
1105        ) -> Option<Vec<Vec<usize>>> {
1106            let left_child_len = self.left_child.schema().fields().len();
1107            let mut left_reqs = vec![];
1108            let mut right_reqs = vec![];
1109            for &out_idx in output_columns {
1110                if out_idx < left_child_len {
1111                    left_reqs.push(out_idx);
1112                } else {
1113                    // Output indices further than the left_child_len
1114                    // comes from right children
1115                    right_reqs.push(out_idx - left_child_len)
1116                }
1117            }
1118            Some(vec![left_reqs, right_reqs])
1119        }
1120
1121        fn supports_limit_pushdown(&self) -> bool {
1122            false // Disallow limit push-down by default
1123        }
1124    }
1125
1126    /// A user-defined node that does NOT implement `necessary_children_exprs`,
1127    /// so the optimizer cannot determine which columns are required from its
1128    /// children and must assume all columns are needed.
1129    #[derive(Debug, Hash, PartialEq, Eq)]
1130    struct OpaqueRequirementsUserDefined {
1131        input: Arc<LogicalPlan>,
1132        schema: DFSchemaRef,
1133    }
1134
1135    // Manual implementation needed because of `schema` field. Comparison excludes this field.
1136    impl PartialOrd for OpaqueRequirementsUserDefined {
1137        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
1138            self.input
1139                .partial_cmp(&other.input)
1140                .filter(|cmp| *cmp != Ordering::Equal || self == other)
1141        }
1142    }
1143
1144    impl UserDefinedLogicalNodeCore for OpaqueRequirementsUserDefined {
1145        fn name(&self) -> &str {
1146            "OpaqueRequirementsUserDefined"
1147        }
1148
1149        fn inputs(&self) -> Vec<&LogicalPlan> {
1150            vec![&self.input]
1151        }
1152
1153        fn schema(&self) -> &DFSchemaRef {
1154            &self.schema
1155        }
1156
1157        fn expressions(&self) -> Vec<Expr> {
1158            vec![]
1159        }
1160
1161        fn with_exprs_and_inputs(
1162            &self,
1163            _exprs: Vec<Expr>,
1164            mut inputs: Vec<LogicalPlan>,
1165        ) -> Result<Self> {
1166            Ok(Self {
1167                input: Arc::new(inputs.swap_remove(0)),
1168                schema: Arc::clone(&self.schema),
1169            })
1170        }
1171
1172        fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result {
1173            write!(f, "OpaqueRequirementsUserDefined")
1174        }
1175    }
1176
1177    #[test]
1178    fn merge_two_projection() -> Result<()> {
1179        let table_scan = test_table_scan()?;
1180        let plan = LogicalPlanBuilder::from(table_scan)
1181            .project(vec![col("a")])?
1182            .project(vec![binary_expr(lit(1), Operator::Plus, col("a"))])?
1183            .build()?;
1184
1185        assert_optimized_plan_equal!(
1186            plan,
1187            @r"
1188        Projection: Int32(1) + test.a
1189          TableScan: test projection=[a]
1190        "
1191        )
1192    }
1193
1194    #[test]
1195    fn merge_three_projection() -> Result<()> {
1196        let table_scan = test_table_scan()?;
1197        let plan = LogicalPlanBuilder::from(table_scan)
1198            .project(vec![col("a"), col("b")])?
1199            .project(vec![col("a")])?
1200            .project(vec![binary_expr(lit(1), Operator::Plus, col("a"))])?
1201            .build()?;
1202
1203        assert_optimized_plan_equal!(
1204            plan,
1205            @r"
1206        Projection: Int32(1) + test.a
1207          TableScan: test projection=[a]
1208        "
1209        )
1210    }
1211
1212    #[test]
1213    fn merge_alias() -> Result<()> {
1214        let table_scan = test_table_scan()?;
1215        let plan = LogicalPlanBuilder::from(table_scan)
1216            .project(vec![col("a")])?
1217            .project(vec![col("a").alias("alias")])?
1218            .build()?;
1219
1220        assert_optimized_plan_equal!(
1221            plan,
1222            @r"
1223        Projection: test.a AS alias
1224          TableScan: test projection=[a]
1225        "
1226        )
1227    }
1228
1229    #[test]
1230    fn merge_nested_alias() -> Result<()> {
1231        let table_scan = test_table_scan()?;
1232        let plan = LogicalPlanBuilder::from(table_scan)
1233            .project(vec![col("a").alias("alias1").alias("alias2")])?
1234            .project(vec![col("alias2").alias("alias")])?
1235            .build()?;
1236
1237        assert_optimized_plan_equal!(
1238            plan,
1239            @r"
1240        Projection: test.a AS alias
1241          TableScan: test projection=[a]
1242        "
1243        )
1244    }
1245
1246    #[test]
1247    fn test_nested_count() -> Result<()> {
1248        let schema = Schema::new(vec![Field::new("foo", DataType::Int32, false)]);
1249
1250        let groups: Vec<Expr> = vec![];
1251
1252        let plan = table_scan(TableReference::none(), &schema, None)
1253            .unwrap()
1254            .aggregate(groups.clone(), vec![count(lit(1))])
1255            .unwrap()
1256            .aggregate(groups, vec![count(lit(1))])
1257            .unwrap()
1258            .build()
1259            .unwrap();
1260
1261        assert_optimized_plan_equal!(
1262            plan,
1263            @r"
1264        Aggregate: groupBy=[[]], aggr=[[count(Int32(1))]]
1265          EmptyRelation: rows=1
1266        "
1267        )
1268    }
1269
1270    #[test]
1271    fn test_neg_push_down() -> Result<()> {
1272        let table_scan = test_table_scan()?;
1273        let plan = LogicalPlanBuilder::from(table_scan)
1274            .project(vec![-col("a")])?
1275            .build()?;
1276
1277        assert_optimized_plan_equal!(
1278            plan,
1279            @r"
1280        Projection: (- test.a)
1281          TableScan: test projection=[a]
1282        "
1283        )
1284    }
1285
1286    #[test]
1287    fn test_is_null() -> Result<()> {
1288        let table_scan = test_table_scan()?;
1289        let plan = LogicalPlanBuilder::from(table_scan)
1290            .project(vec![col("a").is_null()])?
1291            .build()?;
1292
1293        assert_optimized_plan_equal!(
1294            plan,
1295            @r"
1296        Projection: test.a IS NULL
1297          TableScan: test projection=[a]
1298        "
1299        )
1300    }
1301
1302    #[test]
1303    fn test_is_not_null() -> Result<()> {
1304        let table_scan = test_table_scan()?;
1305        let plan = LogicalPlanBuilder::from(table_scan)
1306            .project(vec![col("a").is_not_null()])?
1307            .build()?;
1308
1309        assert_optimized_plan_equal!(
1310            plan,
1311            @r"
1312        Projection: test.a IS NOT NULL
1313          TableScan: test projection=[a]
1314        "
1315        )
1316    }
1317
1318    #[test]
1319    fn test_is_true() -> Result<()> {
1320        let table_scan = test_table_scan()?;
1321        let plan = LogicalPlanBuilder::from(table_scan)
1322            .project(vec![col("a").is_true()])?
1323            .build()?;
1324
1325        assert_optimized_plan_equal!(
1326            plan,
1327            @r"
1328        Projection: test.a IS TRUE
1329          TableScan: test projection=[a]
1330        "
1331        )
1332    }
1333
1334    #[test]
1335    fn test_is_not_true() -> Result<()> {
1336        let table_scan = test_table_scan()?;
1337        let plan = LogicalPlanBuilder::from(table_scan)
1338            .project(vec![col("a").is_not_true()])?
1339            .build()?;
1340
1341        assert_optimized_plan_equal!(
1342            plan,
1343            @r"
1344        Projection: test.a IS NOT TRUE
1345          TableScan: test projection=[a]
1346        "
1347        )
1348    }
1349
1350    #[test]
1351    fn test_is_false() -> Result<()> {
1352        let table_scan = test_table_scan()?;
1353        let plan = LogicalPlanBuilder::from(table_scan)
1354            .project(vec![col("a").is_false()])?
1355            .build()?;
1356
1357        assert_optimized_plan_equal!(
1358            plan,
1359            @r"
1360        Projection: test.a IS FALSE
1361          TableScan: test projection=[a]
1362        "
1363        )
1364    }
1365
1366    #[test]
1367    fn test_is_not_false() -> Result<()> {
1368        let table_scan = test_table_scan()?;
1369        let plan = LogicalPlanBuilder::from(table_scan)
1370            .project(vec![col("a").is_not_false()])?
1371            .build()?;
1372
1373        assert_optimized_plan_equal!(
1374            plan,
1375            @r"
1376        Projection: test.a IS NOT FALSE
1377          TableScan: test projection=[a]
1378        "
1379        )
1380    }
1381
1382    #[test]
1383    fn test_is_unknown() -> Result<()> {
1384        let table_scan = test_table_scan()?;
1385        let plan = LogicalPlanBuilder::from(table_scan)
1386            .project(vec![col("a").is_unknown()])?
1387            .build()?;
1388
1389        assert_optimized_plan_equal!(
1390            plan,
1391            @r"
1392        Projection: test.a IS UNKNOWN
1393          TableScan: test projection=[a]
1394        "
1395        )
1396    }
1397
1398    #[test]
1399    fn test_is_not_unknown() -> Result<()> {
1400        let table_scan = test_table_scan()?;
1401        let plan = LogicalPlanBuilder::from(table_scan)
1402            .project(vec![col("a").is_not_unknown()])?
1403            .build()?;
1404
1405        assert_optimized_plan_equal!(
1406            plan,
1407            @r"
1408        Projection: test.a IS NOT UNKNOWN
1409          TableScan: test projection=[a]
1410        "
1411        )
1412    }
1413
1414    #[test]
1415    fn test_not() -> Result<()> {
1416        let table_scan = test_table_scan()?;
1417        let plan = LogicalPlanBuilder::from(table_scan)
1418            .project(vec![not(col("a"))])?
1419            .build()?;
1420
1421        assert_optimized_plan_equal!(
1422            plan,
1423            @r"
1424        Projection: NOT test.a
1425          TableScan: test projection=[a]
1426        "
1427        )
1428    }
1429
1430    #[test]
1431    fn test_try_cast() -> Result<()> {
1432        let table_scan = test_table_scan()?;
1433        let plan = LogicalPlanBuilder::from(table_scan)
1434            .project(vec![try_cast(col("a"), DataType::Float64)])?
1435            .build()?;
1436
1437        assert_optimized_plan_equal!(
1438            plan,
1439            @r"
1440        Projection: TRY_CAST(test.a AS Float64)
1441          TableScan: test projection=[a]
1442        "
1443        )
1444    }
1445
1446    #[test]
1447    fn test_similar_to() -> Result<()> {
1448        let table_scan = test_table_scan()?;
1449        let expr = Box::new(col("a"));
1450        let pattern = Box::new(lit("[0-9]"));
1451        let similar_to_expr =
1452            Expr::SimilarTo(Like::new(false, expr, pattern, None, false));
1453        let plan = LogicalPlanBuilder::from(table_scan)
1454            .project(vec![similar_to_expr])?
1455            .build()?;
1456
1457        assert_optimized_plan_equal!(
1458            plan,
1459            @r#"
1460        Projection: test.a SIMILAR TO Utf8("[0-9]")
1461          TableScan: test projection=[a]
1462        "#
1463        )
1464    }
1465
1466    #[test]
1467    fn test_between() -> Result<()> {
1468        let table_scan = test_table_scan()?;
1469        let plan = LogicalPlanBuilder::from(table_scan)
1470            .project(vec![col("a").between(lit(1), lit(3))])?
1471            .build()?;
1472
1473        assert_optimized_plan_equal!(
1474            plan,
1475            @r"
1476        Projection: test.a BETWEEN Int32(1) AND Int32(3)
1477          TableScan: test projection=[a]
1478        "
1479        )
1480    }
1481
1482    // Test Case expression
1483    #[test]
1484    fn test_case_merged() -> Result<()> {
1485        let table_scan = test_table_scan()?;
1486        let plan = LogicalPlanBuilder::from(table_scan)
1487            .project(vec![col("a"), lit(0).alias("d")])?
1488            .project(vec![
1489                col("a"),
1490                when(col("a").eq(lit(1)), lit(10))
1491                    .otherwise(col("d"))?
1492                    .alias("d"),
1493            ])?
1494            .build()?;
1495
1496        assert_optimized_plan_equal!(
1497            plan,
1498            @r"
1499        Projection: test.a, CASE WHEN test.a = Int32(1) THEN Int32(10) ELSE Int32(0) END AS d
1500          TableScan: test projection=[a]
1501        "
1502        )
1503    }
1504
1505    // Test outer projection isn't discarded despite the same schema as inner
1506    // https://github.com/apache/datafusion/issues/8942
1507    #[test]
1508    fn test_derived_column() -> Result<()> {
1509        let table_scan = test_table_scan()?;
1510        let plan = LogicalPlanBuilder::from(table_scan)
1511            .project(vec![col("a").add(lit(1)).alias("a"), lit(0).alias("d")])?
1512            .project(vec![
1513                col("a"),
1514                when(col("a").eq(lit(1)), lit(10))
1515                    .otherwise(col("d"))?
1516                    .alias("d"),
1517            ])?
1518            .build()?;
1519
1520        assert_optimized_plan_equal!(
1521            plan,
1522            @r"
1523        Projection: a, CASE WHEN a = Int32(1) THEN Int32(10) ELSE d END AS d
1524          Projection: test.a + Int32(1) AS a, Int32(0) AS d
1525            TableScan: test projection=[a]
1526        "
1527        )
1528    }
1529
1530    // Since only column `a` is referred at the output. Scan should only contain projection=[a].
1531    // User defined node should be able to propagate necessary expressions by its parent to its child.
1532    #[test]
1533    fn test_user_defined_logical_plan_node() -> Result<()> {
1534        let table_scan = test_table_scan()?;
1535        let custom_plan = LogicalPlan::Extension(Extension {
1536            node: Arc::new(NoOpUserDefined::new(
1537                Arc::clone(table_scan.schema()),
1538                Arc::new(table_scan.clone()),
1539            )),
1540        });
1541        let plan = LogicalPlanBuilder::from(custom_plan)
1542            .project(vec![col("a"), lit(0).alias("d")])?
1543            .build()?;
1544
1545        assert_optimized_plan_equal!(
1546            plan,
1547            @r"
1548        Projection: test.a, Int32(0) AS d
1549          NoOpUserDefined
1550            TableScan: test projection=[a]
1551        "
1552        )
1553    }
1554
1555    // Only column `a` is referred at the output. However, User defined node itself uses column `b`
1556    // during its operation. Hence, scan should contain projection=[a, b].
1557    // User defined node should be able to propagate necessary expressions by its parent, as well as its own
1558    // required expressions.
1559    #[test]
1560    fn test_user_defined_logical_plan_node2() -> Result<()> {
1561        let table_scan = test_table_scan()?;
1562        let exprs = vec![Expr::Column(Column::from_qualified_name("b"))];
1563        let custom_plan = LogicalPlan::Extension(Extension {
1564            node: Arc::new(
1565                NoOpUserDefined::new(
1566                    Arc::clone(table_scan.schema()),
1567                    Arc::new(table_scan.clone()),
1568                )
1569                .with_exprs(exprs),
1570            ),
1571        });
1572        let plan = LogicalPlanBuilder::from(custom_plan)
1573            .project(vec![col("a"), lit(0).alias("d")])?
1574            .build()?;
1575
1576        assert_optimized_plan_equal!(
1577            plan,
1578            @r"
1579        Projection: test.a, Int32(0) AS d
1580          NoOpUserDefined
1581            TableScan: test projection=[a, b]
1582        "
1583        )
1584    }
1585
1586    // Only column `a` is referred at the output. However, User defined node itself uses expression `b+c`
1587    // during its operation. Hence, scan should contain projection=[a, b, c].
1588    // User defined node should be able to propagate necessary expressions by its parent, as well as its own
1589    // required expressions. Expressions doesn't have to be just column. Requirements from complex expressions
1590    // should be propagated also.
1591    #[test]
1592    fn test_user_defined_logical_plan_node3() -> Result<()> {
1593        let table_scan = test_table_scan()?;
1594        let left_expr = Expr::Column(Column::from_qualified_name("b"));
1595        let right_expr = Expr::Column(Column::from_qualified_name("c"));
1596        let binary_expr = Expr::BinaryExpr(BinaryExpr::new(
1597            Box::new(left_expr),
1598            Operator::Plus,
1599            Box::new(right_expr),
1600        ));
1601        let exprs = vec![binary_expr];
1602        let custom_plan = LogicalPlan::Extension(Extension {
1603            node: Arc::new(
1604                NoOpUserDefined::new(
1605                    Arc::clone(table_scan.schema()),
1606                    Arc::new(table_scan.clone()),
1607                )
1608                .with_exprs(exprs),
1609            ),
1610        });
1611        let plan = LogicalPlanBuilder::from(custom_plan)
1612            .project(vec![col("a"), lit(0).alias("d")])?
1613            .build()?;
1614
1615        assert_optimized_plan_equal!(
1616            plan,
1617            @r"
1618        Projection: test.a, Int32(0) AS d
1619          NoOpUserDefined
1620            TableScan: test projection=[a, b, c]
1621        "
1622        )
1623    }
1624
1625    // Columns `l.a`, `l.c`, `r.a` is referred at the output.
1626    // User defined node should be able to propagate necessary expressions by its parent, to its children.
1627    // Even if it has multiple children.
1628    // left child should have `projection=[a, c]`, and right side should have `projection=[a]`.
1629    #[test]
1630    fn test_user_defined_logical_plan_node4() -> Result<()> {
1631        let left_table = test_table_scan_with_name("l")?;
1632        let right_table = test_table_scan_with_name("r")?;
1633        let custom_plan = LogicalPlan::Extension(Extension {
1634            node: Arc::new(UserDefinedCrossJoin::new(
1635                Arc::new(left_table),
1636                Arc::new(right_table),
1637            )),
1638        });
1639        let plan = LogicalPlanBuilder::from(custom_plan)
1640            .project(vec![col("l.a"), col("l.c"), col("r.a"), lit(0).alias("d")])?
1641            .build()?;
1642
1643        assert_optimized_plan_equal!(
1644            plan,
1645            @r"
1646        Projection: l.a, l.c, r.a, Int32(0) AS d
1647          UserDefinedCrossJoin
1648            TableScan: l projection=[a, c]
1649            TableScan: r projection=[a]
1650        "
1651        )
1652    }
1653
1654    #[test]
1655    fn aggregate_no_group_by() -> Result<()> {
1656        let table_scan = test_table_scan()?;
1657
1658        let plan = LogicalPlanBuilder::from(table_scan)
1659            .aggregate(Vec::<Expr>::new(), vec![max(col("b"))])?
1660            .build()?;
1661
1662        assert_optimized_plan_equal!(
1663            plan,
1664            @r"
1665        Aggregate: groupBy=[[]], aggr=[[max(test.b)]]
1666          TableScan: test projection=[b]
1667        "
1668        )
1669    }
1670
1671    #[test]
1672    fn aggregate_group_by() -> Result<()> {
1673        let table_scan = test_table_scan()?;
1674
1675        let plan = LogicalPlanBuilder::from(table_scan)
1676            .aggregate(vec![col("c")], vec![max(col("b"))])?
1677            .build()?;
1678
1679        assert_optimized_plan_equal!(
1680            plan,
1681            @r"
1682        Aggregate: groupBy=[[test.c]], aggr=[[max(test.b)]]
1683          TableScan: test projection=[b, c]
1684        "
1685        )
1686    }
1687
1688    #[test]
1689    fn aggregate_group_by_with_table_alias() -> Result<()> {
1690        let table_scan = test_table_scan()?;
1691
1692        let plan = LogicalPlanBuilder::from(table_scan)
1693            .alias("a")?
1694            .aggregate(vec![col("c")], vec![max(col("b"))])?
1695            .build()?;
1696
1697        assert_optimized_plan_equal!(
1698            plan,
1699            @r"
1700        Aggregate: groupBy=[[a.c]], aggr=[[max(a.b)]]
1701          SubqueryAlias: a
1702            TableScan: test projection=[b, c]
1703        "
1704        )
1705    }
1706
1707    #[test]
1708    fn aggregate_no_group_by_with_filter() -> Result<()> {
1709        let table_scan = test_table_scan()?;
1710
1711        let plan = LogicalPlanBuilder::from(table_scan)
1712            .filter(col("c").gt(lit(1)))?
1713            .aggregate(Vec::<Expr>::new(), vec![max(col("b"))])?
1714            .build()?;
1715
1716        assert_optimized_plan_equal!(
1717            plan,
1718            @r"
1719        Aggregate: groupBy=[[]], aggr=[[max(test.b)]]
1720          Projection: test.b
1721            Filter: test.c > Int32(1)
1722              TableScan: test projection=[b, c]
1723        "
1724        )
1725    }
1726
1727    #[test]
1728    fn aggregate_with_periods() -> Result<()> {
1729        let schema = Schema::new(vec![Field::new("tag.one", DataType::Utf8, false)]);
1730
1731        // Build a plan that looks as follows (note "tag.one" is a column named
1732        // "tag.one", not a column named "one" in a table named "tag"):
1733        //
1734        // Projection: tag.one
1735        //   Aggregate: groupBy=[], aggr=[max("tag.one") AS "tag.one"]
1736        //    TableScan
1737        let plan = table_scan(Some("m4"), &schema, None)?
1738            .aggregate(
1739                Vec::<Expr>::new(),
1740                vec![max(col(Column::new_unqualified("tag.one"))).alias("tag.one")],
1741            )?
1742            .project([col(Column::new_unqualified("tag.one"))])?
1743            .build()?;
1744
1745        assert_optimized_plan_equal!(
1746            plan,
1747            @r"
1748        Aggregate: groupBy=[[]], aggr=[[max(m4.tag.one) AS tag.one]]
1749          TableScan: m4 projection=[tag.one]
1750        "
1751        )
1752    }
1753
1754    #[test]
1755    fn redundant_project() -> Result<()> {
1756        let table_scan = test_table_scan()?;
1757
1758        let plan = LogicalPlanBuilder::from(table_scan)
1759            .project(vec![col("a"), col("b"), col("c")])?
1760            .project(vec![col("a"), col("c"), col("b")])?
1761            .build()?;
1762        assert_optimized_plan_equal!(
1763            plan,
1764            @r"
1765        Projection: test.a, test.c, test.b
1766          TableScan: test projection=[a, b, c]
1767        "
1768        )
1769    }
1770
1771    #[test]
1772    fn reorder_scan() -> Result<()> {
1773        let schema = Schema::new(test_table_scan_fields());
1774
1775        let plan = table_scan(Some("test"), &schema, Some(vec![1, 0, 2]))?.build()?;
1776        assert_optimized_plan_equal!(
1777            plan,
1778            @"TableScan: test projection=[b, a, c]"
1779        )
1780    }
1781
1782    #[test]
1783    fn reorder_scan_projection() -> Result<()> {
1784        let schema = Schema::new(test_table_scan_fields());
1785
1786        let plan = table_scan(Some("test"), &schema, Some(vec![1, 0, 2]))?
1787            .project(vec![col("a"), col("b")])?
1788            .build()?;
1789        assert_optimized_plan_equal!(
1790            plan,
1791            @r"
1792        Projection: test.a, test.b
1793          TableScan: test projection=[b, a]
1794        "
1795        )
1796    }
1797
1798    #[test]
1799    fn reorder_projection() -> Result<()> {
1800        let table_scan = test_table_scan()?;
1801
1802        let plan = LogicalPlanBuilder::from(table_scan)
1803            .project(vec![col("c"), col("b"), col("a")])?
1804            .build()?;
1805        assert_optimized_plan_equal!(
1806            plan,
1807            @r"
1808        Projection: test.c, test.b, test.a
1809          TableScan: test projection=[a, b, c]
1810        "
1811        )
1812    }
1813
1814    #[test]
1815    fn noncontinuous_redundant_projection() -> Result<()> {
1816        let table_scan = test_table_scan()?;
1817
1818        let plan = LogicalPlanBuilder::from(table_scan)
1819            .project(vec![col("c"), col("b"), col("a")])?
1820            .filter(col("c").gt(lit(1)))?
1821            .project(vec![col("c"), col("a"), col("b")])?
1822            .filter(col("b").gt(lit(1)))?
1823            .filter(col("a").gt(lit(1)))?
1824            .project(vec![col("a"), col("c"), col("b")])?
1825            .build()?;
1826        assert_optimized_plan_equal!(
1827            plan,
1828            @r"
1829        Projection: test.a, test.c, test.b
1830          Filter: test.a > Int32(1)
1831            Filter: test.b > Int32(1)
1832              Projection: test.c, test.a, test.b
1833                Filter: test.c > Int32(1)
1834                  Projection: test.c, test.b, test.a
1835                    TableScan: test projection=[a, b, c]
1836        "
1837        )
1838    }
1839
1840    #[test]
1841    fn join_schema_trim_full_join_column_projection() -> Result<()> {
1842        let table_scan = test_table_scan()?;
1843
1844        let schema = Schema::new(vec![Field::new("c1", DataType::UInt32, false)]);
1845        let table2_scan = scan_empty(Some("test2"), &schema, None)?.build()?;
1846
1847        let plan = LogicalPlanBuilder::from(table_scan)
1848            .join(table2_scan, JoinType::Left, (vec!["a"], vec!["c1"]), None)?
1849            .project(vec![col("a"), col("b"), col("c1")])?
1850            .build()?;
1851
1852        let optimized_plan = optimize(plan)?;
1853
1854        // make sure projections are pushed down to both table scans
1855        assert_snapshot!(
1856            optimized_plan.clone(),
1857            @r"
1858        Left Join: test.a = test2.c1
1859          TableScan: test projection=[a, b]
1860          TableScan: test2 projection=[c1]
1861        "
1862        );
1863
1864        // make sure schema for join node include both join columns
1865        let optimized_join = optimized_plan;
1866        assert_eq!(
1867            **optimized_join.schema(),
1868            DFSchema::new_with_metadata(
1869                vec![
1870                    (
1871                        Some("test".into()),
1872                        Arc::new(Field::new("a", DataType::UInt32, false))
1873                    ),
1874                    (
1875                        Some("test".into()),
1876                        Arc::new(Field::new("b", DataType::UInt32, false))
1877                    ),
1878                    (
1879                        Some("test2".into()),
1880                        Arc::new(Field::new("c1", DataType::UInt32, true))
1881                    ),
1882                ],
1883                HashMap::new()
1884            )?,
1885        );
1886
1887        Ok(())
1888    }
1889
1890    #[test]
1891    fn join_schema_trim_partial_join_column_projection() -> Result<()> {
1892        // test join column push down without explicit column projections
1893
1894        let table_scan = test_table_scan()?;
1895
1896        let schema = Schema::new(vec![Field::new("c1", DataType::UInt32, false)]);
1897        let table2_scan = scan_empty(Some("test2"), &schema, None)?.build()?;
1898
1899        let plan = LogicalPlanBuilder::from(table_scan)
1900            .join(table2_scan, JoinType::Left, (vec!["a"], vec!["c1"]), None)?
1901            // projecting joined column `a` should push the right side column `c1` projection as
1902            // well into test2 table even though `c1` is not referenced in projection.
1903            .project(vec![col("a"), col("b")])?
1904            .build()?;
1905
1906        let optimized_plan = optimize(plan)?;
1907
1908        // make sure projections are pushed down to both table scans
1909        assert_snapshot!(
1910            optimized_plan.clone(),
1911            @r"
1912        Projection: test.a, test.b
1913          Left Join: test.a = test2.c1
1914            TableScan: test projection=[a, b]
1915            TableScan: test2 projection=[c1]
1916        "
1917        );
1918
1919        // make sure schema for join node include both join columns
1920        let optimized_join = optimized_plan.inputs()[0];
1921        assert_eq!(
1922            **optimized_join.schema(),
1923            DFSchema::new_with_metadata(
1924                vec![
1925                    (
1926                        Some("test".into()),
1927                        Arc::new(Field::new("a", DataType::UInt32, false))
1928                    ),
1929                    (
1930                        Some("test".into()),
1931                        Arc::new(Field::new("b", DataType::UInt32, false))
1932                    ),
1933                    (
1934                        Some("test2".into()),
1935                        Arc::new(Field::new("c1", DataType::UInt32, true))
1936                    ),
1937                ],
1938                HashMap::new()
1939            )?,
1940        );
1941
1942        Ok(())
1943    }
1944
1945    #[test]
1946    fn join_schema_trim_using_join() -> Result<()> {
1947        // shared join columns from using join should be pushed to both sides
1948
1949        let table_scan = test_table_scan()?;
1950
1951        let schema = Schema::new(vec![Field::new("a", DataType::UInt32, false)]);
1952        let table2_scan = scan_empty(Some("test2"), &schema, None)?.build()?;
1953
1954        let plan = LogicalPlanBuilder::from(table_scan)
1955            .join_using(table2_scan, JoinType::Left, vec!["a".into()])?
1956            .project(vec![col("a"), col("b")])?
1957            .build()?;
1958
1959        let optimized_plan = optimize(plan)?;
1960
1961        // make sure projections are pushed down to table scan
1962        assert_snapshot!(
1963            optimized_plan.clone(),
1964            @r"
1965        Projection: test.a, test.b
1966          Left Join: Using test.a = test2.a
1967            TableScan: test projection=[a, b]
1968            TableScan: test2 projection=[a]
1969        "
1970        );
1971
1972        // make sure schema for join node include both join columns
1973        let optimized_join = optimized_plan.inputs()[0];
1974        assert_eq!(
1975            **optimized_join.schema(),
1976            DFSchema::new_with_metadata(
1977                vec![
1978                    (
1979                        Some("test".into()),
1980                        Arc::new(Field::new("a", DataType::UInt32, false))
1981                    ),
1982                    (
1983                        Some("test".into()),
1984                        Arc::new(Field::new("b", DataType::UInt32, false))
1985                    ),
1986                    (
1987                        Some("test2".into()),
1988                        Arc::new(Field::new("a", DataType::UInt32, true))
1989                    ),
1990                ],
1991                HashMap::new()
1992            )?,
1993        );
1994
1995        Ok(())
1996    }
1997
1998    #[test]
1999    fn cast() -> Result<()> {
2000        let table_scan = test_table_scan()?;
2001
2002        let plan = LogicalPlanBuilder::from(table_scan)
2003            .project(vec![Expr::Cast(Cast::new(
2004                Box::new(col("c")),
2005                DataType::Float64,
2006            ))])?
2007            .build()?;
2008
2009        assert_optimized_plan_equal!(
2010            plan,
2011            @r"
2012        Projection: CAST(test.c AS Float64)
2013          TableScan: test projection=[c]
2014        "
2015        )
2016    }
2017
2018    #[test]
2019    fn table_scan_projected_schema() -> Result<()> {
2020        let table_scan = test_table_scan()?;
2021        let plan = LogicalPlanBuilder::from(test_table_scan()?)
2022            .project(vec![col("a"), col("b")])?
2023            .build()?;
2024
2025        assert_eq!(3, table_scan.schema().fields().len());
2026        assert_fields_eq(&table_scan, vec!["a", "b", "c"]);
2027        assert_fields_eq(&plan, vec!["a", "b"]);
2028
2029        assert_optimized_plan_equal!(
2030            plan,
2031            @"TableScan: test projection=[a, b]"
2032        )
2033    }
2034
2035    #[test]
2036    fn table_scan_projected_schema_non_qualified_relation() -> Result<()> {
2037        let table_scan = test_table_scan()?;
2038        let input_schema = table_scan.schema();
2039        assert_eq!(3, input_schema.fields().len());
2040        assert_fields_eq(&table_scan, vec!["a", "b", "c"]);
2041
2042        // Build the LogicalPlan directly (don't use PlanBuilder), so
2043        // that the Column references are unqualified (e.g. their
2044        // relation is `None`). PlanBuilder resolves the expressions
2045        let expr = vec![col("test.a"), col("test.b")];
2046        let plan =
2047            LogicalPlan::Projection(Projection::try_new(expr, Arc::new(table_scan))?);
2048
2049        assert_fields_eq(&plan, vec!["a", "b"]);
2050
2051        assert_optimized_plan_equal!(
2052            plan,
2053            @"TableScan: test projection=[a, b]"
2054        )
2055    }
2056
2057    #[test]
2058    fn table_limit() -> Result<()> {
2059        let table_scan = test_table_scan()?;
2060        assert_eq!(3, table_scan.schema().fields().len());
2061        assert_fields_eq(&table_scan, vec!["a", "b", "c"]);
2062
2063        let plan = LogicalPlanBuilder::from(table_scan)
2064            .project(vec![col("c"), col("a")])?
2065            .limit(0, Some(5))?
2066            .build()?;
2067
2068        assert_fields_eq(&plan, vec!["c", "a"]);
2069
2070        assert_optimized_plan_equal!(
2071            plan,
2072            @r"
2073        Limit: skip=0, fetch=5
2074          Projection: test.c, test.a
2075            TableScan: test projection=[a, c]
2076        "
2077        )
2078    }
2079
2080    #[test]
2081    fn table_scan_without_projection() -> Result<()> {
2082        let table_scan = test_table_scan()?;
2083        let plan = LogicalPlanBuilder::from(table_scan).build()?;
2084        // should expand projection to all columns without projection
2085        assert_optimized_plan_equal!(
2086            plan,
2087            @"TableScan: test projection=[a, b, c]"
2088        )
2089    }
2090
2091    #[test]
2092    fn table_scan_with_literal_projection() -> Result<()> {
2093        let table_scan = test_table_scan()?;
2094        let plan = LogicalPlanBuilder::from(table_scan)
2095            .project(vec![lit(1_i64), lit(2_i64)])?
2096            .build()?;
2097        assert_optimized_plan_equal!(
2098            plan,
2099            @r"
2100        Projection: Int64(1), Int64(2)
2101          TableScan: test projection=[]
2102        "
2103        )
2104    }
2105
2106    /// tests that it removes unused columns in projections
2107    #[test]
2108    fn table_unused_column() -> Result<()> {
2109        let table_scan = test_table_scan()?;
2110        assert_eq!(3, table_scan.schema().fields().len());
2111        assert_fields_eq(&table_scan, vec!["a", "b", "c"]);
2112
2113        // we never use "b" in the first projection => remove it
2114        let plan = LogicalPlanBuilder::from(table_scan)
2115            .project(vec![col("c"), col("a"), col("b")])?
2116            .filter(col("c").gt(lit(1)))?
2117            .aggregate(vec![col("c")], vec![max(col("a"))])?
2118            .build()?;
2119
2120        assert_fields_eq(&plan, vec!["c", "max(test.a)"]);
2121
2122        let plan = optimize(plan).expect("failed to optimize plan");
2123        assert_optimized_plan_equal!(
2124            plan,
2125            @r"
2126        Aggregate: groupBy=[[test.c]], aggr=[[max(test.a)]]
2127          Filter: test.c > Int32(1)
2128            Projection: test.c, test.a
2129              TableScan: test projection=[a, c]
2130        "
2131        )
2132    }
2133
2134    /// tests that it removes un-needed projections
2135    #[test]
2136    fn table_unused_projection() -> Result<()> {
2137        let table_scan = test_table_scan()?;
2138        assert_eq!(3, table_scan.schema().fields().len());
2139        assert_fields_eq(&table_scan, vec!["a", "b", "c"]);
2140
2141        // there is no need for the first projection
2142        let plan = LogicalPlanBuilder::from(table_scan)
2143            .project(vec![col("b")])?
2144            .project(vec![lit(1).alias("a")])?
2145            .build()?;
2146
2147        assert_fields_eq(&plan, vec!["a"]);
2148
2149        assert_optimized_plan_equal!(
2150            plan,
2151            @r"
2152        Projection: Int32(1) AS a
2153          TableScan: test projection=[]
2154        "
2155        )
2156    }
2157
2158    #[test]
2159    fn table_full_filter_pushdown() -> Result<()> {
2160        let schema = Schema::new(test_table_scan_fields());
2161
2162        let table_scan = table_scan_with_filters(
2163            Some("test"),
2164            &schema,
2165            None,
2166            vec![col("b").eq(lit(1))],
2167        )?
2168        .build()?;
2169        assert_eq!(3, table_scan.schema().fields().len());
2170        assert_fields_eq(&table_scan, vec!["a", "b", "c"]);
2171
2172        // there is no need for the first projection
2173        let plan = LogicalPlanBuilder::from(table_scan)
2174            .project(vec![col("b")])?
2175            .project(vec![lit(1).alias("a")])?
2176            .build()?;
2177
2178        assert_fields_eq(&plan, vec!["a"]);
2179
2180        assert_optimized_plan_equal!(
2181            plan,
2182            @r"
2183        Projection: Int32(1) AS a
2184          TableScan: test projection=[], full_filters=[b = Int32(1)]
2185        "
2186        )
2187    }
2188
2189    /// tests that optimizing twice yields same plan
2190    #[test]
2191    fn test_double_optimization() -> Result<()> {
2192        let table_scan = test_table_scan()?;
2193
2194        let plan = LogicalPlanBuilder::from(table_scan)
2195            .project(vec![col("b")])?
2196            .project(vec![lit(1).alias("a")])?
2197            .build()?;
2198
2199        let optimized_plan1 = optimize(plan).expect("failed to optimize plan");
2200        let optimized_plan2 =
2201            optimize(optimized_plan1.clone()).expect("failed to optimize plan");
2202
2203        let formatted_plan1 = format!("{optimized_plan1:?}");
2204        let formatted_plan2 = format!("{optimized_plan2:?}");
2205        assert_eq!(formatted_plan1, formatted_plan2);
2206        Ok(())
2207    }
2208
2209    #[test]
2210    fn test_continue_processing_through_extension() -> Result<()> {
2211        let table_scan = test_table_scan()?;
2212        let plan = LogicalPlanBuilder::from(table_scan.clone())
2213            .project(vec![col("a")])?
2214            .project(vec![col("a")])?
2215            .build()?;
2216        let plan = LogicalPlan::Extension(Extension {
2217            node: Arc::new(OpaqueRequirementsUserDefined {
2218                input: Arc::new(plan),
2219                schema: Arc::clone(table_scan.schema()),
2220            }),
2221        });
2222        let plan = optimize(plan).expect("failed to optimize plan");
2223        assert_optimized_plan_equal!(
2224            plan,
2225            @r"
2226        OpaqueRequirementsUserDefined
2227          TableScan: test projection=[a]
2228        "
2229        )
2230    }
2231
2232    /// tests that it removes an aggregate is never used downstream
2233    #[test]
2234    fn table_unused_aggregate() -> Result<()> {
2235        let table_scan = test_table_scan()?;
2236        assert_eq!(3, table_scan.schema().fields().len());
2237        assert_fields_eq(&table_scan, vec!["a", "b", "c"]);
2238
2239        // we never use "min(b)" => remove it
2240        let plan = LogicalPlanBuilder::from(table_scan)
2241            .aggregate(vec![col("a"), col("c")], vec![max(col("b")), min(col("b"))])?
2242            .filter(col("c").gt(lit(1)))?
2243            .project(vec![col("c"), col("a"), col("max(test.b)")])?
2244            .build()?;
2245
2246        assert_fields_eq(&plan, vec!["c", "a", "max(test.b)"]);
2247
2248        assert_optimized_plan_equal!(
2249            plan,
2250            @r"
2251        Projection: test.c, test.a, max(test.b)
2252          Filter: test.c > Int32(1)
2253            Aggregate: groupBy=[[test.a, test.c]], aggr=[[max(test.b)]]
2254              TableScan: test projection=[a, b, c]
2255        "
2256        )
2257    }
2258
2259    #[test]
2260    fn aggregate_filter_pushdown() -> Result<()> {
2261        let table_scan = test_table_scan()?;
2262        let aggr_with_filter = count_udaf()
2263            .call(vec![col("b")])
2264            .filter(col("c").gt(lit(42)))
2265            .build()?;
2266        let plan = LogicalPlanBuilder::from(table_scan)
2267            .aggregate(
2268                vec![col("a")],
2269                vec![count(col("b")), aggr_with_filter.alias("count2")],
2270            )?
2271            .build()?;
2272
2273        assert_optimized_plan_equal!(
2274            plan,
2275            @r"
2276        Aggregate: groupBy=[[test.a]], aggr=[[count(test.b), count(test.b) FILTER (WHERE test.c > Int32(42)) AS count2]]
2277          TableScan: test projection=[a, b, c]
2278        "
2279        )
2280    }
2281
2282    #[test]
2283    fn pushdown_through_distinct() -> Result<()> {
2284        let table_scan = test_table_scan()?;
2285
2286        let plan = LogicalPlanBuilder::from(table_scan)
2287            .project(vec![col("a"), col("b")])?
2288            .distinct()?
2289            .project(vec![col("a")])?
2290            .build()?;
2291
2292        assert_optimized_plan_equal!(
2293            plan,
2294            @r"
2295        Projection: test.a
2296          Distinct:
2297            TableScan: test projection=[a, b]
2298        "
2299        )
2300    }
2301
2302    #[test]
2303    fn test_window() -> Result<()> {
2304        let table_scan = test_table_scan()?;
2305
2306        let max1 = Expr::from(expr::WindowFunction::new(
2307            WindowFunctionDefinition::AggregateUDF(max_udaf()),
2308            vec![col("test.a")],
2309        ))
2310        .partition_by(vec![col("test.b")])
2311        .build()
2312        .unwrap();
2313
2314        let max2 = Expr::from(expr::WindowFunction::new(
2315            WindowFunctionDefinition::AggregateUDF(max_udaf()),
2316            vec![col("test.b")],
2317        ));
2318        let col1 = col(max1.schema_name().to_string());
2319        let col2 = col(max2.schema_name().to_string());
2320
2321        let plan = LogicalPlanBuilder::from(table_scan)
2322            .window(vec![max1])?
2323            .window(vec![max2])?
2324            .project(vec![col1, col2])?
2325            .build()?;
2326
2327        assert_optimized_plan_equal!(
2328            plan,
2329            @r"
2330        Projection: max(test.a) PARTITION BY [test.b] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING, max(test.b) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
2331          WindowAggr: windowExpr=[[max(test.b) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]]
2332            Projection: test.b, max(test.a) PARTITION BY [test.b] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
2333              WindowAggr: windowExpr=[[max(test.a) PARTITION BY [test.b] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]]
2334                TableScan: test projection=[a, b]
2335        "
2336        )
2337    }
2338
2339    // Regression test for https://github.com/apache/datafusion/issues/20083
2340    // Optimizer must not fail when LeftMark joins from EXISTS OR EXISTS
2341    // feed into a Left join.
2342    #[test]
2343    fn optimize_projections_exists_or_exists_with_outer_join() -> Result<()> {
2344        use datafusion_expr::utils::disjunction;
2345        use datafusion_expr::{exists, out_ref_col};
2346
2347        let table_a = test_table_scan_with_name("a")?;
2348        let table_b = test_table_scan_with_name("b")?;
2349
2350        let sq_a = Arc::new(
2351            LogicalPlanBuilder::from(test_table_scan_with_name("sq_a")?)
2352                .filter(col("sq_a.a").eq(out_ref_col(DataType::UInt32, "a.a")))?
2353                .project(vec![lit(1)])?
2354                .build()?,
2355        );
2356
2357        let sq_b = Arc::new(
2358            LogicalPlanBuilder::from(test_table_scan_with_name("sq_b")?)
2359                .filter(col("sq_b.b").eq(out_ref_col(DataType::UInt32, "a.b")))?
2360                .project(vec![lit(1)])?
2361                .build()?,
2362        );
2363
2364        let plan = LogicalPlanBuilder::from(table_a)
2365            .filter(disjunction(vec![exists(sq_a), exists(sq_b)]).unwrap())?
2366            .join(table_b, JoinType::Left, (vec!["a"], vec!["a"]), None)?
2367            .build()?;
2368
2369        let optimizer = Optimizer::new();
2370        let config = OptimizerContext::new();
2371        optimizer.optimize(plan, &config, observe)?;
2372
2373        Ok(())
2374    }
2375
2376    #[test]
2377    fn optimize_projections_left_mark_join_with_projection() -> Result<()> {
2378        let table_a = test_table_scan_with_name("a")?;
2379        let table_b = test_table_scan_with_name("b")?;
2380        let table_c = test_table_scan_with_name("c")?;
2381
2382        let plan = LogicalPlanBuilder::from(table_a)
2383            .join(table_b, JoinType::LeftMark, (vec!["a"], vec!["a"]), None)?
2384            .project(vec![col("a.a"), col("a.b"), col("a.c")])?
2385            .join(table_c, JoinType::Left, (vec!["a"], vec!["a"]), None)?
2386            .build()?;
2387
2388        assert_optimized_plan_equal!(
2389            plan,
2390            @r"
2391        Left Join: a.a = c.a
2392          Projection: a.a, a.b, a.c
2393            LeftMark Join: a.a = b.a
2394              TableScan: a projection=[a, b, c]
2395              TableScan: b projection=[a]
2396          TableScan: c projection=[a, b, c]
2397        "
2398        )
2399    }
2400
2401    // Stacked filter-less LeftMark joins (from `= ANY` / `<> ALL`) must keep
2402    // each `mark` qualified so they don't collide.
2403    #[test]
2404    fn optimize_projections_stacked_mark_joins_keep_qualified_mark() -> Result<()> {
2405        let person = test_table_scan_with_name("person")?;
2406
2407        let aliased_scan = |table: &str, alias: &str| -> Result<LogicalPlan> {
2408            LogicalPlanBuilder::from(test_table_scan_with_name(table)?)
2409                .project(vec![col(format!("{table}.a"))])?
2410                .alias(alias)?
2411                .build()
2412        };
2413
2414        let plan = LogicalPlanBuilder::from(person)
2415            .join_on(
2416                aliased_scan("s1", "__correlated_sq_1")?,
2417                JoinType::LeftMark,
2418                vec![lit(true)],
2419            )?
2420            .join_on(
2421                aliased_scan("s2", "__correlated_sq_2")?,
2422                JoinType::LeftMark,
2423                vec![lit(true)],
2424            )?
2425            .join_on(
2426                aliased_scan("s3", "__correlated_sq_3")?,
2427                JoinType::LeftMark,
2428                vec![lit(true)],
2429            )?
2430            .filter(
2431                col("__correlated_sq_1.mark")
2432                    .or(col("__correlated_sq_2.mark"))
2433                    .and(not(col("__correlated_sq_3.mark"))),
2434            )?
2435            .project(vec![col("person.a")])?
2436            .build()?;
2437
2438        assert_optimized_plan_equal!(
2439            plan,
2440            @r"
2441        Projection: person.a
2442          Filter: (__correlated_sq_1.mark OR __correlated_sq_2.mark) AND NOT __correlated_sq_3.mark
2443            LeftMark Join:  Filter: Boolean(true)
2444              LeftMark Join:  Filter: Boolean(true)
2445                LeftMark Join:  Filter: Boolean(true)
2446                  TableScan: person projection=[a]
2447                  SubqueryAlias: __correlated_sq_1
2448                    TableScan: s1 projection=[a]
2449                SubqueryAlias: __correlated_sq_2
2450                  TableScan: s2 projection=[a]
2451              SubqueryAlias: __correlated_sq_3
2452                TableScan: s3 projection=[a]
2453        "
2454        )
2455    }
2456
2457    fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
2458
2459    fn optimize(plan: LogicalPlan) -> Result<LogicalPlan> {
2460        let optimizer = Optimizer::with_rules(vec![Arc::new(OptimizeProjections::new())]);
2461        let optimized_plan =
2462            optimizer.optimize(plan, &OptimizerContext::new(), observe)?;
2463        Ok(optimized_plan)
2464    }
2465}