feat: when-then broadcasting

narwhals/_compliant/when_then.py

@@ -154,11 +154,18 @@ def __call__(self, df: EagerDataFrameT, /) -> Sequence[EagerSeriesT]:
         is_expr = self._condition._is_expr
         when: EagerSeriesT = self._condition(df)[0]
         then: EagerSeriesT
+        if (
+            self._condition._metadata is not None
+            and self._condition._metadata.is_scalar_like
+        ):
+            when._broadcast = True
+
         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:

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

@@ -2,7 +2,7 @@
 
 import operator
 from functools import reduce
-from typing import TYPE_CHECKING, Iterable, Sequence, cast
+from typing import TYPE_CHECKING, Any, Iterable, Sequence, cast
 
 import dask.dataframe as dd
 import pandas as pd
@@ -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,25 +285,44 @@ 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)
+        def _aggregates(expr: Any) -> bool:
+            if isinstance(expr, DaskExpr):
+                return expr._metadata is not None and expr._metadata.is_scalar_like
+            elif isinstance(expr, dd.Series):
+                return len(expr) == 1
+            return True
+
+        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]
+        if (
+            _aggregates(self._condition)
+            and _aggregates(then_value)
+            and _aggregates(otherwise_value)
+        ):
+            df = df._with_native(condition.to_frame())
+        elif _aggregates(self._condition):
+            condition = self._condition.broadcast(ExprKind.AGGREGATION)(df)[0]
+
+        (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]
+        return [then_series.where(condition, otherwise_series)]
 
 
 class DaskThen(CompliantThen[DaskLazyFrame, "dx.Series", DaskExpr], DaskExpr): ...

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,
@@ -1449,7 +1448,6 @@ 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:
         return Then(

tests/expr_and_series/when_test.py

@@ -6,7 +6,7 @@
 import pytest
 
 import narwhals as nw
-from narwhals.exceptions import MultiOutputExpressionError, ShapeError
+from narwhals.exceptions import MultiOutputExpressionError
 from tests.utils import Constructor, ConstructorEager, assert_equal_data
 
 if TYPE_CHECKING:
@@ -115,12 +115,6 @@ def test_when_then_otherwise_into_expr(constructor: Constructor) -> None:
     assert_equal_data(result, expected)
 
 
-def test_when_then_invalid(constructor: Constructor) -> None:
-    df = nw.from_native(constructor(data))
-    with pytest.raises(ShapeError):
-        df.select(nw.when(nw.col("a").sum() > 1).then("c"))
-
-
 def test_when_then_otherwise_lit_str(constructor: Constructor) -> None:
     df = nw.from_native(constructor(data))
     result = df.select(nw.when(nw.col("a") > 1).then(nw.col("b")).otherwise(nw.lit("z")))
@@ -144,3 +138,63 @@ 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, 200, [100]),
+        (nw.col("a").sum() == 6, nw.col("a"), 200, [1, 2, 3]),
+        (nw.col("a").sum() == 6, nw.col("a"), nw.col("b"), [1, 2, 3]),
+        (nw.col("a").sum() == 6, 100, nw.col("b"), [100, 100, 100]),
+        (nw.col("a").sum() == 5, 100, 200, [200]),
+        (nw.col("a").sum() == 5, nw.col("a"), 200, [200, 200, 200]),
+        (nw.col("a").sum() == 5, nw.col("a"), nw.col("b"), [4, 5, 6]),
+        (nw.col("a").sum() == 5, 100, nw.col("b"), [4, 5, 6]),
+    ],
+)
+def test_when_then_otherwise_broadcast_select(
+    condition: nw.Expr,
+    then: nw.Expr | int,
+    otherwise: nw.Expr | int,
+    expected: list[int],
+    constructor: Constructor,
+    request: pytest.FixtureRequest,
+) -> None:
+    # lazy backends needs to be able to pushdown the aggregation for broadcasting to work
+    if any(x in str(constructor) for x in ["duckdb", "sqlframe", "pyspark"]) and not (
+        isinstance(then, int) and isinstance(otherwise, int)
+    ):
+        request.applymarker(pytest.mark.xfail)
+    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, 200, [100, 100, 100]),
+        (nw.col("a").sum() == 6, nw.col("a"), 200, [1, 2, 3]),
+        (nw.col("a").sum() == 6, nw.col("a"), nw.col("b"), [1, 2, 3]),
+        (nw.col("a").sum() == 6, 100, nw.col("b"), [100, 100, 100]),
+        (nw.col("a").sum() == 5, 100, 200, [200, 200, 200]),
+        (nw.col("a").sum() == 5, nw.col("a"), 200, [200, 200, 200]),
+        (nw.col("a").sum() == 5, nw.col("a"), nw.col("b"), [4, 5, 6]),
+        (nw.col("a").sum() == 5, 100, nw.col("b"), [4, 5, 6]),
+    ],
+)
+def test_when_then_otherwise_broadcast_with_columns(
+    condition: nw.Expr,
+    then: nw.Expr | int,
+    otherwise: nw.Expr | int,
+    expected: list[int],
+    constructor: Constructor,
+    request: pytest.FixtureRequest,
+) -> None:
+    # lazy backends needs to be able to pushdown the aggregation for broadcasting to work
+    if any(x in str(constructor) for x in ["duckdb", "sqlframe", "pyspark"]):
+        request.applymarker(pytest.mark.xfail)
+    df = nw.from_native(constructor({"a": [1, 2, 3], "b": [4, 5, 6]}))
+    result = df.with_columns(a_when=nw.when(condition).then(then).otherwise(otherwise))
+    assert_equal_data(result.select(nw.col("a_when")), {"a_when": expected})