// SPDX-License-Identifier: Apache-2.0

// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership

// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements

//! [`OptimizeProjections`] identifies and eliminates unused columns

//!

//! # Precedence over `PushDownLeafProjections`

//!

//! `OptimizeProjections` merges adjacent projections, including a recovery

//! projection over an extraction projection that

//! [`PushDownLeafProjections`] creates. To keep the two rules from undoing each

//! other, `PushDownLeafProjections` only creates that pair when the extraction

//! projection moves below the input. See the module documentation of

//! [`extract_leaf_expressions`](crate::extract_leaf_expressions).

//!

//! [`PushDownLeafProjections`]: crate::extract_leaf_expressions::PushDownLeafProjections

mod required_indices;

use crate::optimizer::ApplyOrder;

use crate::{OptimizerConfig, OptimizerRule};

use datafusion_expr::logical_plan::MaterializedCtePlan;

use std::sync::Arc;

use datafusion_common::{

Column, DFSchema, DFSchemaRef, HashMap, JoinType, Result, assert_eq_or_internal_err,

get_required_group_by_exprs_indices, internal_datafusion_err, internal_err,

};

use datafusion_expr::expr::Alias;

use datafusion_expr::utils::passthrough_field_index;

use datafusion_expr::{

Aggregate, Distinct, EmptyRelation, Expr, Projection, TableScanBuilder, Unnest,

Window, logical_plan::LogicalPlan,

};

use crate::optimize_projections::required_indices::RequiredIndices;

use crate::utils::NamePreserver;

use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeContainer};

/// Optimizer rule to prune unnecessary columns from intermediate schemas

/// inside the [`LogicalPlan`]. This rule:

/// - Removes unnecessary columns that do not appear at the output and/or are

/// not used during any computation step.

/// - Adds projections to decrease table column size before operators that

/// benefit from a smaller memory footprint at its input.

/// - Removes unnecessary [`LogicalPlan::Projection`]s from the [`LogicalPlan`].

///

/// `OptimizeProjections` is an optimizer rule that identifies and eliminates

/// columns from a logical plan that are not used by downstream operations.

/// This can improve query performance and reduce unnecessary data processing.

///

/// The rule analyzes the input logical plan, determines the necessary column

/// indices, and then removes any unnecessary columns. It also removes any

/// unnecessary projections from the plan tree.

///

/// ## Schema, Field Properties, and Metadata Handling

///

/// The `OptimizeProjections` rule preserves schema and field metadata in most optimization scenarios:

///

/// **Schema-level metadata preservation by plan type**:

/// - **Window and Aggregate plans**: Schema metadata is preserved

/// - **Projection plans**: Schema metadata is preserved per [`projection_schema`](datafusion_expr::logical_plan::projection_schema).

/// - **Other logical plans**: Schema metadata is preserved unless [`LogicalPlan::recompute_schema`]

/// is called on plan types that drop metadata

///

/// **Field-level properties and metadata**: Individual field properties are preserved when fields

/// are retained in the optimized plan, determined by [`exprlist_to_fields`](datafusion_expr::utils::exprlist_to_fields)

/// and [`ExprSchemable::to_field`](datafusion_expr::expr_schema::ExprSchemable::to_field).

///

/// **Field precedence**: When the same field appears multiple times, the optimizer

/// maintains one occurrence and removes duplicates (refer to `RequiredIndices::compact()`),

/// preserving the properties and metadata of that occurrence.

#[derive(Default, Debug)]

pub struct OptimizeProjections {}

impl OptimizeProjections {

#[expect(missing_docs)]

pub fn new() -> Self {

Self {}

}

}

impl OptimizerRule for OptimizeProjections {

fn name(&self) -> &str {

"optimize_projections"

}

fn apply_order(&self) -> Option<ApplyOrder> {

None

}

fn rewrite(

&self,

plan: LogicalPlan,

config: &dyn OptimizerConfig,

) -> Result<Transformed<LogicalPlan>> {

// All output fields are necessary:

let indices = RequiredIndices::new_for_all_exprs(&plan);

optimize_projections(plan, config, indices)

}

}

/// Removes unnecessary columns (e.g. columns that do not appear in the output

/// schema and/or are not used during any computation step such as expression

/// evaluation) from the logical plan and its inputs.

///

/// # Parameters

///

/// - `plan`: A reference to the input `LogicalPlan` to optimize.

/// - `config`: A reference to the optimizer configuration.

/// - `indices`: A slice of column indices that represent the necessary column

/// indices for downstream (parent) plan nodes.

///

/// # Returns

///

/// A `Result` object with the following semantics:

///

/// - `Ok(Some(LogicalPlan))`: An optimized `LogicalPlan` without unnecessary

/// columns.

/// - `Ok(None)`: Signal that the given logical plan did not require any change.

/// - `Err(error)`: An error occurred during the optimization process.

#[cfg_attr(feature = "recursive_protection", recursive::recursive)]

fn optimize_projections(

plan: LogicalPlan,

config: &dyn OptimizerConfig,

indices: RequiredIndices,

) -> Result<Transformed<LogicalPlan>> {

// Recursively rewrite any nodes that may be able to avoid computation given

// their parents' required indices.

match plan {

LogicalPlan::Projection(proj) => {

return merge_consecutive_projections(proj)?

.transform_data(|proj| {

rewrite_projection_given_requirements(proj, config, &indices)

})?

.transform_data(|plan| optimize_subqueries(plan, config));

}

LogicalPlan::Aggregate(aggregate) => {

// Split parent requirements to GROUP BY and aggregate sections:

let n_group_exprs = aggregate.group_expr_len()?;

// Offset aggregate indices so that they point to valid indices at

// `aggregate.aggr_expr`:

let (group_by_reqs, aggregate_reqs) = indices.split_off(n_group_exprs);

// Get absolutely necessary GROUP BY fields.

//

// When the input has no functional dependencies, we can

// short-circuit this analysis.

let new_group_bys = if aggregate

.input

.schema()

.functional_dependencies()

.is_empty()

{

aggregate.group_expr

} else {

let group_by_input_indices = aggregate

.group_expr

.iter()

.map(|group_by_expr| {

passthrough_field_index(group_by_expr, aggregate.input.schema())

})

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

if let Some(simplest_groupby_indices) =

get_required_group_by_exprs_indices(

aggregate.input.schema(),

&group_by_input_indices,

)

{

// Some of the fields in the GROUP BY may be required by

// the parent even if these fields are unnecessary in

// terms of functional dependency.

group_by_reqs

.append(&simplest_groupby_indices)

.get_at_indices(&aggregate.group_expr)

} else {

aggregate.group_expr

}

};

// Only use the absolutely necessary aggregate expressions required

// by the parent:

let new_aggr_expr = aggregate_reqs.get_at_indices(&aggregate.aggr_expr);

if new_group_bys.is_empty() && new_aggr_expr.is_empty() {

// Global aggregation with no aggregate functions always produces 1 row and no columns.

return Ok(Transformed::yes(LogicalPlan::EmptyRelation(

EmptyRelation {

produce_one_row: true,

schema: Arc::new(DFSchema::empty()),

},

)));

}

let all_exprs_iter = new_group_bys.iter().chain(new_aggr_expr.iter());

let schema = aggregate.input.schema();

let necessary_indices =

RequiredIndices::new().with_exprs(schema, all_exprs_iter);

let necessary_exprs = necessary_indices.get_required_exprs(schema);

return optimize_projections(

Arc::unwrap_or_clone(aggregate.input),

config,

necessary_indices,

)?

.transform_data(|aggregate_input| {

// Simplify the input of the aggregation by adding a projection so

// that its input only contains absolutely necessary columns for

// the aggregate expressions. Note that necessary_indices refer to

// fields in `aggregate.input.schema()`.

add_projection_on_top_if_helpful(aggregate_input, necessary_exprs)

})?

.map_data(|aggregate_input| {

// Create a new aggregate plan with the updated input and only the

// absolutely necessary fields:

Aggregate::try_new(

Arc::new(aggregate_input),

new_group_bys,

new_aggr_expr,

)

.map(LogicalPlan::Aggregate)

})?

.transform_data(|plan| optimize_subqueries(plan, config));

}

LogicalPlan::Window(window) => {

let input_schema = Arc::clone(window.input.schema());

// Split parent requirements to child and window expression sections:

let n_input_fields = input_schema.fields().len();

// Offset window expression indices so that they point to valid

// indices at `window.window_expr`:

let (child_reqs, window_reqs) = indices.split_off(n_input_fields);

// Only use window expressions that are absolutely necessary according

// to parent requirements:

let new_window_expr = window_reqs.get_at_indices(&window.window_expr);

// Get all the required column indices at the input, either by the

// parent or window expression requirements.

let required_indices = child_reqs.with_exprs(&input_schema, &new_window_expr);

return optimize_projections(

Arc::unwrap_or_clone(window.input),

config,

required_indices.clone(),

)?

.transform_data(|window_child| {

if new_window_expr.is_empty() {

// When no window expression is necessary, use the input directly:

Ok(Transformed::no(window_child))

} else {

// Calculate required expressions at the input of the window.

// Please note that we use `input_schema`, because `required_indices`

// refers to that schema

let required_exprs =

required_indices.get_required_exprs(&input_schema);

let window_child =

add_projection_on_top_if_helpful(window_child, required_exprs)?

.data;

Window::try_new(new_window_expr, Arc::new(window_child))

.map(LogicalPlan::Window)

.map(Transformed::yes)

}

})?

.transform_data(|plan| optimize_subqueries(plan, config));

}

LogicalPlan::TableScan(table_scan) => {

// Get indices referred to in the original (schema with all fields)

// given projected indices.

let projection = match &table_scan.projection {

Some(projection) => indices.into_mapped_indices(|idx| projection[idx]),

None => indices.into_inner(),

};

let new_scan = TableScanBuilder::from(table_scan)

.with_projection(Some(projection))

.build()?;

return Transformed::yes(LogicalPlan::TableScan(new_scan))

.transform_data(|plan| optimize_subqueries(plan, config));

}

// Other node types are handled below

_ => {}

}

// For other plan node types, calculate indices for columns they use and

// try to rewrite their children

let mut child_required_indices: Vec<RequiredIndices> = match &plan {

LogicalPlan::Sort(_)

| LogicalPlan::Filter(_)

| LogicalPlan::Repartition(_)

| LogicalPlan::Union(_)

| LogicalPlan::SubqueryAlias(_)

| LogicalPlan::Distinct(Distinct::On(_)) => {

// Pass index requirements from the parent as well as column indices

// that appear in this plan's expressions to its child. All these

// operators benefit from "small" inputs, so the projection_beneficial

// flag is `true`.

plan.inputs()

.into_iter()

.map(|input| {

indices

.clone()

.with_projection_beneficial()

.with_plan_exprs(&plan, input.schema())

})

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

}

LogicalPlan::Limit(_) => {

// Pass index requirements from the parent as well as column indices

// that appear in this plan's expressions to its child. These operators

// do not benefit from "small" inputs, so the projection_beneficial

// flag is `false`.

plan.inputs()

.into_iter()

.map(|input| indices.clone().with_plan_exprs(&plan, input.schema()))

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

}

LogicalPlan::Copy(_)

| LogicalPlan::Ddl(_)

| LogicalPlan::Dml(_)

| LogicalPlan::Explain(_)

| LogicalPlan::Analyze(_)

| LogicalPlan::Subquery(_)

| LogicalPlan::Statement(_)

| LogicalPlan::Distinct(Distinct::All(_)) => {

// These plans require all their fields, and their children should

// be treated as final plans -- otherwise, we may have schema a

// mismatch.

// TODO: For some subquery variants (e.g. a subquery arising from an

// EXISTS expression), we may not need to require all indices.

plan.inputs()

.into_iter()

.map(RequiredIndices::new_for_all_exprs)

.collect()

}

LogicalPlan::Extension(extension) => {

if let Some(necessary_children_indices) =

extension.node.necessary_children_exprs(indices.indices())

{

let children = extension.node.inputs();

assert_eq_or_internal_err!(

children.len(),

necessary_children_indices.len(),

"Inconsistent length between children and necessary children indices. \

Make sure `.necessary_children_exprs` implementation of the \

`UserDefinedLogicalNode` is consistent with actual children length \

for the node."

);

children

.into_iter()

.zip(necessary_children_indices)

.map(|(child, necessary_indices)| {

RequiredIndices::new_from_indices(necessary_indices)

.with_plan_exprs(&plan, child.schema())

})

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

} else {

// Requirements from parent cannot be routed down to user defined logical plan safely

// Assume it requires all input exprs here

plan.inputs()

.into_iter()

.map(RequiredIndices::new_for_all_exprs)

.collect()

}

}

LogicalPlan::MaterializedCte(MaterializedCtePlan::Cte(cte)) => {

// The scans are leaves whose schema is the schema of the body, so

// the body keeps every column. The output is the output of the

// continuation, so the continuation needs what the parent needs.

vec![

RequiredIndices::new_for_all_exprs(&cte.cte),

RequiredIndices::new_from_indices(indices.into_inner()),

]

}

LogicalPlan::EmptyRelation(_)

| LogicalPlan::Values(_)

| LogicalPlan::MaterializedCte(MaterializedCtePlan::Scan(_))

| LogicalPlan::Show(_) => {

// These operators have no inputs, so stop the optimization process.

return Ok(Transformed::no(plan));

}

LogicalPlan::RecursiveQuery(_) => {

// optimize the static and recursive terms: treat each recursive CTE term like a

// standalone subquery: optimize its internals, but do not push parent required indices

// through the RecursiveQuery boundary, as this can otherwise lead to bugs

// (see: https://github.com/apache/datafusion/issues/22249)

return plan.map_children(|c| {

let indices = RequiredIndices::new_for_all_exprs(&c);

optimize_projections(c, config, indices)

});

}

LogicalPlan::Join(join) => {

let left_len = join.left.schema().fields().len();

let right_len = join.right.schema().fields().len();

let (left_req_indices, right_req_indices) =

split_join_requirements(left_len, right_len, indices, &join.join_type);

let mut left_indices =

left_req_indices.with_plan_exprs(&plan, join.left.schema())?;

let mut right_indices =

right_req_indices.with_plan_exprs(&plan, join.right.schema())?;

// Ensure an empty mark join still has a column to qualify mark

match join.join_type {

JoinType::LeftMark if right_indices.indices().is_empty() => {

right_indices = right_indices.append(&[0]);

}

JoinType::RightMark if left_indices.indices().is_empty() => {

left_indices = left_indices.append(&[0]);

}

_ => {}

}

// Joins benefit from "small" input tables (lower memory usage).

// Therefore, each child benefits from projection:

vec![

left_indices.with_projection_beneficial(),

right_indices.with_projection_beneficial(),

]

}

LogicalPlan::AsOfJoin(join) => {

let left_len = join.left.schema().fields().len();

let mut left_required = Vec::new();

let mut right_required = Vec::new();

for index in indices.indices() {

if *index < left_len {

left_required.push(*index);

} else {

right_required.push(*index - left_len);

}

}

let left_indices = RequiredIndices::new_from_indices(left_required)

.with_plan_exprs(&plan, join.left.schema())?;

let right_indices = RequiredIndices::new_from_indices(right_required)

.with_plan_exprs(&plan, join.right.schema())?;

vec![

left_indices.with_projection_beneficial(),

right_indices.with_projection_beneficial(),

]

}

// these nodes are explicitly rewritten in the match statement above

LogicalPlan::Projection(_)

| LogicalPlan::Aggregate(_)

| LogicalPlan::Window(_)

| LogicalPlan::TableScan(_) => {

return internal_err!(

"OptimizeProjection: should have handled in the match statement above"

);

}

LogicalPlan::Unnest(Unnest {

input,

dependency_indices,

..

}) => {

// at least provide the indices for the exec-columns as a starting point

let required_indices =

RequiredIndices::new().with_plan_exprs(&plan, input.schema())?;

// Add additional required indices from the parent

let mut additional_necessary_child_indices =

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

indices.indices().iter().for_each(|idx| {

if let Some(index) = dependency_indices.get(*idx) {

additional_necessary_child_indices.push(*index);

}

});

vec![required_indices.append(&additional_necessary_child_indices)]

}

};

// Required indices are currently ordered (child0, child1, ...)

// but the loop pops off the last element, so we need to reverse the order

child_required_indices.reverse();

assert_eq_or_internal_err!(

child_required_indices.len(),

plan.inputs().len(),

"OptimizeProjection: child_required_indices length mismatch with plan inputs"

);

// Rewrite children of the plan

let transformed_plan = plan.map_children(|child| {

let required_indices = child_required_indices.pop().ok_or_else(|| {

internal_datafusion_err!(

"Unexpected number of required_indices in OptimizeProjections rule"

)

})?;

let projection_beneficial = required_indices.projection_beneficial();

let project_exprs = required_indices.get_required_exprs(child.schema());

optimize_projections(child, config, required_indices)?.transform_data(

|new_input| {

if projection_beneficial {

add_projection_on_top_if_helpful(new_input, project_exprs)

} else {

Ok(Transformed::no(new_input))

}

},

)

})?;

let transformed_plan =

transformed_plan.transform_data(|plan| optimize_subqueries(plan, config))?;

// If any of the children are transformed, we need to potentially update the plan's schema

if transformed_plan.transformed {

transformed_plan.map_data(|plan| plan.recompute_schema())

} else {

Ok(transformed_plan)

}

}

/// Optimizes uncorrelated subquery plans embedded in expressions of the given

/// plan node (e.g., `Expr::ScalarSubquery`). `map_children` only visits direct

/// plan inputs, so subqueries must be handled separately.

fn optimize_subqueries(

plan: LogicalPlan,

config: &dyn OptimizerConfig,

) -> Result<Transformed<LogicalPlan>> {

plan.map_uncorrelated_subqueries(|subquery_plan| {

let indices = RequiredIndices::new_for_all_exprs(&subquery_plan);

optimize_projections(subquery_plan, config, indices)

})

}

/// Given a projection `proj`, this function attempts to merge it with a previous

/// projection if it exists and if merging is beneficial. Merging is considered

/// beneficial when expressions in the current projection are non-trivial and

/// appear more than once in its input fields. This can act as a caching mechanism

/// for non-trivial computations.

///

/// ## Metadata Handling During Projection Merging

///

/// **Alias metadata preservation**: When merging projections, alias metadata from both

/// the current and previous projections is carefully preserved. The presence of metadata

/// precludes alias trimming.

///

/// **Schema, Fields, and metadata**: If a projection is rewritten, the schema and metadata

/// are preserved. Individual field properties and metadata flows through expression rewriting

/// and are preserved when fields are referenced in the merged projection.

/// Refer to [`projection_schema`](datafusion_expr::logical_plan::projection_schema)

/// for more details.

///

/// # Parameters

///

/// * `proj` - A reference to the `Projection` to be merged.

///

/// # Returns

///

/// A `Result` object with the following semantics:

///

/// - `Ok(Some(Projection))`: Merge was beneficial and successful. Contains the

/// merged projection.

/// - `Ok(None)`: Signals that merge is not beneficial (and has not taken place).

/// - `Err(error)`: An error occurred during the function call.

fn merge_consecutive_projections(proj: Projection) -> Result<Transformed<Projection>> {

// Collapse the whole chain in one pass; otherwise an N-deep chain needs

// N outer optimizer passes to fully fold.

let mut current = proj;

let mut transformed_any = false;

loop {

let Transformed {

data, transformed, ..

} = merge_consecutive_projections_one_level(current)?;

current = data;

if !transformed {

break;

}

transformed_any = true;

}

Ok(if transformed_any {

Transformed::yes(current)

} else {

Transformed::no(current)

})

}

fn merge_consecutive_projections_one_level(

proj: Projection,

) -> Result<Transformed<Projection>> {

let Projection {

expr,

input,

schema,

..

} = proj;

let LogicalPlan::Projection(prev_projection) = input.as_ref() else {

return Projection::try_new_with_schema(expr, input, schema).map(Transformed::no);

};

// Count usages (referrals) of each projection expression in its input fields:

let mut column_referral_map = HashMap::<&Column, usize>::new();

for expr in &expr {

expr.add_column_ref_counts(&mut column_referral_map);

}

// If an expression is non-trivial (KeepInPlace) and appears more than once, do not merge

// them as consecutive projections will benefit from a compute-once approach.

// For details, see: https://github.com/apache/datafusion/issues/8296

if column_referral_map.into_iter().any(|(col, usage)| {

usage > 1

&& !prev_projection.expr[prev_projection.schema.index_of_column(col).unwrap()]

.placement()

.should_push_to_leaves()

}) {

// no change

return Projection::try_new_with_schema(expr, input, schema).map(Transformed::no);

}

let LogicalPlan::Projection(prev_projection) = Arc::unwrap_or_clone(input) else {

// We know it is a `LogicalPlan::Projection` from check above

unreachable!();

};

// Try to rewrite the expressions in the current projection using the

// previous projection as input:

let name_preserver = NamePreserver::new_for_projection();

let mut original_names = Vec::with_capacity(expr.len());

let new_exprs = expr.map_elements(|expr| {

original_names.push(name_preserver.save(&expr));

// do not rewrite top level Aliases (rewriter will remove all aliases within exprs)

match expr {

Expr::Alias(Alias {

expr,

relation,

name,

metadata,

}) => rewrite_expr(*expr, &prev_projection).map(|result| {

result.update_data(|expr| {

// After substitution, the inner expression may now have the

// same schema_name as the alias (e.g. when an extraction

// alias like `__extracted_1 AS f(x)` is resolved back to

// `f(x)`). Wrapping in a redundant self-alias causes a

// cosmetic `f(x) AS f(x)` due to Display vs schema_name

// formatting differences. Drop the alias when it matches.

if metadata.is_none() && expr.schema_name().to_string() == name {

expr

} else {

Expr::Alias(Alias {

expr: Box::new(expr),

relation,

name,

metadata,

})

}

})

}),

e => rewrite_expr(e, &prev_projection),

}

})?;

// if the expressions could be rewritten, create a new projection with the

// new expressions

if new_exprs.transformed {

// Add any needed aliases back to the expressions

let new_exprs = new_exprs

.data

.into_iter()

.zip(original_names)

.map(|(expr, original_name)| original_name.restore(expr))

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

Projection::try_new(new_exprs, prev_projection.input).map(Transformed::yes)

} else {

// not rewritten, so put the projection back together

let input = Arc::new(LogicalPlan::Projection(prev_projection));

Projection::try_new_with_schema(new_exprs.data, input, schema)

.map(Transformed::no)

}

}

/// Rewrites a projection expression using the projection before it (i.e. its input)

/// This is a subroutine to the `merge_consecutive_projections` function.

///

/// # Parameters

///

/// * `expr` - A reference to the expression to rewrite.

/// * `input` - A reference to the input of the projection expression (itself

/// a projection).

///

/// # Returns

///

/// A `Result` object with the following semantics:

///

/// - `Ok(Some(Expr))`: Rewrite was successful. Contains the rewritten result.

/// - `Ok(None)`: Signals that `expr` can not be rewritten.

/// - `Err(error)`: An error occurred during the function call.

///

/// # Notes

/// This rewrite also removes any unnecessary layers of aliasing. "Unnecessary" is

/// defined as not contributing new information, such as metadata.

///

/// Without trimming, we can end up with unnecessary indirections inside expressions

/// during projection merges.

///

/// Consider:

///

/// ```text

/// Projection(a1 + b1 as sum1)

/// --Projection(a as a1, b as b1)

/// ----Source(a, b)

/// ```

///

/// After merge, we want to produce:

///

/// ```text

/// Projection(a + b as sum1)

/// --Source(a, b)

/// ```

///

/// Without trimming, we would end up with:

///

/// ```text

/// Projection((a as a1 + b as b1) as sum1)

/// --Source(a, b)

/// ```

fn rewrite_expr(expr: Expr, input: &Projection) -> Result<Transformed<Expr>> {

expr.transform_up(|expr| {

match expr {

// remove any intermediate aliases if they do not carry metadata

Expr::Alias(alias) => {

match alias

.metadata

.as_ref()

.map(|h| h.is_empty())

.unwrap_or(true)

{

true => Ok(Transformed::yes(*alias.expr)),

false => Ok(Transformed::no(Expr::Alias(alias))),

}

}

Expr::Column(col) => {

// Find index of column:

let idx = input.schema.index_of_column(&col)?;

// get the corresponding unaliased input expression

//

// For example:

// * the input projection is [`a + b` as c, `d + e` as f]

// * the current column is an expression "f"

//

// return the expression `d + e` (not `d + e` as f)

let input_expr = input.expr[idx].clone().unalias_nested().data;

Ok(Transformed::yes(input_expr))

}

// Unsupported type for consecutive projection merge analysis.

_ => Ok(Transformed::no(expr)),

}

})

}

/// Splits requirement indices for a join into left and right children based on

/// the join type.

///

/// This function takes the length of the left child, a slice of requirement

/// indices, and the type of join (e.g. `INNER`, `LEFT`, `RIGHT`) as arguments.

/// Depending on the join type, it divides the requirement indices into those

/// that apply to the left child and those that apply to the right child.

///

/// - For `INNER`, `LEFT`, `RIGHT`, `FULL`, `LEFTMARK`, and `RIGHTMARK` joins,

/// the requirements are split between left and right children. The right

/// child indices are adjusted to point to valid positions within the right

/// child by subtracting the length of the left child.

///

/// - For `LEFT ANTI`, `LEFT SEMI`, `RIGHT SEMI` and `RIGHT ANTI` joins, all

/// requirements are re-routed to either the left child or the right child

/// directly, depending on the join type.

///

/// # Parameters

///

/// * `left_len` - The length of the left child.

/// * `right_len` - The length of the right child.

/// * `indices` - A slice of requirement indices.

/// * `join_type` - The type of join (e.g. `INNER`, `LEFT`, `RIGHT`).

///

/// # Returns

///

/// A tuple containing two vectors of `usize` indices: The first vector represents

/// the requirements for the left child, and the second vector represents the

/// requirements for the right child. The indices are appropriately split and

/// adjusted based on the join type.

fn split_join_requirements(

left_len: usize,

right_len: usize,

indices: RequiredIndices,

join_type: &JoinType,

) -> (RequiredIndices, RequiredIndices) {

match join_type {

// In these cases requirements are split between left/right children:

JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => {

// Decrease right side indices by `left_len` so that they point to valid

// positions within the right child:

indices.split_off(left_len)

}

JoinType::LeftMark => {

// LeftMark output: [left_cols(0..left_len), mark]

// The mark column is synthetic (produced by the join itself),

// so discard it and route only to the left child.

let (left_indices, _mark) = indices.split_off(left_len);

(left_indices, RequiredIndices::new())

}

JoinType::RightMark => {

// Same as LeftMark, but for the right child.

let (right_indices, _mark) = indices.split_off(right_len);

(RequiredIndices::new(), right_indices)

}

// All requirements can be re-routed to left child directly.

JoinType::LeftAnti | JoinType::LeftSemi => (indices, RequiredIndices::new()),

// All requirements can be re-routed to right side directly.

// No need to change index, join schema is right child schema.

JoinType::RightSemi | JoinType::RightAnti => (RequiredIndices::new(), indices),

}

}

/// Adds a projection on top of a logical plan if doing so reduces the number

/// of columns for the parent operator.

///

/// This function takes a `LogicalPlan` and a list of projection expressions.

/// If the projection is beneficial (it reduces the number of columns in the

/// plan) a new `LogicalPlan` with the projection is created and returned, along

/// with a `true` flag. If the projection doesn't reduce the number of columns,

/// the original plan is returned with a `false` flag.

///

/// # Parameters

///

/// * `plan` - The input `LogicalPlan` to potentially add a projection to.

/// * `project_exprs` - A list of expressions for the projection.

///

/// # Returns

///

/// A `Transformed` indicating if a projection was added

fn add_projection_on_top_if_helpful(

plan: LogicalPlan,

project_exprs: Vec<Expr>,

) -> Result<Transformed<LogicalPlan>> {

// Make sure projection decreases the number of columns, otherwise it is unnecessary.

if project_exprs.len() >= plan.schema().fields().len() {

Ok(Transformed::no(plan))

} else {

Projection::try_new(project_exprs, Arc::new(plan))

.map(LogicalPlan::Projection)

.map(Transformed::yes)

}

}

/// Rewrite the given projection according to the fields required by its

/// ancestors.

///

/// # Parameters

///

/// * `proj` - A reference to the original projection to rewrite.

/// * `config` - A reference to the optimizer configuration.

/// * `indices` - A slice of indices representing the columns required by the

/// ancestors of the given projection.

///

/// # Returns

///

/// A `Result` object with the following semantics:

///

/// - `Ok(Some(LogicalPlan))`: Contains the rewritten projection

/// - `Ok(None)`: No rewrite necessary.

/// - `Err(error)`: An error occurred during the function call.

fn rewrite_projection_given_requirements(

proj: Projection,

config: &dyn OptimizerConfig,

indices: &RequiredIndices,

) -> Result<Transformed<LogicalPlan>> {

let Projection {

expr,

input,

schema,

..

} = proj;

let exprs_used = indices.get_at_indices(&expr);

let required_indices =

RequiredIndices::new().with_exprs(input.schema(), exprs_used.iter());

let old_input_schema = Arc::clone(input.schema());

let read_indices = required_indices.indices().to_vec();

// rewrite the children projection, and if they are changed rewrite the

// projection down

optimize_projections(Arc::unwrap_or_clone(input), config, required_indices)?

.transform_data(|input| {

if is_projection_unnecessary(&input, &exprs_used)? {

Ok(Transformed::yes(input))

} else if exprs_used.iter().any(contains_subquery)

|| !read_fields_unchanged(

&old_input_schema,

input.schema(),

&read_indices,

)

{

// A subquery is optimized on its own, and rewriting the input

// can change the fields this projection reads, so the stored

// schema may be stale. Recompute it rather than reuse it.

Projection::try_new(exprs_used, Arc::new(input))

.map(LogicalPlan::Projection)

.map(Transformed::yes)

} else {

// The retained expressions are a subset of the original ones and

// read unchanged input fields, so their output fields are

// unchanged too. Select them from the existing schema instead of

// recomputing each one via `Expr::to_field` (see

// `projection_schema`), which costs O(exprs * schema_width).

// `merge_consecutive_projections` reuses the schema the same way.

let projected_schema =

project_schema_by_indices(&schema, indices.indices())?;

Projection::try_new_with_schema(

exprs_used,

Arc::new(input),

projected_schema,

)

.map(LogicalPlan::Projection)

.map(Transformed::yes)

}

})

}

/// Returns true if the input fields at `read_indices` of `old` are unchanged in

/// `new`, the input after it was rewritten for those indices.

///

/// The rewritten input either keeps every field of `old` or only the fields at

/// `read_indices`, in order. Any other shape is treated as changed.

fn read_fields_unchanged(old: &DFSchema, new: &DFSchema, read_indices: &[usize]) -> bool {

if new.fields().len() == old.fields().len() {

read_indices.iter().all(|&i| old.field(i) == new.field(i))

} else if new.fields().len() == read_indices.len() {

read_indices

.iter()

.enumerate()

.all(|(j, &i)| old.field(i) == new.field(j))

} else {

false

}

}

/// Returns true if `expr` contains a subquery expression.

fn contains_subquery(expr: &Expr) -> bool {

expr.exists(|e| {

Ok(matches!(

e,

Expr::Exists(_)

| Expr::InSubquery(_)

| Expr::SetComparison(_)

| Expr::ScalarSubquery(_)

))

})

.expect("exists closure is infallible")

}

/// Builds the output schema of a projection that keeps only the fields at

/// `indices` (a sorted, deduplicated subset produced by [`RequiredIndices`]) of

/// `schema`, reusing the already-computed fields instead of recomputing each

/// field's type/nullability from the expressions.

///

/// Pruning unreferenced sibling columns cannot change the retained fields, so

/// as long as `schema` is current for the projection's expressions and input,

/// the sliced fields are identical to the ones [`projection_schema`] would

/// recompute, at O(k) instead of O(exprs * schema_width). Functional

/// dependencies can be weaker than recomputed ones, never stronger.

///

/// [`projection_schema`]: datafusion_expr::logical_plan::projection_schema

fn project_schema_by_indices(

schema: &DFSchemaRef,

indices: &[usize],

) -> Result<DFSchemaRef> {

// Nothing pruned: the output schema is unchanged, reuse it as-is.

if indices.len() == schema.fields().len() {

return Ok(Arc::clone(schema));

}

let qualified_fields = indices

.iter()

.map(|&i| {

let (qualifier, field) = schema.qualified_field(i);

(qualifier.cloned(), Arc::clone(field))

})

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

let func_deps = schema

.functional_dependencies()

.project_functional_dependencies(indices, indices.len());

let projected =

DFSchema::new_with_metadata(qualified_fields, schema.metadata().clone())?

.with_functional_dependencies(func_deps)?;

Ok(Arc::new(projected))

}

/// Projection is unnecessary, when

/// - input schema of the projection, output schema of the projection are same, and

/// - all projection expressions are either Column or Literal

pub fn is_projection_unnecessary(

input: &LogicalPlan,

proj_exprs: &[Expr],

) -> Result<bool> {

// First check if the number of expressions is equal to the number of fields in the input schema.

if proj_exprs.len() != input.schema().fields().len() {

return Ok(false);

}

Ok(input.schema().iter().zip(proj_exprs.iter()).all(