diff --git a/narwhals/_compliant/when_then.py b/narwhals/_compliant/when_then.py index a2a2b4fe78..1de91f9369 100644 --- a/narwhals/_compliant/when_then.py +++ b/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] diff --git a/narwhals/_dask/expr.py b/narwhals/_dask/expr.py index 193c0db949..4167e03b0f 100644 --- a/narwhals/_dask/expr.py +++ b/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, diff --git a/narwhals/_dask/namespace.py b/narwhals/_dask/namespace.py index 5e9150f244..8f004acfbf 100644 --- a/narwhals/_dask/namespace.py +++ b/narwhals/_dask/namespace.py @@ -18,6 +18,7 @@ validate_comparand, ) from narwhals._expression_parsing import ( + ExprKind, combine_alias_output_names, combine_evaluate_output_names, ) @@ -283,23 +284,36 @@ def _then(self) -> type[DaskThen]: return DaskThen def __call__(self, df: DaskLazyFrame) -> Sequence[dx.Series]: - condition = self._condition(df)[0] + 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 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) + condition = self._condition(df)[0] + # re-evaluate DataFrame if the condition aggregates to force + # then/otherwise to be evaluated against the aggregated frame + assert self._condition._metadata is not None # noqa: S101 + if 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 if self._otherwise_value is None: - return [then_series.where(condition)] - - 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, then_series) = align_series_full_broadcast( + df, condition, then_value + ) + validate_comparand(condition, then_series) + return [then_series.where(condition)] # pyright: ignore[reportArgumentType] + (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] diff --git a/narwhals/_duckdb/namespace.py b/narwhals/_duckdb/namespace.py index f6b697c319..7397e76495 100644 --- a/narwhals/_duckdb/namespace.py +++ b/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): ... diff --git a/narwhals/_spark_like/namespace.py b/narwhals/_spark_like/namespace.py index 818c7ebb57..f0f629cb46 100644 --- a/narwhals/_spark_like/namespace.py +++ b/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 diff --git a/narwhals/dependencies.py b/narwhals/dependencies.py index 9d4731ab6c..d775677d22 100644 --- a/narwhals/dependencies.py +++ b/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) diff --git a/narwhals/functions.py b/narwhals/functions.py index 8f70aa95b6..78a4f49ea7 100644 --- a/narwhals/functions.py +++ b/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 and not kind.is_scalar_like: + msg = ( + "If you pass a scalar-like predicate to `nw.when`, then " + "the `then` value must also be scalar-like." + ) + raise ShapeError(msg) + return Then( lambda plx: apply_n_ary_operation( plx, @@ -1473,11 +1479,21 @@ 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 and not is_scalar_like(kind): + msg = ( + "If you pass a scalar-like predicate to `nw.when`, then " + "the `otherwise` value must also be scalar-like." + ) + 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] diff --git a/tests/expr_and_series/when_test.py b/tests/expr_and_series/when_test.py index faa4715502..53cd3ad4b3 100644 --- a/tests/expr_and_series/when_test.py +++ b/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})