feat: Implement chained when-then-otherwise 🥳

Related (#668 (comment))

- This would be how we should model it in *actual narwhals*
- Almost identical to the `rust` version
  - See this in particular (https://github.com/pola-rs/polars/blob/b9dd8cdbd6e6ec8373110536955ed5940b9460ec/crates/polars-plan/src/dsl/arity.rs#L89-L130)

Author: dangotbanned

narwhals/_plan/demo.py

@@ -9,11 +9,12 @@
     expr_parsing as parse,
     functions as F,  # noqa: N812
 )
-from narwhals._plan.common import ExprIR, IntoExpr, is_non_nested_literal
+from narwhals._plan.common import ExprIR, IntoExpr, is_expr, is_non_nested_literal
 from narwhals._plan.dummy import DummySeries
 from narwhals._plan.expr import All, Column, Columns, IndexColumns, Len, Nth
 from narwhals._plan.literal import ScalarLiteral, SeriesLiteral
 from narwhals._plan.strings import ConcatHorizontal
+from narwhals._plan.when_then import When
 from narwhals.dtypes import DType
 from narwhals.exceptions import OrderDependentExprError
 from narwhals.utils import Version, flatten
@@ -131,6 +132,32 @@ def concat_str(
     )
 
 
+def when(*predicates: IntoExpr | t.Iterable[IntoExpr]) -> When:
+    """Start a `when-then-otherwise` expression.
+
+    Examples:
+        >>> from narwhals._plan import demo as nwd
+
+        >>> when_then_many = (
+        ...     nwd.when(nwd.col("x") == "a")
+        ...     .then(1)
+        ...     .when(nwd.col("x") == "b")
+        ...     .then(2)
+        ...     .when(nwd.col("x") == "c")
+        ...     .then(3)
+        ...     .otherwise(4)
+        ... )
+        >>> when_then_many
+        Narwhals DummyExpr (main):
+        .when([(col('x')) == (lit(str: a))]).then(lit(int: 1)).otherwise(.when([(col('x')) == (lit(str: b))]).then(lit(int: 2)).otherwise(.when([(col('x')) == (lit(str: c))]).then(lit(int: 3)).otherwise(lit(int: 4))))
+    """
+    if builtins.len(predicates) == 1 and is_expr(predicates[0]):
+        expr = predicates[0]
+    else:
+        expr = all_horizontal(*predicates)
+    return When._from_expr(expr)
+
+
 def _is_order_enforcing_previous(obj: t.Any) -> TypeIs[SortBy]:
     """In theory, we could add other nodes to this check."""
     from narwhals._plan.expr import SortBy

narwhals/_plan/expr.py

@@ -13,7 +13,7 @@
 import typing as t
 
 from narwhals._plan.aggregation import Agg, OrderableAgg
-from narwhals._plan.common import ExprIR, SelectorIR
+from narwhals._plan.common import ExprIR, SelectorIR, _field_str
 from narwhals._plan.name import KeepName, RenameAlias
 from narwhals._plan.typing import (
     FunctionT,
@@ -475,3 +475,23 @@ class Ternary(ExprIR):
 
     Deferring this for now.
     """
+
+    __slots__ = ("falsy", "predicate", "truthy")
+
+    predicate: ExprIR
+    truthy: ExprIR
+    falsy: ExprIR
+
+    def __str__(self) -> str:
+        # NOTE: Default slot ordering made it difficult to read
+        fields = (
+            _field_str("predicate", self.predicate),
+            _field_str("truthy", self.truthy),
+            _field_str("falsy", self.falsy),
+        )
+        return f"{type(self).__name__}({', '.join(fields)})"
+
+    def __repr__(self) -> str:
+        return (
+            f".when({self.predicate!r}).then({self.truthy!r}).otherwise({self.falsy!r})"
+        )

narwhals/_plan/when_then.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from narwhals._plan.common import Immutable
+from narwhals._plan.expr_parsing import parse_into_expr_ir
+
+if TYPE_CHECKING:
+    from narwhals._plan.common import ExprIR, IntoExpr, Seq
+    from narwhals._plan.dummy import DummyExpr
+    from narwhals._plan.expr import Ternary
+
+
+class When(Immutable):
+    __slots__ = ("condition",)
+
+    condition: ExprIR
+
+    def then(self, expr: IntoExpr, /) -> Then:
+        return Then(condition=self.condition, statement=parse_into_expr_ir(expr))
+
+    @staticmethod
+    def _from_expr(expr: DummyExpr, /) -> When:
+        return When(condition=expr._ir)
+
+
+class Then(Immutable):
+    __slots__ = ("condition", "statement")
+
+    condition: ExprIR
+    statement: ExprIR
+
+    def when(self, condition: IntoExpr, /) -> ChainedWhen:
+        return ChainedWhen(
+            conditions=(self.condition, parse_into_expr_ir(condition)),
+            statements=(self.statement,),
+        )
+
+    def otherwise(self, statement: IntoExpr, /) -> DummyExpr:
+        return ternary_expr(
+            self.condition, self.condition, parse_into_expr_ir(statement)
+        ).to_narwhals()
+
+
+class ChainedWhen(Immutable):
+    __slots__ = ("conditions", "statements")
+
+    conditions: Seq[ExprIR]
+    statements: Seq[ExprIR]
+
+    def then(self, statement: IntoExpr, /) -> ChainedThen:
+        return ChainedThen(
+            conditions=self.conditions,
+            statements=(*self.statements, parse_into_expr_ir(statement)),
+        )
+
+
+class ChainedThen(Immutable):
+    """https://github.com/pola-rs/polars/blob/b9dd8cdbd6e6ec8373110536955ed5940b9460ec/crates/polars-plan/src/dsl/arity.rs#L89-L130."""
+
+    __slots__ = ("conditions", "statements")
+
+    conditions: Seq[ExprIR]
+    statements: Seq[ExprIR]
+
+    def when(self, condition: IntoExpr, /) -> ChainedWhen:
+        return ChainedWhen(
+            conditions=(*self.conditions, parse_into_expr_ir(condition)),
+            statements=self.statements,
+        )
+
+    def otherwise(self, statement: IntoExpr, /) -> DummyExpr:
+        otherwise = parse_into_expr_ir(statement)
+        it_conditions = reversed(self.conditions)
+        it_statements = reversed(self.statements)
+        for e in it_conditions:
+            otherwise = ternary_expr(e, next(it_statements), otherwise)
+        return otherwise.to_narwhals()
+
+
+def ternary_expr(predicate: ExprIR, truthy: ExprIR, falsy: ExprIR, /) -> Ternary:
+    from narwhals._plan.expr import Ternary
+
+    return Ternary(predicate=predicate, truthy=truthy, falsy=falsy)