From f611196373e51b8dda018e2fedb8cbfa3da8726b Mon Sep 17 00:00:00 2001
From: Jack Wright <56345+ayax79@users.noreply.github.com>
Date: Thu, 5 Sep 2024 09:06:37 -0700
Subject: [PATCH] Expression support for `polars str-slice` (#13783)

# Description
Provides expression support for `polars str-slice`:
<img width="893" alt="Screenshot 2024-09-04 at 18 03 05"
src="https://github.com/user-attachments/assets/d8a8a2a7-53cf-4c3a-ae7a-dfdaf48a05ee">

# User-Facing Changes
- `polars str-slice` can now be used as an expression
---
 .../src/dataframe/series/string/str_slice.rs  | 81 ++++++++++++++++---
 1 file changed, 71 insertions(+), 10 deletions(-)

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<Example> {
         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<PipelineData, LabeledError> {
-        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<PipelineData, ShellError> {
+    let start = lit(call.req::<i64>(0)?);
+    let length: Expr = call
+        .get_flag::<i64>("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<PipelineData, ShellError> {
     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 {