feat: when-then broadcasting

narwhals/_compliant/when_then.py

@@ -21,6 +21,7 @@
     from typing_extensions import Self, TypeAlias
 
     from narwhals._compliant.typing import EvalSeries, ScalarKwargs
+    from narwhals._compliant.window import WindowInputs
     from narwhals._utils import Implementation, Version, _FullContext
     from narwhals.typing import NonNestedLiteral
 
@@ -50,6 +51,9 @@ class CompliantWhen(Protocol38[FrameT, SeriesT, ExprT]):
     @property
     def _then(self) -> type[CompliantThen[FrameT, SeriesT, ExprT]]: ...
     def __call__(self, compliant_frame: FrameT, /) -> Sequence[SeriesT]: ...
+    def _window_function(
+        self, compliant_frame: FrameT, window_inputs: WindowInputs[Any]
+    ) -> Sequence[SeriesT]: ...
 
     def then(
         self, value: IntoExpr[SeriesT, ExprT], /
@@ -124,9 +128,7 @@ def from_when(
         obj = cls.__new__(cls)
         obj._call = when
 
-        # This may require more complicated logic if we want to push down the
-        # `over`: https://github.com/narwhals-dev/narwhals/issues/2652.
-        obj._window_function = None
+        obj._window_function = when._window_function
 
         obj._when_value = when
         obj._depth = 0
@@ -154,11 +156,13 @@ def __call__(self, df: EagerDataFrameT, /) -> Sequence[EagerSeriesT]:
         is_expr = self._condition._is_expr
         when: EagerSeriesT = self._condition(df)[0]
         then: EagerSeriesT
+
         if is_expr(self._then_value):
             then = self._then_value(df)[0]
         else:
             then = when.alias("literal")._from_scalar(self._then_value)
             then._broadcast = True
+
         if is_expr(self._otherwise_value):
             otherwise = self._otherwise_value(df)[0]
         elif self._otherwise_value is not None:
@@ -175,7 +179,6 @@ class LazyWhen(
 ):
     when: Callable[..., NativeExprT]
     lit: Callable[..., NativeExprT]
-    _window_function: WindowFunction[CompliantLazyFrameT, NativeExprT] | None
 
     def __call__(self, df: CompliantLazyFrameT) -> Sequence[NativeExprT]:
         is_expr = self._condition._is_expr
@@ -197,13 +200,33 @@ def from_expr(cls, condition: LazyExprT, /, *, context: _FullContext) -> Self:
         obj = cls.__new__(cls)
         obj._condition = condition
 
-        # This may require more complicated logic if we want to push down the
-        # `over`: https://github.com/narwhals-dev/narwhals/issues/2652.
-        obj._window_function = None
-
         obj._then_value = None
         obj._otherwise_value = None
         obj._implementation = context._implementation
         obj._backend_version = context._backend_version
         obj._version = context._version
         return obj
+
+    def _window_function(
+        self, df: CompliantLazyFrameT, window_inputs: WindowInputs[NativeExprT]
+    ) -> Sequence[NativeExprT]:
+        is_expr = self._condition._is_expr
+        condition = self._condition.window_function(df, window_inputs)[0]
+        then_ = self._then_value
+        then = (
+            then_.window_function(df, window_inputs)[0]
+            if is_expr(then_)
+            else self.lit(then_)
+        )
+
+        other_ = self._otherwise_value
+        if other_ is None:
+            result = self.when(condition, then)
+        else:
+            other = (
+                other_.window_function(df, window_inputs)[0]
+                if is_expr(other_)
+                else self.lit(other_)
+            )
+            result = self.when(condition, then).otherwise(other)  # type: ignore  # noqa: PGH003
+        return [result]

narwhals/_dask/expr.py

@@ -81,7 +81,9 @@ def __narwhals_namespace__(self) -> DaskNamespace:  # pragma: no cover
 
     def broadcast(self, kind: Literal[ExprKind.AGGREGATION, ExprKind.LITERAL]) -> Self:
         def func(df: DaskLazyFrame) -> list[dx.Series]:
-            return [result[0] for result in self(df)]
+            # result.loc[0][0] is a workaround for dask~<=2024.10.0/dask_expr~<=1.1.16
+            #   that raised a KeyErrror for result[0] during collection.
+            return [result.loc[0][0] for result in self(df)]
 
         return self.__class__(
             func,

narwhals/_dask/namespace.py

@@ -18,10 +18,12 @@
     validate_comparand,
 )
 from narwhals._expression_parsing import (
+    ExprKind,
     combine_alias_output_names,
     combine_evaluate_output_names,
 )
 from narwhals._utils import Implementation
+from narwhals.dependencies import get_dask_expr
 
 if TYPE_CHECKING:
     import dask.dataframe.dask_expr as dx
@@ -283,23 +285,33 @@ def _then(self) -> type[DaskThen]:
         return DaskThen
 
     def __call__(self, df: DaskLazyFrame) -> Sequence[dx.Series]:
-        condition = self._condition(df)[0]
-
-        if isinstance(self._then_value, DaskExpr):
-            then_value = self._then_value(df)[0]
-        else:
-            then_value = self._then_value
-        (then_series,) = align_series_full_broadcast(df, then_value)
-        validate_comparand(condition, then_series)
+        then_value = (
+            self._then_value(df)[0]
+            if isinstance(self._then_value, DaskExpr)
+            else self._then_value
+        )
+        otherwise_value = (
+            self._otherwise_value(df)[0]
+            if isinstance(self._otherwise_value, DaskExpr)
+            else self._otherwise_value
+        )
 
         if self._otherwise_value is None:
-            return [then_series.where(condition)]
+            otherwise_value = get_dask_expr()._expr.Where._defaults["other"]
 
-        if isinstance(self._otherwise_value, DaskExpr):
-            otherwise_value = self._otherwise_value(df)[0]
-        else:
-            return [then_series.where(condition, self._otherwise_value)]  # pyright: ignore[reportArgumentType]
-        (otherwise_series,) = align_series_full_broadcast(df, otherwise_value)
+        condition = self._condition(df)[0]
+        # re-evaluate DataFrame if the condition aggregates to force
+        #   then/otherwise to be evaluated against the aggregated frame
+        if self._condition._metadata is None or self._condition._metadata.is_scalar_like:
+            new_df = df._with_native(condition.to_frame())
+            condition = self._condition.broadcast(ExprKind.AGGREGATION)(df)[0]
+            df = new_df
+
+        (condition, then_series, otherwise_series) = align_series_full_broadcast(
+            df, condition, then_value, otherwise_value
+        )
+
+        validate_comparand(condition, then_series)
         validate_comparand(condition, otherwise_series)
         return [then_series.where(condition, otherwise_series)]  # pyright: ignore[reportArgumentType]
 

narwhals/_duckdb/namespace.py

@@ -199,5 +199,12 @@ def __call__(self, df: DuckDBLazyFrame) -> Sequence[Expression]:
         self.lit = lit
         return super().__call__(df)
 
+    def _window_function(
+        self, df: DuckDBLazyFrame, window_inputs: DuckDBWindowInputs
+    ) -> Sequence[Expression]:
+        self.when = when
+        self.lit = lit
+        return super()._window_function(df, window_inputs)
+
 
 class DuckDBThen(LazyThen["DuckDBLazyFrame", Expression, DuckDBExpr], DuckDBExpr): ...

narwhals/_spark_like/namespace.py

@@ -280,6 +280,13 @@ def __call__(self, df: SparkLikeLazyFrame) -> Sequence[Column]:
         self.lit = df._F.lit
         return super().__call__(df)
 
+    def _window_function(
+        self, df: SparkLikeLazyFrame, window_inputs: SparkWindowInputs
+    ) -> Sequence[Column]:
+        self.when = df._F.when
+        self.lit = df._F.lit
+        return super()._window_function(df, window_inputs)
+
 
 class SparkLikeThen(
     LazyThen[SparkLikeLazyFrame, "Column", SparkLikeExpr], SparkLikeExpr

narwhals/dependencies.py

@@ -100,6 +100,8 @@ def get_ibis() -> Any:
 
 def get_dask_expr() -> Any:  # pragma: no cover
     """Get dask_expr module (if already imported - else return None)."""
+    if (dd := get_dask_dataframe()) is not None and hasattr(dd, "dask_expr"):
+        return dd.dask_expr
     return sys.modules.get("dask_expr", None)
 
 

narwhals/functions.py

@@ -9,7 +9,6 @@
     ExprKind,
     ExprMetadata,
     apply_n_ary_operation,
-    check_expressions_preserve_length,
     combine_metadata,
     extract_compliant,
     is_scalar_like,
@@ -32,7 +31,7 @@
     is_numpy_array_2d,
     is_pyarrow_table,
 )
-from narwhals.exceptions import InvalidOperationError
+from narwhals.exceptions import InvalidOperationError, ShapeError
 from narwhals.expr import Expr
 from narwhals.translate import from_native, to_native
 
@@ -1449,9 +1448,16 @@ def max_horizontal(*exprs: IntoExpr | Iterable[IntoExpr]) -> Expr:
 class When:
     def __init__(self, *predicates: IntoExpr | Iterable[IntoExpr]) -> None:
         self._predicate = all_horizontal(*flatten(predicates))
-        check_expressions_preserve_length(self._predicate, function_name="when")
 
     def then(self, value: IntoExpr | NonNestedLiteral | _1DArray) -> Then:
+        kind = ExprKind.from_into_expr(value, str_as_lit=False)
+        if (
+            self._predicate._metadata.is_scalar_like is True
+            and kind.is_scalar_like is False
+        ):
+            msg = "When produced a scalar-like result and Then did not"
+            raise ShapeError(msg)
+
         return Then(
             lambda plx: apply_n_ary_operation(
                 plx,
@@ -1473,11 +1479,18 @@ def then(self, value: IntoExpr | NonNestedLiteral | _1DArray) -> Then:
 class Then(Expr):
     def otherwise(self, value: IntoExpr | NonNestedLiteral | _1DArray) -> Expr:
         kind = ExprKind.from_into_expr(value, str_as_lit=False)
+        if self._metadata.is_scalar_like is True and is_scalar_like(kind) is False:
+            msg = "When/Then produced a scalar-like result and otherwise did not"
+            raise ShapeError(msg)
 
         def func(plx: CompliantNamespace[Any, Any]) -> CompliantExpr[Any, Any]:
             compliant_expr = self._to_compliant_expr(plx)
             compliant_value = extract_compliant(plx, value, str_as_lit=False)
-            if is_scalar_like(kind) and is_compliant_expr(compliant_value):
+            if (
+                not self._metadata.is_scalar_like
+                and is_scalar_like(kind)
+                and is_compliant_expr(compliant_value)
+            ):
                 compliant_value = compliant_value.broadcast(kind)
             return compliant_expr.otherwise(compliant_value)  # type: ignore[attr-defined, no-any-return]
 

tests/expr_and_series/when_test.py

@@ -120,6 +120,9 @@ def test_when_then_invalid(constructor: Constructor) -> None:
     with pytest.raises(ShapeError):
         df.select(nw.when(nw.col("a").sum() > 1).then("c"))
 
+    with pytest.raises(ShapeError):
+        df.select(nw.when(nw.col("a").sum() > 1).then(1).otherwise("c"))
+
 
 def test_when_then_otherwise_lit_str(constructor: Constructor) -> None:
     df = nw.from_native(constructor(data))
@@ -144,3 +147,58 @@ def test_when_then_otherwise_multi_output(constructor: Constructor) -> None:
         df.select(x1=nw.when(nw.all() > 1).then(nw.col("a", "b")))
     with pytest.raises(MultiOutputExpressionError):
         df.select(x1=nw.when(nw.all() > 1).then(nw.lit(1)).otherwise(nw.all()))
+
+
+@pytest.mark.parametrize(
+    ("condition", "then", "otherwise", "expected"),
+    [
+        (nw.col("a").sum() == 6, 100, None, [100]),
+        (nw.col("a").sum() == 6, 100, 200, [100]),
+        (nw.col("a").sum() == 6, nw.col("a").sum(), 200, [6]),
+        (nw.col("a").sum() == 6, 100, nw.col("b").sum(), [100]),
+        (nw.col("a").sum() == 6, nw.col("a").sum(), nw.col("b").sum(), [6]),
+        (nw.col("a").sum() == 5, 100, None, [None]),
+        (nw.col("a").sum() == 5, 100, 200, [200]),
+        (nw.col("a").sum() == 5, nw.col("a").sum(), 200, [200]),
+        (nw.col("a").sum() == 5, 100, nw.col("b").sum(), [15]),
+        (nw.col("a").sum() == 5, nw.col("a").sum(), nw.col("b").sum(), [15]),
+    ],
+)
+def test_when_then_otherwise_aggregate_select(
+    condition: nw.Expr,
+    then: nw.Expr | int,
+    otherwise: nw.Expr | int,
+    expected: list[int],
+    constructor: Constructor,
+) -> None:
+    df = nw.from_native(constructor({"a": [1, 2, 3], "b": [4, 5, 6]}))
+    result = df.select(a_when=nw.when(condition).then(then).otherwise(otherwise))
+    assert_equal_data(result, {"a_when": expected})
+
+
+@pytest.mark.parametrize(
+    ("condition", "then", "otherwise", "expected"),
+    [
+        (nw.col("a").sum() == 6, 100, None, [100, 100, 100]),
+        (nw.col("a").sum() == 6, 100, 200, [100, 100, 100]),
+        (nw.col("a").sum() == 6, nw.col("a").sum(), 200, [6, 6, 6]),
+        (nw.col("a").sum() == 6, 100, nw.col("b").sum(), [100, 100, 100]),
+        (nw.col("a").sum() == 6, nw.col("a").sum(), nw.col("b").sum(), [6, 6, 6]),
+        (nw.col("a").sum() == 5, 100, None, [None, None, None]),
+        (nw.col("a").sum() == 5, 100, 200, [200, 200, 200]),
+        (nw.col("a").sum() == 5, nw.col("a").sum(), 200, [200, 200, 200]),
+        (nw.col("a").sum() == 5, 100, nw.col("b").sum(), [15, 15, 15]),
+        (nw.col("a").sum() == 5, nw.col("a").sum(), nw.col("b").sum(), [15, 15, 15]),
+    ],
+)
+def test_when_then_otherwise_aggregate_with_columns(
+    condition: nw.Expr,
+    then: nw.Expr | int,
+    otherwise: nw.Expr | int,
+    expected: list[int],
+    constructor: Constructor,
+) -> None:
+    df = nw.from_native(constructor({"a": [1, 2, 3], "b": [4, 5, 6]}))
+    expr = nw.when(condition).then(then).otherwise(otherwise)
+    result = df.with_columns(a_when=expr)
+    assert_equal_data(result.select(nw.col("a_when")), {"a_when": expected})