feat: str.split (#1932) (#2054)

skritsotalakis · Stelios Kritsotalakis · dangotbanned · web-flow · commit 72dea2a54cfb · 2025-02-27T22:07:47.000Z
---------

Co-authored-by: Stelios Kritsotalakis &lt;kstelios@DESKTOP-D65QO0G&gt;
Co-authored-by: Dan Redding &lt;125183946+dangotbanned@users.noreply.github.com&gt;
diff --git a/docs/api-reference/expr_str.md b/docs/api-reference/expr_str.md
@@ -11,6 +11,7 @@
         - replace
         - replace_all
         - slice
+        - split
         - starts_with
         - strip_chars
         - tail
diff --git a/docs/api-reference/series_str.md b/docs/api-reference/series_str.md
@@ -11,6 +11,7 @@
         - replace
         - replace_all
         - slice
+        - split
         - starts_with
         - strip_chars
         - tail
diff --git a/narwhals/_arrow/expr_str.py b/narwhals/_arrow/expr_str.py
@@ -67,6 +67,11 @@ def slice(self: Self, offset: int, length: int | None) -> ArrowExpr:
             self._compliant_expr, "str", "slice", offset=offset, length=length
         )
 
+    def split(self: Self, by: str) -> ArrowExpr:
+        return reuse_series_namespace_implementation(
+            self._compliant_expr, "str", "split", by=by
+        )
+
     def to_datetime(self: Self, format: str | None) -> ArrowExpr:  # noqa: A002
         return reuse_series_namespace_implementation(
             self._compliant_expr, "str", "to_datetime", format=format
diff --git a/narwhals/_arrow/series_str.py b/narwhals/_arrow/series_str.py
@@ -72,6 +72,10 @@ def slice(self: Self, offset: int, length: int | None) -> ArrowSeries:
             )
         )
 
+    def split(self: Self, by: str) -> ArrowSeries:
+        split_series = pc.split_pattern(self._compliant_series._native_series, by)  # type: ignore[call-overload]
+        return self._compliant_series._from_native_series(split_series)
+
     def to_datetime(self: Self, format: str | None) -> ArrowSeries:  # noqa: A002
         native = self._compliant_series._native_series
         format = parse_datetime_format(native) if format is None else format
diff --git a/narwhals/_dask/expr_str.py b/narwhals/_dask/expr_str.py
@@ -81,6 +81,13 @@ def slice(self: Self, offset: int, length: int | None) -> DaskExpr:
             length=length,
         )
 
+    def split(self: Self, by: str) -> DaskExpr:
+        return self._compliant_expr._from_call(
+            lambda _input, by: _input.str.split(pat=by),
+            "split",
+            by=by,
+        )
+
     def to_datetime(self: Self, format: str | None) -> DaskExpr:  # noqa: A002
         return self._compliant_expr._from_call(
             lambda _input, format: dd.to_datetime(_input, format=format),  # noqa: A006
diff --git a/narwhals/_duckdb/expr_str.py b/narwhals/_duckdb/expr_str.py
@@ -54,6 +54,12 @@ def func(_input: duckdb.Expression) -> duckdb.Expression:
 
         return self._compliant_expr._from_call(func, "slice")
 
+    def split(self: Self, by: str) -> DuckDBExpr:
+        return self._compliant_expr._from_call(
+            lambda _input: FunctionExpression("str_split", _input, lit(by)),
+            "split",
+        )
+
     def len_chars(self: Self) -> DuckDBExpr:
         return self._compliant_expr._from_call(
             lambda _input: FunctionExpression("length", _input), "len_chars"
diff --git a/narwhals/_pandas_like/expr_str.py b/narwhals/_pandas_like/expr_str.py
@@ -93,6 +93,11 @@ def slice(self: Self, offset: int, length: int | None) -> PandasLikeExpr:
             self._compliant_expr, "str", "slice", offset=offset, length=length
         )
 
+    def split(self: Self, by: str) -> PandasLikeExpr:
+        return reuse_series_namespace_implementation(
+            self._compliant_expr, "str", "split", by=by
+        )
+
     def to_datetime(self: Self, format: str | None) -> PandasLikeExpr:  # noqa: A002
         return reuse_series_namespace_implementation(
             self._compliant_expr,
diff --git a/narwhals/_pandas_like/series_str.py b/narwhals/_pandas_like/series_str.py
@@ -2,7 +2,9 @@
 
 from typing import TYPE_CHECKING
 
+from narwhals._pandas_like.utils import get_dtype_backend
 from narwhals._pandas_like.utils import to_datetime
+from narwhals.utils import Implementation
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -61,6 +63,27 @@ def slice(self: Self, offset: int, length: int | None) -> PandasLikeSeries:
             self._compliant_series._native_series.str.slice(start=offset, stop=stop),
         )
 
+    def split(self: Self, by: str) -> PandasLikeSeries:
+        if (
+            self._compliant_series._implementation is not Implementation.CUDF
+        ):  # pragma: no cover
+            dtype_backend = get_dtype_backend(
+                self._compliant_series._native_series.dtype,
+                self._compliant_series._implementation,
+            )
+            if dtype_backend != "pyarrow":
+                msg = (
+                    "This operation requires a pyarrow-backed series. "
+                    "Please refer to https://narwhals-dev.github.io/narwhals/api-reference/narwhals/#narwhals.maybe_convert_dtypes "
+                    "and ensure you are using dtype_backend='pyarrow'. "
+                    "Additionally, make sure you have pandas version 1.5+ and pyarrow installed. "
+                )
+                raise TypeError(msg)
+
+        return self._compliant_series._from_native_series(
+            self._compliant_series._native_series.str.split(pat=by),
+        )
+
     def to_datetime(self: Self, format: str | None) -> PandasLikeSeries:  # noqa: A002
         return self._compliant_series._from_native_series(
             to_datetime(self._compliant_series._implementation)(
diff --git a/narwhals/_spark_like/expr_str.py b/narwhals/_spark_like/expr_str.py
@@ -90,6 +90,12 @@ def func(_input: Column) -> Column:
 
         return self._compliant_expr._from_call(func, "slice")
 
+    def split(self: Self, by: str) -> SparkLikeExpr:
+        return self._compliant_expr._from_call(
+            lambda _input: self._compliant_expr._F.split(_input, by),
+            "split",
+        )
+
     def to_uppercase(self: Self) -> SparkLikeExpr:
         return self._compliant_expr._from_call(
             self._compliant_expr._F.upper, "to_uppercase"
diff --git a/narwhals/expr_str.py b/narwhals/expr_str.py
@@ -271,6 +271,40 @@ def slice(self: Self, offset: int, length: int | None = None) -> ExprT:
             self._expr._metadata,
         )
 
+    def split(self: Self, by: str) -> ExprT:
+        r"""Split the string values of an expression by a substring.
+
+        Arguments:
+            by: Substring to split by.
+
+        Returns:
+            A new expression.
+
+        Examples:
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> df_native = pl.DataFrame({"s": ["foo bar", "foo_bar"]})
+            >>> df = nw.from_native(df_native)
+            >>> df.with_columns(nw.col("s").str.split("_").alias("s_split"))
+            ┌────────────────────────────┐
+            |     Narwhals DataFrame     |
+            |----------------------------|
+            |shape: (2, 2)               |
+            |┌─────────┬────────────────┐|
+            |│ s       ┆ s_split        │|
+            |│ ---     ┆ ---            │|
+            |│ str     ┆ list[str]      │|
+            |╞═════════╪════════════════╡|
+            |│ foo bar ┆ ["foo bar"]    │|
+            |│ foo_bar ┆ ["foo", "bar"] │|
+            |└─────────┴────────────────┘|
+            └────────────────────────────┘
+        """
+        return self._expr.__class__(
+            lambda plx: self._expr._to_compliant_expr(plx).str.split(by=by),
+            self._expr._metadata,
+        )
+
     def head(self: Self, n: int = 5) -> ExprT:
         r"""Take the first n elements of each string.
 
diff --git a/narwhals/series_str.py b/narwhals/series_str.py
@@ -236,6 +236,32 @@ def slice(self: Self, offset: int, length: int | None = None) -> SeriesT:
             )
         )
 
+    def split(self: Self, by: str) -> SeriesT:
+        r"""Split the string values of a Series by a substring.
+
+        Arguments:
+            by: Substring to split by.
+
+        Returns:
+            A new Series containing lists of strings.
+
+        Examples:
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> s_native = pl.Series(["foo bar", "foo_bar"])
+            >>> s = nw.from_native(s_native, series_only=True)
+            >>> s.str.split("_").to_native()  # doctest: +NORMALIZE_WHITESPACE
+            shape: (2,)
+            Series: '' [list[str]]
+            [
+                    ["foo bar"]
+                    ["foo", "bar"]
+            ]
+        """
+        return self._narwhals_series._from_compliant_series(
+            self._narwhals_series._compliant_series.str.split(by=by)
+        )
+
     def head(self: Self, n: int = 5) -> SeriesT:
         r"""Take the first n elements of each string.
 
diff --git a/tests/expr_and_series/str/split_test.py b/tests/expr_and_series/str/split_test.py
@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+import re
+from typing import Any
+
+import pytest
+
+import narwhals.stable.v1 as nw
+from tests.utils import Constructor
+from tests.utils import ConstructorEager
+from tests.utils import assert_equal_data
+
+data = {"s": ["foo bar", "foo_bar", "foo_bar_baz", "foo,bar"]}
+
+
+@pytest.mark.parametrize(
+    ("by", "expected"),
+    [
+        ("_", {"s": [["foo bar"], ["foo", "bar"], ["foo", "bar", "baz"], ["foo,bar"]]}),
+        (
+            ",",
+            {"s": [["foo bar"], ["foo_bar"], ["foo_bar_baz"], ["foo", "bar"]]},
+        ),
+    ],
+)
+def test_str_split(
+    constructor: Constructor,
+    by: str,
+    expected: Any,
+) -> None:
+    if "cudf" not in str(constructor) and (
+        constructor.__name__.startswith("pandas")
+        and "pyarrow" not in constructor.__name__
+    ):
+        df = nw.from_native(constructor(data))
+        msg = re.escape("This operation requires a pyarrow-backed series. ")
+        with pytest.raises(TypeError, match=msg):
+            df.select(nw.col("s").str.split(by=by))
+        return
+    df = nw.from_native(constructor(data))
+    result_frame = df.select(nw.col("s").str.split(by=by))
+    assert_equal_data(result_frame, expected)
+
+
+@pytest.mark.parametrize(
+    ("by", "expected"),
+    [
+        ("_", {"s": [["foo bar"], ["foo", "bar"], ["foo", "bar", "baz"], ["foo,bar"]]}),
+        (
+            ",",
+            {"s": [["foo bar"], ["foo_bar"], ["foo_bar_baz"], ["foo", "bar"]]},
+        ),
+    ],
+)
+def test_str_split_series(
+    constructor_eager: ConstructorEager,
+    by: str,
+    expected: Any,
+) -> None:
+    if "cudf" not in str(constructor_eager) and (
+        constructor_eager.__name__.startswith("pandas")
+        and "pyarrow" not in constructor_eager.__name__
+    ):
+        df = nw.from_native(constructor_eager(data), eager_only=True)
+        msg = re.escape("This operation requires a pyarrow-backed series. ")
+        with pytest.raises(TypeError, match=msg):
+            df["s"].str.split(by=by)
+        return
+    df = nw.from_native(constructor_eager(data), eager_only=True)
+    result_series = df["s"].str.split(by=by)
+    assert_equal_data({"s": result_series}, expected)
diff --git a/tests/utils.py b/tests/utils.py
@@ -131,6 +131,10 @@ def assert_equal_data(result: Any, expected: Mapping[str, Any]) -> None:
                 are_equivalent_values = lhs is None or math.isnan(lhs)
             elif lhs is None:
                 are_equivalent_values = rhs is None
+            elif isinstance(lhs, list) and isinstance(rhs, list):
+                are_equivalent_values = all(
+                    left_side == right_side for left_side, right_side in zip(lhs, rhs)
+                )
             elif pd.isna(lhs):
                 are_equivalent_values = pd.isna(rhs)
             else:

Original file line number	Diff line number	Diff line change
`@@ -72,6 +72,10 @@ def slice(self: Self, offset: int, length: int \| None) -> ArrowSeries:`
`72`	`72`	`)`
`73`	`73`	`)`
`74`	`74`
	`75`	`+ def split(self: Self, by: str) -> ArrowSeries:`
	`76`	`+ split_series = pc.split_pattern(self._compliant_series._native_series, by) # type: ignore[call-overload]`
	`77`	`+ return self._compliant_series._from_native_series(split_series)`
	`78`	`+`
`75`	`79`	`def to_datetime(self: Self, format: str \| None) -> ArrowSeries: # noqa: A002`
`76`	`80`	`native = self._compliant_series._native_series`
`77`	`81`	`format = parse_datetime_format(native) if format is None else format`