Repository: narwhals-dev/narwhals (python)

Issue:
[Enh]: str zfill
### We would like to learn about your use case. For example, if this feature is needed  to adopt Narwhals in an open source project, could you please enter the link to it below?


I would like to use narwhals for the latent-calendar project. However, it uses zfill

https://github.com/williambdean/latent-calendar/blob/07ffd832188098af21a469470bb988560d7b6e80/latent_calendar/transformers.py#L155

### Please describe the purpose of the new feature or describe the problem to solve.

pad a string

### Suggest a solution if possible.

use the pandas, polars, spark, etc functionality

### If you have tried alternatives, please describe them below.

looked through the docs

### Additional information that may help us understand your needs.

_No response_

Pull request:
feat: Add support for `Expr.str.zfill`
<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

This adds `zfill` for both Series and Expr

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues

- Related issue #\<issue number\>
- Closes #2596

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added
- [x] Documented the changes

## If you have comments or can explain your changes, please do so below


Patch:
diff --git a/docs/api-reference/expr_str.md b/docs/api-reference/expr_str.md
index bb6abb7e..9fd31704 100644
--- a/docs/api-reference/expr_str.md
+++ b/docs/api-reference/expr_str.md
@@ -18,5 +18,6 @@
         - to_datetime
         - to_lowercase
         - to_uppercase
+        - zfill
       show_source: false
       show_bases: false
diff --git a/docs/api-reference/series_str.md b/docs/api-reference/series_str.md
index 42810161..d2d70e69 100644
--- a/docs/api-reference/series_str.md
+++ b/docs/api-reference/series_str.md
@@ -18,5 +18,6 @@
         - to_datetime
         - to_lowercase
         - to_uppercase
+        - zfill
       show_source: false
       show_bases: false
diff --git a/narwhals/_arrow/series_str.py b/narwhals/_arrow/series_str.py
index 64dcce8c..3731e078 100644
--- a/narwhals/_arrow/series_str.py
+++ b/narwhals/_arrow/series_str.py
@@ -3,12 +3,14 @@ from __future__ import annotations
 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 @@ class ArrowSeriesStringNamespace(ArrowSeriesNamespace):
 
     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)
diff --git a/narwhals/_compliant/any_namespace.py b/narwhals/_compliant/any_namespace.py
index 3365d25b..93fc5323 100644
--- a/narwhals/_compliant/any_namespace.py
+++ b/narwhals/_compliant/any_namespace.py
@@ -79,6 +79,7 @@ class StringNamespace(_StoresCompliant[CompliantT_co], Protocol[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]):
diff --git a/narwhals/_compliant/expr.py b/narwhals/_compliant/expr.py
index 1ea3fcd8..14c70415 100644
--- a/narwhals/_compliant/expr.py
+++ b/narwhals/_compliant/expr.py
@@ -1127,6 +1127,9 @@ class EagerExprStringNamespace(
     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]
diff --git a/narwhals/_dask/expr_str.py b/narwhals/_dask/expr_str.py
index b770b53d..938b0844 100644
--- a/narwhals/_dask/expr_str.py
+++ b/narwhals/_dask/expr_str.py
@@ -96,3 +96,8 @@ class DaskExprStringNamespace:
         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
+        )
diff --git a/narwhals/_duckdb/expr_str.py b/narwhals/_duckdb/expr_str.py
index 6a9d8c2e..e3969eda 100644
--- a/narwhals/_duckdb/expr_str.py
+++ b/narwhals/_duckdb/expr_str.py
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
 
 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 @@ class DuckDBExprStringNamespace:
             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()
diff --git a/narwhals/_ibis/expr_str.py b/narwhals/_ibis/expr_str.py
index 1c0d6e59..6602995e 100644
--- a/narwhals/_ibis/expr_str.py
+++ b/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 @@ class IbisExprStringNamespace:
         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()
diff --git a/narwhals/_pandas_like/series_str.py b/narwhals/_pandas_like/series_str.py
index c4bef09e..8a8d6081 100644
--- a/narwhals/_pandas_like/series_str.py
+++ b/narwhals/_pandas_like/series_str.py
@@ -77,3 +77,6 @@ class PandasLikeSeriesStringNamespace(
 
     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))
diff --git a/narwhals/_polars/expr.py b/narwhals/_polars/expr.py
index 426a01dc..261e2fff 100644
--- a/narwhals/_polars/expr.py
+++ b/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)
diff --git a/narwhals/_polars/series.py b/narwhals/_polars/series.py
index 3c9cd8ba..adf1ab72 100644
--- a/narwhals/_polars/series.py
+++ b/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)
diff --git a/narwhals/_spark_like/expr_str.py b/narwhals/_spark_like/expr_str.py
index 7c659525..4d64d653 100644
--- a/narwhals/_spark_like/expr_str.py
+++ b/narwhals/_spark_like/expr_str.py
@@ -113,3 +113,30 @@ class SparkLikeExprStringNamespace:
         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)
diff --git a/narwhals/expr_str.py b/narwhals/expr_str.py
index e598ff78..d6b5115e 100644
--- a/narwhals/expr_str.py
+++ b/narwhals/expr_str.py
@@ -447,3 +447,34 @@ class ExprStringNamespace(Generic[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)
+        )
diff --git a/narwhals/series_str.py b/narwhals/series_str.py
index 82d50f5a..e2b149af 100644
--- a/narwhals/series_str.py
+++ b/narwhals/series_str.py
@@ -398,3 +398,28 @@ class SeriesStringNamespace(Generic[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)
+        )


Test patch:
diff --git a/tests/expr_and_series/str/zfill_test.py b/tests/expr_and_series/str/zfill_test.py
new file mode 100644
index 00000000..8e4eccf2
--- /dev/null
+++ b/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)