diff --git a/crates/nu_plugin_polars/src/dataframe/series/string/str_slice.rs b/crates/nu_plugin_polars/src/dataframe/series/string/str_slice.rs index d81500a7cd..f51b42c3d0 100644 --- a/crates/nu_plugin_polars/src/dataframe/series/string/str_slice.rs +++ b/crates/nu_plugin_polars/src/dataframe/series/string/str_slice.rs @@ -1,4 +1,9 @@ -use crate::{values::CustomValueSupport, PolarsPlugin}; +use crate::{ + values::{ + cant_convert_err, CustomValueSupport, NuExpression, PolarsPluginObject, PolarsPluginType, + }, + PolarsPlugin, +}; use super::super::super::values::{Column, NuDataFrame}; @@ -8,7 +13,7 @@ use nu_protocol::{ Value, }; use polars::{ - prelude::{IntoSeries, NamedFrom, StringNameSpaceImpl}, + prelude::{lit, Expr, IntoSeries, NamedFrom, Null, StringNameSpaceImpl}, series::Series, }; @@ -30,15 +35,40 @@ impl PluginCommand for StrSlice { Signature::build(self.name()) .required("start", SyntaxShape::Int, "start of slice") .named("length", SyntaxShape::Int, "optional length", Some('l')) - .input_output_type( - Type::Custom("dataframe".into()), - Type::Custom("dataframe".into()), - ) + .input_output_types(vec![ + ( + Type::Custom("expression".into()), + Type::Custom("expression".into()), + ), + ( + Type::Custom("dataframe".into()), + Type::Custom("dataframe".into()), + ), + ]) .category(Category::Custom("dataframe".into())) } fn examples(&self) -> Vec { vec![ + Example { + description: "Creates slices from the strings in a specified column", + example: "[[a]; [abcded] [abc321] [abc123]] | polars into-df | polars select (polars col a | polars str-slice 1 --length 2) | polars collect", + result: Some( + NuDataFrame::try_from_columns( + vec![Column::new( + "a".to_string(), + vec![ + Value::test_string("bc"), + Value::test_string("bc"), + Value::test_string("bc"), + ], + )], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, Example { description: "Creates slices from the strings", example: "[abcded abc321 abc123] | polars into-df | polars str-slice 1 --length 2", @@ -87,15 +117,47 @@ impl PluginCommand for StrSlice { call: &EvaluatedCall, input: PipelineData, ) -> Result { - command(plugin, engine, call, input).map_err(LabeledError::from) + let value = input.into_value(call.head)?; + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuDataFrame(df) => command_df(plugin, engine, call, df), + PolarsPluginObject::NuLazyFrame(lazy) => { + command_df(plugin, engine, call, lazy.collect(call.head)?) + } + PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr), + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + ], + )), + } + .map_err(LabeledError::from) } } -fn command( +fn command_expr( plugin: &PolarsPlugin, engine: &EngineInterface, call: &EvaluatedCall, - input: PipelineData, + expr: NuExpression, +) -> Result { + let start = lit(call.req::(0)?); + let length: Expr = call + .get_flag::("length")? + .map(lit) + .unwrap_or(lit(Null {})); + + let res: NuExpression = expr.into_polars().str().slice(start, length).into(); + res.to_pipeline_data(plugin, engine, call.head) +} + +fn command_df( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + df: NuDataFrame, ) -> Result { let start: i64 = call.req(0)?; let start = Series::new("", &[start]); @@ -106,7 +168,6 @@ fn command( None => Series::new_null("", 1), }; - let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?; let series = df.as_series(call.head)?; let chunked = series.str().map_err(|e| ShellError::GenericError {