diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/groupby.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/groupby.rs index 24372e47aa..8a98864655 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/groupby.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/groupby.rs @@ -39,9 +39,10 @@ impl PluginCommand for ToLazyGroupBy { } fn examples(&self) -> Vec { - vec![Example { - description: "Group by and perform an aggregation", - example: r#"[[a b]; [1 2] [1 4] [2 6] [2 4]] + vec![ + Example { + description: "Group by and perform an aggregation", + example: r#"[[a b]; [1 2] [1 4] [2 6] [2 4]] | polars into-lazy | polars group-by a | polars agg [ @@ -51,19 +52,45 @@ impl PluginCommand for ToLazyGroupBy { ] | polars collect | polars sort-by a"#, - result: Some( - NuDataFrame::from( - df!( - "a" => &[1i64, 2], - "b_min" => &[2i64, 4], - "b_max" => &[4i64, 6], - "b_sum" => &[6i64, 10], + result: Some( + NuDataFrame::from( + df!( + "a" => &[1i64, 2], + "b_min" => &[2i64, 4], + "b_max" => &[4i64, 6], + "b_sum" => &[6i64, 10], + ) + .expect("should not fail"), ) - .expect("should not fail"), - ) - .into_value(Span::test_data()), - ), - }] + .into_value(Span::test_data()), + ), + }, + Example { + description: "Group by an expression and perform an aggregation", + example: r#"[[a b]; [2025-04-01 1] [2025-04-02 2] [2025-04-03 3] [2025-04-04 4]] + | polars into-lazy + | polars group-by (polars col a | polars get-day | $in mod 2) + | polars agg [ + (polars col b | polars min | polars as "b_min") + (polars col b | polars max | polars as "b_max") + (polars col b | polars sum | polars as "b_sum") + ] + | polars collect + | polars sort-by a"#, + result: Some( + NuDataFrame::from( + df!( + "a" => &[0i64, 1], + "b_min" => &[2i64, 1], + "b_max" => &[4i64, 3], + "b_sum" => &[6i64, 4], + ) + .expect("should not fail"), + ) + .into_value(Span::test_data()), + ), + }, + ] } fn run( @@ -79,11 +106,11 @@ impl PluginCommand for ToLazyGroupBy { if expressions .iter() - .any(|expr| !matches!(expr, Expr::Column(..))) + .any(|expr| matches!(expr, Expr::Agg(..) | Expr::Window { .. })) { let value: Value = call.req(0)?; Err(ShellError::IncompatibleParametersSingle { - msg: "Expected only Col expressions".into(), + msg: "Cannot group by an aggregation or window expression".into(), span: value.span(), })?; }