feat(rust): add str_slice method to StringNameSpace (#8427)

pola-rs · Apr 22, 2023 · 71dd8f1 · 71dd8f1
1 parent 77a6465
commit 71dd8f1
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 10 deletions.
diff --git a/polars/polars-lazy/polars-plan/src/dsl/function_expr/mod.rs b/polars/polars-lazy/polars-plan/src/dsl/function_expr/mod.rs
@@ -515,6 +515,7 @@ impl From<StringFunction> for SpecialEq<Arc<dyn SeriesUdf>> {
             RStrip(matches) => map!(strings::rstrip, matches.as_deref()),
             #[cfg(feature = "string_from_radix")]
             FromRadix(radix, strict) => map!(strings::from_radix, radix, strict),
+            Slice(start, length) => map!(strings::str_slice, start, length),
         }
     }
 }

diff --git a/polars/polars-lazy/polars-plan/src/dsl/function_expr/strings.rs b/polars/polars-lazy/polars-plan/src/dsl/function_expr/strings.rs
@@ -62,6 +62,7 @@ pub enum StringFunction {
     LStrip(Option<String>),
     #[cfg(feature = "string_from_radix")]
     FromRadix(u32, bool),
+    Slice(i64, Option<u64>),
 }
 
 impl StringFunction {
@@ -82,7 +83,7 @@ impl StringFunction {
             ConcatVertical(_) | ConcatHorizontal(_) => mapper.with_dtype(DataType::Utf8),
             #[cfg(feature = "regex")]
             Replace { .. } => mapper.with_dtype(DataType::Utf8),
-            Uppercase | Lowercase | Strip(_) | LStrip(_) | RStrip(_) => {
+            Uppercase | Lowercase | Strip(_) | LStrip(_) | RStrip(_) | Slice(_, _) => {
                 mapper.with_dtype(DataType::Utf8)
             }
             #[cfg(feature = "string_from_radix")]
@@ -123,6 +124,7 @@ impl Display for StringFunction {
             StringFunction::RStrip(_) => "rstrip",
             #[cfg(feature = "string_from_radix")]
             StringFunction::FromRadix { .. } => "from_radix",
+            StringFunction::Slice(_, _) => "str_slice",
         };
 
         write!(f, "str.{s}")
@@ -596,3 +598,7 @@ pub(super) fn from_radix(s: &Series, radix: u32, strict: bool) -> PolarsResult<S
     let ca = s.utf8()?;
     ca.parse_int(radix, strict).map(|ok| ok.into_series())
 }
+pub(super) fn str_slice(s: &Series, start: i64, length: Option<u64>) -> PolarsResult<Series> {
+    let ca = s.utf8()?;
+    ca.str_slice(start, length).map(|ca| ca.into_series())
+}
diff --git a/polars/polars-lazy/polars-plan/src/dsl/string.rs b/polars/polars-lazy/polars-plan/src/dsl/string.rs
@@ -406,4 +406,12 @@ impl StringNameSpace {
                 radix, strict,
             )))
     }
+
+    /// Slice the string values.
+    pub fn str_slice(self, start: i64, length: Option<u64>) -> Expr {
+        self.0
+            .map_private(FunctionExpr::StringExpr(StringFunction::Slice(
+                start, length,
+            )))
+    }
 }
diff --git a/py-polars/src/lazy/dsl.rs b/py-polars/src/lazy/dsl.rs
@@ -654,15 +654,7 @@ impl PyExpr {
     }
 
     pub fn str_slice(&self, start: i64, length: Option<u64>) -> PyExpr {
-        let function = move |s: Series| {
-            let ca = s.utf8()?;
-            Ok(Some(ca.str_slice(start, length)?.into_series()))
-        };
-        self.clone()
-            .inner
-            .map(function, GetOutput::from_type(DataType::Utf8))
-            .with_fmt("str.slice")
-            .into()
+        self.inner.clone().str().str_slice(start, length).into()
     }
 
     pub fn str_to_uppercase(&self) -> PyExpr {