From aa710eeb9aac2d2a686c0679eb2efb7bee715761 Mon Sep 17 00:00:00 2001
From: Jack Wright <56345+ayax79@users.noreply.github.com>
Date: Fri, 13 Jun 2025 12:10:29 -0700
Subject: [PATCH] Add groupby support for `polars last` (#15953)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

# 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>
---
 .../src/dataframe/command/data/last.rs        | 38 ++++++++++++++++++-
 1 file changed, 37 insertions(+), 1 deletion(-)

diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/last.rs b/crates/nu_plugin_polars/src/dataframe/command/data/last.rs
index b1a7f6b0c3..6e1f9c53b7 100644
--- a/crates/nu_plugin_polars/src/dataframe/command/data/last.rs
+++ b/crates/nu_plugin_polars/src/dataframe/command/data/last.rs
@@ -1,6 +1,6 @@
 use crate::{
     PolarsPlugin,
-    values::{Column, CustomValueSupport, NuLazyFrame, PolarsPluginObject},
+    values::{Column, CustomValueSupport, NuLazyFrame, NuLazyGroupBy, PolarsPluginObject},
 };
 
 use crate::values::{NuDataFrame, NuExpression, utils::DEFAULT_ROWS};
@@ -9,6 +9,7 @@ use nu_protocol::{
     Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
     Value,
 };
+use polars::df;
 
 #[derive(Clone)]
 pub struct LastDF;
@@ -62,6 +63,24 @@ impl PluginCommand for LastDF {
                 example: "polars col a | polars last",
                 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) => {
                 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 = expr.into_polars().last().into();
@@ -121,6 +143,20 @@ fn command_lazy(
     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)]
 mod test {
     use crate::test::test_polars_plugin_command;