extern/nushell
1
0
Fork 1
mirror of https://github.com/nushell/nushell.git synced 2025-05-31 07:08:22 +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

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

@ -40,7 +40,8 @@ impl PluginCommand for LazyAggregate {
} }
fn examples(&self) -> Vec<Example> { fn examples(&self) -> Vec<Example> {
vec![Example { vec![
Example {
description: "Group by and perform an aggregation", description: "Group by and perform an aggregation",
example: r#"[[a b]; [1 2] [1 4] [2 6] [2 4]] example: r#"[[a b]; [1 2] [1 4] [2 6] [2 4]]
| polars into-lazy | polars into-lazy
@ -77,7 +78,46 @@ impl PluginCommand for LazyAggregate {
.expect("simple df for test should not fail") .expect("simple df for test should not fail")
.into_value(Span::test_data()), .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( fn run(

47
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> { fn examples(&self) -> Vec<Example> {
vec![Example { vec![
Example {
description: "Select a column from the dataframe", description: "Select a column from the dataframe",
example: "[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars select a", example: "[[a b]; [6 2] [4 2] [2 2]] | polars into-df | polars select a",
result: Some( result: Some(
NuDataFrame::try_from_columns( NuDataFrame::try_from_columns(
vec![Column::new( vec![
Column::new(
"a".to_string(), "a".to_string(),
vec![Value::test_int(6), Value::test_int(4), Value::test_int(2)], 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, None,
) )
.expect("simple df for test should not fail") .expect("simple df for test should not fail")
.into_value(Span::test_data()), .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( 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) * 2 | polars as "c")
((polars col a) * 3 | polars as "d") ((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"#, | polars collect"#,
result: Some( result: Some(
NuDataFrame::try_from_columns( 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) Value::Custom { .. } => NuExpression::try_from_value(plugin, &value)
.map(NuExpression::into_polars) .map(NuExpression::into_polars)
.map(ExtractedExpr::Single), .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 Value::List { vals, .. } => vals
.into_iter() .into_iter()
.map(|x| Self::extract_exprs(plugin, x)) .map(|x| Self::extract_exprs(plugin, x))