split exact expression (#2872)

pola-rs · Mar 11, 2022 · 10948a5 · 10948a5
1 parent 0038bd2
commit 10948a5
Show file tree

Hide file tree

Showing 7 changed files with 227 additions and 0 deletions.
diff --git a/polars/polars-lazy/src/dsl/string.rs b/polars/polars-lazy/src/dsl/string.rs
@@ -1,5 +1,6 @@
 use super::*;
 use polars_arrow::array::ValueSize;
+use polars_arrow::export::arrow::array::{MutableArray, MutableUtf8Array};
 
 /// Specialized expressions for [`Series`] of [`DataType::Utf8`].
 pub struct StringNameSpace(pub(crate) Expr);
@@ -109,6 +110,104 @@ impl StringNameSpace {
             .with_fmt("str.split")
     }
 
+    pub fn split_exact(self, by: &str, n: usize) -> Expr {
+        let by = by.to_string();
+
+        let function = move |s: Series| {
+            let ca = s.utf8()?;
+
+            let mut arrs = (0..n)
+                .map(|_| MutableUtf8Array::<i64>::with_capacity(ca.len()))
+                .collect::<Vec<_>>();
+
+            ca.into_iter().for_each(|opt_s| match opt_s {
+                None => {
+                    for arr in &mut arrs {
+                        arr.push_null()
+                    }
+                }
+                Some(s) => {
+                    let mut arr_iter = arrs.iter_mut();
+                    let split_iter = s.split(&by);
+                    (split_iter)
+                        .zip(&mut arr_iter)
+                        .for_each(|(splitted, arr)| arr.push(Some(splitted)));
+                    // fill the remaining with null
+                    for arr in arr_iter {
+                        arr.push_null()
+                    }
+                }
+            });
+            let fields = arrs
+                .into_iter()
+                .enumerate()
+                .map(|(i, arr)| {
+                    Series::try_from((format!("field_{i}").as_str(), arr.into_arc())).unwrap()
+                })
+                .collect::<Vec<_>>();
+            Ok(StructChunked::new(ca.name(), &fields)?.into_series())
+        };
+        self.0
+            .map(
+                function,
+                GetOutput::from_type(DataType::Struct(
+                    (0..n)
+                        .map(|i| Field::new(&format!("field_{i}"), DataType::Utf8))
+                        .collect(),
+                )),
+            )
+            .with_fmt("str.split_exact")
+    }
+
+    pub fn split_exact_inclusive(self, by: &str, n: usize) -> Expr {
+        let by = by.to_string();
+
+        let function = move |s: Series| {
+            let ca = s.utf8()?;
+
+            let mut arrs = (0..n)
+                .map(|_| MutableUtf8Array::<i64>::with_capacity(ca.len()))
+                .collect::<Vec<_>>();
+
+            ca.into_iter().for_each(|opt_s| match opt_s {
+                None => {
+                    for arr in &mut arrs {
+                        arr.push_null()
+                    }
+                }
+                Some(s) => {
+                    let mut arr_iter = arrs.iter_mut();
+                    let split_iter = s.split_inclusive(&by);
+                    (split_iter)
+                        .zip(&mut arr_iter)
+                        .for_each(|(splitted, arr)| arr.push(Some(splitted)));
+                    // fill the remaining with null
+                    for arr in arr_iter {
+                        arr.push_null()
+                    }
+                }
+            });
+            let fields = arrs
+                .into_iter()
+                .enumerate()
+                .map(|(i, arr)| {
+                    Series::try_from((format!("field_{i}").as_str(), arr.into_arc())).unwrap()
+                })
+                .collect::<Vec<_>>();
+            Ok(StructChunked::new(ca.name(), &fields)?.into_series())
+        };
+        self.0
+            .map(
+                function,
+                GetOutput::from_type(DataType::Struct(
+                    (0..n)
+                        .map(|i| Field::new(&format!("field_{i}"), DataType::Utf8))
+                        .collect(),
+                )),
+            )
+            .with_fmt("str.split_exact")
+    }
+
     /// Split the string by a substring.
     pub fn split_inclusive(self, by: &str) -> Expr {
         let by = by.to_string();

diff --git a/py-polars/docs/source/reference/expression.rst b/py-polars/docs/source/reference/expression.rst
@@ -296,6 +296,7 @@ The following methods are available under the `Expr.str` attribute.
     ExprStringNameSpace.json_path_match
     ExprStringNameSpace.extract
     ExprStringNameSpace.split
+    ExprStringNameSpace.split_exact
     ExprStringNameSpace.replace
     ExprStringNameSpace.replace_all
     ExprStringNameSpace.slice

diff --git a/py-polars/docs/source/reference/series.rst b/py-polars/docs/source/reference/series.rst
@@ -252,6 +252,7 @@ The following methods are available under the `Series.str` attribute.
     StringNameSpace.json_path_match
     StringNameSpace.extract
     StringNameSpace.split
+    StringNameSpace.split_exact
     StringNameSpace.replace
     StringNameSpace.replace_all
     StringNameSpace.to_lowercase

diff --git a/py-polars/polars/internals/expr.py b/py-polars/polars/internals/expr.py
@@ -3587,6 +3587,52 @@ def split(self, by: str, inclusive: bool = False) -> Expr:
             return wrap_expr(self._pyexpr.str_split_inclusive(by))
         return wrap_expr(self._pyexpr.str_split(by))
 
+    def split_exact(self, by: str, n: int, inclusive: bool = False) -> Expr:
+        """
+        Split the string by a substring into a struct of `n` fields.
+        The return type will by of type Struct<Utf8>
+
+        If it cannot make `n` splits, the remaiming field elements will be null
+
+        Parameters
+        ----------
+        by
+            substring
+        n
+            Number of splits to make
+        inclusive
+            Include the split character/string in the results
+
+        Examples
+        --------
+
+        >>> (
+        ...     pl.DataFrame({"x": ["a_1", None, "c", "d_4"]}).select(
+        ...         [
+        ...             pl.col("x").str.split_exact("_", 2).alias("fields"),
+        ...         ]
+        ...     )
+        ... )
+        shape: (4, 1)
+        ┌───────────────────────────────────────────┐
+        │ fields                                    │
+        │ ---                                       │
+        │ struct[2]{'field_0': str, 'field_1': str} │
+        ╞═══════════════════════════════════════════╡
+        │ {"a","1"}                                 │
+        ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+        │ {null,null}                               │
+        ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+        │ {"c",null}                                │
+        ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+        │ {"d","4"}                                 │
+        └───────────────────────────────────────────┘
+
+        """
+        if inclusive:
+            return wrap_expr(self._pyexpr.str_split_exact_inclusive(by, n))
+        return wrap_expr(self._pyexpr.str_split_exact(by, n))
+
     def replace(self, pattern: str, value: str) -> Expr:
         """
         Replace first regex match with a string value.

diff --git a/py-polars/polars/internals/series.py b/py-polars/polars/internals/series.py
@@ -3868,6 +3868,55 @@ def split(self, by: str, inclusive: bool = False) -> Series:
         s = wrap_s(self._s)
         return s.to_frame().select(pli.col(s.name).str.split(by, inclusive)).to_series()
 
+    def split_exact(self, by: str, n: int, inclusive: bool = False) -> Series:
+        """
+        Split the string by a substring into a struct of `n` fields.
+        The return type will by of type Struct<Utf8>
+
+        If it cannot make `n` splits, the remaiming field elements will be null
+
+        Parameters
+        ----------
+        by
+            substring
+        n
+            Number of splits to make
+        inclusive
+            Include the split character/string in the results
+
+        Examples
+        --------
+
+        >>> (
+        ...     pl.DataFrame({"x": ["a_1", None, "c", "d_4"]}).select(
+        ...         [
+        ...             pl.col("x").str.split_exact("_", 2).alias("fields"),
+        ...         ]
+        ...     )
+        ... )
+        shape: (4, 1)
+        ┌───────────────────────────────────────────┐
+        │ fields                                    │
+        │ ---                                       │
+        │ struct[2]{'field_0': str, 'field_1': str} │
+        ╞═══════════════════════════════════════════╡
+        │ {"a","1"}                                 │
+        ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+        │ {null,null}                               │
+        ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+        │ {"c",null}                                │
+        ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+        │ {"d","4"}                                 │
+        └───────────────────────────────────────────┘
+
+        """
+        s = wrap_s(self._s)
+        return (
+            s.to_frame()
+            .select(pli.col(s.name).str.split_exact(by, n, inclusive))
+            .to_series()
+        )
+
     def replace(self, pattern: str, value: str) -> Series:
         """
         Replace first regex match with a string value.

diff --git a/py-polars/src/lazy/dsl.rs b/py-polars/src/lazy/dsl.rs
@@ -592,6 +592,13 @@ impl PyExpr {
         self.inner.clone().str().split_inclusive(by).into()
     }
 
+    pub fn str_split_exact(&self, by: &str, n: usize) -> PyExpr {
+        self.inner.clone().str().split_exact(by, n).into()
+    }
+    pub fn str_split_exact_inclusive(&self, by: &str, n: usize) -> PyExpr {
+        self.inner.clone().str().split_exact_inclusive(by, n).into()
+    }
+
     pub fn arr_lengths(&self) -> PyExpr {
         self.inner.clone().arr().lengths().into()
     }

diff --git a/py-polars/tests/test_exprs.py b/py-polars/tests/test_exprs.py
@@ -122,3 +122,27 @@ def test_wildcard_expansion() -> None:
     assert df.select(
         pl.concat_str(pl.all()).str.to_lowercase()
     ).to_series().to_list() == ["xs", "yo", "zs"]
+
+
+def test_split_exact() -> None:
+    df = pl.DataFrame(dict(x=["a_a", None, "b", "c_c"]))
+    out = df.select([pl.col("x").str.split_exact("_", 3, inclusive=False)]).unnest("x")
+
+    expected = pl.DataFrame(
+        {
+            "field_0": ["a", None, "b", "c"],
+            "field_1": ["a", None, None, "c"],
+            "field_2": [None, None, None, None],
+        }
+    )
+
+    assert out.frame_equal(expected)
+
+    out = df.select([pl.col("x").str.split_exact("_", 2, inclusive=True)]).unnest("x")
+
+    expected = pl.DataFrame(
+        {"field_0": ["a_", None, "b", "c_"], "field_1": ["a", None, None, "c"]}
+    )
+    assert out.frame_equal(expected)
+    assert df["x"].str.split_exact("_", 1).dtype == pl.Struct
+    assert df["x"].str.split_exact("_", 1, inclusive=False).dtype == pl.Struct