// 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

//! Two-pass optimizer pipeline that pushes cheap expressions (like struct field

//! access `user['status']`) closer to data sources, enabling early data reduction

//! and source-level optimizations (e.g., Parquet column pruning). See

//! [`ExtractLeafExpressions`] (pass 1) and [`PushDownLeafProjections`] (pass 2).

//!

//! # Precedence over [`PushDownFilter`]

//!

//! [`PushDownLeafProjections`] and [`PushDownFilter`] both move nodes towards

//! the leaves, and for an adjacent filter and *pure extraction projection* (a

//! projection whose expressions are only `__datafusion_extracted_N` aliases and

//! pass-through columns) they want the opposite order:

//!

//! ```text

//! Filter: t.date = '2025-01-03' <-- (A)

//! Projection: get_field(t.ids, 'id1') AS __datafusion_extracted_1, t.date <-- (B)

//! TableScan: t

//! ```

//!

//! `PushDownFilter` wants (A) below (B). `PushDownLeafProjections` wants (B)

//! below (A). Before <https://github.com/apache/datafusion/issues/14540> was

//! decided the two rules undid each other on every optimizer pass, and the rule

//! that runs later in the list won.

//!

//! **Invariant: pure extraction projections win.** `PushDownFilter` does not

//! move a filter below a pure extraction projection, so the plan above is the

//! final plan. `PushDownLeafProjections` keeps moving such a projection through

//! a filter, which is how the projection reaches the scan.

//!

//! The reason is that the extraction projection is the node a source absorbs.

//! A Parquet scan merges it into the file projection and reads only the struct

//! leaf. With the opposite precedence the scan reads the whole struct whenever

//! `datafusion.execution.parquet.pushdown_filters` is `false`, which is the

//! default. The filter loses nothing by staying one node higher, because

//! `PushDownFilter` records the predicate in

//! [`TableScan::filters`](datafusion_expr::logical_plan::TableScan) in the pass

//! that runs before the extraction projection exists, so row group pruning and

//! source level filtering still happen.

//!

//! # Yields to [`OptimizeProjections`]

//!

//! When the extraction projection cannot move below the input of a projection

//! (for example a `TableScan` or an `Aggregate`), [`PushDownLeafProjections`]

//! leaves the projection unchanged. Splitting it in place would give a

//! recovery projection over an extraction projection on the same input, and

//! `OptimizeProjections` merges those two back into the original projection.

//! The two rules would then undo each other on every optimizer pass. The split

//! has no use there: the source absorbs the leaf expressions of the original

//! projection just as well.

//!

//! [`PushDownFilter`]: crate::push_down_filter::PushDownFilter

//! [`OptimizeProjections`]: crate::optimize_projections::OptimizeProjections

use indexmap::{IndexMap, IndexSet};

use std::collections::{BTreeSet, HashMap};

use std::sync::Arc;

use datafusion_common::alias::AliasGenerator;

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

use datafusion_common::{Column, DFSchema, Result, qualified_name};

use datafusion_expr::logical_plan::LogicalPlan;

use datafusion_expr::{Expr, ExpressionPlacement, Projection};

use crate::optimizer::ApplyOrder;

use crate::push_down_filter::replace_cols_by_name;

use crate::utils::{ColumnReference, has_all_column_refs, schema_columns};

use crate::{OptimizerConfig, OptimizerRule};

/// Prefix for aliases generated by the extraction optimizer passes.

///

/// This prefix is **reserved for internal optimizer use**. User-defined aliases

/// starting with this prefix may be misidentified as optimizer-generated

/// extraction aliases, leading to unexpected behavior. Do not use this prefix

/// in user queries.

const EXTRACTED_EXPR_PREFIX: &str = "__datafusion_extracted";

/// Returns `true` if any sub-expression in `exprs` has

/// [`ExpressionPlacement::MoveTowardsLeafNodes`] placement.

///

/// This is a lightweight pre-check that short-circuits as soon as one

/// extractable expression is found, avoiding the expensive allocations

/// (column HashSets, extractors, expression rewrites) that the full

/// extraction pipeline requires.

fn has_extractable_expr(exprs: &[Expr]) -> bool {

exprs.iter().any(|expr| {

expr.exists(|e| Ok(e.placement() == ExpressionPlacement::MoveTowardsLeafNodes))

.unwrap_or(false)

})

}

/// Returns the flat names of `plan`'s output columns whose defining expression

/// is volatile (e.g. `random()`).

///

/// Only a [`LogicalPlan::Projection`] can define such a column: anywhere else

/// the value has already been materialized by the projection that produced it,

/// so referencing the column again does not re-evaluate anything.

fn volatile_output_columns(plan: &LogicalPlan) -> BTreeSet<String> {

let LogicalPlan::Projection(projection) = plan else {

return BTreeSet::new();

};

projection

.schema

.iter()

.zip(projection.expr.iter())

.filter(|(_, expr)| expr.is_volatile())

.map(|((qualifier, field), _)| Column::from((qualifier, field)).flat_name())

.collect()

}

/// Returns `true` if building an extraction projection for `exprs` on top of

/// `input` would duplicate a volatile computation.

///

/// When `input` is already a projection, [`build_extraction_projection_impl`]

/// *merges* into it: every column reference in an extracted expression is

/// replaced by that column's defining expression (see

/// [`build_projection_replace_map`]). Inlining a volatile definition makes the

/// merged projection evaluate it a second, independent time, so the extracted

/// value no longer matches the column it was derived from:

///

/// ```text

/// Projection: s, get_field(s, 'a') AS field

/// Projection: named_struct('a', random()) AS s

/// ```

///

/// would merge into a single projection computing `random()` twice, and

/// `field` would then differ from `s['a']` on every row. Callers skip the

/// extraction instead.

fn would_duplicate_volatile<'a>(

exprs: impl IntoIterator<Item = &'a Expr>,

input: &LogicalPlan,

) -> bool {

let volatile = volatile_output_columns(input);

if volatile.is_empty() {

return false;

}

exprs.into_iter().any(|expr| {

expr.column_refs()

.iter()

.any(|col| volatile.contains(&col.flat_name()))

})

}

/// Extracts `MoveTowardsLeafNodes` sub-expressions from non-projection nodes

/// into **extraction projections** (pass 1 of 2).

///

/// This handles Filter, Sort, Limit, Aggregate, and Join nodes. For Projection

/// nodes, extraction and pushdown are handled by [`PushDownLeafProjections`].

///

/// # Key Concepts

///

/// **Extraction projection**: a projection inserted *below* a node that

/// pre-computes a cheap expression and exposes it under an alias

/// (`__datafusion_extracted_N`). The parent node then references the alias

/// instead of the original expression.

///

/// **Recovery projection**: a projection inserted *above* a node to restore

/// the original output schema when extraction changes it.

/// Schema-preserving nodes (Filter, Sort, Limit) gain extra columns from

/// the extraction projection that bubble up; the recovery projection selects

/// only the original columns to hide the extras.

///

/// # Example

///

/// Given a filter with a struct field access:

///

/// ```text

/// Filter: user['status'] = 'active'

/// TableScan: t [id, user]

/// ```

///

/// This rule:

/// 1. Inserts an **extraction projection** below the filter:

/// 2. Adds a **recovery projection** above to hide the extra column:

///

/// ```text

/// Projection: id, user <-- recovery projection

/// Filter: __datafusion_extracted_1 = 'active'

/// Projection: user['status'] AS __datafusion_extracted_1, id, user <-- extraction projection

/// TableScan: t [id, user]

/// ```

///

/// **Important:** The `PushDownFilter` rule is aware of projections created by this rule

/// and will not push filters through them. It uses `ExpressionPlacement` to detect

/// `MoveTowardsLeafNodes` expressions and skip filter pushdown past them.

#[derive(Default, Debug)]

pub struct ExtractLeafExpressions {}

impl ExtractLeafExpressions {

/// Create a new [`ExtractLeafExpressions`]

pub fn new() -> Self {

Self {}

}

}

impl OptimizerRule for ExtractLeafExpressions {

fn name(&self) -> &str {

"extract_leaf_expressions"

}

fn rewrite(

&self,

plan: LogicalPlan,

config: &dyn OptimizerConfig,

) -> Result<Transformed<LogicalPlan>> {

if !config.options().optimizer.enable_leaf_expression_pushdown {

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

}

let alias_generator = config.alias_generator();

// Advance the alias generator past any user-provided __datafusion_extracted_N

// aliases to prevent collisions when generating new extraction aliases.

advance_generator_past_existing(&plan, alias_generator)?;

plan.transform_down_with_subqueries(|plan| {

extract_from_plan(plan, alias_generator)

})

}

}

/// Scans the plan for pre-existing `__datafusion_extracted_N` aliases and

/// advances the generator counter past them to avoid collisions with

/// user-provided aliases.

///

/// Subquery plans nested inside expressions are scanned as well: extraction

/// rewrites with `transform_down_with_subqueries`, so it can generate aliases

/// *inside* a subquery and would otherwise collide with a user alias there.

fn advance_generator_past_existing(

plan: &LogicalPlan,

alias_generator: &AliasGenerator,

) -> Result<()> {

plan.apply_with_subqueries(|plan| {

plan.apply_expressions(|expr| {

expr.apply(|e| {

if let Expr::Alias(alias) = e

&& let Some(id) = alias

.name

.strip_prefix(EXTRACTED_EXPR_PREFIX)

.and_then(|s| s.strip_prefix('_'))

.and_then(|s| s.parse().ok())

{

alias_generator.update_min_id(id);

}

Ok(TreeNodeRecursion::Continue)

})

})

})

.map(|_| ())

}

/// Extracts `MoveTowardsLeafNodes` sub-expressions from a plan node.

///

/// Works for any number of inputs (0, 1, 2, …N). For multi-input nodes

/// like Join, each extracted sub-expression is routed to the correct input

/// by checking which input's schema contains all of the expression's column

/// references.

fn extract_from_plan(

plan: LogicalPlan,

alias_generator: &Arc<AliasGenerator>,

) -> Result<Transformed<LogicalPlan>> {

// Only extract from plan types whose output schema is predictable after

// expression rewriting. Nodes like Window derive column names from

// their expressions, so rewriting `get_field` inside a window function

// changes the output schema and breaks the recovery projection.

if !matches!(

&plan,

LogicalPlan::Aggregate(_)

| LogicalPlan::Filter(_)

| LogicalPlan::Sort(_)

| LogicalPlan::Limit(_)

| LogicalPlan::Join(_)

) {

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

}

let inputs = plan.inputs();

if inputs.is_empty() {

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

}

// Fast pre-check: skip all allocations if no extractable expressions exist

let node_exprs = plan.expressions();

if !has_extractable_expr(&node_exprs) {

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

}

// The extraction projection is merged into an input that is already a

// projection, which inlines the referenced columns' definitions. Skip the

// extraction when that would duplicate a volatile computation.

if inputs

.iter()

.any(|input| would_duplicate_volatile(node_exprs.iter(), input))

{

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

}

// Save original output schema before any transformation

let original_schema = Arc::clone(plan.schema());

// Build per-input schemas from borrowed inputs (before plan is consumed

// by map_expressions). We only need schemas and column sets for routing;

// the actual inputs are cloned later only if extraction succeeds.

let input_schemas: Vec<Arc<DFSchema>> =

inputs.iter().map(|i| Arc::clone(i.schema())).collect();

// Build per-input extractors

let mut extractors: Vec<LeafExpressionExtractor> = input_schemas

.iter()

.map(|schema| LeafExpressionExtractor::new(schema.as_ref(), alias_generator))

.collect();

// Build per-input column sets for routing expressions to the correct input

let input_column_sets: Vec<std::collections::HashSet<ColumnReference>> =

input_schemas

.iter()

.map(|schema| schema_columns(schema.as_ref()))

.collect();

// Transform expressions via map_expressions with routing

let transformed = plan.map_expressions(|expr| {

routing_extract(expr, &mut extractors, &input_column_sets)

})?;

// If no expressions were rewritten, nothing was extracted

if !transformed.transformed {

return Ok(transformed);

}

// Clone inputs now that we know extraction succeeded. Wrap in Arc

// upfront since build_extraction_projection expects &Arc<LogicalPlan>.

let owned_inputs: Vec<Arc<LogicalPlan>> = transformed

.data

.inputs()

.into_iter()

.map(|i| Arc::new(i.clone()))

.collect();

// Build per-input extraction projections (None means no extractions for that input)

let new_inputs: Vec<LogicalPlan> = owned_inputs

.into_iter()

.zip(extractors.iter())

.map(|(input_arc, extractor)| {

match extractor.build_extraction_projection(&input_arc)? {

Some(plan) => Ok(plan),

// No extractions for this input — recover the LogicalPlan

// without cloning (refcount is 1 since build returned None).

None => Ok(Arc::unwrap_or_clone(input_arc)),

}

})

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

// Rebuild the plan keeping its rewritten expressions but replacing

// inputs with the new extraction projections.

let new_plan = transformed

.data

.with_new_exprs(transformed.data.expressions(), new_inputs)?;

// Add recovery projection if the output schema changed

let recovered = build_recovery_projection(original_schema.as_ref(), new_plan)?;

Ok(Transformed::yes(recovered))

}

/// Given an expression, returns the index of the input whose columns fully

/// cover the expression's column references.

/// Returns `None` if the expression references columns from multiple inputs

/// or if multiple inputs match (ambiguous, e.g. unqualified columns present

/// in both sides of a join).

fn find_owning_input(

expr: &Expr,

input_column_sets: &[std::collections::HashSet<ColumnReference>],

) -> Option<usize> {

let mut found = None;

for (idx, cols) in input_column_sets.iter().enumerate() {

if has_all_column_refs(expr, cols) {

if found.is_some() {

// Ambiguous — multiple inputs match

return None;

}

found = Some(idx);

}

}

found

}

/// Walks an expression tree top-down, extracting `MoveTowardsLeafNodes`

/// sub-expressions and routing each to the correct per-input extractor.

fn routing_extract(

expr: Expr,

extractors: &mut [LeafExpressionExtractor],

input_column_sets: &[std::collections::HashSet<ColumnReference>],

) -> Result<Transformed<Expr>> {

expr.transform_down(|e| {

// Skip expressions already aliased with extracted expression pattern

if let Expr::Alias(alias) = &e

&& alias.name.starts_with(EXTRACTED_EXPR_PREFIX)

{

return Ok(Transformed {

data: e,

transformed: false,

tnr: TreeNodeRecursion::Jump,

});

}

// Don't extract Alias nodes directly — preserve the alias and let

// transform_down recurse into the inner expression

if matches!(&e, Expr::Alias(_)) {

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

}

match e.placement() {

ExpressionPlacement::MoveTowardsLeafNodes => {

if let Some(idx) = find_owning_input(&e, input_column_sets) {

let col_ref = extractors[idx].add_extracted(e)?;

Ok(Transformed::yes(col_ref))

} else {

// References columns from multiple inputs — cannot extract

Ok(Transformed::no(e))

}

}

ExpressionPlacement::Column => {

// Track columns that the parent node references so the

// extraction projection includes them as pass-through.

// Without this, the extraction projection would only

// contain __datafusion_extracted_N aliases, and the parent couldn't

// resolve its other column references.

if let Expr::Column(col) = &e

&& let Some(idx) = find_owning_input(&e, input_column_sets)

{

extractors[idx].columns_needed.insert(col.clone());

}

Ok(Transformed::no(e))

}

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

}

})

}

/// Rewrites extraction pairs and column references from one qualifier

/// space to another.

///

/// Builds a replacement map by zipping `from_schema` (whose qualifiers

/// currently appear in `pairs` / `columns`) with `to_schema` (the

/// qualifiers we want), then applies `replace_cols_by_name`.

///

/// Used for SubqueryAlias (alias-space -> input-space) and Union

/// (union output-space -> per-branch input-space).

fn remap_pairs_and_columns(

pairs: &[(Expr, String)],

columns: &IndexSet<Column>,

from_schema: &DFSchema,

to_schema: &DFSchema,

) -> Result<ExtractionTarget> {

let mut replace_map = HashMap::with_capacity(from_schema.fields().len());

for ((from_q, from_f), (to_q, to_f)) in from_schema.iter().zip(to_schema.iter()) {

replace_map.insert(

qualified_name(from_q, from_f.name()),

Expr::Column(Column::new(to_q.cloned(), to_f.name())),

);

}

let remapped_pairs: Vec<(Expr, String)> = pairs

.iter()

.map(|(expr, alias)| {

Ok((

replace_cols_by_name(expr.clone(), &replace_map)?,

alias.clone(),

))

})

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

let remapped_columns: IndexSet<Column> = columns

.iter()

.filter_map(|col| {

let rewritten =

replace_cols_by_name(Expr::Column(col.clone()), &replace_map).ok()?;

if let Expr::Column(c) = rewritten {

Some(c)

} else {

Some(col.clone())

}

})

.collect();

Ok(ExtractionTarget {

pairs: remapped_pairs,

columns: remapped_columns,

})

}

// =============================================================================

// Helper Types & Functions for Extraction Targeting

// =============================================================================

/// A bundle of extraction pairs (expression + alias) and standalone columns

/// that need to be pushed through a plan node.

struct ExtractionTarget {

/// Extracted expressions paired with their generated aliases.

pairs: Vec<(Expr, String)>,

/// Standalone column references needed by the parent node.

columns: IndexSet<Column>,

}

/// Build a replacement map from a projection: output_column_name -> underlying_expr.

///

/// This is used to resolve column references through a renaming projection.

/// For example, if a projection has `user AS x`, this maps `x` -> `col("user")`.

fn build_projection_replace_map(projection: &Projection) -> HashMap<String, Expr> {

projection

.schema

.iter()

.zip(projection.expr.iter())

.map(|((qualifier, field), expr)| {

let key = Column::from((qualifier, field)).flat_name();

(key, expr.clone().unalias())

})

.collect()

}

/// Build a recovery projection to restore the original output schema.

///

/// After extraction, a node's output schema may differ from the original:

///

/// - **Schema-preserving nodes** (Filter/Sort/Limit): the extraction projection

/// below adds extra `__datafusion_extracted_N` columns that bubble up through

/// the node. Recovery selects only the original columns to hide the extras.

/// ```text

/// Original schema: [id, user]

/// After extraction: [__datafusion_extracted_1, id, user] ← extra column leaked through

/// Recovery: SELECT id, user FROM ... ← hides __datafusion_extracted_1

/// ```

///

/// - **Schema-defining nodes** (Aggregate): same number of columns but names

/// may differ because extracted aliases replaced the original expressions.

/// Recovery maps positionally, aliasing where names changed.

/// ```text

/// Original: [SUM(user['balance'])]

/// After: [SUM(__datafusion_extracted_1)] ← name changed

/// Recovery: SUM(__datafusion_extracted_1) AS "SUM(user['balance'])"

/// ```

///

/// - **Schemas identical** → no recovery projection needed.

fn build_recovery_projection(

original_schema: &DFSchema,

input: LogicalPlan,

) -> Result<LogicalPlan> {

let new_schema = input.schema();

let orig_len = original_schema.fields().len();

let new_len = new_schema.fields().len();

if orig_len == new_len {

// Same number of fields — check if schemas are identical

let schemas_match = original_schema.iter().zip(new_schema.iter()).all(

|((orig_q, orig_f), (new_q, new_f))| {

orig_f.name() == new_f.name() && orig_q == new_q

},

);

if schemas_match {

return Ok(input);

}

// Schema-defining nodes (Aggregate, Join): names may differ at some

// positions because extracted aliases replaced the original expressions.

// Map positionally, aliasing where the name changed.

//

// Invariant: `with_new_exprs` on all supported node types (Aggregate,

// Filter, Sort, Limit, Join) preserves column order, so positional

// mapping is safe here.

debug_assert!(

orig_len == new_len,

"build_recovery_projection: positional mapping requires same field count, \

got original={orig_len} vs new={new_len}"

);

let mut proj_exprs = Vec::with_capacity(orig_len);

for (i, (orig_qualifier, orig_field)) in original_schema.iter().enumerate() {

let (new_qualifier, new_field) = new_schema.qualified_field(i);

if orig_field.name() == new_field.name() && orig_qualifier == new_qualifier {

proj_exprs.push(Expr::from((orig_qualifier, orig_field)));

} else {

let new_col = Expr::Column(Column::from((new_qualifier, new_field)));

proj_exprs.push(

new_col.alias_qualified(orig_qualifier.cloned(), orig_field.name()),

);

}

}

let projection = Projection::try_new(proj_exprs, Arc::new(input))?;

Ok(LogicalPlan::Projection(projection))

} else {

// Schema-preserving nodes: new schema has extra extraction columns.

// Original columns still exist by name; select them to hide extras.

let col_exprs: Vec<Expr> = original_schema.iter().map(Expr::from).collect();

let projection = Projection::try_new(col_exprs, Arc::new(input))?;

Ok(LogicalPlan::Projection(projection))

}

}

/// Collects `MoveTowardsLeafNodes` sub-expressions found during expression

/// tree traversal and can build an extraction projection from them.

///

/// # Example

///

/// Given `Filter: user['status'] = 'active' AND user['name'] IS NOT NULL`:

/// - `add_extracted(user['status'])` → stores it, returns `col("__datafusion_extracted_1")`

/// - `add_extracted(user['name'])` → stores it, returns `col("__datafusion_extracted_2")`

/// - `build_extraction_projection()` produces:

/// `Projection: user['status'] AS __datafusion_extracted_1, user['name'] AS __datafusion_extracted_2, <all input columns>`

struct LeafExpressionExtractor<'a> {

/// Extracted expressions: maps expression -> alias

extracted: IndexMap<Expr, String>,

/// Columns referenced by extracted expressions or the parent node,

/// included as pass-through in the extraction projection.

columns_needed: IndexSet<Column>,

/// Input schema

input_schema: &'a DFSchema,

/// Alias generator

alias_generator: &'a Arc<AliasGenerator>,

}

impl<'a> LeafExpressionExtractor<'a> {

fn new(input_schema: &'a DFSchema, alias_generator: &'a Arc<AliasGenerator>) -> Self {

Self {

extracted: IndexMap::new(),

columns_needed: IndexSet::new(),

input_schema,

alias_generator,

}

}

/// Adds an expression to extracted set, returns column reference.

fn add_extracted(&mut self, expr: Expr) -> Result<Expr> {

// Deduplication: reuse existing alias if same expression

if let Some(alias) = self.extracted.get(&expr) {

return Ok(Expr::Column(Column::new_unqualified(alias)));

}

// Track columns referenced by this expression

for col in expr.column_refs() {

self.columns_needed.insert(col.clone());

}

// Generate unique alias

let alias = self.alias_generator.next(EXTRACTED_EXPR_PREFIX);

self.extracted.insert(expr, alias.clone());

Ok(Expr::Column(Column::new_unqualified(&alias)))

}

/// Builds an extraction projection above the given input, or merges into

/// it if the input is already a projection. Delegates to

/// [`build_extraction_projection_impl`].

///

/// Returns `None` if there are no extractions.

fn build_extraction_projection(

&self,

input: &Arc<LogicalPlan>,

) -> Result<Option<LogicalPlan>> {

if self.extracted.is_empty() {

return Ok(None);

}

let pairs: Vec<(Expr, String)> = self

.extracted

.iter()

.map(|(e, a)| (e.clone(), a.clone()))

.collect();

let proj = build_extraction_projection_impl(

&pairs,

&self.columns_needed,

input,

self.input_schema,

)?;

Ok(Some(LogicalPlan::Projection(proj)))

}

}

/// Resolves `col` against `schema` and returns it with the qualifier that

/// `schema` gives the field. Returns `None` if `schema` does not hold `col`,

/// or if the name is ambiguous in `schema`.

///

/// Two spellings of the same input column (a bare `c` and a qualified `t.c`)

/// resolve to the same `Column`, so they compare equal. A column pushed into

/// a projection this way uses the name that the input gives it.

fn resolve_against(schema: &DFSchema, col: &Column) -> Option<Column> {

schema

.qualified_field_from_column(col)

.ok()

.map(Column::from)

}

/// Build an extraction projection above the target node (shared by both passes).

///

/// If the target is an existing projection, merges into it. This requires

/// resolving column references through the projection's rename mapping:

/// if the projection has `user AS u`, and an extracted expression references

/// `u['name']`, we must rewrite it to `user['name']` since the merged

/// projection reads from the same input as the original.

///

/// Deduplicates by resolved expression equality and adds pass-through

/// columns as needed. Otherwise builds a fresh projection with extracted

/// expressions + ALL input schema columns.

fn build_extraction_projection_impl(

extracted_exprs: &[(Expr, String)],

columns_needed: &IndexSet<Column>,

target: &Arc<LogicalPlan>,

target_schema: &DFSchema,

) -> Result<Projection> {

if let LogicalPlan::Projection(existing) = target.as_ref() {

// Merge into existing projection

let mut proj_exprs = existing.expr.clone();

// Build a map of existing expressions (by Expr equality) to their aliases

let existing_extractions: IndexMap<Expr, String> = existing

.expr

.iter()

.filter_map(|e| {

if let Expr::Alias(alias) = e

&& alias.name.starts_with(EXTRACTED_EXPR_PREFIX)

{

return Some((*alias.expr.clone(), alias.name.clone()));

}

None

})

.collect();

// Resolve column references through the projection's rename mapping

let replace_map = build_projection_replace_map(existing);

// Add new extracted expressions, resolving column refs through the projection

for (expr, alias) in extracted_exprs {

let resolved = replace_cols_by_name(expr.clone().alias(alias), &replace_map)?;

let resolved_inner = if let Expr::Alias(a) = &resolved {

a.expr.as_ref()

} else {

&resolved

};

if let Some(existing_alias) = existing_extractions.get(resolved_inner) {

// Same expression already extracted under a different alias —

// add the expression with the new alias so both names are

// available in the output. We can't reference the existing alias

// as a column within the same projection, so we duplicate the

// computation.

if existing_alias != alias {

proj_exprs.push(resolved);

}

} else {

proj_exprs.push(resolved);

}

}

// Add any new pass-through columns that aren't already in the projection.

// We check against existing.input.schema() (the projection's source) rather

// than target_schema (the projection's output) because columns produced

// by alias expressions (e.g., CSE's __common_expr_N) exist in the output but

// not the input, and cannot be added as pass-through Column references.

//

// Resolve the pass-through columns of `existing` (including same-name

// aliases such as `t.c AS c`) and the columns in `columns_needed`

// against the input schema before comparing them. Otherwise a bare `c`

// and a qualified `t.c` do not match, and the merged projection holds

// both, which `Projection::try_new` rejects as ambiguous.

let input_schema = existing.input.schema();

let existing_cols: IndexSet<Column> = existing

.expr

.iter()

.filter_map(|e| resolve_against(input_schema, passthrough_column(e)?))

.collect();

// The names the merged projection already produces. The merge keeps

// every expression of `existing`, so the parent can still read each of

// these names, and a pass-through column that carries one of them makes

// the output schema ambiguous.

let output_names: std::collections::HashSet<&str> = existing

.schema

.fields()

.iter()

.map(|f| f.name().as_str())

.collect();

for col in columns_needed {

// The projection produces this name, so the parent reads it from the

// projection's own output. Do not resolve the name through the rename

// map: a rename such as `t.a AS b` is a computed output, and its

// input column `t.a` beside the output field `a` of a second rename

// gives an ambiguous schema (issue #25446).

if output_names.contains(col.name.as_str()) {

continue;

}

let col_expr = Expr::Column(col.clone());

let resolved = replace_cols_by_name(col_expr, &replace_map)?;

if let Expr::Column(resolved_col) = &resolved

&& let Some(input_col) = resolve_against(input_schema, resolved_col)

&& !existing_cols.contains(&input_col)

&& !output_names.contains(input_col.name.as_str())

{

proj_exprs.push(Expr::Column(input_col));

}

// Skip a column that resolves to a non-column expr: the existing

// projection already computes it. Also skip a column that is not in

// the input or is ambiguous there: no single spelling is correct.

}

Projection::try_new(proj_exprs, Arc::clone(&existing.input))

} else {

// Build new projection with extracted expressions + all input columns

let mut proj_exprs =

Vec::with_capacity(extracted_exprs.len() + target_schema.fields().len());

for (expr, alias) in extracted_exprs {

proj_exprs.push(expr.clone().alias(alias));

}

for (qualifier, field) in target_schema.iter() {

proj_exprs.push(Expr::from((qualifier, field)));

}

Projection::try_new(proj_exprs, Arc::clone(target))

}

}

// =============================================================================

// Pass 2: PushDownLeafProjections

// =============================================================================

/// Pushes extraction projections down through schema-preserving nodes towards

/// leaf nodes (pass 2 of 2, after [`ExtractLeafExpressions`]).

///

/// Handles two types of projections:

/// - **Pure extraction projections** (all `__datafusion_extracted` aliases + columns):

/// pushes through Filter/Sort/Limit, merges into existing projections, or routes

/// into multi-input node inputs (Join, SubqueryAlias, etc.)

/// - **Mixed projections** (user projections containing `MoveTowardsLeafNodes`

/// sub-expressions): splits into a recovery projection + extraction projection,

/// then pushes the extraction projection down.

///

/// # Example: Pushing through a Filter

///

/// After pass 1, the extraction projection sits directly below the filter:

/// ```text

/// Projection: id, user <-- recovery

/// Filter: __datafusion_extracted_1 = 'active'

/// Projection: user['status'] AS __datafusion_extracted_1, id, user <-- extraction

/// TableScan: t [id, user]

/// ```

///

/// Pass 2 pushes the extraction projection through the recovery and filter,

/// and a subsequent `OptimizeProjections` pass removes the (now-redundant)

/// recovery projection:

/// ```text

/// Filter: __datafusion_extracted_1 = 'active'

/// Projection: user['status'] AS __datafusion_extracted_1, id, user <-- extraction (pushed down)

/// TableScan: t [id, user]

/// ```

#[derive(Default, Debug)]

pub struct PushDownLeafProjections {}

impl PushDownLeafProjections {

pub fn new() -> Self {

Self {}

}

}

impl OptimizerRule for PushDownLeafProjections {

fn name(&self) -> &str {

"push_down_leaf_projections"

}

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

Some(ApplyOrder::TopDown)

}

fn rewrite(

&self,

plan: LogicalPlan,

config: &dyn OptimizerConfig,

) -> Result<Transformed<LogicalPlan>> {

if !config.options().optimizer.enable_leaf_expression_pushdown {

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

}

let alias_generator = config.alias_generator();

match try_push_input(&plan, alias_generator)? {

Some(new_plan) => Ok(Transformed::yes(new_plan)),

None => Ok(Transformed::no(plan)),

}

}

}

/// Attempts to push a projection's extractable expressions further down.

///

/// Returns `Some(new_subtree)` if the projection was pushed down or merged,

/// `None` if there is nothing to push or the projection sits above a barrier.

fn try_push_input(

input: &LogicalPlan,

alias_generator: &Arc<AliasGenerator>,

) -> Result<Option<LogicalPlan>> {

let LogicalPlan::Projection(proj) = input else {

return Ok(None);

};

split_and_push_projection(proj, alias_generator)

}

/// If this is a passthrough column. I.e,

///

/// - A bare column expression

/// - A trivial alias rename (i.e, `p.column AS column`)

fn passthrough_column(expr: &Expr) -> Option<&Column> {

match expr {

Expr::Column(col) => Some(col),

Expr::Alias(alias) => match alias.expr.as_ref() {

Expr::Column(col) if col.name == alias.name => Some(col),

_ => None,

},

_ => None,

}

}

/// Returns `true` if merging the extracted `pairs` (and the pass-through

/// `columns_needed`) into `child` would inline a `KeepInPlace` definition into

/// more than one reference site — e.g. a struct-returning UDF `f(c)` behind both

/// `f(c)['a']` and `f(c)['b']`, or behind `f(c)['a']` and a bare `f(c)`.

///

/// `CommonSubexprEliminate` hoists such an expression into a shared column so it

/// runs once; [`build_extraction_projection_impl`] then resolves that column

/// back to its definition and re-duplicates it. Mirrors the

/// `optimize_projections` merge guard (#8296): a compute-once expression used

/// more than once stays in place.

fn merge_would_duplicate_kept_expr(

pairs: &[(Expr, String)],

columns_needed: &IndexSet<Column>,

child: &Projection,

) -> bool {

// Columns whose `KeepInPlace` definition must stay put; inlining one would

// duplicate a compute-once expression. Plain columns, pushable exprs, and

// literals are all cheap to duplicate, so only `KeepInPlace` counts.

let kept_columns: std::collections::HashSet<String> = child

.schema

.iter()

.zip(child.expr.iter())

.filter_map(|((qualifier, field), expr)| {

if matches!(expr.placement(), ExpressionPlacement::KeepInPlace) {

Some(Column::from((qualifier, field)).flat_name())

} else {

None

}

})

.collect();

if kept_columns.is_empty() {

return false;

}

// Count every site that references a kept column. A pair that references the

// same kept column twice still inlines one copy, so count each pair at most

// once per column.

let mut usage: HashMap<String, usize> = HashMap::new();

let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();

for (expr, _alias) in pairs {

seen.clear();

for col in expr.column_refs() {

let name = col.flat_name();

if kept_columns.contains(&name) && seen.insert(name.clone()) {

*usage.entry(name).or_insert(0) += 1;

}

}

}

// Pass-through / pre-existing-alias inputs are additional reference sites:

// a bare `f(c)` alongside `f(c)['a']` also pins the compute-once column, so

// counting only `pairs` under-counted and re-duplicated it (#23655).

for col in columns_needed {

let name = col.flat_name();

if kept_columns.contains(&name) {

*usage.entry(name).or_insert(0) += 1;

}

}

usage.values().any(|&count| count > 1)

}

/// Splits a projection into extractable pieces, pushes them towards leaf

/// nodes, and adds a recovery projection if needed.

///

/// Handles both:

/// - **Pure extraction projections** (all `__datafusion_extracted` aliases + columns)

/// - **Mixed projections** (containing `MoveTowardsLeafNodes` sub-expressions)

///

/// Returns `Some(new_subtree)` if extractions were pushed down,

/// `None` if there is nothing to extract or push.

///

/// # Example: Mixed Projection

///

/// ```text

/// Input plan:

/// Projection: user['name'] IS NOT NULL AS has_name, id

/// Filter: ...

/// TableScan

///

/// Phase 1 (Split):

/// extraction_pairs: [(user['name'], "__datafusion_extracted_1")]

/// recovery_exprs: [__datafusion_extracted_1 IS NOT NULL AS has_name, id]

///

/// Phase 2 (Push):

/// Push extraction projection through Filter toward TableScan

///

/// Phase 3 (Recovery):

/// Projection: __datafusion_extracted_1 IS NOT NULL AS has_name, id <-- recovery

/// Filter: ...

/// Projection: user['name'] AS __datafusion_extracted_1, id <-- extraction (pushed)

/// TableScan