Expression support polars replace and polars replace-all (#13726)

# Description Adds the ability for `polars replace` and `polars replace-all` to work as expressions. # User-Facing Changes - `polars replace` can be used with polars expressions - `polars replace-all` can be used with polars expressions
nushell · Aug 29, 2024 · 0119534 · 0119534
1 parent 84e1ac2
commit 0119534
Show file tree

Hide file tree

Showing 3 changed files with 208 additions and 61 deletions.
diff --git a/crates/nu_plugin_polars/src/dataframe/series/string/replace.rs b/crates/nu_plugin_polars/src/dataframe/series/string/replace.rs
@@ -1,4 +1,10 @@
-use crate::{missing_flag_error, values::CustomValueSupport, PolarsPlugin};
+use crate::{
+    missing_flag_error,
+    values::{
+        cant_convert_err, CustomValueSupport, NuExpression, PolarsPluginObject, PolarsPluginType,
+    },
+    PolarsPlugin,
+};
 
 use super::super::super::values::{Column, NuDataFrame};
 
@@ -7,7 +13,7 @@ use nu_protocol::{
     Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
     Value,
 };
-use polars::prelude::{IntoSeries, StringNameSpaceImpl};
+use polars::prelude::{lit, IntoSeries, StringNameSpaceImpl};
 
 #[derive(Clone)]
 pub struct Replace;
@@ -37,33 +43,61 @@ impl PluginCommand for Replace {
                 "replacing string",
                 Some('r'),
             )
-            .input_output_type(
-                Type::Custom("dataframe".into()),
-                Type::Custom("dataframe".into()),
-            )
+            .input_output_types(vec![
+                (
+                    Type::Custom("dataframe".into()),
+                    Type::Custom("dataframe".into()),
+                ),
+                (
+                    Type::Custom("expression".into()),
+                    Type::Custom("expression".into()),
+                ),
+            ])
             .category(Category::Custom("dataframe".into()))
     }
 
     fn examples(&self) -> Vec<Example> {
-        vec![Example {
-            description: "Replaces string",
-            example: "[abc abc abc] | polars into-df | polars replace --pattern ab --replace AB",
-            result: Some(
-                NuDataFrame::try_from_columns(
-                    vec![Column::new(
-                        "0".to_string(),
-                        vec![
-                            Value::test_string("ABc"),
-                            Value::test_string("ABc"),
-                            Value::test_string("ABc"),
-                        ],
-                    )],
-                    None,
-                )
-                .expect("simple df for test should not fail")
-                .into_value(Span::test_data()),
-            ),
-        }]
+        vec![
+            Example {
+                description: "Replaces string in column",
+                example:
+                    "[[a]; [abc] [abcabc]] | polars into-df | polars select (polars col a | polars replace --pattern ab --replace AB) | polars collect",
+                result: Some(
+                    NuDataFrame::try_from_columns(
+                        vec![Column::new(
+                            "a".to_string(),
+                            vec![
+                                Value::test_string("ABc"),
+                                Value::test_string("ABcabc"),
+                            ],
+                        )],
+                        None,
+                    )
+                    .expect("simple df for test should not fail")
+                    .into_value(Span::test_data()),
+                ),
+            },
+            Example {
+                description: "Replaces string",
+                example:
+                    "[abc abc abc] | polars into-df | polars replace --pattern ab --replace AB",
+                result: Some(
+                    NuDataFrame::try_from_columns(
+                        vec![Column::new(
+                            "0".to_string(),
+                            vec![
+                                Value::test_string("ABc"),
+                                Value::test_string("ABc"),
+                                Value::test_string("ABc"),
+                            ],
+                        )],
+                        None,
+                    )
+                    .expect("simple df for test should not fail")
+                    .into_value(Span::test_data()),
+                ),
+            },
+        ]
     }
 
     fn run(
@@ -73,15 +107,56 @@ impl PluginCommand for Replace {
         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,
+    expr: NuExpression,
+) -> Result<PipelineData, ShellError> {
+    let pattern: String = call
+        .get_flag("pattern")?
+        .ok_or_else(|| missing_flag_error("pattern", call.head))?;
+    let pattern = lit(pattern);
+
+    let replace: String = call
+        .get_flag("replace")?
+        .ok_or_else(|| missing_flag_error("replace", call.head))?;
+    let replace = lit(replace);
+
+    let res: NuExpression = expr
+        .into_polars()
+        .str()
+        .replace(pattern, replace, false)
+        .into();
+
+    res.to_pipeline_data(plugin, engine, call.head)
+}
+
+fn command_df(
     plugin: &PolarsPlugin,
     engine: &EngineInterface,
     call: &EvaluatedCall,
-    input: PipelineData,
+    df: NuDataFrame,
 ) -> Result<PipelineData, ShellError> {
     let pattern: String = call
         .get_flag("pattern")?
@@ -90,7 +165,6 @@ fn command(
         .get_flag("replace")?
         .ok_or_else(|| missing_flag_error("replace", call.head))?;
 
-    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 {
         error: "Error conversion to string".into(),

diff --git a/crates/nu_plugin_polars/src/dataframe/series/string/replace_all.rs b/crates/nu_plugin_polars/src/dataframe/series/string/replace_all.rs
@@ -1,4 +1,10 @@
-use crate::{missing_flag_error, values::CustomValueSupport, PolarsPlugin};
+use crate::{
+    missing_flag_error,
+    values::{
+        cant_convert_err, CustomValueSupport, NuExpression, PolarsPluginObject, PolarsPluginType,
+    },
+    PolarsPlugin,
+};
 
 use super::super::super::values::{Column, NuDataFrame};
 
@@ -7,7 +13,7 @@ use nu_protocol::{
     Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
     Value,
 };
-use polars::prelude::{IntoSeries, StringNameSpaceImpl};
+use polars::prelude::{lit, IntoSeries, StringNameSpaceImpl};
 
 #[derive(Clone)]
 pub struct ReplaceAll;
@@ -37,34 +43,62 @@ impl PluginCommand for ReplaceAll {
                 "replacing string",
                 Some('r'),
             )
-            .input_output_type(
-                Type::Custom("dataframe".into()),
-                Type::Custom("dataframe".into()),
-            )
+            .input_output_types(vec![
+                (
+                    Type::Custom("dataframe".into()),
+                    Type::Custom("dataframe".into()),
+                ),
+                (
+                    Type::Custom("expression".into()),
+                    Type::Custom("expression".into()),
+                ),
+            ])
             .category(Category::Custom("dataframe".into()))
     }
 
     fn examples(&self) -> Vec<Example> {
-        vec![Example {
-            description: "Replaces string",
-            example:
-                "[abac abac abac] | polars into-df | polars replace-all --pattern a --replace A",
-            result: Some(
-                NuDataFrame::try_from_columns(
-                    vec![Column::new(
-                        "0".to_string(),
-                        vec![
-                            Value::test_string("AbAc"),
-                            Value::test_string("AbAc"),
-                            Value::test_string("AbAc"),
-                        ],
-                    )],
-                    None,
-                )
-                .expect("simple df for test should not fail")
-                .into_value(Span::test_data()),
-            ),
-        }]
+        vec![
+            Example {
+                description: "Replaces string in a column",
+                example:
+                    "[[a]; [abac] [abac] [abac]] | polars into-df | polars select (polars col a | polars replace-all --pattern a --replace A) | polars collect",
+                result: Some(
+                    NuDataFrame::try_from_columns(
+                        vec![Column::new(
+                            "a".to_string(),
+                            vec![
+                                Value::test_string("AbAc"),
+                                Value::test_string("AbAc"),
+                                Value::test_string("AbAc"),
+                            ],
+                        )],
+                        None,
+                    )
+                    .expect("simple df for test should not fail")
+                    .into_value(Span::test_data()),
+                ),
+            },
+            Example {
+                description: "Replaces string",
+                example:
+                    "[abac abac abac] | polars into-df | polars replace-all --pattern a --replace A",
+                result: Some(
+                    NuDataFrame::try_from_columns(
+                        vec![Column::new(
+                            "0".to_string(),
+                            vec![
+                                Value::test_string("AbAc"),
+                                Value::test_string("AbAc"),
+                                Value::test_string("AbAc"),
+                            ],
+                        )],
+                        None,
+                    )
+                    .expect("simple df for test should not fail")
+                    .into_value(Span::test_data()),
+                ),
+            },
+        ]
     }
 
     fn run(
@@ -74,15 +108,55 @@ impl PluginCommand for ReplaceAll {
         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_state: &EngineInterface,
+    call: &EvaluatedCall,
+    expr: NuExpression,
+) -> Result<PipelineData, ShellError> {
+    let pattern: String = call
+        .get_flag("pattern")?
+        .ok_or_else(|| missing_flag_error("pattern", call.head))?;
+    let pattern = lit(pattern);
+    let replace: String = call
+        .get_flag("replace")?
+        .ok_or_else(|| missing_flag_error("replace", call.head))?;
+    let replace = lit(replace);
+
+    let res: NuExpression = expr
+        .into_polars()
+        .str()
+        .replace_all(pattern, replace, false)
+        .into();
+
+    res.to_pipeline_data(plugin, engine_state, call.head)
+}
+
+fn command_df(
     plugin: &PolarsPlugin,
     engine_state: &EngineInterface,
     call: &EvaluatedCall,
-    input: PipelineData,
+    df: NuDataFrame,
 ) -> Result<PipelineData, ShellError> {
     let pattern: String = call
         .get_flag("pattern")?
@@ -91,7 +165,6 @@ fn command(
         .get_flag("replace")?
         .ok_or_else(|| missing_flag_error("replace", call.head))?;
 
-    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 {
         error: "Error conversion to string".into(),

diff --git a/crates/nu_plugin_polars/src/dataframe/values/mod.rs b/crates/nu_plugin_polars/src/dataframe/values/mod.rs
@@ -185,8 +185,8 @@ impl CustomValueType {
     }
 }
 
-pub fn cant_convert_err(value: &Value, types: &[PolarsPluginType]) -> ShellError {
-    let type_string = types
+pub fn cant_convert_err(value: &Value, expected_types: &[PolarsPluginType]) -> ShellError {
+    let type_string = expected_types
         .iter()
         .map(ToString::to_string)
         .collect::<Vec<String>>()