Repository: pola-rs/polars
File: crates/polars-plan/src/plans/optimizer/projection_pushdown/generic.rs (Rust)

Code before:
use polars_utils::idx_vec::UnitVec;

use super::*;

#[allow(clippy::too_many_arguments)]
pub(super) fn process_generic(
    proj_pd: &mut ProjectionPushDown,
    lp: IR,
    ctx: ProjectionContext,
    lp_arena: &mut Arena<IR>,
    expr_arena: &mut Arena<AExpr>,
) -> PolarsResult<IR> {
    let inputs = lp.get_inputs();
    let input_count = inputs.len();

    // let mut first_schema = None;
    // let mut names = None;

    let new_inputs = inputs
        .into_iter()
        .map(|node| {
            proj_pd.is_count_star = false;

            let alp = lp_arena.take(node);
            let mut alp = proj_pd.push_down(alp, ctx.clone(), lp_arena, expr_arena)?;

            // double projection can mess up the schema ordering
            // here we ensure the ordering is maintained.
            //
            // Consider this query
            // df1 => a, b
            // df2 => a, b
            //
            // df3 = df1.join(df2, on = a, b)
            //
            // concat([df1, df3]).select(a)
            //
            // schema after projection pd
            // df3 => a, b
            // df1 => a
            // so we ensure we do the 'a' projection again before we concatenate
            if !ctx.acc_projections.is_empty() && input_count > 1 {
                alp = IRBuilder::from_lp(alp, expr_arena, lp_arena)
                    .project_simple_nodes(ctx.acc_projections.iter().map(|e| e.0))
                    .unwrap()
                    .build();
            }

Review comment (suggestion):
This smells to me. `proj_pd` only contains `is_count_star`. If we do this, why is `proj_pd` even an argument to `process_generic`?

I'm not too sure what this does but I think removing the `proj_pd` argument and just using `ProjectionPushdown::new()` here makes more sense.

@ritchie46 wrote this (three years ago) so perhaps he has a better idea, but I wouldn't blame him for forgetting why this is the case.

Code after:
use polars_utils::idx_vec::UnitVec;

use super::*;

#[allow(clippy::too_many_arguments)]
pub(super) fn process_generic(
    proj_pd: &mut ProjectionPushDown,
    lp: IR,
    ctx: ProjectionContext,
    lp_arena: &mut Arena<IR>,
    expr_arena: &mut Arena<AExpr>,
    reset_count_star: bool,
) -> PolarsResult<IR> {
    let inputs = lp.get_inputs();
    let input_count = inputs.len();

    // let mut first_schema = None;
    // let mut names = None;

    let new_inputs = inputs
        .into_iter()
        .map(|node| {
            if reset_count_star {
                proj_pd.is_count_star = false;
            }
            let alp = lp_arena.take(node);
            let mut alp = proj_pd.push_down(alp, ctx.clone(), lp_arena, expr_arena)?;

            // double projection can mess up the schema ordering
            // here we ensure the ordering is maintained.
            //
            // Consider this query
            // df1 => a, b
            // df2 => a, b
            //
            // df3 = df1.join(df2, on = a, b)
            //
            // concat([df1, df3]).select(a)
            //
            // schema after projection pd
            // df3 => a, b
            // df1 => a
            // so we ensure we do the 'a' projection again before we concatenate
            if !ctx.acc_projections.is_empty() && input_count > 1 {
                alp = IRBuilder::from_lp(alp, expr_arena, lp_arena)
                    .project_simple_nodes(ctx.acc_projections.iter().map(|e| e.0))
                    .unwrap()