extern/nushell
1
0
Fork 1
mirror of https://github.com/nushell/nushell.git synced 2025-04-16 01:08:21 +02:00
Code Issues Packages Projects Releases Wiki Activity

polars: extend NuExpression::extract_exprs to handle records

Browse Source
This commit is contained in:
pyz 2025-04-11 17:52:43 -04:00
parent 885b87a842
commit 467f193c34
4 changed files with 170 additions and 44 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

98
crates/nu_plugin_polars/src/dataframe/command/aggregation/aggregate.rs
View File

@ -40,9 +40,10 @@ impl PluginCommand for LazyAggregate {
}
fn examples(&self) -> Vec<Example> {
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(

71
crates/nu_plugin_polars/src/dataframe/command/data/select.rs
View File

@ -37,21 +37,62 @@ impl PluginCommand for LazySelect {
}
fn examples(&self) -> Vec<Example> {
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(

35
crates/nu_plugin_polars/src/dataframe/command/data/with_column.rs
View File

@ -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(

10
crates/nu_plugin_polars/src/dataframe/values/nu_expression/mod.rs
View File

@ -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::<Result<Vec<ExtractedExpr>, ShellError>>()
.map(ExtractedExpr::List),
Value::List { vals, .. } => vals
.into_iter()
.map(|x| Self::extract_exprs(plugin, x))