feat: add when-then-otherwise expression

docs/api-reference/narwhals.md

@@ -27,6 +27,7 @@ Here are the top-level functions available in Narwhals.
         - narwhalify
         - sum
         - sum_horizontal
+        - when
         - show_versions
         - to_native
       show_source: false

narwhals/__init__.py

@@ -33,6 +33,7 @@
 from narwhals.expr import min
 from narwhals.expr import sum
 from narwhals.expr import sum_horizontal
+from narwhals.expr import when
 from narwhals.functions import concat
 from narwhals.functions import from_dict
 from narwhals.functions import get_level
@@ -73,6 +74,7 @@
     "mean",
     "sum",
     "sum_horizontal",
+    "when",
     "DataFrame",
     "LazyFrame",
     "Series",

narwhals/_pandas_like/namespace.py

@@ -250,3 +250,116 @@ def concat(
                 backend_version=self._backend_version,
             )
         raise NotImplementedError
+
+    def when(
+        self,
+        *predicates: IntoPandasLikeExpr,
+    ) -> PandasWhen:
+        plx = self.__class__(self._implementation, self._backend_version)
+        if predicates:
+            condition = plx.all_horizontal(*predicates)
+        else:
+            msg = "at least one predicate needs to be provided"
+            raise TypeError(msg)
+
+        return PandasWhen(condition, self._implementation, self._backend_version)
+
+
+def _when_then_value_arg_process(
+    plx: PandasLikeNamespace,
+    value: PandasLikeExpr | PandasLikeSeries | Any,
+    *,
+    shape: tuple[int],
+) -> PandasLikeExpr:
+    from narwhals.dependencies import get_numpy
+
+    # NumPy is a required dependency of pandas
+    np = get_numpy()
+
+    if isinstance(value, PandasLikeExpr):
+        return value
+    elif isinstance(value, PandasLikeSeries):
+        return plx._create_expr_from_series(value)
+    elif (np := get_numpy()) is not None and isinstance(value, np.ndarray):
+        return plx._create_expr_from_series(plx._create_compliant_series(value))
+    else:
+        return plx._create_expr_from_series(
+            plx._create_compliant_series(np.full(shape[0], value))
+        )
+
+
+class PandasWhen:
+    def __init__(
+        self,
+        condition: PandasLikeExpr,
+        implementation: Implementation,
+        backend_version: tuple[int, ...],
+        then_value: Any = None,
+        otherise_value: Any = None,
+    ) -> None:
+        self._implementation = implementation
+        self._backend_version = backend_version
+        self._condition = condition
+        self._then_value = then_value
+        self._otherwise_value = otherise_value
+
+    def __call__(self, df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
+        from narwhals._pandas_like.namespace import PandasLikeNamespace
+
+        plx = PandasLikeNamespace(
+            implementation=self._implementation, backend_version=self._backend_version
+        )
+
+        condition = self._condition._call(df)[0]
+
+        value_series = _when_then_value_arg_process(
+            plx, self._then_value, shape=condition.shape
+        )._call(df)[0]
+        otherwise_series = _when_then_value_arg_process(
+            plx, self._otherwise_value, shape=condition.shape
+        )._call(df)[0]
+
+        return [value_series.zip_with(condition, otherwise_series)]
+
+    def then(self, value: PandasLikeExpr | PandasLikeSeries | Any) -> PandasThen:
+        self._then_value = value
+
+        return PandasThen(
+            self,
+            depth=0,
+            function_name="whenthen",
+            root_names=None,
+            output_names=None,
+            implementation=self._implementation,
+            backend_version=self._backend_version,
+        )
+
+
+class PandasThen(PandasLikeExpr):
+    def __init__(
+        self,
+        call: PandasWhen,
+        *,
+        depth: int,
+        function_name: str,
+        root_names: list[str] | None,
+        output_names: list[str] | None,
+        implementation: Implementation,
+        backend_version: tuple[int, ...],
+    ) -> None:
+        self._implementation = implementation
+        self._backend_version = backend_version
+
+        self._call = call
+        self._depth = depth
+        self._function_name = function_name
+        self._root_names = root_names
+        self._output_names = output_names
+
+    def otherwise(self, value: PandasLikeExpr | PandasLikeSeries | Any) -> PandasLikeExpr:
+        # type ignore because we are setting the `_call` attribute to a
+        # callable object of type `PandasWhen`, base class has the attribute as
+        # only a `Callable`
+        self._call._otherwise_value = value  # type: ignore[attr-defined]
+        self._function_name = "whenotherwise"
+        return self

narwhals/expr.py

@@ -3643,6 +3643,76 @@ def sum_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr:
     )
 
 
+class When:
+    def __init__(self, *predicates: IntoExpr | Iterable[IntoExpr]) -> None:
+        self._predicates = flatten([predicates])
+
+    def _extract_predicates(self, plx: Any) -> Any:
+        return [extract_compliant(plx, v) for v in self._predicates]
+
+    def then(self, value: Any) -> Then:
+        return Then(
+            lambda plx: plx.when(*self._extract_predicates(plx)).then(
+                extract_compliant(plx, value)
+            )
+        )
+
+
+class Then(Expr):
+    def __init__(self, call: Callable[[Any], Any]) -> None:
+        self._call = call
+
+    def otherwise(self, value: Any) -> Expr:
+        return Expr(lambda plx: self._call(plx).otherwise(extract_compliant(plx, value)))
+
+
+def when(*predicates: IntoExpr | Iterable[IntoExpr]) -> When:
+    """
+    Start a `when-then-otherwise` expression.
+    Expression similar to an `if-else` statement in Python. Always initiated by a `pl.when(<condition>).then(<value if condition>)`., and optionally followed by chaining one or more `.when(<condition>).then(<value>)` statements.
+    Chained when-then operations should be read as Python `if, elif, ... elif` blocks, not as `if, if, ... if`, i.e. the first condition that evaluates to `True` will be picked.
+    If none of the conditions are `True`, an optional `.otherwise(<value if all statements are false>)` can be appended at the end. If not appended, and none of the conditions are `True`, `None` will be returned.
+
+    Arguments:
+        predicates: Condition(s) that must be met in order to apply the subsequent statement. Accepts one or more boolean expressions, which are implicitly combined with `&`. String input is parsed as a column name.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import polars as pl
+        >>> import narwhals as nw
+        >>> df_pl = pl.DataFrame({"a": [1, 2, 3], "b": [5, 10, 15]})
+        >>> df_pd = pd.DataFrame({"a": [1, 2, 3], "b": [5, 10, 15]})
+
+        We define a dataframe-agnostic function:
+
+        >>> @nw.narwhalify
+        ... def func(df_any):
+        ...     return df_any.with_columns(
+        ...         nw.when(nw.col("a") < 3).then(5).otherwise(6).alias("a_when")
+        ...     )
+
+        We can then pass either pandas or polars to `func`:
+
+        >>> func(df_pd)
+           a   b  a_when
+        0  1   5       5
+        1  2  10       5
+        2  3  15       6
+        >>> func(df_pl)
+        shape: (3, 3)
+        ┌─────┬─────┬────────┐
+        │ a   ┆ b   ┆ a_when │
+        │ --- ┆ --- ┆ ---    │
+        │ i64 ┆ i64 ┆ i32    │
+        ╞═════╪═════╪════════╡
+        │ 1   ┆ 5   ┆ 5      │
+        │ 2   ┆ 10  ┆ 5      │
+        │ 3   ┆ 15  ┆ 6      │
+        └─────┴─────┴────────┘
+    """
+    return When(*predicates)
+
+
 def all_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr:
     r"""
     Compute the bitwise AND horizontally across columns.

narwhals/stable/v1.py

@@ -34,6 +34,9 @@
 from narwhals.dtypes import UInt64
 from narwhals.dtypes import Unknown
 from narwhals.expr import Expr as NwExpr
+from narwhals.expr import Then as NwThen
+from narwhals.expr import When as NwWhen
+from narwhals.expr import when as nw_when
 from narwhals.functions import concat
 from narwhals.functions import show_versions
 from narwhals.schema import Schema as NwSchema
@@ -1469,6 +1472,71 @@ def get_level(
     return nw.get_level(obj)
 
 
+class When(NwWhen):
+    @classmethod
+    def from_when(cls, when: NwWhen) -> Self:
+        return cls(*when._predicates)
+
+    def then(self, value: Any) -> Then:
+        return Then.from_then(super().then(value))
+
+
+class Then(NwThen, Expr):
+    @classmethod
+    def from_then(cls, then: NwThen) -> Self:
+        return cls(then._call)
+
+    def otherwise(self, value: Any) -> Expr:
+        return _stableify(super().otherwise(value))
+
+
+def when(*predicates: IntoExpr | Iterable[IntoExpr]) -> When:
+    """
+    Start a `when-then-otherwise` expression.
+    Expression similar to an `if-else` statement in Python. Always initiated by a `pl.when(<condition>).then(<value if condition>)`., and optionally followed by chaining one or more `.when(<condition>).then(<value>)` statements.
+    Chained when-then operations should be read as Python `if, elif, ... elif` blocks, not as `if, if, ... if`, i.e. the first condition that evaluates to `True` will be picked.
+    If none of the conditions are `True`, an optional `.otherwise(<value if all statements are false>)` can be appended at the end. If not appended, and none of the conditions are `True`, `None` will be returned.
+
+    Arguments:
+        predicates: Condition(s) that must be met in order to apply the subsequent statement. Accepts one or more boolean expressions, which are implicitly combined with `&`. String input is parsed as a column name.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import polars as pl
+        >>> import narwhals.stable.v1 as nw
+        >>> df_pl = pl.DataFrame({"a": [1, 2, 3], "b": [5, 10, 15]})
+        >>> df_pd = pd.DataFrame({"a": [1, 2, 3], "b": [5, 10, 15]})
+
+        We define a dataframe-agnostic function:
+
+        >>> @nw.narwhalify
+        ... def func(df_any):
+        ...     return df_any.with_columns(
+        ...         nw.when(nw.col("a") < 3).then(5).otherwise(6).alias("a_when")
+        ...     )
+
+        We can then pass either pandas or polars to `func`:
+
+        >>> func(df_pd)
+           a   b  a_when
+        0  1   5       5
+        1  2  10       5
+        2  3  15       6
+        >>> func(df_pl)
+        shape: (3, 3)
+        ┌─────┬─────┬────────┐
+        │ a   ┆ b   ┆ a_when │
+        │ --- ┆ --- ┆ ---    │
+        │ i64 ┆ i64 ┆ i32    │
+        ╞═════╪═════╪════════╡
+        │ 1   ┆ 5   ┆ 5      │
+        │ 2   ┆ 10  ┆ 5      │
+        │ 3   ┆ 15  ┆ 6      │
+        └─────┴─────┴────────┘
+    """
+    return When.from_when(nw_when(*predicates))
+
+
 def from_dict(
     data: dict[str, Any],
     schema: dict[str, DType] | Schema | None = None,
@@ -1545,6 +1613,7 @@ def from_dict(
     "mean",
     "sum",
     "sum_horizontal",
+    "when",
     "DataFrame",
     "LazyFrame",
     "Series",