From 8e4489e364dd3e0ce89e1989ee4c35aa8ffcb413 Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Tue, 10 Jun 2025 13:17:48 -0700 Subject: [PATCH 01/13] add when_then_broadcast tests --- tests/expr_and_series/when_test.py | 68 +++++++++++++++++++++++++++--- 1 file changed, 61 insertions(+), 7 deletions(-) diff --git a/tests/expr_and_series/when_test.py b/tests/expr_and_series/when_test.py index faa4715502..7fffb75178 100644 --- a/tests/expr_and_series/when_test.py +++ b/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}) From 51ad4d6cc3e5dff63a50be9e3e2fb826543473a8 Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Tue, 10 Jun 2025 13:18:29 -0700 Subject: [PATCH 02/13] remove when-then expression_preserve_length check --- narwhals/functions.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/narwhals/functions.py b/narwhals/functions.py index 8f70aa95b6..8a93e0bf49 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, @@ -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( From 2bd2940cf4fbd30cc168fbf6430c52836ba001bd Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Tue, 10 Jun 2025 13:18:58 -0700 Subject: [PATCH 03/13] enh when-then eager backends --- narwhals/_compliant/when_then.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/narwhals/_compliant/when_then.py b/narwhals/_compliant/when_then.py index a2a2b4fe78..3c012f776d 100644 --- a/narwhals/_compliant/when_then.py +++ b/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: From 3f6934295f418867a4db9607070fd60e1b75bd10 Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Tue, 10 Jun 2025 13:37:17 -0700 Subject: [PATCH 04/13] enh when-then broadcast dask backend --- narwhals/_dask/namespace.py | 54 ++++++++++++++++++++++++++----------- 1 file changed, 38 insertions(+), 16 deletions(-) diff --git a/narwhals/_dask/namespace.py b/narwhals/_dask/namespace.py index 5e9150f244..c47266883c 100644 --- a/narwhals/_dask/namespace.py +++ b/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,6 +18,7 @@ validate_comparand, ) from narwhals._expression_parsing import ( + ExprKind, combine_alias_output_names, combine_evaluate_output_names, ) @@ -283,25 +284,46 @@ 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)] + from dask.dataframe.dask_expr._expr import Where - 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) + otherwise_value = Where._defaults["other"] + + 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): ... From 2d4aa4aba548e8a52466bd8823d3e623d3e416f3 Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Tue, 10 Jun 2025 14:41:42 -0700 Subject: [PATCH 05/13] _dask.expr.DaskExpr.broadcast workaround --- narwhals/_dask/expr.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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, From 2f9edfc97cbfaa843dcc5c23a30ef19976ebe36c Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Tue, 10 Jun 2025 14:46:43 -0700 Subject: [PATCH 06/13] fix get_dask_Expr import --- narwhals/_dask/namespace.py | 5 ++--- narwhals/dependencies.py | 2 ++ 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/narwhals/_dask/namespace.py b/narwhals/_dask/namespace.py index c47266883c..2457ddfe92 100644 --- a/narwhals/_dask/namespace.py +++ b/narwhals/_dask/namespace.py @@ -23,6 +23,7 @@ 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 @@ -303,9 +304,7 @@ def _aggregates(expr: Any) -> bool: ) if self._otherwise_value is None: - from dask.dataframe.dask_expr._expr import Where - - otherwise_value = Where._defaults["other"] + otherwise_value = get_dask_expr()._expr.Where._defaults["other"] condition = self._condition(df)[0] if ( 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) From f02c7f0342cd9cd2d73adae30cca99694802ab06 Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Wed, 11 Jun 2025 20:14:55 -0700 Subject: [PATCH 07/13] when/then catch anticipated shape errors early --- narwhals/functions.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/narwhals/functions.py b/narwhals/functions.py index 8a93e0bf49..8b10bb4ec8 100644 --- a/narwhals/functions.py +++ b/narwhals/functions.py @@ -31,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 @@ -1450,6 +1450,14 @@ def __init__(self, *predicates: IntoExpr | Iterable[IntoExpr]) -> None: self._predicate = all_horizontal(*flatten(predicates)) 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, @@ -1471,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] From bd5dce3d1288c37de31c39b8a96c3ba227110ecb Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Wed, 11 Jun 2025 20:15:29 -0700 Subject: [PATCH 08/13] add when/then shape and aggregation/broadcast tests --- tests/expr_and_series/when_test.py | 52 +++++++++++++++--------------- 1 file changed, 26 insertions(+), 26 deletions(-) diff --git a/tests/expr_and_series/when_test.py b/tests/expr_and_series/when_test.py index 7fffb75178..a695733229 100644 --- a/tests/expr_and_series/when_test.py +++ b/tests/expr_and_series/when_test.py @@ -6,7 +6,7 @@ import pytest import narwhals as nw -from narwhals.exceptions import MultiOutputExpressionError +from narwhals.exceptions import MultiOutputExpressionError, ShapeError from tests.utils import Constructor, ConstructorEager, assert_equal_data if TYPE_CHECKING: @@ -115,6 +115,15 @@ 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")) + + 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)) result = df.select(nw.when(nw.col("a") > 1).then(nw.col("b")).otherwise(nw.lit("z"))) @@ -144,28 +153,22 @@ def test_when_then_otherwise_multi_output(constructor: Constructor) -> None: ("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() == 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, 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]), + (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_broadcast_select( +def test_when_then_otherwise_aggregate_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}) @@ -175,26 +178,23 @@ def test_when_then_otherwise_broadcast_select( ("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() == 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, 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]), + (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_broadcast_with_columns( +def test_when_then_otherwise_aggregate_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)) + 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}) From 283a017346cd4d365d04cb956bd37a095aed3143 Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Wed, 11 Jun 2025 20:16:31 -0700 Subject: [PATCH 09/13] swap the dask when/then broadcasting hack for a cleaner hack --- narwhals/_dask/namespace.py | 23 +++++++---------------- 1 file changed, 7 insertions(+), 16 deletions(-) diff --git a/narwhals/_dask/namespace.py b/narwhals/_dask/namespace.py index 2457ddfe92..f257f5134c 100644 --- a/narwhals/_dask/namespace.py +++ b/narwhals/_dask/namespace.py @@ -2,7 +2,7 @@ import operator from functools import reduce -from typing import TYPE_CHECKING, Any, Iterable, Sequence, cast +from typing import TYPE_CHECKING, Iterable, Sequence, cast import dask.dataframe as dd import pandas as pd @@ -285,13 +285,6 @@ def _then(self) -> type[DaskThen]: return DaskThen def __call__(self, df: DaskLazyFrame) -> Sequence[dx.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) @@ -307,14 +300,12 @@ def _aggregates(expr: Any) -> bool: otherwise_value = get_dask_expr()._expr.Where._defaults["other"] 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): + # 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 @@ -322,7 +313,7 @@ def _aggregates(expr: Any) -> bool: validate_comparand(condition, then_series) validate_comparand(condition, otherwise_series) - return [then_series.where(condition, otherwise_series)] + return [then_series.where(condition, otherwise_series)] # pyright: ignore[reportArgumentType] class DaskThen(CompliantThen[DaskLazyFrame, "dx.Series", DaskExpr], DaskExpr): ... From 83e188b14bd35ffc7adf739834f6ba3556af2303 Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Wed, 11 Jun 2025 20:19:23 -0700 Subject: [PATCH 10/13] when/then over pushdown in lazy backends --- narwhals/_compliant/when_then.py | 42 +++++++++++++++++++++---------- narwhals/_duckdb/namespace.py | 7 ++++++ narwhals/_spark_like/namespace.py | 7 ++++++ 3 files changed, 43 insertions(+), 13 deletions(-) diff --git a/narwhals/_compliant/when_then.py b/narwhals/_compliant/when_then.py index 3c012f776d..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,6 @@ 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] @@ -182,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 @@ -204,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/_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 From f4a7fba319c7de62460f6cd9ef9cdd40ca18fa88 Mon Sep 17 00:00:00 2001 From: Cameron Riddell Date: Wed, 11 Jun 2025 20:35:29 -0700 Subject: [PATCH 11/13] cov when/then aggregation with None test --- tests/expr_and_series/when_test.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/expr_and_series/when_test.py b/tests/expr_and_series/when_test.py index a695733229..53cd3ad4b3 100644 --- a/tests/expr_and_series/when_test.py +++ b/tests/expr_and_series/when_test.py @@ -152,10 +152,12 @@ def test_when_then_otherwise_multi_output(constructor: Constructor) -> None: @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]), @@ -177,10 +179,12 @@ def test_when_then_otherwise_aggregate_select( @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]), From 4156549de193c5fdebf68926e6477a284688da90 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Thu, 12 Jun 2025 10:39:48 +0100 Subject: [PATCH 12/13] dont use Dask private methods --- narwhals/_dask/namespace.py | 14 ++++++++------ narwhals/functions.py | 17 ++++++++++------- 2 files changed, 18 insertions(+), 13 deletions(-) diff --git a/narwhals/_dask/namespace.py b/narwhals/_dask/namespace.py index f257f5134c..8f004acfbf 100644 --- a/narwhals/_dask/namespace.py +++ b/narwhals/_dask/namespace.py @@ -23,7 +23,6 @@ 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 @@ -296,21 +295,24 @@ def __call__(self, df: DaskLazyFrame) -> Sequence[dx.Series]: else self._otherwise_value ) - if self._otherwise_value is None: - otherwise_value = get_dask_expr()._expr.Where._defaults["other"] - 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: + 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: + (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/functions.py b/narwhals/functions.py index 8b10bb4ec8..c7bb214c86 100644 --- a/narwhals/functions.py +++ b/narwhals/functions.py @@ -1451,11 +1451,11 @@ def __init__(self, *predicates: IntoExpr | Iterable[IntoExpr]) -> None: 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" + 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( @@ -1479,8 +1479,11 @@ 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" + 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 `other` value must also be scalar-like." + ) raise ShapeError(msg) def func(plx: CompliantNamespace[Any, Any]) -> CompliantExpr[Any, Any]: From fd16347de51c03e919ac35488cbfae12c128f8f0 Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Thu, 12 Jun 2025 10:42:08 +0100 Subject: [PATCH 13/13] typo --- narwhals/functions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/narwhals/functions.py b/narwhals/functions.py index c7bb214c86..78a4f49ea7 100644 --- a/narwhals/functions.py +++ b/narwhals/functions.py @@ -1482,7 +1482,7 @@ def otherwise(self, value: IntoExpr | NonNestedLiteral | _1DArray) -> Expr: 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 `other` value must also be scalar-like." + "the `otherwise` value must also be scalar-like." ) raise ShapeError(msg)