Repository: pola-rs/polars
File: crates/polars-expr/src/expressions/apply.rs (Rust)

Code before:
                },
            }
        }
    }

    fn to_field(&self, input_schema: &Schema) -> PolarsResult<Field> {
        self.expr.to_field(input_schema)
    }
    fn as_partitioned_aggregator(&self) -> Option<&dyn PartitionedAggregation> {
        if self.inputs.len() == 1 && self.flags.is_elementwise() {
            Some(self)
        } else {
            None
        }
    }
    fn is_scalar(&self) -> bool {
        self.flags.returns_scalar()
            || (self.function_operates_on_scalar && self.flags.is_length_preserving())
    }
}

impl PartitionedAggregation for ApplyExpr {
    fn evaluate_partitioned(
        &self,
        df: &DataFrame,
        groups: &GroupPositions,
        state: &ExecutionState,
    ) -> PolarsResult<Column> {
        let a = self.inputs[0].as_partitioned_aggregator().unwrap();
        let s = a.evaluate_partitioned(df, groups, state)?;

        if self.flags.contains(FunctionFlags::ALLOW_RENAME) {
            self.eval_and_flatten(&mut [s])
        } else {
            let in_name = s.name().clone();
            Ok(self.eval_and_flatten(&mut [s])?.with_name(in_name))
        }
    }

    fn finalize(
        &self,
        partitioned: Column,
        _groups: &GroupPositions,
        _state: &ExecutionState,
    ) -> PolarsResult<Column> {
        Ok(partitioned)
    }
}


Review comment (performance):
I think what you want here is:

```rust
if all_literal {
    self.apply_all_literal(acs)
} else if has_agg_scalar && (has_agg_list || has_not_agg) {
    // Not compatible
    self.apply_multiple_group_aware(acs, df)
} else if has_not_agg_with_rolling_groups && (?) {
    // Compatible but calling aggregated() is too expensive
    self.apply_multiple_group_aware(acs, df)
} else {
    self.apply_multiple_elementwise(acs) // determine may diverge in there
}
```

Code after:
                    for ac in &acs {
                        match ac.state {
                            AggState::AggregatedList(_) => {
                                has_agg_list = true;
                            },
                            AggState::AggregatedScalar(_) => has_agg_scalar = true,
                            AggState::NotAggregated(_) => {
                                has_not_agg = true;
                                if let Some(p) = previous {
                                    not_agg_groups_may_diverge |=
                                        !std::ptr::eq(p.groups.as_ref(), ac.groups.as_ref());
                                }
                                previous = Some(ac);
                                if ac.groups.is_overlapping() {
                                    has_not_agg_with_overlapping_groups = true;
                                }
                            },
                            _ => {},
                        }
                    }

                    let all_literal = !(has_agg_list || has_agg_scalar || has_not_agg);
                    let elementwise_must_aggregate =
                        has_not_agg && (has_agg_list || not_agg_groups_may_diverge);

                    if all_literal {
                        // Fast path
                        self.apply_all_literal_elementwise(acs)
                    } else if has_agg_scalar && (has_agg_list || has_not_agg) {
                        // Not compatible
                        self.apply_multiple_group_aware(acs, df)
                    } else if elementwise_must_aggregate && has_not_agg_with_overlapping_groups {
                        // Compatible but calling aggregated() is too expensive
                        self.apply_multiple_group_aware(acs, df)
                    } else {
                        self.apply_multiple_elementwise(acs, elementwise_must_aggregate)
                    }
                },
            }
        }
    }

    fn to_field(&self, _input_schema: &Schema) -> PolarsResult<Field> {
        Ok(self.output_field.clone())
    }
    fn as_partitioned_aggregator(&self) -> Option<&dyn PartitionedAggregation> {
        if self.inputs.len() == 1 && self.flags.is_elementwise() {
            Some(self)
        } else {
            None
        }