datafusion_optimizer/

extract_leaf_expressions.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! 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).

use indexmap::{IndexMap, IndexSet};
use std::collections::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::has_all_column_refs;
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)
    })
}

/// 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 current plan node's expressions for pre-existing
/// `__datafusion_extracted_N` aliases and advances the generator
/// counter past them to avoid collisions with user-provided aliases.
fn advance_generator_past_existing(
    plan: &LogicalPlan,
    alias_generator: &AliasGenerator,
) -> Result<()> {
    plan.apply(|plan| {
        plan.expressions().iter().try_for_each(|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)
            })?;
            Ok::<(), datafusion_common::error::DataFusionError>(())
        })?;
        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
    if !has_extractable_expr(&plan.expressions()) {
        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<Column>> = 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::try_unwrap(input_arc).unwrap_or_else(|arc| (*arc).clone()))
                }
            }
        })
        .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<Column>],
) -> 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<Column>],
) -> 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)),
        }
    })
}

/// Returns all columns in the schema (both qualified and unqualified forms)
fn schema_columns(schema: &DFSchema) -> std::collections::HashSet<Column> {
    schema
        .iter()
        .flat_map(|(qualifier, field)| {
            [
                Column::new(qualifier.cloned(), field.name()),
                Column::new_unqualified(field.name()),
            ]
        })
        .collect()
}

/// 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::new();
    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)))
    }
}

/// 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.
        let existing_cols: IndexSet<Column> = existing
            .expr
            .iter()
            .filter_map(|e| {
                if let Expr::Column(c) = e {
                    Some(c.clone())
                } else {
                    None
                }
            })
            .collect();

        let input_schema = existing.input.schema();
        for col in columns_needed {
            let col_expr = Expr::Column(col.clone());
            let resolved = replace_cols_by_name(col_expr, &replace_map)?;
            if let Expr::Column(resolved_col) = &resolved
                && !existing_cols.contains(resolved_col)
                && input_schema.has_column(resolved_col)
            {
                proj_exprs.push(Expr::Column(resolved_col.clone()));
            }
            // If resolved to non-column expr, it's already computed by existing projection
        }

        Projection::try_new(proj_exprs, Arc::clone(&existing.input))
    } else {
        // Build new projection with extracted expressions + all input columns
        let mut proj_exprs = Vec::new();
        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)
}

/// 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
/// ```
fn split_and_push_projection(
    proj: &Projection,
    alias_generator: &Arc<AliasGenerator>,
) -> Result<Option<LogicalPlan>> {
    // Fast pre-check: skip if there are no pre-existing extracted aliases
    // and no new extractable expressions.
    let has_existing_extracted = proj.expr.iter().any(|e| {
        matches!(e, Expr::Alias(alias) if alias.name.starts_with(EXTRACTED_EXPR_PREFIX))
    });
    if !has_existing_extracted && !has_extractable_expr(&proj.expr) {
        return Ok(None);
    }

    let input = &proj.input;
    let input_schema = input.schema();

    // ── Phase 1: Split ──────────────────────────────────────────────────
    // For each projection expression, collect extraction pairs and build
    // recovery expressions.
    //
    // Pre-existing `__datafusion_extracted` aliases are inserted into the
    // extractor's `IndexMap` with the **full** `Expr::Alias(…)` as the key,
    // so the alias name participates in equality. This prevents collisions
    // when CSE rewrites produce the same inner expression under different
    // alias names (e.g. `__common_expr_4 AS __datafusion_extracted_1` and
    // `__common_expr_4 AS __datafusion_extracted_3`). New extractions from
    // `routing_extract` use bare (non-Alias) keys and get normal dedup.
    //
    // When building the final `extraction_pairs`, the Alias wrapper is
    // stripped so consumers see the usual `(inner_expr, alias_name)` tuples.

    let mut extractors = vec![LeafExpressionExtractor::new(
        input_schema.as_ref(),
        alias_generator,
    )];
    let input_column_sets = vec![schema_columns(input_schema.as_ref())];

    let original_schema = proj.schema.as_ref();
    let mut recovery_exprs: Vec<Expr> = Vec::with_capacity(proj.expr.len());
    let mut needs_recovery = false;
    let mut has_new_extractions = false;
    let mut proj_exprs_captured: usize = 0;
    // Track standalone column expressions (Case B) to detect column refs
    // from extracted aliases (Case A) that aren't also standalone expressions.
    let mut standalone_columns: IndexSet<Column> = IndexSet::new();

    for (expr, (qualifier, field)) in proj.expr.iter().zip(original_schema.iter()) {
        if let Expr::Alias(alias) = expr
            && alias.name.starts_with(EXTRACTED_EXPR_PREFIX)
        {
            // Insert the full Alias expression as the key so that
            // distinct alias names don't collide in the IndexMap.
            let alias_name = alias.name.clone();

            for col_ref in alias.expr.column_refs() {
                extractors[0].columns_needed.insert(col_ref.clone());
            }

            extractors[0]
                .extracted
                .insert(expr.clone(), alias_name.clone());
            recovery_exprs.push(Expr::Column(Column::new_unqualified(&alias_name)));
            proj_exprs_captured += 1;
        } else if let Expr::Column(col) = expr {
            // Plain column pass-through — track it in the extractor
            extractors[0].columns_needed.insert(col.clone());
            standalone_columns.insert(col.clone());
            recovery_exprs.push(expr.clone());
            proj_exprs_captured += 1;
        } else {
            // Everything else: run through routing_extract
            let transformed =
                routing_extract(expr.clone(), &mut extractors, &input_column_sets)?;
            if transformed.transformed {
                has_new_extractions = true;
            }
            let transformed_expr = transformed.data;

            // Build recovery expression, aliasing back to original name if needed
            let original_name = field.name();
            let needs_alias = if let Expr::Column(col) = &transformed_expr {
                col.name.as_str() != original_name
            } else {
                let expr_name = transformed_expr.schema_name().to_string();
                original_name != &expr_name
            };
            let recovery_expr = if needs_alias {
                needs_recovery = true;
                transformed_expr
                    .clone()
                    .alias_qualified(qualifier.cloned(), original_name)
            } else {
                transformed_expr.clone()
            };

            // If the expression was transformed (i.e., has extracted sub-parts),
            // it differs from what the pushed projection outputs → needs recovery.
            // Also, any non-column, non-__datafusion_extracted expression needs recovery
            // because the pushed extraction projection won't output it directly.
            if transformed.transformed || !matches!(expr, Expr::Column(_)) {
                needs_recovery = true;
            }

            recovery_exprs.push(recovery_expr);
        }
    }

    // Build extraction_pairs, stripping the Alias wrapper from pre-existing
    // entries (they used the full Alias as the map key to avoid dedup).
    let extractor = &extractors[0];
    let extraction_pairs: Vec<(Expr, String)> = extractor
        .extracted
        .iter()
        .map(|(e, a)| match e {
            Expr::Alias(alias) => (*alias.expr.clone(), a.clone()),
            _ => (e.clone(), a.clone()),
        })
        .collect();
    let columns_needed = &extractor.columns_needed;

    // If no extractions found, nothing to do
    if extraction_pairs.is_empty() {
        return Ok(None);
    }

    // If columns_needed has entries that aren't standalone projection columns
    // (i.e., they came from column refs inside extracted aliases), a merge
    // into an inner projection will widen the schema with those extra columns,
    // requiring a recovery projection to restore the original schema.
    if columns_needed
        .iter()
        .any(|c| !standalone_columns.contains(c))
    {
        needs_recovery = true;
    }

    // ── Phase 2: Push down ──────────────────────────────────────────────
    let proj_input = Arc::clone(&proj.input);
    let pushed = push_extraction_pairs(
        &extraction_pairs,
        columns_needed,
        proj,
        &proj_input,
        alias_generator,
        proj_exprs_captured,
    )?;

    // ── Phase 3: Recovery ───────────────────────────────────────────────
    // Determine the base plan: either the pushed result or an in-place extraction.
    let base_plan = match pushed {
        Some(plan) => plan,
        None => {
            if !has_new_extractions {
                // Only pre-existing __datafusion_extracted aliases and columns, no new
                // extractions from routing_extract. The original projection is
                // already an extraction projection that couldn't be pushed
                // further. Return None.
                return Ok(None);
            }
            // Build extraction projection in-place (couldn't push down)
            let input_arc = Arc::clone(input);
            let extraction = build_extraction_projection_impl(
                &extraction_pairs,
                columns_needed,
                &input_arc,
                input_schema.as_ref(),
            )?;
            LogicalPlan::Projection(extraction)
        }
    };

    // Wrap with recovery projection if the output schema changed
    if needs_recovery {
        let recovery = LogicalPlan::Projection(Projection::try_new(
            recovery_exprs,
            Arc::new(base_plan),
        )?);
        Ok(Some(recovery))
    } else {
        Ok(Some(base_plan))
    }
}

/// Returns true if the plan is a Projection where ALL expressions are either
/// `Alias(EXTRACTED_EXPR_PREFIX, ...)` or `Column`, with at least one extraction.
/// Such projections can safely be pushed further without re-extraction.
fn is_pure_extraction_projection(plan: &LogicalPlan) -> bool {
    let LogicalPlan::Projection(proj) = plan else {
        return false;
    };
    let mut has_extraction = false;
    for expr in &proj.expr {
        match expr {
            Expr::Alias(alias) if alias.name.starts_with(EXTRACTED_EXPR_PREFIX) => {
                has_extraction = true;
            }
            Expr::Column(_) => {}
            _ => return false,
        }
    }
    has_extraction
}

/// Pushes extraction pairs down through the projection's input node,
/// dispatching to the appropriate handler based on the input node type.
fn push_extraction_pairs(
    pairs: &[(Expr, String)],
    columns_needed: &IndexSet<Column>,
    proj: &Projection,
    proj_input: &Arc<LogicalPlan>,
    alias_generator: &Arc<AliasGenerator>,
    proj_exprs_captured: usize,
) -> Result<Option<LogicalPlan>> {
    match proj_input.as_ref() {
        // Merge into existing projection, then try to push the result further down.
        // Only merge when every expression in the outer projection is fully
        // captured as either an extraction pair (Case A: __datafusion_extracted
        // alias) or a plain column (Case B). Uncaptured expressions (e.g.
        // `col AS __common_expr_1` from CSE, or complex expressions with
        // extracted sub-parts) would be lost during the merge.
        LogicalPlan::Projection(_) if proj_exprs_captured == proj.expr.len() => {
            let target_schema = Arc::clone(proj_input.schema());
            let merged = build_extraction_projection_impl(
                pairs,
                columns_needed,
                proj_input,
                target_schema.as_ref(),
            )?;
            let merged_plan = LogicalPlan::Projection(merged);

            // After merging, try to push the result further down, but ONLY
            // if the merged result is still a pure extraction projection
            // (all __datafusion_extracted aliases + columns). If the merge inherited
            // bare MoveTowardsLeafNodes expressions from the inner projection,
            // pushing would re-extract them into new aliases and fail when
            // the (None, true) fallback can't find the original aliases.
            // This handles: Extraction → Recovery(cols) → Filter → ... → TableScan
            // by pushing through the recovery projection AND the filter in one pass.
            if is_pure_extraction_projection(&merged_plan)
                && let Some(pushed) = try_push_input(&merged_plan, alias_generator)?
            {
                return Ok(Some(pushed));
            }
            Ok(Some(merged_plan))
        }
        // Generic: handles Filter/Sort/Limit (via recursion),
        // SubqueryAlias (with qualifier remap in try_push_into_inputs),
        // Join, and anything else.
        // Safely bails out for nodes that don't pass through extracted
        // columns (Aggregate, Window) via the output schema check.
        _ => try_push_into_inputs(
            pairs,
            columns_needed,
            proj_input.as_ref(),
            alias_generator,
        ),
    }
}

/// Routes extraction pairs and columns to the appropriate inputs.
///
/// - **Union**: broadcasts to every input via [`remap_pairs_and_columns`].
/// - **Other nodes**: routes each expression to the one input that owns
///   all of its column references (via [`find_owning_input`]).
///
/// Returns `None` if any expression can't be routed or no input has pairs.
fn route_to_inputs(
    pairs: &[(Expr, String)],
    columns: &IndexSet<Column>,
    node: &LogicalPlan,
    input_column_sets: &[std::collections::HashSet<Column>],
    input_schemas: &[Arc<DFSchema>],
) -> Result<Option<Vec<ExtractionTarget>>> {
    let num_inputs = input_schemas.len();
    let mut per_input: Vec<ExtractionTarget> = (0..num_inputs)
        .map(|_| ExtractionTarget {
            pairs: vec![],
            columns: IndexSet::new(),
        })
        .collect();

    if matches!(node, LogicalPlan::Union(_)) {
        // Union output schema and each input schema have the same fields by
        // index but may differ in qualifiers (e.g. output `s` vs input
        // `simple_struct.s`). Remap pairs/columns to each input's space.
        let union_schema = node.schema();
        for (idx, input_schema) in input_schemas.iter().enumerate() {
            per_input[idx] =
                remap_pairs_and_columns(pairs, columns, union_schema, input_schema)?;
        }
    } else {
        for (expr, alias) in pairs {
            match find_owning_input(expr, input_column_sets) {
                Some(idx) => per_input[idx].pairs.push((expr.clone(), alias.clone())),
                None => return Ok(None), // Cross-input expression — bail out
            }
        }
        for col in columns {
            let col_expr = Expr::Column(col.clone());
            match find_owning_input(&col_expr, input_column_sets) {
                Some(idx) => {
                    per_input[idx].columns.insert(col.clone());
                }
                None => return Ok(None), // Ambiguous column — bail out
            }
        }
    }

    // Check at least one input has extractions to push
    if per_input.iter().all(|t| t.pairs.is_empty()) {
        return Ok(None);
    }

    Ok(Some(per_input))
}

/// Pushes extraction expressions into a node's inputs by routing each
/// expression to the input that owns all of its column references.
///
/// Works for any number of inputs (1, 2, …N). For single-input nodes,
/// all expressions trivially route to that input. For multi-input nodes
/// (Join, etc.), each expression is routed to the side that owns its columns.
///
/// Returns `Some(new_node)` if all expressions could be routed AND the
/// rebuilt node's output schema contains all extracted aliases.
/// Returns `None` if any expression references columns from multiple inputs
/// or the node doesn't pass through the extracted columns.
///
/// # Example: Join with expressions from both sides
///
/// ```text
/// Extraction projection above a Join:
///   Projection: left.user['name'] AS __datafusion_extracted_1, right.order['total'] AS __datafusion_extracted_2, ...
///     Join: left.id = right.user_id
///       TableScan: left [id, user]
///       TableScan: right [user_id, order]
///
/// After routing each expression to its owning input:
///   Join: left.id = right.user_id
///     Projection: user['name'] AS __datafusion_extracted_1, id, user              <-- left-side extraction
///       TableScan: left [id, user]
///     Projection: order['total'] AS __datafusion_extracted_2, user_id, order      <-- right-side extraction
///       TableScan: right [user_id, order]
/// ```
fn try_push_into_inputs(
    pairs: &[(Expr, String)],
    columns_needed: &IndexSet<Column>,
    node: &LogicalPlan,
    alias_generator: &Arc<AliasGenerator>,
) -> Result<Option<LogicalPlan>> {
    let inputs = node.inputs();
    if inputs.is_empty() {
        return Ok(None);
    }

    // SubqueryAlias remaps qualifiers between input and output.
    // Rewrite pairs/columns from alias-space to input-space before routing.
    let remapped = if let LogicalPlan::SubqueryAlias(sa) = node {
        remap_pairs_and_columns(pairs, columns_needed, &sa.schema, sa.input.schema())?
    } else {
        ExtractionTarget {
            pairs: pairs.to_vec(),
            columns: columns_needed.clone(),
        }
    };
    let pairs = &remapped.pairs[..];
    let columns_needed = &remapped.columns;

    // Build per-input schemas and column sets for routing
    let input_schemas: Vec<Arc<DFSchema>> =
        inputs.iter().map(|i| Arc::clone(i.schema())).collect();
    let input_column_sets: Vec<std::collections::HashSet<Column>> =
        input_schemas.iter().map(|s| schema_columns(s)).collect();

    // Route pairs and columns to the appropriate inputs
    let per_input = match route_to_inputs(
        pairs,
        columns_needed,
        node,
        &input_column_sets,
        &input_schemas,
    )? {
        Some(routed) => routed,
        None => return Ok(None),
    };

    let num_inputs = inputs.len();

    // Build per-input extraction projections and push them as far as possible
    // immediately. This is critical because map_children preserves cached schemas,
    // so if the TopDown pass later pushes a child further (changing its output
    // schema), the parent node's schema becomes stale.
    let mut new_inputs: Vec<LogicalPlan> = Vec::with_capacity(num_inputs);
    for (idx, input) in inputs.into_iter().enumerate() {
        if per_input[idx].pairs.is_empty() {
            new_inputs.push(input.clone());
        } else {
            let input_arc = Arc::new(input.clone());
            let target_schema = Arc::clone(input.schema());
            let proj = build_extraction_projection_impl(
                &per_input[idx].pairs,
                &per_input[idx].columns,
                &input_arc,
                target_schema.as_ref(),
            )?;
            // Verify all requested aliases appear in the projection's output.
            // A merge may deduplicate if the same expression already exists
            // under a different alias, leaving the requested alias missing.
            let proj_schema = proj.schema.as_ref();
            for (_expr, alias) in &per_input[idx].pairs {
                if !proj_schema.fields().iter().any(|f| f.name() == alias) {
                    return Ok(None);
                }
            }
            let proj_plan = LogicalPlan::Projection(proj);
            // Try to push the extraction projection further down within
            // this input (e.g., through Filter → existing extraction projection).
            // This ensures the input's output schema is stable and won't change
            // when the TopDown pass later visits children.
            match try_push_input(&proj_plan, alias_generator)? {
                Some(pushed) => new_inputs.push(pushed),
                None => new_inputs.push(proj_plan),
            }
        }
    }

    // Rebuild the node with new inputs
    let new_node = node.with_new_exprs(node.expressions(), new_inputs)?;

    // Safety check: verify all extracted aliases appear in the rebuilt
    // node's output schema. Nodes like Aggregate define their own output
    // and won't pass through extracted columns — bail out for those.
    let output_schema = new_node.schema();
    for (_expr, alias) in pairs {
        if !output_schema.fields().iter().any(|f| f.name() == alias) {
            return Ok(None);
        }
    }

    Ok(Some(new_node))
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use super::*;
    use crate::optimize_projections::OptimizeProjections;
    use crate::test::udfs::PlacementTestUDF;
    use crate::test::*;
    use crate::{Optimizer, OptimizerContext};
    use datafusion_common::Result;
    use datafusion_expr::expr::ScalarFunction;
    use datafusion_expr::{Expr, ExpressionPlacement};
    use datafusion_expr::{
        ScalarUDF, col, lit, logical_plan::builder::LogicalPlanBuilder,
    };

    fn leaf_udf(expr: Expr, name: &str) -> Expr {
        Expr::ScalarFunction(ScalarFunction::new_udf(
            Arc::new(ScalarUDF::new_from_impl(
                PlacementTestUDF::new()
                    .with_placement(ExpressionPlacement::MoveTowardsLeafNodes),
            )),
            vec![expr, lit(name)],
        ))
    }

    // =========================================================================
    // Combined optimization stage formatter
    // =========================================================================

    /// Runs all 4 optimization stages and returns a single formatted string.
    /// Stages that produce the same plan as the previous stage show
    /// "(same as <previous>)" to reduce noise.
    ///
    /// Stages:
    /// 1. **Original** - OptimizeProjections only (baseline)
    /// 2. **After Extraction** - + ExtractLeafExpressions
    /// 3. **After Pushdown** - + PushDownLeafProjections
    /// 4. **Optimized** - + final OptimizeProjections
    fn format_optimization_stages(plan: &LogicalPlan) -> Result<String> {
        let run = |rules: Vec<Arc<dyn OptimizerRule + Send + Sync>>| -> Result<String> {
            let ctx = OptimizerContext::new().with_max_passes(1);
            let optimizer = Optimizer::with_rules(rules);
            let optimized = optimizer.optimize(plan.clone(), &ctx, |_, _| {})?;
            Ok(format!("{optimized}"))
        };

        let original = run(vec![Arc::new(OptimizeProjections::new())])?;

        let after_extract = run(vec![
            Arc::new(OptimizeProjections::new()),
            Arc::new(ExtractLeafExpressions::new()),
        ])?;

        let after_pushdown = run(vec![
            Arc::new(OptimizeProjections::new()),
            Arc::new(ExtractLeafExpressions::new()),
            Arc::new(PushDownLeafProjections::new()),
        ])?;

        let optimized = run(vec![
            Arc::new(OptimizeProjections::new()),
            Arc::new(ExtractLeafExpressions::new()),
            Arc::new(PushDownLeafProjections::new()),
            Arc::new(OptimizeProjections::new()),
        ])?;

        let mut out = format!("## Original Plan\n{original}");

        out.push_str("\n\n## After Extraction\n");
        if after_extract == original {
            out.push_str("(same as original)");
        } else {
            out.push_str(&after_extract);
        }

        out.push_str("\n\n## After Pushdown\n");
        if after_pushdown == after_extract {
            out.push_str("(same as after extraction)");
        } else {
            out.push_str(&after_pushdown);
        }

        out.push_str("\n\n## Optimized\n");
        if optimized == after_pushdown {
            out.push_str("(same as after pushdown)");
        } else {
            out.push_str(&optimized);
        }

        Ok(out)
    }

    /// Assert all optimization stages for a plan in a single insta snapshot.
    macro_rules! assert_stages {
        ($plan:expr, @ $expected:literal $(,)?) => {{
            let result = format_optimization_stages(&$plan)?;
            insta::assert_snapshot!(result, @ $expected);
            Ok::<(), datafusion_common::DataFusionError>(())
        }};
    }

    #[test]
    fn test_extract_from_filter() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan.clone())
            .filter(leaf_udf(col("user"), "status").eq(lit("active")))?
            .select(vec![
                table_scan
                    .schema()
                    .index_of_column_by_name(None, "id")
                    .unwrap(),
            ])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: test.id
          Filter: leaf_udf(test.user, Utf8("status")) = Utf8("active")
            TableScan: test projection=[id, user]

        ## After Extraction
        Projection: test.id
          Projection: test.id, test.user
            Filter: __datafusion_extracted_1 = Utf8("active")
              Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, test.user
                TableScan: test projection=[id, user]

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        Projection: test.id
          Filter: __datafusion_extracted_1 = Utf8("active")
            Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id
              TableScan: test projection=[id, user]
        "#)
    }

    #[test]
    fn test_no_extraction_for_column() -> Result<()> {
        let table_scan = test_table_scan()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .filter(col("a").eq(lit(1)))?
            .build()?;

        assert_stages!(plan, @"
        ## Original Plan
        Filter: test.a = Int32(1)
          TableScan: test projection=[a, b, c]

        ## After Extraction
        (same as original)

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        (same as after pushdown)
        ")
    }

    #[test]
    fn test_extract_from_projection() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .project(vec![leaf_udf(col("user"), "name")])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: leaf_udf(test.user, Utf8("name"))
          TableScan: test projection=[user]

        ## After Extraction
        (same as original)

        ## After Pushdown
        Projection: __datafusion_extracted_1 AS leaf_udf(test.user,Utf8("name"))
          Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.user
            TableScan: test projection=[user]

        ## Optimized
        Projection: leaf_udf(test.user, Utf8("name"))
          TableScan: test projection=[user]
        "#)
    }

    #[test]
    fn test_extract_from_projection_with_subexpression() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .project(vec![
                leaf_udf(col("user"), "name")
                    .is_not_null()
                    .alias("has_name"),
            ])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: leaf_udf(test.user, Utf8("name")) IS NOT NULL AS has_name
          TableScan: test projection=[user]

        ## After Extraction
        (same as original)

        ## After Pushdown
        Projection: __datafusion_extracted_1 IS NOT NULL AS has_name
          Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.user
            TableScan: test projection=[user]

        ## Optimized
        Projection: leaf_udf(test.user, Utf8("name")) IS NOT NULL AS has_name
          TableScan: test projection=[user]
        "#)
    }

    #[test]
    fn test_projection_no_extraction_for_column() -> Result<()> {
        let table_scan = test_table_scan()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .project(vec![col("a"), col("b")])?
            .build()?;

        assert_stages!(plan, @"
        ## Original Plan
        TableScan: test projection=[a, b]

        ## After Extraction
        (same as original)

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        (same as after pushdown)
        ")
    }

    #[test]
    fn test_filter_with_deduplication() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let field_access = leaf_udf(col("user"), "name");
        // Filter with the same expression used twice
        let plan = LogicalPlanBuilder::from(table_scan)
            .filter(
                field_access
                    .clone()
                    .is_not_null()
                    .and(field_access.is_null()),
            )?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Filter: leaf_udf(test.user, Utf8("name")) IS NOT NULL AND leaf_udf(test.user, Utf8("name")) IS NULL
          TableScan: test projection=[id, user]

        ## After Extraction
        Projection: test.id, test.user
          Filter: __datafusion_extracted_1 IS NOT NULL AND __datafusion_extracted_1 IS NULL
            Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.id, test.user
              TableScan: test projection=[id, user]

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        (same as after pushdown)
        "#)
    }

    #[test]
    fn test_already_leaf_expression_in_filter() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .filter(leaf_udf(col("user"), "name").eq(lit("test")))?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Filter: leaf_udf(test.user, Utf8("name")) = Utf8("test")
          TableScan: test projection=[id, user]

        ## After Extraction
        Projection: test.id, test.user
          Filter: __datafusion_extracted_1 = Utf8("test")
            Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.id, test.user
              TableScan: test projection=[id, user]

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        (same as after pushdown)
        "#)
    }

    #[test]
    fn test_extract_from_aggregate_group_by() -> Result<()> {
        use datafusion_expr::test::function_stub::count;

        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .aggregate(vec![leaf_udf(col("user"), "status")], vec![count(lit(1))])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Aggregate: groupBy=[[leaf_udf(test.user, Utf8("status"))]], aggr=[[COUNT(Int32(1))]]
          TableScan: test projection=[user]

        ## After Extraction
        Projection: __datafusion_extracted_1 AS leaf_udf(test.user,Utf8("status")), COUNT(Int32(1))
          Aggregate: groupBy=[[__datafusion_extracted_1]], aggr=[[COUNT(Int32(1))]]
            Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.user
              TableScan: test projection=[user]

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        Projection: __datafusion_extracted_1 AS leaf_udf(test.user,Utf8("status")), COUNT(Int32(1))
          Aggregate: groupBy=[[__datafusion_extracted_1]], aggr=[[COUNT(Int32(1))]]
            Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1
              TableScan: test projection=[user]
        "#)
    }

    #[test]
    fn test_extract_from_aggregate_args() -> Result<()> {
        use datafusion_expr::test::function_stub::count;

        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .aggregate(
                vec![col("user")],
                vec![count(leaf_udf(col("user"), "value"))],
            )?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Aggregate: groupBy=[[test.user]], aggr=[[COUNT(leaf_udf(test.user, Utf8("value")))]]
          TableScan: test projection=[user]

        ## After Extraction
        Projection: test.user, COUNT(__datafusion_extracted_1) AS COUNT(leaf_udf(test.user,Utf8("value")))
          Aggregate: groupBy=[[test.user]], aggr=[[COUNT(__datafusion_extracted_1)]]
            Projection: leaf_udf(test.user, Utf8("value")) AS __datafusion_extracted_1, test.user
              TableScan: test projection=[user]

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        (same as after pushdown)
        "#)
    }

    #[test]
    fn test_projection_with_filter_combined() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .filter(leaf_udf(col("user"), "status").eq(lit("active")))?
            .project(vec![leaf_udf(col("user"), "name")])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: leaf_udf(test.user, Utf8("name"))
          Filter: leaf_udf(test.user, Utf8("status")) = Utf8("active")
            TableScan: test projection=[user]

        ## After Extraction
        Projection: leaf_udf(test.user, Utf8("name"))
          Projection: test.user
            Filter: __datafusion_extracted_1 = Utf8("active")
              Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.user
                TableScan: test projection=[user]

        ## After Pushdown
        Projection: __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("name"))
          Filter: __datafusion_extracted_1 = Utf8("active")
            Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.user, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2
              TableScan: test projection=[user]

        ## Optimized
        Projection: __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("name"))
          Filter: __datafusion_extracted_1 = Utf8("active")
            Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2
              TableScan: test projection=[user]
        "#)
    }

    #[test]
    fn test_projection_preserves_alias() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .project(vec![leaf_udf(col("user"), "name").alias("username")])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: leaf_udf(test.user, Utf8("name")) AS username
          TableScan: test projection=[user]

        ## After Extraction
        (same as original)

        ## After Pushdown
        Projection: __datafusion_extracted_1 AS username
          Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.user
            TableScan: test projection=[user]

        ## Optimized
        Projection: leaf_udf(test.user, Utf8("name")) AS username
          TableScan: test projection=[user]
        "#)
    }

    /// Test: Projection with different field than Filter
    /// SELECT id, s['label'] FROM t WHERE s['value'] > 150
    /// Both s['label'] and s['value'] should be in a single extraction projection.
    #[test]
    fn test_projection_different_field_from_filter() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .filter(leaf_udf(col("user"), "value").gt(lit(150)))?
            .project(vec![col("user"), leaf_udf(col("user"), "label")])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: test.user, leaf_udf(test.user, Utf8("label"))
          Filter: leaf_udf(test.user, Utf8("value")) > Int32(150)
            TableScan: test projection=[user]

        ## After Extraction
        Projection: test.user, leaf_udf(test.user, Utf8("label"))
          Projection: test.user
            Filter: __datafusion_extracted_1 > Int32(150)
              Projection: leaf_udf(test.user, Utf8("value")) AS __datafusion_extracted_1, test.user
                TableScan: test projection=[user]

        ## After Pushdown
        Projection: test.user, __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("label"))
          Filter: __datafusion_extracted_1 > Int32(150)
            Projection: leaf_udf(test.user, Utf8("value")) AS __datafusion_extracted_1, test.user, leaf_udf(test.user, Utf8("label")) AS __datafusion_extracted_2
              TableScan: test projection=[user]

        ## Optimized
        (same as after pushdown)
        "#)
    }

    #[test]
    fn test_projection_deduplication() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let field = leaf_udf(col("user"), "name");
        let plan = LogicalPlanBuilder::from(table_scan)
            .project(vec![field.clone(), field.clone().alias("name2")])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: leaf_udf(test.user, Utf8("name")), leaf_udf(test.user, Utf8("name")) AS name2
          TableScan: test projection=[user]

        ## After Extraction
        (same as original)

        ## After Pushdown
        Projection: __datafusion_extracted_1 AS leaf_udf(test.user,Utf8("name")), __datafusion_extracted_1 AS name2
          Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.user
            TableScan: test projection=[user]

        ## Optimized
        Projection: leaf_udf(test.user, Utf8("name")), leaf_udf(test.user, Utf8("name")) AS name2
          TableScan: test projection=[user]
        "#)
    }

    // =========================================================================
    // Additional tests for code coverage
    // =========================================================================

    /// Extractions push through Sort nodes to reach the TableScan.
    #[test]
    fn test_extract_through_sort() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .sort(vec![col("user").sort(true, true)])?
            .project(vec![leaf_udf(col("user"), "name")])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: leaf_udf(test.user, Utf8("name"))
          Sort: test.user ASC NULLS FIRST
            TableScan: test projection=[user]

        ## After Extraction
        (same as original)

        ## After Pushdown
        Projection: __datafusion_extracted_1 AS leaf_udf(test.user,Utf8("name"))
          Sort: test.user ASC NULLS FIRST
            Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.user
              TableScan: test projection=[user]

        ## Optimized
        (same as after pushdown)
        "#)
    }

    /// Extractions push through Limit nodes to reach the TableScan.
    #[test]
    fn test_extract_through_limit() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .limit(0, Some(10))?
            .project(vec![leaf_udf(col("user"), "name")])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: leaf_udf(test.user, Utf8("name"))
          Limit: skip=0, fetch=10
            TableScan: test projection=[user]

        ## After Extraction
        (same as original)

        ## After Pushdown
        Projection: __datafusion_extracted_1 AS leaf_udf(test.user,Utf8("name"))
          Limit: skip=0, fetch=10
            Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.user
              TableScan: test projection=[user]

        ## Optimized
        Projection: __datafusion_extracted_1 AS leaf_udf(test.user,Utf8("name"))
          Limit: skip=0, fetch=10
            Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1
              TableScan: test projection=[user]
        "#)
    }

    /// Aliased aggregate functions like count(...).alias("cnt") are handled.
    #[test]
    fn test_extract_from_aliased_aggregate() -> Result<()> {
        use datafusion_expr::test::function_stub::count;

        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .aggregate(
                vec![col("user")],
                vec![count(leaf_udf(col("user"), "value")).alias("cnt")],
            )?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Aggregate: groupBy=[[test.user]], aggr=[[COUNT(leaf_udf(test.user, Utf8("value"))) AS cnt]]
          TableScan: test projection=[user]

        ## After Extraction
        Aggregate: groupBy=[[test.user]], aggr=[[COUNT(__datafusion_extracted_1) AS cnt]]
          Projection: leaf_udf(test.user, Utf8("value")) AS __datafusion_extracted_1, test.user
            TableScan: test projection=[user]

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        (same as after pushdown)
        "#)
    }

    /// Aggregates with no MoveTowardsLeafNodes expressions return unchanged.
    #[test]
    fn test_aggregate_no_extraction() -> Result<()> {
        use datafusion_expr::test::function_stub::count;

        let table_scan = test_table_scan()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .aggregate(vec![col("a")], vec![count(col("b"))])?
            .build()?;

        assert_stages!(plan, @"
        ## Original Plan
        Aggregate: groupBy=[[test.a]], aggr=[[COUNT(test.b)]]
          TableScan: test projection=[a, b]

        ## After Extraction
        (same as original)

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        (same as after pushdown)
        ")
    }

    /// Projections containing extracted expression aliases are skipped (already extracted).
    #[test]
    fn test_skip_extracted_projection() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .project(vec![
                leaf_udf(col("user"), "name").alias("__datafusion_extracted_manual"),
                col("user"),
            ])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_manual, test.user
          TableScan: test projection=[user]

        ## After Extraction
        (same as original)

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        (same as after pushdown)
        "#)
    }

    /// Multiple extractions merge into a single extracted expression projection.
    #[test]
    fn test_merge_into_existing_extracted_projection() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .filter(leaf_udf(col("user"), "status").eq(lit("active")))?
            .filter(leaf_udf(col("user"), "name").is_not_null())?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Filter: leaf_udf(test.user, Utf8("name")) IS NOT NULL
          Filter: leaf_udf(test.user, Utf8("status")) = Utf8("active")
            TableScan: test projection=[id, user]

        ## After Extraction
        Projection: test.id, test.user
          Filter: __datafusion_extracted_1 IS NOT NULL
            Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.id, test.user
              Projection: test.id, test.user
                Filter: __datafusion_extracted_2 = Utf8("active")
                  Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_2, test.id, test.user
                    TableScan: test projection=[id, user]

        ## After Pushdown
        Projection: test.id, test.user
          Filter: __datafusion_extracted_1 IS NOT NULL
            Filter: __datafusion_extracted_2 = Utf8("active")
              Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_2, test.id, test.user, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1
                TableScan: test projection=[id, user]

        ## Optimized
        Projection: test.id, test.user
          Filter: __datafusion_extracted_1 IS NOT NULL
            Projection: test.id, test.user, __datafusion_extracted_1
              Filter: __datafusion_extracted_2 = Utf8("active")
                Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_2, test.id, test.user, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1
                  TableScan: test projection=[id, user]
        "#)
    }

    /// Extractions push through passthrough projections (columns only).
    #[test]
    fn test_extract_through_passthrough_projection() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .project(vec![col("user")])?
            .project(vec![leaf_udf(col("user"), "name")])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: leaf_udf(test.user, Utf8("name"))
          TableScan: test projection=[user]

        ## After Extraction
        (same as original)

        ## After Pushdown
        Projection: __datafusion_extracted_1 AS leaf_udf(test.user,Utf8("name"))
          Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.user
            TableScan: test projection=[user]

        ## Optimized
        Projection: leaf_udf(test.user, Utf8("name"))
          TableScan: test projection=[user]
        "#)
    }

    /// Projections with aliased columns (nothing to extract) return unchanged.
    #[test]
    fn test_projection_early_return_no_extraction() -> Result<()> {
        let table_scan = test_table_scan()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .project(vec![col("a").alias("x"), col("b")])?
            .build()?;

        assert_stages!(plan, @"
        ## Original Plan
        Projection: test.a AS x, test.b
          TableScan: test projection=[a, b]

        ## After Extraction
        (same as original)

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        (same as after pushdown)
        ")
    }

    /// Projections with arithmetic expressions but no MoveTowardsLeafNodes return unchanged.
    #[test]
    fn test_projection_with_arithmetic_no_extraction() -> Result<()> {
        let table_scan = test_table_scan()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .project(vec![(col("a") + col("b")).alias("sum")])?
            .build()?;

        assert_stages!(plan, @"
        ## Original Plan
        Projection: test.a + test.b AS sum
          TableScan: test projection=[a, b]

        ## After Extraction
        (same as original)

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        (same as after pushdown)
        ")
    }

    /// Aggregate extractions merge into existing extracted projection created by Filter.
    #[test]
    fn test_aggregate_merge_into_extracted_projection() -> Result<()> {
        use datafusion_expr::test::function_stub::count;

        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .filter(leaf_udf(col("user"), "status").eq(lit("active")))?
            .aggregate(vec![leaf_udf(col("user"), "name")], vec![count(lit(1))])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Aggregate: groupBy=[[leaf_udf(test.user, Utf8("name"))]], aggr=[[COUNT(Int32(1))]]
          Filter: leaf_udf(test.user, Utf8("status")) = Utf8("active")
            TableScan: test projection=[user]

        ## After Extraction
        Projection: __datafusion_extracted_1 AS leaf_udf(test.user,Utf8("name")), COUNT(Int32(1))
          Aggregate: groupBy=[[__datafusion_extracted_1]], aggr=[[COUNT(Int32(1))]]
            Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.user
              Projection: test.user
                Filter: __datafusion_extracted_2 = Utf8("active")
                  Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_2, test.user
                    TableScan: test projection=[user]

        ## After Pushdown
        Projection: __datafusion_extracted_1 AS leaf_udf(test.user,Utf8("name")), COUNT(Int32(1))
          Aggregate: groupBy=[[__datafusion_extracted_1]], aggr=[[COUNT(Int32(1))]]
            Filter: __datafusion_extracted_2 = Utf8("active")
              Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_2, test.user, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1
                TableScan: test projection=[user]

        ## Optimized
        Projection: __datafusion_extracted_1 AS leaf_udf(test.user,Utf8("name")), COUNT(Int32(1))
          Aggregate: groupBy=[[__datafusion_extracted_1]], aggr=[[COUNT(Int32(1))]]
            Projection: __datafusion_extracted_1
              Filter: __datafusion_extracted_2 = Utf8("active")
                Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_2, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1
                  TableScan: test projection=[user]
        "#)
    }

    /// Projection containing a MoveTowardsLeafNodes sub-expression above an
    /// Aggregate. Aggregate blocks pushdown, so the (None, true) recovery
    /// fallback path fires: in-place extraction + recovery projection.
    #[test]
    fn test_projection_with_leaf_expr_above_aggregate() -> Result<()> {
        use datafusion_expr::test::function_stub::count;

        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .aggregate(vec![col("user")], vec![count(lit(1))])?
            .project(vec![
                leaf_udf(col("user"), "name")
                    .is_not_null()
                    .alias("has_name"),
                col("COUNT(Int32(1))"),
            ])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: leaf_udf(test.user, Utf8("name")) IS NOT NULL AS has_name, COUNT(Int32(1))
          Aggregate: groupBy=[[test.user]], aggr=[[COUNT(Int32(1))]]
            TableScan: test projection=[user]

        ## After Extraction
        (same as original)

        ## After Pushdown
        Projection: __datafusion_extracted_1 IS NOT NULL AS has_name, COUNT(Int32(1))
          Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.user, COUNT(Int32(1))
            Aggregate: groupBy=[[test.user]], aggr=[[COUNT(Int32(1))]]
              TableScan: test projection=[user]

        ## Optimized
        Projection: leaf_udf(test.user, Utf8("name")) IS NOT NULL AS has_name, COUNT(Int32(1))
          Aggregate: groupBy=[[test.user]], aggr=[[COUNT(Int32(1))]]
            TableScan: test projection=[user]
        "#)
    }

    /// Merging adds new pass-through columns not in the existing extracted projection.
    #[test]
    fn test_merge_with_new_columns() -> Result<()> {
        let table_scan = test_table_scan()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .filter(leaf_udf(col("a"), "x").eq(lit(1)))?
            .filter(leaf_udf(col("b"), "y").eq(lit(2)))?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Filter: leaf_udf(test.b, Utf8("y")) = Int32(2)
          Filter: leaf_udf(test.a, Utf8("x")) = Int32(1)
            TableScan: test projection=[a, b, c]

        ## After Extraction
        Projection: test.a, test.b, test.c
          Filter: __datafusion_extracted_1 = Int32(2)
            Projection: leaf_udf(test.b, Utf8("y")) AS __datafusion_extracted_1, test.a, test.b, test.c
              Projection: test.a, test.b, test.c
                Filter: __datafusion_extracted_2 = Int32(1)
                  Projection: leaf_udf(test.a, Utf8("x")) AS __datafusion_extracted_2, test.a, test.b, test.c
                    TableScan: test projection=[a, b, c]

        ## After Pushdown
        Projection: test.a, test.b, test.c
          Filter: __datafusion_extracted_1 = Int32(2)
            Filter: __datafusion_extracted_2 = Int32(1)
              Projection: leaf_udf(test.a, Utf8("x")) AS __datafusion_extracted_2, test.a, test.b, test.c, leaf_udf(test.b, Utf8("y")) AS __datafusion_extracted_1
                TableScan: test projection=[a, b, c]

        ## Optimized
        Projection: test.a, test.b, test.c
          Filter: __datafusion_extracted_1 = Int32(2)
            Projection: test.a, test.b, test.c, __datafusion_extracted_1
              Filter: __datafusion_extracted_2 = Int32(1)
                Projection: leaf_udf(test.a, Utf8("x")) AS __datafusion_extracted_2, test.a, test.b, test.c, leaf_udf(test.b, Utf8("y")) AS __datafusion_extracted_1
                  TableScan: test projection=[a, b, c]
        "#)
    }

    // =========================================================================
    // Join extraction tests
    // =========================================================================

    /// Create a second table scan with struct field for join tests
    fn test_table_scan_with_struct_named(name: &str) -> Result<LogicalPlan> {
        use arrow::datatypes::Schema;
        let schema = Schema::new(test_table_scan_with_struct_fields());
        datafusion_expr::logical_plan::table_scan(Some(name), &schema, None)?.build()
    }

    /// Extraction from equijoin keys (`on` expressions).
    #[test]
    fn test_extract_from_join_on() -> Result<()> {
        use datafusion_expr::JoinType;

        let left = test_table_scan_with_struct()?;
        let right = test_table_scan_with_struct_named("right")?;

        let plan = LogicalPlanBuilder::from(left)
            .join_with_expr_keys(
                right,
                JoinType::Inner,
                (
                    vec![leaf_udf(col("user"), "id")],
                    vec![leaf_udf(col("user"), "id")],
                ),
                None,
            )?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Inner Join: leaf_udf(test.user, Utf8("id")) = leaf_udf(right.user, Utf8("id"))
          TableScan: test projection=[id, user]
          TableScan: right projection=[id, user]

        ## After Extraction
        Projection: test.id, test.user, right.id, right.user
          Inner Join: __datafusion_extracted_1 = __datafusion_extracted_2
            Projection: leaf_udf(test.user, Utf8("id")) AS __datafusion_extracted_1, test.id, test.user
              TableScan: test projection=[id, user]
            Projection: leaf_udf(right.user, Utf8("id")) AS __datafusion_extracted_2, right.id, right.user
              TableScan: right projection=[id, user]

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        (same as after pushdown)
        "#)
    }

    /// Extraction from non-equi join filter.
    #[test]
    fn test_extract_from_join_filter() -> Result<()> {
        use datafusion_expr::JoinType;

        let left = test_table_scan_with_struct()?;
        let right = test_table_scan_with_struct_named("right")?;

        let plan = LogicalPlanBuilder::from(left)
            .join_on(
                right,
                JoinType::Inner,
                vec![
                    col("test.user").eq(col("right.user")),
                    leaf_udf(col("test.user"), "status").eq(lit("active")),
                ],
            )?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Inner Join:  Filter: test.user = right.user AND leaf_udf(test.user, Utf8("status")) = Utf8("active")
          TableScan: test projection=[id, user]
          TableScan: right projection=[id, user]

        ## After Extraction
        Projection: test.id, test.user, right.id, right.user
          Inner Join:  Filter: test.user = right.user AND __datafusion_extracted_1 = Utf8("active")
            Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, test.user
              TableScan: test projection=[id, user]
            TableScan: right projection=[id, user]

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        (same as after pushdown)
        "#)
    }

    /// Extraction from both left and right sides of a join.
    #[test]
    fn test_extract_from_join_both_sides() -> Result<()> {
        use datafusion_expr::JoinType;

        let left = test_table_scan_with_struct()?;
        let right = test_table_scan_with_struct_named("right")?;

        let plan = LogicalPlanBuilder::from(left)
            .join_on(
                right,
                JoinType::Inner,
                vec![
                    col("test.user").eq(col("right.user")),
                    leaf_udf(col("test.user"), "status").eq(lit("active")),
                    leaf_udf(col("right.user"), "role").eq(lit("admin")),
                ],
            )?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Inner Join:  Filter: test.user = right.user AND leaf_udf(test.user, Utf8("status")) = Utf8("active") AND leaf_udf(right.user, Utf8("role")) = Utf8("admin")
          TableScan: test projection=[id, user]
          TableScan: right projection=[id, user]

        ## After Extraction
        Projection: test.id, test.user, right.id, right.user
          Inner Join:  Filter: test.user = right.user AND __datafusion_extracted_1 = Utf8("active") AND __datafusion_extracted_2 = Utf8("admin")
            Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, test.user
              TableScan: test projection=[id, user]
            Projection: leaf_udf(right.user, Utf8("role")) AS __datafusion_extracted_2, right.id, right.user
              TableScan: right projection=[id, user]

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        (same as after pushdown)
        "#)
    }

    /// Join with no MoveTowardsLeafNodes expressions returns unchanged.
    #[test]
    fn test_extract_from_join_no_extraction() -> Result<()> {
        use datafusion_expr::JoinType;

        let left = test_table_scan()?;
        let right = test_table_scan_with_name("right")?;

        let plan = LogicalPlanBuilder::from(left)
            .join(right, JoinType::Inner, (vec!["a"], vec!["a"]), None)?
            .build()?;

        assert_stages!(plan, @"
        ## Original Plan
        Inner Join: test.a = right.a
          TableScan: test projection=[a, b, c]
          TableScan: right projection=[a, b, c]

        ## After Extraction
        (same as original)

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        (same as after pushdown)
        ")
    }

    /// Join followed by filter with extraction.
    #[test]
    fn test_extract_from_filter_above_join() -> Result<()> {
        use datafusion_expr::JoinType;

        let left = test_table_scan_with_struct()?;
        let right = test_table_scan_with_struct_named("right")?;

        let plan = LogicalPlanBuilder::from(left)
            .join_with_expr_keys(
                right,
                JoinType::Inner,
                (
                    vec![leaf_udf(col("user"), "id")],
                    vec![leaf_udf(col("user"), "id")],
                ),
                None,
            )?
            .filter(leaf_udf(col("test.user"), "status").eq(lit("active")))?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Filter: leaf_udf(test.user, Utf8("status")) = Utf8("active")
          Inner Join: leaf_udf(test.user, Utf8("id")) = leaf_udf(right.user, Utf8("id"))
            TableScan: test projection=[id, user]
            TableScan: right projection=[id, user]

        ## After Extraction
        Projection: test.id, test.user, right.id, right.user
          Filter: __datafusion_extracted_1 = Utf8("active")
            Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, test.user, right.id, right.user
              Projection: test.id, test.user, right.id, right.user
                Inner Join: __datafusion_extracted_2 = __datafusion_extracted_3
                  Projection: leaf_udf(test.user, Utf8("id")) AS __datafusion_extracted_2, test.id, test.user
                    TableScan: test projection=[id, user]
                  Projection: leaf_udf(right.user, Utf8("id")) AS __datafusion_extracted_3, right.id, right.user
                    TableScan: right projection=[id, user]

        ## After Pushdown
        Projection: test.id, test.user, right.id, right.user
          Filter: __datafusion_extracted_1 = Utf8("active")
            Inner Join: __datafusion_extracted_2 = __datafusion_extracted_3
              Projection: leaf_udf(test.user, Utf8("id")) AS __datafusion_extracted_2, test.id, test.user, leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1
                TableScan: test projection=[id, user]
              Projection: leaf_udf(right.user, Utf8("id")) AS __datafusion_extracted_3, right.id, right.user
                TableScan: right projection=[id, user]

        ## Optimized
        Projection: test.id, test.user, right.id, right.user
          Filter: __datafusion_extracted_1 = Utf8("active")
            Projection: test.id, test.user, __datafusion_extracted_1, right.id, right.user
              Inner Join: __datafusion_extracted_2 = __datafusion_extracted_3
                Projection: leaf_udf(test.user, Utf8("id")) AS __datafusion_extracted_2, test.id, test.user, leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1
                  TableScan: test projection=[id, user]
                Projection: leaf_udf(right.user, Utf8("id")) AS __datafusion_extracted_3, right.id, right.user
                  TableScan: right projection=[id, user]
        "#)
    }

    /// Extraction projection (get_field in SELECT) above a Join pushes into
    /// the correct input side.
    #[test]
    fn test_extract_projection_above_join() -> Result<()> {
        use datafusion_expr::JoinType;

        let left = test_table_scan_with_struct()?;
        let right = test_table_scan_with_struct_named("right")?;

        let plan = LogicalPlanBuilder::from(left)
            .join(right, JoinType::Inner, (vec!["id"], vec!["id"]), None)?
            .project(vec![
                leaf_udf(col("test.user"), "status"),
                leaf_udf(col("right.user"), "role"),
            ])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: leaf_udf(test.user, Utf8("status")), leaf_udf(right.user, Utf8("role"))
          Inner Join: test.id = right.id
            TableScan: test projection=[id, user]
            TableScan: right projection=[id, user]

        ## After Extraction
        (same as original)

        ## After Pushdown
        Projection: __datafusion_extracted_1 AS leaf_udf(test.user,Utf8("status")), __datafusion_extracted_2 AS leaf_udf(right.user,Utf8("role"))
          Inner Join: test.id = right.id
            Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, test.user
              TableScan: test projection=[id, user]
            Projection: leaf_udf(right.user, Utf8("role")) AS __datafusion_extracted_2, right.id, right.user
              TableScan: right projection=[id, user]

        ## Optimized
        Projection: __datafusion_extracted_1 AS leaf_udf(test.user,Utf8("status")), __datafusion_extracted_2 AS leaf_udf(right.user,Utf8("role"))
          Inner Join: test.id = right.id
            Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id
              TableScan: test projection=[id, user]
            Projection: leaf_udf(right.user, Utf8("role")) AS __datafusion_extracted_2, right.id
              TableScan: right projection=[id, user]
        "#)
    }

    /// Join where both sides have same-named columns: a qualified reference
    /// to the right side must be routed to the right input, not the left.
    #[test]
    fn test_extract_from_join_qualified_right_side() -> Result<()> {
        use datafusion_expr::JoinType;

        let left = test_table_scan_with_struct()?;
        let right = test_table_scan_with_struct_named("right")?;

        // Filter references right.user explicitly — must route to right side
        let plan = LogicalPlanBuilder::from(left)
            .join_on(
                right,
                JoinType::Inner,
                vec![
                    col("test.id").eq(col("right.id")),
                    leaf_udf(col("right.user"), "status").eq(lit("active")),
                ],
            )?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Inner Join:  Filter: test.id = right.id AND leaf_udf(right.user, Utf8("status")) = Utf8("active")
          TableScan: test projection=[id, user]
          TableScan: right projection=[id, user]

        ## After Extraction
        Projection: test.id, test.user, right.id, right.user
          Inner Join:  Filter: test.id = right.id AND __datafusion_extracted_1 = Utf8("active")
            TableScan: test projection=[id, user]
            Projection: leaf_udf(right.user, Utf8("status")) AS __datafusion_extracted_1, right.id, right.user
              TableScan: right projection=[id, user]

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        (same as after pushdown)
        "#)
    }

    /// When both inputs contain the same unqualified column, an unqualified
    /// column reference is ambiguous and `find_owning_input` must return
    /// `None` rather than always returning 0 (the left side).
    #[test]
    fn test_find_owning_input_ambiguous_unqualified_column() {
        use std::collections::HashSet;

        // Simulate schema_columns output for two sides of a join where both
        // have a "user" column — each set contains the qualified and
        // unqualified form.
        let left_cols: HashSet<Column> = [
            Column::new(Some("test"), "user"),
            Column::new_unqualified("user"),
        ]
        .into_iter()
        .collect();

        let right_cols: HashSet<Column> = [
            Column::new(Some("right"), "user"),
            Column::new_unqualified("user"),
        ]
        .into_iter()
        .collect();

        let input_column_sets = vec![left_cols, right_cols];

        // Unqualified "user" matches both sets — must return None (ambiguous)
        let unqualified = Expr::Column(Column::new_unqualified("user"));
        assert_eq!(find_owning_input(&unqualified, &input_column_sets), None);

        // Qualified "right.user" matches only the right set — must return Some(1)
        let qualified_right = Expr::Column(Column::new(Some("right"), "user"));
        assert_eq!(
            find_owning_input(&qualified_right, &input_column_sets),
            Some(1)
        );

        // Qualified "test.user" matches only the left set — must return Some(0)
        let qualified_left = Expr::Column(Column::new(Some("test"), "user"));
        assert_eq!(
            find_owning_input(&qualified_left, &input_column_sets),
            Some(0)
        );
    }

    /// Two leaf_udf expressions from different sides of a Join in a Filter.
    /// Each is routed to its respective input side independently.
    #[test]
    fn test_extract_from_join_cross_input_expression() -> Result<()> {
        let left = test_table_scan_with_struct()?;
        let right = test_table_scan_with_struct_named("right")?;

        let plan = LogicalPlanBuilder::from(left)
            .join_on(
                right,
                datafusion_expr::JoinType::Inner,
                vec![col("test.id").eq(col("right.id"))],
            )?
            .filter(
                leaf_udf(col("test.user"), "status")
                    .eq(leaf_udf(col("right.user"), "status")),
            )?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Filter: leaf_udf(test.user, Utf8("status")) = leaf_udf(right.user, Utf8("status"))
          Inner Join:  Filter: test.id = right.id
            TableScan: test projection=[id, user]
            TableScan: right projection=[id, user]

        ## After Extraction
        Projection: test.id, test.user, right.id, right.user
          Filter: __datafusion_extracted_1 = __datafusion_extracted_2
            Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, leaf_udf(right.user, Utf8("status")) AS __datafusion_extracted_2, test.id, test.user, right.id, right.user
              Inner Join:  Filter: test.id = right.id
                TableScan: test projection=[id, user]
                TableScan: right projection=[id, user]

        ## After Pushdown
        Projection: test.id, test.user, right.id, right.user
          Filter: __datafusion_extracted_1 = __datafusion_extracted_2
            Inner Join:  Filter: test.id = right.id
              Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, test.user
                TableScan: test projection=[id, user]
              Projection: leaf_udf(right.user, Utf8("status")) AS __datafusion_extracted_2, right.id, right.user
                TableScan: right projection=[id, user]

        ## Optimized
        (same as after pushdown)
        "#)
    }

    // =========================================================================
    // Column-rename through intermediate node tests
    // =========================================================================

    /// Projection with leaf expr above Filter above renaming Projection.
    #[test]
    fn test_extract_through_filter_with_column_rename() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .project(vec![col("user").alias("x")])?
            .filter(col("x").is_not_null())?
            .project(vec![leaf_udf(col("x"), "a")])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: leaf_udf(x, Utf8("a"))
          Filter: x IS NOT NULL
            Projection: test.user AS x
              TableScan: test projection=[user]

        ## After Extraction
        (same as original)

        ## After Pushdown
        Projection: __datafusion_extracted_1 AS leaf_udf(x,Utf8("a"))
          Filter: x IS NOT NULL
            Projection: test.user AS x, leaf_udf(test.user, Utf8("a")) AS __datafusion_extracted_1, test.user
              TableScan: test projection=[user]

        ## Optimized
        Projection: __datafusion_extracted_1 AS leaf_udf(x,Utf8("a"))
          Filter: x IS NOT NULL
            Projection: test.user AS x, leaf_udf(test.user, Utf8("a")) AS __datafusion_extracted_1
              TableScan: test projection=[user]
        "#)
    }

    /// Same as above but with a partial extraction (leaf + arithmetic).
    #[test]
    fn test_extract_partial_through_filter_with_column_rename() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .project(vec![col("user").alias("x")])?
            .filter(col("x").is_not_null())?
            .project(vec![leaf_udf(col("x"), "a").is_not_null()])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: leaf_udf(x, Utf8("a")) IS NOT NULL
          Filter: x IS NOT NULL
            Projection: test.user AS x
              TableScan: test projection=[user]

        ## After Extraction
        (same as original)

        ## After Pushdown
        Projection: __datafusion_extracted_1 IS NOT NULL AS leaf_udf(x,Utf8("a")) IS NOT NULL
          Filter: x IS NOT NULL
            Projection: test.user AS x, leaf_udf(test.user, Utf8("a")) AS __datafusion_extracted_1, test.user
              TableScan: test projection=[user]

        ## Optimized
        Projection: __datafusion_extracted_1 IS NOT NULL AS leaf_udf(x,Utf8("a")) IS NOT NULL
          Filter: x IS NOT NULL
            Projection: test.user AS x, leaf_udf(test.user, Utf8("a")) AS __datafusion_extracted_1
              TableScan: test projection=[user]
        "#)
    }

    /// Tests merge_into_extracted_projection path through a renaming projection.
    #[test]
    fn test_extract_from_filter_above_renaming_projection() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .project(vec![col("user").alias("x")])?
            .filter(leaf_udf(col("x"), "a").eq(lit("active")))?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Filter: leaf_udf(x, Utf8("a")) = Utf8("active")
          Projection: test.user AS x
            TableScan: test projection=[user]

        ## After Extraction
        Projection: x
          Filter: __datafusion_extracted_1 = Utf8("active")
            Projection: test.user AS x, leaf_udf(test.user, Utf8("a")) AS __datafusion_extracted_1, test.user
              TableScan: test projection=[user]

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        Projection: x
          Filter: __datafusion_extracted_1 = Utf8("active")
            Projection: test.user AS x, leaf_udf(test.user, Utf8("a")) AS __datafusion_extracted_1
              TableScan: test projection=[user]
        "#)
    }

    // =========================================================================
    // SubqueryAlias extraction tests
    // =========================================================================

    /// Extraction projection pushes through SubqueryAlias.
    #[test]
    fn test_extract_through_subquery_alias() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .alias("sub")?
            .project(vec![leaf_udf(col("sub.user"), "name")])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: leaf_udf(sub.user, Utf8("name"))
          SubqueryAlias: sub
            TableScan: test projection=[user]

        ## After Extraction
        (same as original)

        ## After Pushdown
        Projection: __datafusion_extracted_1 AS leaf_udf(sub.user,Utf8("name"))
          SubqueryAlias: sub
            Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.user
              TableScan: test projection=[user]

        ## Optimized
        Projection: __datafusion_extracted_1 AS leaf_udf(sub.user,Utf8("name"))
          SubqueryAlias: sub
            Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1
              TableScan: test projection=[user]
        "#)
    }

    /// Extraction projection pushes through SubqueryAlias + Filter.
    #[test]
    fn test_extract_through_subquery_alias_with_filter() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .alias("sub")?
            .filter(leaf_udf(col("sub.user"), "status").eq(lit("active")))?
            .project(vec![leaf_udf(col("sub.user"), "name")])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: leaf_udf(sub.user, Utf8("name"))
          Filter: leaf_udf(sub.user, Utf8("status")) = Utf8("active")
            SubqueryAlias: sub
              TableScan: test projection=[user]

        ## After Extraction
        Projection: leaf_udf(sub.user, Utf8("name"))
          Projection: sub.user
            Filter: __datafusion_extracted_1 = Utf8("active")
              Projection: leaf_udf(sub.user, Utf8("status")) AS __datafusion_extracted_1, sub.user
                SubqueryAlias: sub
                  TableScan: test projection=[user]

        ## After Pushdown
        Projection: __datafusion_extracted_2 AS leaf_udf(sub.user,Utf8("name"))
          Filter: __datafusion_extracted_1 = Utf8("active")
            SubqueryAlias: sub
              Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2, test.user
                TableScan: test projection=[user]

        ## Optimized
        Projection: __datafusion_extracted_2 AS leaf_udf(sub.user,Utf8("name"))
          Filter: __datafusion_extracted_1 = Utf8("active")
            SubqueryAlias: sub
              Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2
                TableScan: test projection=[user]
        "#)
    }

    /// Two layers of SubqueryAlias: extraction pushes through both.
    #[test]
    fn test_extract_through_nested_subquery_alias() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .alias("inner_sub")?
            .alias("outer_sub")?
            .project(vec![leaf_udf(col("outer_sub.user"), "name")])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: leaf_udf(outer_sub.user, Utf8("name"))
          SubqueryAlias: outer_sub
            SubqueryAlias: inner_sub
              TableScan: test projection=[user]

        ## After Extraction
        (same as original)

        ## After Pushdown
        Projection: __datafusion_extracted_1 AS leaf_udf(outer_sub.user,Utf8("name"))
          SubqueryAlias: outer_sub
            SubqueryAlias: inner_sub
              Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.user
                TableScan: test projection=[user]

        ## Optimized
        Projection: __datafusion_extracted_1 AS leaf_udf(outer_sub.user,Utf8("name"))
          SubqueryAlias: outer_sub
            SubqueryAlias: inner_sub
              Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1
                TableScan: test projection=[user]
        "#)
    }

    /// Plain columns through SubqueryAlias -- no extraction needed.
    #[test]
    fn test_subquery_alias_no_extraction() -> Result<()> {
        let table_scan = test_table_scan()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .alias("sub")?
            .project(vec![col("sub.a"), col("sub.b")])?
            .build()?;

        assert_stages!(plan, @"
        ## Original Plan
        SubqueryAlias: sub
          TableScan: test projection=[a, b]

        ## After Extraction
        (same as original)

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        (same as after pushdown)
        ")
    }

    /// Two UDFs with the same `name()` but different concrete types should NOT be
    /// deduplicated -- they are semantically different expressions that happen to
    /// collide on `schema_name()`.
    #[test]
    fn test_different_udfs_same_schema_name_not_deduplicated() -> Result<()> {
        let udf_a = Arc::new(ScalarUDF::new_from_impl(
            PlacementTestUDF::new()
                .with_placement(ExpressionPlacement::MoveTowardsLeafNodes)
                .with_id(1),
        ));
        let udf_b = Arc::new(ScalarUDF::new_from_impl(
            PlacementTestUDF::new()
                .with_placement(ExpressionPlacement::MoveTowardsLeafNodes)
                .with_id(2),
        ));

        let expr_a = Expr::ScalarFunction(ScalarFunction::new_udf(
            udf_a,
            vec![col("user"), lit("field")],
        ));
        let expr_b = Expr::ScalarFunction(ScalarFunction::new_udf(
            udf_b,
            vec![col("user"), lit("field")],
        ));

        // Verify preconditions: same schema_name but different Expr
        assert_eq!(
            expr_a.schema_name().to_string(),
            expr_b.schema_name().to_string(),
            "Both expressions should have the same schema_name"
        );
        assert_ne!(
            expr_a, expr_b,
            "Expressions should NOT be equal (different UDF instances)"
        );

        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan.clone())
            .filter(expr_a.clone().eq(lit("a")).and(expr_b.clone().eq(lit("b"))))?
            .select(vec![
                table_scan
                    .schema()
                    .index_of_column_by_name(None, "id")
                    .unwrap(),
            ])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: test.id
          Filter: leaf_udf(test.user, Utf8("field")) = Utf8("a") AND leaf_udf(test.user, Utf8("field")) = Utf8("b")
            TableScan: test projection=[id, user]

        ## After Extraction
        Projection: test.id
          Projection: test.id, test.user
            Filter: __datafusion_extracted_1 = Utf8("a") AND __datafusion_extracted_2 = Utf8("b")
              Projection: leaf_udf(test.user, Utf8("field")) AS __datafusion_extracted_1, leaf_udf(test.user, Utf8("field")) AS __datafusion_extracted_2, test.id, test.user
                TableScan: test projection=[id, user]

        ## After Pushdown
        (same as after extraction)

        ## Optimized
        Projection: test.id
          Filter: __datafusion_extracted_1 = Utf8("a") AND __datafusion_extracted_2 = Utf8("b")
            Projection: leaf_udf(test.user, Utf8("field")) AS __datafusion_extracted_1, leaf_udf(test.user, Utf8("field")) AS __datafusion_extracted_2, test.id
              TableScan: test projection=[id, user]
        "#)
    }

    // =========================================================================
    // Filter pushdown interaction tests
    // =========================================================================

    /// Extraction pushdown through a filter that already had its own
    /// `leaf_udf` extracted.
    #[test]
    fn test_extraction_pushdown_through_filter_with_extracted_predicate() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .filter(leaf_udf(col("user"), "status").eq(lit("active")))?
            .project(vec![col("id"), leaf_udf(col("user"), "name")])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: test.id, leaf_udf(test.user, Utf8("name"))
          Filter: leaf_udf(test.user, Utf8("status")) = Utf8("active")
            TableScan: test projection=[id, user]

        ## After Extraction
        Projection: test.id, leaf_udf(test.user, Utf8("name"))
          Projection: test.id, test.user
            Filter: __datafusion_extracted_1 = Utf8("active")
              Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, test.user
                TableScan: test projection=[id, user]

        ## After Pushdown
        Projection: test.id, __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("name"))
          Filter: __datafusion_extracted_1 = Utf8("active")
            Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, test.user, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2
              TableScan: test projection=[id, user]

        ## Optimized
        Projection: test.id, __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("name"))
          Filter: __datafusion_extracted_1 = Utf8("active")
            Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2
              TableScan: test projection=[id, user]
        "#)
    }

    /// Same expression in filter predicate and projection output.
    #[test]
    fn test_extraction_pushdown_same_expr_in_filter_and_projection() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let field_expr = leaf_udf(col("user"), "status");
        let plan = LogicalPlanBuilder::from(table_scan)
            .filter(field_expr.clone().gt(lit(5)))?
            .project(vec![col("id"), field_expr])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: test.id, leaf_udf(test.user, Utf8("status"))
          Filter: leaf_udf(test.user, Utf8("status")) > Int32(5)
            TableScan: test projection=[id, user]

        ## After Extraction
        Projection: test.id, leaf_udf(test.user, Utf8("status"))
          Projection: test.id, test.user
            Filter: __datafusion_extracted_1 > Int32(5)
              Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, test.user
                TableScan: test projection=[id, user]

        ## After Pushdown
        Projection: test.id, __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("status"))
          Filter: __datafusion_extracted_1 > Int32(5)
            Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, test.user, leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_2
              TableScan: test projection=[id, user]

        ## Optimized
        Projection: test.id, __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("status"))
          Filter: __datafusion_extracted_1 > Int32(5)
            Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_2
              TableScan: test projection=[id, user]
        "#)
    }

    /// Left join with a `leaf_udf` filter on the right side AND
    /// the projection also selects `leaf_udf` from the right side.
    #[test]
    fn test_left_join_with_filter_and_projection_extraction() -> Result<()> {
        use datafusion_expr::JoinType;

        let left = test_table_scan_with_struct()?;
        let right = test_table_scan_with_struct_named("right")?;

        let plan = LogicalPlanBuilder::from(left)
            .join_on(
                right,
                JoinType::Left,
                vec![
                    col("test.id").eq(col("right.id")),
                    leaf_udf(col("right.user"), "status").gt(lit(5)),
                ],
            )?
            .project(vec![
                col("test.id"),
                leaf_udf(col("test.user"), "name"),
                leaf_udf(col("right.user"), "status"),
            ])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: test.id, leaf_udf(test.user, Utf8("name")), leaf_udf(right.user, Utf8("status"))
          Left Join:  Filter: test.id = right.id AND leaf_udf(right.user, Utf8("status")) > Int32(5)
            TableScan: test projection=[id, user]
            TableScan: right projection=[id, user]

        ## After Extraction
        Projection: test.id, leaf_udf(test.user, Utf8("name")), leaf_udf(right.user, Utf8("status"))
          Projection: test.id, test.user, right.id, right.user
            Left Join:  Filter: test.id = right.id AND __datafusion_extracted_1 > Int32(5)
              TableScan: test projection=[id, user]
              Projection: leaf_udf(right.user, Utf8("status")) AS __datafusion_extracted_1, right.id, right.user
                TableScan: right projection=[id, user]

        ## After Pushdown
        Projection: test.id, __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("name")), __datafusion_extracted_3 AS leaf_udf(right.user,Utf8("status"))
          Left Join:  Filter: test.id = right.id AND __datafusion_extracted_1 > Int32(5)
            Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2, test.id, test.user
              TableScan: test projection=[id, user]
            Projection: leaf_udf(right.user, Utf8("status")) AS __datafusion_extracted_1, right.id, right.user, leaf_udf(right.user, Utf8("status")) AS __datafusion_extracted_3
              TableScan: right projection=[id, user]

        ## Optimized
        Projection: test.id, __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("name")), __datafusion_extracted_3 AS leaf_udf(right.user,Utf8("status"))
          Left Join:  Filter: test.id = right.id AND __datafusion_extracted_1 > Int32(5)
            Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2, test.id
              TableScan: test projection=[id, user]
            Projection: leaf_udf(right.user, Utf8("status")) AS __datafusion_extracted_1, right.id, leaf_udf(right.user, Utf8("status")) AS __datafusion_extracted_3
              TableScan: right projection=[id, user]
        "#)
    }

    /// Extraction projection pushed through a filter whose predicate
    /// references a different extracted expression.
    #[test]
    fn test_pure_extraction_proj_push_through_filter() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;
        let plan = LogicalPlanBuilder::from(table_scan)
            .filter(leaf_udf(col("user"), "status").gt(lit(5)))?
            .project(vec![
                col("id"),
                leaf_udf(col("user"), "name"),
                leaf_udf(col("user"), "status"),
            ])?
            .build()?;

        assert_stages!(plan, @r#"
        ## Original Plan
        Projection: test.id, leaf_udf(test.user, Utf8("name")), leaf_udf(test.user, Utf8("status"))
          Filter: leaf_udf(test.user, Utf8("status")) > Int32(5)
            TableScan: test projection=[id, user]

        ## After Extraction
        Projection: test.id, leaf_udf(test.user, Utf8("name")), leaf_udf(test.user, Utf8("status"))
          Projection: test.id, test.user
            Filter: __datafusion_extracted_1 > Int32(5)
              Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, test.user
                TableScan: test projection=[id, user]

        ## After Pushdown
        Projection: test.id, __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("name")), __datafusion_extracted_3 AS leaf_udf(test.user,Utf8("status"))
          Filter: __datafusion_extracted_1 > Int32(5)
            Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, test.user, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2, leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_3
              TableScan: test projection=[id, user]

        ## Optimized
        Projection: test.id, __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("name")), __datafusion_extracted_3 AS leaf_udf(test.user,Utf8("status"))
          Filter: __datafusion_extracted_1 > Int32(5)
            Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2, leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_3
              TableScan: test projection=[id, user]
        "#)
    }

    /// When an extraction projection's __extracted alias references a column
    /// (e.g. `user`) that is NOT a standalone expression in the projection,
    /// the merge into the inner projection should still succeed.
    #[test]
    fn test_merge_extraction_into_projection_with_column_ref_inflation() -> Result<()> {
        let table_scan = test_table_scan_with_struct()?;

        // Inner projection (simulates a trimmed projection)
        let inner = LogicalPlanBuilder::from(table_scan)
            .project(vec![col("user"), col("id")])?
            .build()?;

        // Outer projection: __extracted alias + id (but NOT user as standalone).
        // The alias references `user` internally, inflating columns_needed.
        let plan = LogicalPlanBuilder::from(inner)
            .project(vec![
                leaf_udf(col("user"), "status")
                    .alias(format!("{EXTRACTED_EXPR_PREFIX}_1")),
                col("id"),
            ])?
            .build()?;

        // Run only PushDownLeafProjections
        let ctx = OptimizerContext::new().with_max_passes(1);
        let optimizer =
            Optimizer::with_rules(vec![Arc::new(PushDownLeafProjections::new())]);
        let result = optimizer.optimize(plan, &ctx, |_, _| {})?;

        // With the fix: merge succeeds → extraction merged into inner projection.
        // Without the fix: merge rejected → two separate projections remain.
        insta::assert_snapshot!(format!("{result}"), @r#"
        Projection: __datafusion_extracted_1, test.id
          Projection: test.user, test.id, leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1
            TableScan: test
        "#);

        Ok(())
    }
}