From d8255040f1b53689a799f8d18bc1445f503bc9e8 Mon Sep 17 00:00:00 2001
From: Jack Wright <56345+ayax79@users.noreply.github.com>
Date: Tue, 8 Jul 2025 10:24:47 -0700
Subject: [PATCH] Added flag limit for `polars arg-sort` (#16132)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

# Description
Exposes the polars sort option limit for the arg-sort command.

```nu
> ❯ : [1 2 2 3 3] | polars into-df | polars arg-sort --limit 2
╭───┬──────────╮
│ # │ arg_sort │
├───┼──────────┤
│ 0 │        0 │
│ 1 │        1 │
╰───┴──────────╯
```


# User-Facing Changes
- The `--limit` flag is now available for `polars arg-sort`

Co-authored-by: Jack Wright <jack.wright@nike.com>
---
 .../src/dataframe/command/index/arg_sort.rs   | 28 +++++++++++++++++--
 1 file changed, 25 insertions(+), 3 deletions(-)

diff --git a/crates/nu_plugin_polars/src/dataframe/command/index/arg_sort.rs b/crates/nu_plugin_polars/src/dataframe/command/index/arg_sort.rs
index e1c1214641..1414744665 100644
--- a/crates/nu_plugin_polars/src/dataframe/command/index/arg_sort.rs
+++ b/crates/nu_plugin_polars/src/dataframe/command/index/arg_sort.rs
@@ -4,7 +4,8 @@ use super::super::super::values::{Column, NuDataFrame};
 
 use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
 use nu_protocol::{
-    Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, Value,
+    Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
+    Value,
 };
 use polars::prelude::{IntoSeries, SortOptions};
 
@@ -30,6 +31,12 @@ impl PluginCommand for ArgSort {
         Signature::build(self.name())
             .switch("reverse", "reverse order", Some('r'))
             .switch("nulls-last", "nulls ordered last", Some('n'))
+            .named(
+                "limit",
+                SyntaxShape::Int,
+                "Limit a sort output, this is for optimization purposes and might be ignored.",
+                Some('l'),
+            )
             .switch(
                 "maintain-order",
                 "maintain order on sorted items",
@@ -86,6 +93,21 @@ impl PluginCommand for ArgSort {
                     .into_value(Span::test_data()),
                 ),
             },
+            Example {
+                description: "Returns indexes for a sorted series",
+                example: "[1 2 2 3 3] | polars into-df | polars arg-sort --limit 2",
+                result: Some(
+                    NuDataFrame::try_from_columns(
+                        vec![Column::new(
+                            "arg_sort".to_string(),
+                            vec![Value::test_int(0), Value::test_int(1)],
+                        )],
+                        None,
+                    )
+                    .expect("simple df for test should not fail")
+                    .into_value(Span::test_data()),
+                ),
+            },
         ]
     }
 
@@ -106,6 +128,7 @@ fn command(
     call: &EvaluatedCall,
     input: PipelineData,
 ) -> Result<PipelineData, ShellError> {
+    let limit: Option<u64> = call.get_flag("limit")?;
     let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?;
 
     let sort_options = SortOptions {
@@ -113,8 +136,7 @@ fn command(
         nulls_last: call.has_flag("nulls-last")?,
         multithreaded: true,
         maintain_order: call.has_flag("maintain-order")?,
-        // todo - expose limit
-        limit: None,
+        limit,
     };
 
     let mut res = df