feat: Add support for Series|Expr.str.zfill (#2598)

* get pandas zfill to work

* add to the api-reference

* add tests for the zfill

* implement for other backends

* add the series tests

* add additional test cases

* add additional test cases

* implement for duckdb

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* support polars like zfill

* make test cases act like polars

* implement for ibis

* implement for spark_like

* add attempt at arrow

* add skip reason

Co-authored-by: Francesco Bruzzesi <[email protected]>

* consolidate on inner function name

* add to docstring

* feat: Get `pyarrow` working

Man they really make the simple things complex 😅]

* fix(typing): `ibis` casts

* fix(typing): `pyspark` ignore

* add dask implementation

* add exceptions and skips

* adapt pandas / stdlib zfill behavior for all

* correct the example in the docstring

* skip based on pandas version

* change the condition

* add expr_str example as well

* fixing the import

* implement for older polars versions

* add skipif for second test :tear:

* fix the doctests

* refactor logic to look like _spark_like

* pyarrow variable naming, typing and some misc

* pyarrow use repeat

* remove warnings

* chore(typing): Add missing `zfill`

Would've been caught by #2500

* Update narwhals/_ibis/expr_str.py

Co-authored-by: Dan Redding <[email protected]>

* Update narwhals/_ibis/expr_str.py

Co-authored-by: Dan Redding <[email protected]>

* Update narwhals/_arrow/series_str.py

Co-authored-by: Dan Redding <[email protected]>

* perf: Reuse `PolarsExpr` impl in `PolarsSeries`

Resolves #2598 (comment)

* use xfail instead

* rename helper function

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Francesco Bruzzesi <[email protected]>
Co-authored-by: dangotbanned <[email protected]>
Co-authored-by: FBruzzesi <[email protected]>

Author: 4 people

docs/api-reference/expr_str.md

@@ -18,5 +18,6 @@
         - to_datetime
         - to_lowercase
         - to_uppercase
+        - zfill
       show_source: false
       show_bases: false

docs/api-reference/series_str.md

@@ -18,5 +18,6 @@
         - to_datetime
         - to_lowercase
         - to_uppercase
+        - zfill
       show_source: false
       show_bases: false

narwhals/_arrow/series_str.py

@@ -3,12 +3,14 @@
 import string
 from typing import TYPE_CHECKING
 
+import pyarrow as pa
 import pyarrow.compute as pc
 
 from narwhals._arrow.utils import ArrowSeriesNamespace, lit, parse_datetime_format
 
 if TYPE_CHECKING:
     from narwhals._arrow.series import ArrowSeries
+    from narwhals._arrow.typing import Incomplete
 
 
 class ArrowSeriesStringNamespace(ArrowSeriesNamespace):
@@ -60,3 +62,36 @@ def to_uppercase(self) -> ArrowSeries:
 
     def to_lowercase(self) -> ArrowSeries:
         return self.with_native(pc.utf8_lower(self.native))
+
+    def zfill(self, width: int) -> ArrowSeries:
+        binary_join: Incomplete = pc.binary_join_element_wise
+        native = self.native
+        hyphen, plus = lit("-"), lit("+")
+        first_char, remaining_chars = self.slice(0, 1).native, self.slice(1, None).native
+
+        # Conditions
+        less_than_width = pc.less(pc.utf8_length(native), lit(width))
+        starts_with_hyphen = pc.equal(first_char, hyphen)
+        starts_with_plus = pc.equal(first_char, plus)
+
+        conditions = pc.make_struct(
+            pc.and_(starts_with_hyphen, less_than_width),
+            pc.and_(starts_with_plus, less_than_width),
+            less_than_width,
+        )
+
+        # Cases
+        padded_remaining_chars = pc.utf8_lpad(remaining_chars, width - 1, padding="0")
+
+        result = pc.case_when(
+            conditions,
+            binary_join(
+                pa.repeat(hyphen, len(native)), padded_remaining_chars, ""
+            ),  # starts with hyphen and less than width
+            binary_join(
+                pa.repeat(plus, len(native)), padded_remaining_chars, ""
+            ),  # starts with plus and less than width
+            pc.utf8_lpad(native, width=width, padding="0"),  # less than width
+            native,
+        )
+        return self.with_native(result)

narwhals/_compliant/any_namespace.py

@@ -79,6 +79,7 @@ def split(self, by: str) -> CompliantT_co: ...
     def to_datetime(self, format: str | None) -> CompliantT_co: ...
     def to_lowercase(self) -> CompliantT_co: ...
     def to_uppercase(self) -> CompliantT_co: ...
+    def zfill(self, width: int) -> CompliantT_co: ...
 
 
 class StructNamespace(_StoresCompliant[CompliantT_co], Protocol[CompliantT_co]):

narwhals/_compliant/expr.py

@@ -1127,6 +1127,9 @@ def to_lowercase(self) -> EagerExprT:
     def to_uppercase(self) -> EagerExprT:
         return self.compliant._reuse_series_namespace("str", "to_uppercase")
 
+    def zfill(self, width: int) -> EagerExprT:
+        return self.compliant._reuse_series_namespace("str", "zfill", width=width)
+
 
 class EagerExprStructNamespace(
     EagerExprNamespace[EagerExprT], StructNamespace[EagerExprT], Generic[EagerExprT]

narwhals/_dask/expr_str.py

@@ -96,3 +96,8 @@ def to_lowercase(self) -> DaskExpr:
         return self._compliant_expr._with_callable(
             lambda expr: expr.str.lower(), "to_lowercase"
         )
+
+    def zfill(self, width: int) -> DaskExpr:
+        return self._compliant_expr._with_callable(
+            lambda expr, width: expr.str.zfill(width), "zfill", width=width
+        )

narwhals/_duckdb/expr_str.py

@@ -4,7 +4,7 @@
 
 from duckdb import FunctionExpression
 
-from narwhals._duckdb.utils import lit
+from narwhals._duckdb.utils import lit, when
 from narwhals._utils import not_implemented
 
 if TYPE_CHECKING:
@@ -100,4 +100,31 @@ def to_datetime(self, format: str | None) -> DuckDBExpr:
             lambda expr: FunctionExpression("strptime", expr, lit(format))
         )
 
+    def zfill(self, width: int) -> DuckDBExpr:
+        # DuckDB does not have a built-in zfill function, so we need to implement it manually
+        # using string manipulation functions.
+
+        def func(expr: Expression) -> Expression:
+            less_than_width = FunctionExpression("length", expr) < lit(width)
+            zero, hyphen, plus = lit("0"), lit("-"), lit("+")
+
+            starts_with_minus = FunctionExpression("starts_with", expr, hyphen)
+            starts_with_plus = FunctionExpression("starts_with", expr, plus)
+            substring = FunctionExpression("substr", expr, lit(2))
+            padded_substring = FunctionExpression("lpad", substring, lit(width - 1), zero)
+            return (
+                when(
+                    starts_with_minus & less_than_width,
+                    FunctionExpression("concat", hyphen, padded_substring),
+                )
+                .when(
+                    starts_with_plus & less_than_width,
+                    FunctionExpression("concat", plus, padded_substring),
+                )
+                .when(less_than_width, FunctionExpression("lpad", expr, lit(width), zero))
+                .otherwise(expr)
+            )
+
+        return self._compliant_expr._with_callable(func)
+
     replace = not_implemented()

narwhals/_ibis/expr_str.py

@@ -1,14 +1,15 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any, Callable, cast
 
+import ibis
+import ibis.expr.types as ir
 from ibis.expr.datatypes import Timestamp
 
+from narwhals._ibis.utils import lit
 from narwhals._utils import _is_naive_format, not_implemented
 
 if TYPE_CHECKING:
-    import ibis.expr.types as ir
-
     from narwhals._ibis.expr import IbisExpr
 
 
@@ -100,4 +101,23 @@ def to_datetime(self, format: str | None) -> IbisExpr:
         fn = self._to_datetime_naive if _is_naive_format(format) else self._to_datetime
         return self._compliant_expr._with_callable(fn(format))
 
+    def zfill(self, width: int) -> IbisExpr:
+        def func(expr: ir.StringColumn) -> ir.Value:
+            length = expr.length()
+            less_than_width = length < lit(width)
+            zero, hyphen, plus = "0", "-", "+"
+            starts_with_minus = expr.startswith(hyphen)
+            starts_with_plus = expr.startswith(plus)
+            one = cast("ir.IntegerScalar", lit(1))
+            sub_length = cast("ir.IntegerValue", length - one)
+            substring = expr.substr(one, sub_length).lpad(width - 1, zero)
+            return ibis.cases(
+                (starts_with_minus & less_than_width, (substring.lpad(width, hyphen))),
+                (starts_with_plus & less_than_width, (substring.lpad(width, plus))),
+                (less_than_width, expr.lpad(width, zero)),
+                else_=expr,
+            )
+
+        return self._compliant_expr._with_callable(func)
+
     replace = not_implemented()

narwhals/_pandas_like/series_str.py

@@ -77,3 +77,6 @@ def to_uppercase(self) -> PandasLikeSeries:
 
     def to_lowercase(self) -> PandasLikeSeries:
         return self.with_native(self.native.str.lower())
+
+    def zfill(self, width: int) -> PandasLikeSeries:
+        return self.with_native(self.native.str.zfill(width))

narwhals/_polars/expr.py

@@ -339,6 +339,27 @@ class PolarsExprStringNamespace:
     def __init__(self, expr: PolarsExpr) -> None:
         self._compliant_expr = expr
 
+    def zfill(self, width: int) -> PolarsExpr:
+        native_expr = self._compliant_expr.native
+        native_result = native_expr.str.zfill(width)
+
+        if self._compliant_expr._backend_version <= (1, 30, 0):
+            length = native_expr.str.len_chars()
+            less_than_width = length < width
+            plus = "+"
+            starts_with_plus = native_expr.str.starts_with(plus)
+            native_result = (
+                pl.when(starts_with_plus & less_than_width)
+                .then(
+                    native_expr.str.slice(1, length)
+                    .str.zfill(width - 1)
+                    .str.pad_start(width, plus)
+                )
+                .otherwise(native_result)
+            )
+
+        return self._compliant_expr._with_native(native_result)
+
     def __getattr__(self, attr: str) -> Callable[[Any], PolarsExpr]:
         def func(*args: Any, **kwargs: Any) -> PolarsExpr:
             pos, kwds = extract_args_kwargs(args, kwargs)

narwhals/_polars/series.py

@@ -686,6 +686,12 @@ class PolarsSeriesStringNamespace:
     def __init__(self, series: PolarsSeries) -> None:
         self._compliant_series = series
 
+    def zfill(self, width: int) -> PolarsSeries:
+        series = self._compliant_series
+        name = series.name
+        ns = series.__narwhals_namespace__()
+        return series.to_frame().select(ns.col(name).str.zfill(width)).get_column(name)
+
     def __getattr__(self, attr: str) -> Any:
         def func(*args: Any, **kwargs: Any) -> Any:
             pos, kwds = extract_args_kwargs(args, kwargs)

narwhals/_spark_like/expr_str.py

@@ -113,3 +113,30 @@ def to_datetime(self, format: str | None) -> SparkLikeExpr:
         return self._compliant_expr._with_callable(
             lambda expr: function(F.replace(expr, F.lit("T"), F.lit(" ")))
         )
+
+    def zfill(self, width: int) -> SparkLikeExpr:
+        def func(expr: Column) -> Column:
+            F = self._compliant_expr._F  # noqa: N806
+
+            length = F.length(expr)
+            less_than_width = length < width
+            hyphen, plus = F.lit("-"), F.lit("+")
+            starts_with_minus = F.startswith(expr, hyphen)
+            starts_with_plus = F.startswith(expr, plus)
+            sub_length = length - F.lit(1)
+            # NOTE: `len` annotated as `int`, but `Column.substr` accepts `int | Column`
+            substring = F.substring(expr, 2, sub_length)  # pyright: ignore[reportArgumentType]
+            padded_substring = F.lpad(substring, width - 1, "0")
+            return (
+                F.when(
+                    starts_with_minus & less_than_width,
+                    F.concat(hyphen, padded_substring),
+                )
+                .when(
+                    starts_with_plus & less_than_width, F.concat(plus, padded_substring)
+                )
+                .when(less_than_width, F.lpad(expr, width, "0"))
+                .otherwise(expr)
+            )
+
+        return self._compliant_expr._with_callable(func)

narwhals/expr_str.py

@@ -447,3 +447,34 @@ def to_lowercase(self) -> ExprT:
         return self._expr._with_elementwise_op(
             lambda plx: self._expr._to_compliant_expr(plx).str.to_lowercase()
         )
+
+    def zfill(self, width: int) -> ExprT:
+        """Transform string to zero-padded variant.
+
+        Arguments:
+            width: The desired length of the string after padding. If the length of the
+                string is greater than `width`, no padding is applied.
+                If `width` is less than 0, no padding is applied.
+
+        Returns:
+            A new expression.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import narwhals as nw
+            >>> df_native = pd.DataFrame({"digits": ["+1", "-1", "1", None]})
+            >>> df = nw.from_native(df_native)
+            >>> df.with_columns(zfill_col=nw.col("digits").str.zfill(3))
+            ┌──────────────────┐
+            |Narwhals DataFrame|
+            |------------------|
+            |  digits zfill_col|
+            |0     +1       +01|
+            |1     -1       -01|
+            |2      1       001|
+            |3   None      None|
+            └──────────────────┘
+        """
+        return self._expr._with_elementwise_op(
+            lambda plx: self._expr._to_compliant_expr(plx).str.zfill(width)
+        )

narwhals/series_str.py

@@ -398,3 +398,28 @@ def to_datetime(self, format: str | None = None) -> SeriesT:
         return self._narwhals_series._with_compliant(
             self._narwhals_series._compliant_series.str.to_datetime(format=format)
         )
+
+    def zfill(self, width: int) -> SeriesT:
+        r"""Pad strings with zeros on the left.
+
+        Arguments:
+            width: The target width of the string. If the string is shorter than this width, it will be padded with zeros on the left.
+
+        Returns:
+            A new Series with strings padded with zeros on the left.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import narwhals as nw
+            >>> s_native = pd.Series(["+1", "-23", "456", "123456"])
+            >>> s = nw.from_native(s_native, series_only=True)
+            >>> s.str.zfill(5).to_native()
+            0     +0001
+            1     -0023
+            2     00456
+            3    123456
+            dtype: object
+        """
+        return self._narwhals_series._with_compliant(
+            self._narwhals_series._compliant_series.str.zfill(width)
+        )

tests/expr_and_series/str/zfill_test.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import pytest
+
+import narwhals as nw
+from tests.utils import PANDAS_VERSION, Constructor, ConstructorEager, assert_equal_data
+
+data = {"a": ["-1", "+1", "1", "12", "123", "99999", "+9999", None]}
+expected = {"a": ["-01", "+01", "001", "012", "123", "99999", "+9999", None]}
+
+
+def uses_pyarrow_backend(constructor: Constructor | ConstructorEager) -> bool:
+    return constructor.__name__ in {
+        "pandas_pyarrow_constructor",
+        "modin_pyarrow_constructor",
+    }
+
+
+@pytest.mark.skipif(PANDAS_VERSION < (1, 5), reason="different zfill behavior")
+def test_str_zfill(request: pytest.FixtureRequest, constructor: Constructor) -> None:
+    if uses_pyarrow_backend(constructor):
+        reason = (
+            "pandas with pyarrow backend doesn't support str.zfill, see "
+            "https://github.com/pandas-dev/pandas/issues/61485"
+        )
+        request.applymarker(pytest.mark.xfail(reason=reason))
+
+    df = nw.from_native(constructor(data))
+    result = df.select(nw.col("a").str.zfill(3))
+    assert_equal_data(result, expected)
+
+
+@pytest.mark.skipif(PANDAS_VERSION < (1, 5), reason="different zfill behavior")
+def test_str_zfill_series(
+    request: pytest.FixtureRequest, constructor_eager: ConstructorEager
+) -> None:
+    if uses_pyarrow_backend(constructor_eager):
+        reason = (
+            "pandas with pyarrow backend doesn't support str.zfill, see "
+            "https://github.com/pandas-dev/pandas/issues/61485"
+        )
+        request.applymarker(pytest.mark.xfail(reason=reason))
+
+    df = nw.from_native(constructor_eager(data), eager_only=True)
+    result = df["a"].str.zfill(3)
+    assert_equal_data({"a": result}, expected)