Function unproject_sort_expr

datafusion/sql/src/unparser/utils.rs:306–353

Transforms all Column expressions in a sort expression into the actual expression from aggregation or projection if found. This is required because if an ORDER BY expression is present in an Aggregate or Select, it is replaced with a Column expression (e.g., "sum(catalog_returns.cr_net_loss)"). We need to transform it back to the actual expression, such as sum("catalog_returns"."cr_net_loss").

(
    mut sort_expr: SortExpr,
    agg: Option<&Aggregate>,
    input: &LogicalPlan,
)

/// with a Column expression (e.g., "sum(catalog_returns.cr_net_loss)"). We need to transform it back to
/// the actual expression, such as sum("catalog_returns"."cr_net_loss").
pub(crate) fn unproject_sort_expr(
    mut sort_expr: SortExpr,
    agg: Option<&Aggregate>,
    input: &LogicalPlan,
) -> Result<SortExpr> {
    sort_expr.expr = sort_expr
        .expr
        .transform(|sub_expr| {
            match sub_expr {
                // Remove alias if present, because ORDER BY cannot use aliases
                Expr::Alias(alias) => Ok(Transformed::yes(*alias.expr)),
                Expr::Column(col) => {
                    if col.relation.is_some() {
                        return Ok(Transformed::no(Expr::Column(col)));
                    }

                    // In case of aggregation there could be columns containing aggregation functions we need to unproject
                    if let Some(agg) = agg
                        && agg.schema.is_column_from_schema(&col)
                    {
                        return Ok(Transformed::yes(unproject_agg_exprs(
                            Expr::Column(col),
                            agg,
                            None,
                        )?));
                    }

                    // If SELECT and ORDER BY contain the same expression with a scalar function, the ORDER BY expression will
                    // be replaced by a Column expression (e.g., "substr(customer.c_last_name, Int64(0), Int64(5))"), and we need
                    // to transform it back to the actual expression.
                    if let LogicalPlan::Projection(Projection { expr, schema, .. }) =
                        input
                        && let Ok(idx) = schema.index_of_column(&col)
                        && let Some(Expr::ScalarFunction(scalar_fn)) = expr.get(idx)
                    {
                        return Ok(Transformed::yes(Expr::ScalarFunction(
                            scalar_fn.clone(),
                        )));
                    }

                    Ok(Transformed::no(Expr::Column(col)))
                }
                _ => Ok(Transformed::no(sub_expr)),
            }
        })
        .map(|e| e.data)?;
    Ok(sort_expr)
}

/// Iterates through the children of a [LogicalPlan] to find a TableScan node before encountering
/// a Projection or any unexpected node that indicates the presence of a Projection (SELECT) in the plan.