Function extract_from_plan

datafusion/optimizer/src/extract_leaf_expressions.rs:171–265

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.

(
    plan: LogicalPlan,
    alias_generator: &Arc<AliasGenerator>,
)

/// 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<ColumnReference>> =
        input_schemas
            .iter()
            .map(|schema| schema_columns(schema.as_ref()))
            .collect();

    // Transform expressions via map_expressions with routing
    let transformed = plan.map_expressions(|expr| {
        routing_extract(expr, &mut extractors, &input_column_sets)
    })?;

    // If no expressions were rewritten, nothing was extracted
    if !transformed.transformed {