diff --git a/crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregate.rs b/crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregate.rs index 73ff561a28..64ff8952c5 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregate.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregate.rs @@ -40,9 +40,10 @@ impl PluginCommand for LazyAggregate { } 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 [ @@ -52,32 +53,71 @@ impl PluginCommand for LazyAggregate { ] | polars collect | polars sort-by a"#, - result: Some( - NuDataFrame::try_from_columns( - vec![ - Column::new( - "a".to_string(), - vec![Value::test_int(1), Value::test_int(2)], - ), - Column::new( - "b_min".to_string(), - vec![Value::test_int(2), Value::test_int(4)], - ), - Column::new( - "b_max".to_string(), - vec![Value::test_int(4), Value::test_int(6)], - ), - Column::new( - "b_sum".to_string(), - vec![Value::test_int(6), Value::test_int(10)], - ), - ], - None, - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }] + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new( + "a".to_string(), + vec![Value::test_int(1), Value::test_int(2)], + ), + Column::new( + "b_min".to_string(), + vec![Value::test_int(2), Value::test_int(4)], + ), + Column::new( + "b_max".to_string(), + vec![Value::test_int(4), Value::test_int(6)], + ), + Column::new( + "b_sum".to_string(), + vec![Value::test_int(6), Value::test_int(10)], + ), + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Group by and perform an aggregation using a record", + example: r#"[[a b]; [1 2] [1 4] [2 6] [2 4]] + | polars into-lazy + | polars group-by a + | polars agg { + b_min: (polars col b | polars min) + b_max: (polars col b | polars max) + b_sum: (polars col b | polars sum) + } + | polars collect + | polars sort-by a"#, + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new( + "a".to_string(), + vec![Value::test_int(1), Value::test_int(2)], + ), + Column::new( + "b_min".to_string(), + vec![Value::test_int(2), Value::test_int(4)], + ), + Column::new( + "b_max".to_string(), + vec![Value::test_int(4), Value::test_int(6)], + ), + Column::new( + "b_sum".to_string(), + vec![Value::test_int(6), Value::test_int(10)], + ), + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + ] } fn run( diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/select.rs b/crates/nu_plugin_polars/src/dataframe/command/data/select.rs index ac6c2683d7..8c9a135f4b 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/select.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/select.rs @@ -37,21 +37,62 @@ impl PluginCommand for LazySelect { } fn examples(&self) -> Vec { - vec![Example { - description: "Select a column from the dataframe", - example: "[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars select a", - result: Some( - NuDataFrame::try_from_columns( - vec![Column::new( - "a".to_string(), - vec![Value::test_int(6), Value::test_int(4), Value::test_int(2)], - )], - None, - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }] + vec![ + Example { + description: "Select a column from the dataframe", + example: "[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars select a", + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new( + "a".to_string(), + vec![Value::test_int(6), Value::test_int(4), Value::test_int(2)], + ) + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Select a column from a dataframe using a record", + example: "[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars select {c: ((polars col a) * 2)}", + result: Some( + NuDataFrame::try_from_columns( + vec![Column::new( + "c".to_string(), + vec![Value::test_int(12), Value::test_int(8), Value::test_int(4)], + )], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Select a column from a dataframe using a mix of expressions and record of expressions", + example: "[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars select a b {c: ((polars col a) * 2)}", + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new( + "a".to_string(), + vec![Value::test_int(6), Value::test_int(4), Value::test_int(2)]), + Column::new( + "b".to_string(), + vec![Value::test_int(2), Value::test_int(2), Value::test_int(2)]), + Column::new( + "c".to_string(), + vec![Value::test_int(12), Value::test_int(8), Value::test_int(4)]) + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + ] } fn run( diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/with_column.rs b/crates/nu_plugin_polars/src/dataframe/command/data/with_column.rs index 637caa1776..fba10c02a0 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/with_column.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/with_column.rs @@ -76,6 +76,41 @@ impl PluginCommand for WithColumn { ((polars col a) * 2 | polars as "c") ((polars col a) * 3 | polars as "d") ] + | polars collect"#, + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new( + "a".to_string(), + vec![Value::test_int(1), Value::test_int(3)], + ), + Column::new( + "b".to_string(), + vec![Value::test_int(2), Value::test_int(4)], + ), + Column::new( + "c".to_string(), + vec![Value::test_int(2), Value::test_int(6)], + ), + Column::new( + "d".to_string(), + vec![Value::test_int(3), Value::test_int(9)], + ), + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Add series to the dataframe using a record", + example: r#"[[a b]; [1 2] [3 4]] + | polars into-lazy + | polars with-column { + c: ((polars col a) * 2) + d: ((polars col a) * 3) + } | polars collect"#, result: Some( NuDataFrame::try_from_columns( diff --git a/crates/nu_plugin_polars/src/dataframe/values/nu_expression/mod.rs b/crates/nu_plugin_polars/src/dataframe/values/nu_expression/mod.rs index 7b60f9b1e4..dcc6b18435 100644 --- a/crates/nu_plugin_polars/src/dataframe/values/nu_expression/mod.rs +++ b/crates/nu_plugin_polars/src/dataframe/values/nu_expression/mod.rs @@ -126,6 +126,16 @@ impl ExtractedExpr { Value::Custom { .. } => NuExpression::try_from_value(plugin, &value) .map(NuExpression::into_polars) .map(ExtractedExpr::Single), + Value::Record { val, .. } => val + .iter() + .map(|(key, value)| { + NuExpression::try_from_value(plugin, value) + .map(NuExpression::into_polars) + .map(|expr| expr.alias(key)) + .map(ExtractedExpr::Single) + }) + .collect::, ShellError>>() + .map(ExtractedExpr::List), Value::List { vals, .. } => vals .into_iter() .map(|x| Self::extract_exprs(plugin, x))