Skip to content

Commit 672e4c6

Browse files
authored
refactor: Small drive-by's (#20772)
1 parent 6753bb6 commit 672e4c6

File tree

3 files changed

+11
-7
lines changed

3 files changed

+11
-7
lines changed

crates/polars-plan/src/plans/aexpr/mod.rs

+1
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ pub enum IRAggExpr {
4747
method: QuantileMethod,
4848
},
4949
Sum(Node),
50+
// include_nulls
5051
Count(Node, bool),
5152
Std(Node, u8),
5253
Var(Node, u8),

crates/polars-plan/src/plans/aexpr/traverse.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use super::*;
33
impl AExpr {
44
/// Push the inputs of this node to the given container, in reverse order.
55
/// This ensures the primary node responsible for the name is pushed last.
6-
pub(crate) fn inputs_rev<E>(&self, container: &mut E)
6+
pub fn inputs_rev<E>(&self, container: &mut E)
77
where
88
E: Extend<Node>,
99
{

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

+9-6
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ use polars_plan::prelude::GroupbyOptions;
1212
use polars_utils::arena::{Arena, Node};
1313
use polars_utils::itertools::Itertools;
1414
use polars_utils::pl_str::PlSmallStr;
15+
use recursive::recursive;
1516
use slotmap::SlotMap;
1617

1718
use super::lower_expr::lower_exprs;
@@ -71,6 +72,7 @@ fn build_group_by_fallback(
7172
///
7273
/// Such an expression is defined as the elementwise combination of scalar
7374
/// aggregations of elementwise combinations of the input columns / scalar literals.
75+
#[recursive]
7476
fn try_lower_elementwise_scalar_agg_expr(
7577
expr: Node,
7678
inside_agg: bool,
@@ -178,7 +180,10 @@ fn try_lower_elementwise_scalar_agg_expr(
178180
},
179181

180182
AExpr::Agg(agg) => {
181-
let orig_agg = agg.clone();
183+
// Nested aggregates not supported.
184+
if inside_agg {
185+
return None;
186+
}
182187
match agg {
183188
IRAggExpr::Min { input, .. }
184189
| IRAggExpr::Max { input, .. }
@@ -188,10 +193,7 @@ fn try_lower_elementwise_scalar_agg_expr(
188193
| IRAggExpr::Sum(input)
189194
| IRAggExpr::Var(input, ..)
190195
| IRAggExpr::Std(input, ..) => {
191-
// Nested aggregates not supported.
192-
if inside_agg {
193-
return None;
194-
}
196+
let orig_agg = agg.clone();
195197
// Lower and replace input.
196198
let trans_input = lower_rec!(*input, true)?;
197199
let mut trans_agg = orig_agg;
@@ -311,7 +313,8 @@ fn try_build_streaming_group_by(
311313
&mut trans_agg_exprs,
312314
&trans_input_cols,
313315
)?;
314-
trans_output_exprs.push(ExprIR::new(trans_node, agg.output_name_inner().clone()));
316+
let output_name = OutputName::Alias(agg.output_name().clone());
317+
trans_output_exprs.push(ExprIR::new(trans_node, output_name));
315318
}
316319

317320
let input_schema = &phys_sm[trans_input.node].output_schema;

0 commit comments

Comments
 (0)