perf: Improve streaming groupby CSE (#23092)

Author: orlp

crates/polars-stream/src/physical_plan/lower_group_by.rs

@@ -77,6 +77,7 @@ fn build_group_by_fallback(
 /// Such an expression is defined as the elementwise combination of scalar
 /// aggregations of elementwise combinations of the input columns / scalar literals.
 #[recursive]
+#[allow(clippy::too_many_arguments)]
 fn try_lower_elementwise_scalar_agg_expr(
     expr: Node,
     outer_name: Option<PlSmallStr>,
@@ -85,6 +86,7 @@ fn try_lower_elementwise_scalar_agg_expr(
     expr_arena: &mut Arena<AExpr>,
     agg_exprs: &mut Vec<ExprIR>,
     uniq_input_exprs: &mut PlIndexMap<u32, PlSmallStr>,
+    uniq_agg_exprs: &mut PlIndexMap<u32, PlSmallStr>,
 ) -> Option<Node> {
     // Helper macro to simplify recursive calls.
     macro_rules! lower_rec {
@@ -97,6 +99,7 @@ fn try_lower_elementwise_scalar_agg_expr(
                 expr_arena,
                 agg_exprs,
                 uniq_input_exprs,
+                uniq_agg_exprs,
             )
         };
     }
@@ -219,34 +222,44 @@ fn try_lower_elementwise_scalar_agg_expr(
                 | IRAggExpr::Var(input, ..)
                 | IRAggExpr::Std(input, ..)
                 | IRAggExpr::Count(input, ..) => {
-                    if is_input_independent(*input, expr_arena, expr_cache) {
+                    let agg = agg.clone();
+                    let input = *input;
+                    if is_input_independent(input, expr_arena, expr_cache) {
                         // TODO: we could simply return expr here, but we first need an is_scalar function, because if
                         // it is not a scalar we need to return expr.implode().
                         return None;
                     }
 
-                    if !is_elementwise_rec_cached(*input, expr_arena, expr_cache) {
+                    if !is_elementwise_rec_cached(input, expr_arena, expr_cache) {
                         return None;
                     }
 
-                    let mut trans_agg = agg.clone();
-                    let input_id = expr_merger.get_uniq_id(*input).unwrap();
-                    let input_col = uniq_input_exprs
-                        .entry(input_id)
-                        .or_insert_with(unique_column_name)
+                    let agg_id = expr_merger.get_uniq_id(expr).unwrap();
+                    let name = uniq_agg_exprs
+                        .entry(agg_id)
+                        .or_insert_with(|| {
+                            let mut trans_agg = agg;
+                            let input_id = expr_merger.get_uniq_id(input).unwrap();
+                            let input_col = uniq_input_exprs
+                                .entry(input_id)
+                                .or_insert_with(unique_column_name)
+                                .clone();
+                            let input_col_node = expr_arena.add(AExpr::Column(input_col.clone()));
+                            trans_agg.set_input(input_col_node);
+                            let trans_agg_node = expr_arena.add(AExpr::Agg(trans_agg));
+
+                            // Add to aggregation expressions and replace with a reference to its output.
+                            let agg_expr = if let Some(name) = outer_name {
+                                ExprIR::new(trans_agg_node, OutputName::Alias(name))
+                            } else {
+                                ExprIR::new(trans_agg_node, OutputName::Alias(unique_column_name()))
+                            };
+                            agg_exprs.push(agg_expr.clone());
+                            agg_expr.output_name().clone()
+                        })
                         .clone();
-                    let input_col_node = expr_arena.add(AExpr::Column(input_col.clone()));
-                    trans_agg.set_input(input_col_node);
-                    let trans_agg_node = expr_arena.add(AExpr::Agg(trans_agg));
-
-                    // Add to aggregation expressions and replace with a reference to its output.
-                    let agg_expr = if let Some(name) = outer_name {
-                        ExprIR::new(trans_agg_node, OutputName::Alias(name))
-                    } else {
-                        ExprIR::new(trans_agg_node, OutputName::Alias(unique_column_name()))
-                    };
-                    let result_node = expr_arena.add(AExpr::Column(agg_expr.output_name().clone()));
-                    agg_exprs.push(agg_expr);
+
+                    let result_node = expr_arena.add(AExpr::Column(name));
                     Some(result_node)
                 },
                 IRAggExpr::Median(..)
@@ -338,6 +351,8 @@ fn try_build_streaming_group_by(
         let trans_output_node = expr_arena.add(AExpr::Column(uniq_name));
         trans_output_exprs.push(ExprIR::new(trans_output_node, output_name));
     }
+
+    let mut uniq_agg_exprs = PlIndexMap::new();
     for agg in aggs {
         let trans_node = try_lower_elementwise_scalar_agg_expr(
             agg.node(),
@@ -347,6 +362,7 @@ fn try_build_streaming_group_by(
             expr_arena,
             &mut trans_agg_exprs,
             &mut uniq_input_exprs,
+            &mut uniq_agg_exprs,
         )?;
         let output_name = OutputName::Alias(agg.output_name().clone());
         trans_output_exprs.push(ExprIR::new(trans_node, output_name));