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 { + let limit: Option = 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