Function update_expr

datafusion/physical-expr/src/projection.rs:938–1010

The function projects / unprojects an expression with respect to set of projection expressions. See also [`ProjectionExprs::unproject_expr`] and [`ProjectionExprs::project_expr`] 1) When `unproject` is `true`: Rewrites an expression with respect to the projection expressions, effectively "unprojecting" it to reference the original input columns. For example, given the expressions `a@1 + b@2` a

(
    expr: &Arc<dyn PhysicalExpr>,
    projected_exprs: &[ProjectionExpr],
    unproject: bool,
)

///
/// In this case, `a@3` references index 3, which is out of bounds for `projected_exprs` (which has length 2).
pub fn update_expr(
    expr: &Arc<dyn PhysicalExpr>,
    projected_exprs: &[ProjectionExpr],
    unproject: bool,
) -> Result<Option<Arc<dyn PhysicalExpr>>> {
    #[derive(Debug, PartialEq)]
    enum RewriteState {
        /// The expression is unchanged.
        Unchanged,
        /// Some part of the expression has been rewritten
        RewrittenValid,
        /// Some part of the expression has been rewritten, but some column
        /// references could not be.
        RewrittenInvalid,
    }

    let mut state = RewriteState::Unchanged;

    let new_expr = Arc::clone(expr)
        .transform_up(|expr| {
            if state == RewriteState::RewrittenInvalid {
                return Ok(Transformed::no(expr));
            }

            let Some(column) = expr.downcast_ref::<Column>() else {
                return Ok(Transformed::no(expr));
            };
            if unproject {
                state = RewriteState::RewrittenValid;
                // Update the index of `column`:
                let projected_expr = projected_exprs.get(column.index()).ok_or_else(|| {
                    internal_datafusion_err!(
                        "Column index {} out of bounds for projected expressions of length {}",
                        column.index(),
                        projected_exprs.len()
                    )
                })?;
                Ok(Transformed::yes(Arc::clone(&projected_expr.expr)))
            } else {
                // default to invalid, in case we can't find the relevant column
                state = RewriteState::RewrittenInvalid;
                // Determine how to update `column` to accommodate `projected_exprs`
                projected_exprs
                    .iter()
                    .enumerate()
                    .find_map(|(index, proj_expr)| {
                        proj_expr.expr.downcast_ref::<Column>().and_then(
                            |projected_column| {
                                (column.name().eq(projected_column.name())
                                    && column.index() == projected_column.index())
                                .then(|| {
                                    state = RewriteState::RewrittenValid;
                                    Arc::new(Column::new(&proj_expr.alias, index)) as _
                                })
                            },
                        )
                    })
                    .map_or_else(