Added flag limit for polars arg-sort (#16132)

# 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>
This commit is contained in:
Jack Wright
2025-07-08 10:24:47 -07:00
committed by GitHub
parent 4da755895d
commit d8255040f1

View File

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