Add groupby support for polars last (#15953)

# Description
Allows `polars last` to be used with group-by
```nu
> ❯ : [[a b c d]; [1 0.5 true Apple] [2 0.5 true Orange] [2 4 true Apple] [3 10 false Apple] [4 13 false Banana] [5 14 true Banana]] | polars into-df -s {a: u8, b: f32, c: bool, d: str} | polars group-by d | polars last | polars sort-by [a] | polars collect
╭───┬────────┬───┬───────┬───────╮
│ # │   d    │ a │   b   │   c   │
├───┼────────┼───┼───────┼───────┤
│ 0 │ Orange │ 2 │  0.50 │ true  │
│ 1 │ Apple  │ 3 │ 10.00 │ false │
│ 2 │ Banana │ 5 │ 14.00 │ true  │
╰───┴────────┴───┴───────┴───────╯
```

# User-Facing Changes
- `polars last` can now be used with group-by expressions

Co-authored-by: Jack Wright <jack.wright@nike.com>
This commit is contained in:
Jack Wright 2025-06-13 12:10:29 -07:00 committed by GitHub
parent 91e843a6d4
commit aa710eeb9a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -1,6 +1,6 @@
use crate::{ use crate::{
PolarsPlugin, PolarsPlugin,
values::{Column, CustomValueSupport, NuLazyFrame, PolarsPluginObject}, values::{Column, CustomValueSupport, NuLazyFrame, NuLazyGroupBy, PolarsPluginObject},
}; };
use crate::values::{NuDataFrame, NuExpression, utils::DEFAULT_ROWS}; use crate::values::{NuDataFrame, NuExpression, utils::DEFAULT_ROWS};
@ -9,6 +9,7 @@ use nu_protocol::{
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type, Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
Value, Value,
}; };
use polars::df;
#[derive(Clone)] #[derive(Clone)]
pub struct LastDF; pub struct LastDF;
@ -62,6 +63,24 @@ impl PluginCommand for LastDF {
example: "polars col a | polars last", example: "polars col a | polars last",
result: None, result: None,
}, },
Example {
description: "Aggregate the last values in the group.",
example: "[[a b c d]; [1 0.5 true Apple] [2 0.5 true Orange] [2 4 true Apple] [3 10 false Apple] [4 13 false Banana] [5 14 true Banana]] | polars into-df -s {a: u8, b: f32, c: bool, d: str} | polars group-by d | polars last | polars sort-by [a] | polars collect",
result: Some(
NuDataFrame::new(
false,
df!(
"d" => &["Orange", "Apple", "Banana"],
"a" => &[2, 3, 5],
"b" => &[0.50, 10.0, 14.0],
"c" => &[true, false, true],
)
.expect("dataframe creation should succeed"),
)
.into_value(Span::test_data()),
),
},
] ]
} }
@ -81,6 +100,9 @@ impl PluginCommand for LastDF {
PolarsPluginObject::NuLazyFrame(lazy) => { PolarsPluginObject::NuLazyFrame(lazy) => {
command_lazy(plugin, engine, call, lazy).map_err(|e| e.into()) command_lazy(plugin, engine, call, lazy).map_err(|e| e.into())
} }
PolarsPluginObject::NuLazyGroupBy(groupby) => {
command_groupby(plugin, engine, call, groupby).map_err(|e| e.into())
}
_ => { _ => {
let expr = NuExpression::try_from_value(plugin, &value)?; let expr = NuExpression::try_from_value(plugin, &value)?;
let expr: NuExpression = expr.into_polars().last().into(); let expr: NuExpression = expr.into_polars().last().into();
@ -121,6 +143,20 @@ fn command_lazy(
res.to_pipeline_data(plugin, engine, call.head) res.to_pipeline_data(plugin, engine, call.head)
} }
fn command_groupby(
plugin: &PolarsPlugin,
engine: &EngineInterface,
call: &EvaluatedCall,
groupby: NuLazyGroupBy,
) -> Result<PipelineData, ShellError> {
let rows: Option<usize> = call.opt(0)?;
let rows = rows.unwrap_or(1);
let res = groupby.to_polars().tail(Some(rows));
let res: NuLazyFrame = res.into();
res.to_pipeline_data(plugin, engine, call.head)
}
#[cfg(test)] #[cfg(test)]
mod test { mod test {
use crate::test::test_polars_plugin_command; use crate::test::test_polars_plugin_command;