fix n arg in split_exact (#2879)

pola-rs · Mar 11, 2022 · 64be773 · 64be773
1 parent 255fb9a
commit 64be773
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 8 deletions.
diff --git a/polars/polars-lazy/src/dsl/string.rs b/polars/polars-lazy/src/dsl/string.rs
@@ -116,7 +116,7 @@ impl StringNameSpace {
         let function = move |s: Series| {
             let ca = s.utf8()?;
 
-            let mut arrs = (0..n)
+            let mut arrs = (0..n + 1)
                 .map(|_| MutableUtf8Array::<i64>::with_capacity(ca.len()))
                 .collect::<Vec<_>>();
 
@@ -151,7 +151,7 @@ impl StringNameSpace {
             .map(
                 function,
                 GetOutput::from_type(DataType::Struct(
-                    (0..n)
+                    (0..n + 1)
                         .map(|i| Field::new(&format!("field_{i}"), DataType::Utf8))
                         .collect(),
                 )),
@@ -165,7 +165,7 @@ impl StringNameSpace {
         let function = move |s: Series| {
             let ca = s.utf8()?;
 
-            let mut arrs = (0..n)
+            let mut arrs = (0..n + 1)
                 .map(|_| MutableUtf8Array::<i64>::with_capacity(ca.len()))
                 .collect::<Vec<_>>();
 
@@ -200,7 +200,7 @@ impl StringNameSpace {
             .map(
                 function,
                 GetOutput::from_type(DataType::Struct(
-                    (0..n)
+                    (0..n + 1)
                         .map(|i| Field::new(&format!("field_{i}"), DataType::Utf8))
                         .collect(),
                 )),

diff --git a/py-polars/polars/internals/expr.py b/py-polars/polars/internals/expr.py
@@ -3609,7 +3609,7 @@ def split_exact(self, by: str, n: int, inclusive: bool = False) -> Expr:
         >>> (
         ...     pl.DataFrame({"x": ["a_1", None, "c", "d_4"]}).select(
         ...         [
-        ...             pl.col("x").str.split_exact("_", 2).alias("fields"),
+        ...             pl.col("x").str.split_exact("_", 1).alias("fields"),
         ...         ]
         ...     )
         ... )

diff --git a/py-polars/polars/internals/series.py b/py-polars/polars/internals/series.py
@@ -3890,7 +3890,7 @@ def split_exact(self, by: str, n: int, inclusive: bool = False) -> Series:
         >>> (
         ...     pl.DataFrame({"x": ["a_1", None, "c", "d_4"]}).select(
         ...         [
-        ...             pl.col("x").str.split_exact("_", 2).alias("fields"),
+        ...             pl.col("x").str.split_exact("_", 1).alias("fields"),
         ...         ]
         ...     )
         ... )

diff --git a/py-polars/tests/test_exprs.py b/py-polars/tests/test_exprs.py
@@ -126,7 +126,7 @@ def test_wildcard_expansion() -> None:
 
 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")
+    out = df.select([pl.col("x").str.split_exact("_", 2, inclusive=False)]).unnest("x")
 
     expected = pl.DataFrame(
         {
@@ -138,7 +138,7 @@ def test_split_exact() -> None:
 
     assert out.frame_equal(expected)
 
-    out = df.select([pl.col("x").str.split_exact("_", 2, inclusive=True)]).unnest("x")
+    out = df.select([pl.col("x").str.split_exact("_", 1, inclusive=True)]).unnest("x")
 
     expected = pl.DataFrame(
         {"field_0": ["a_", None, "b", "c_"], "field_1": ["a", None, None, "c"]}