feat: Adds `Compliant(When|Then)` #2261

dangotbanned · 2025-03-21T12:20:49Z

What type of PR is this? (check all applicable)

Related issues

Checklist

Code follows style guide (ruff)
Tests added
Documented the changes

Tasks

- Very barebones to start with - Mostly just defining the current impl in a generic way - There are some quirks I'd like to iron out before merging

dangotbanned · 2025-03-21T15:15:31Z

narwhals/_compliant/when_then.py

+    def otherwise(
+        self, value: CompliantExprT | CompliantSeriesOrNativeExprT_co | _Scalar, /
+    ) -> CompliantExprT:
+        self._call._otherwise_value = value
+        self._function_name = "whenotherwise"
+        return cast("CompliantExprT", self)


I've been thinking about this for a while.

@MarcoGorelli would you be able to help me understand what is going on here?

For reference, I'm just mirroring what is in (Arrow|Pandas)Then

narwhals/narwhals/_arrow/namespace.py

Lines 388 to 391 in 29c0f6c

def otherwise(self: Self, value: ArrowExpr | ArrowSeries | _Scalar) -> ArrowExpr:

self._call._otherwise_value = value

self._function_name = "whenotherwise"

return self

Side note for DuckDBThen

I think that is typing issue would be fixed by annotating DuckDB._call: DuckDBWhen

narwhals/narwhals/_duckdb/namespace.py

Lines 320 to 326 in 29c0f6c

def otherwise(self: Self, value: DuckDBExpr | Any) -> DuckDBExpr:

# type ignore because we are setting the `_call` attribute to a

# callable object of type `DuckDBWhen`, base class has the attribute as

# only a `Callable`

self._call._otherwise_value = value # type: ignore[attr-defined]

self._function_name = "whenotherwise"

return self

1. Mutability

*Then.otherwise mutates both *When and *Then.

I expected that would cause an issue for something like:

Code block

import pandas as pd import polars as pl import narwhals as nw df = pl.DataFrame({"a": ["A", "B", "A"], "b": [1, 2, 3], "c": [9, 2, 4], "d": [8, 7, 8]}) nw_pd = nw.from_native(nw.from_native(df).to_pandas()) when = nw.when(nw.col("a") == "B") when_then = when.then("b") when_then_otherwise = when.then("b").otherwise("c") nw_pd.with_columns( when_then.alias("e"), when_then_otherwise.alias("f"), when_then.otherwise("d").alias("g"), when.then(nw.col("c") * 10).otherwise(nw.col("c") * 100).alias("h"), )

But it all worked 🎉!

Output

┌─────────────────────────────┐ | Narwhals DataFrame | |-----------------------------| | a b c d e f g h| |0 A 1 9 8 NaN 9 8 900| |1 B 2 2 7 2.0 2 2 20| |2 A 3 4 8 NaN 4 8 400| └─────────────────────────────┘

However, AFAICT we're depending on the implementation details of nw.(When|Then) to make this safe.

My best guess is we get new function(s) (and therefore Compliant* objects) on every narwhals-level method call?

nw.(When|Then)

narwhals/narwhals/functions.py

Lines 1618 to 1650 in 29c0f6c

class When:

def __init__(self: Self, *predicates: IntoExpr | Iterable[IntoExpr]) -> None:

self._predicate = all_horizontal(*flatten(predicates))

check_expressions_preserve_length(self._predicate, function_name="when")

def then(self: Self, value: IntoExpr | Any) -> Then:

return Then(

lambda plx: apply_n_ary_operation(

plx,

lambda *args: plx.when(args[0]).then(args[1]),

self._predicate,

value,

str_as_lit=False,

),

combine_metadata(self._predicate, value, str_as_lit=False),

)

class Then(Expr):

def otherwise(self: Self, value: IntoExpr | Any) -> Expr:

kind = infer_kind(value, str_as_lit=False)

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):

compliant_value = compliant_value.broadcast(kind)

return compliant_expr.otherwise(compliant_value) # type: ignore[attr-defined, no-any-return]

return Expr(

func,

combine_metadata(self, value, str_as_lit=False),

)

Suggestions

If we change nothing

It might be helpful to add some tests for this kind of thing to avoid unintentionally regressing.
The tests don't appear to cover it

https://github.com/narwhals-dev/narwhals/blob/4bb8c21bdf0d5e5eeb573d2c1e4ec8689d0893d2/tests/expr_and_series/when_test.py

The example code I've given was arbitrary, but the pattern of reusing parts of a when-then-otherwise is something I've done pretty frequently in altair:

https://altair-viz.github.io/gallery/dot_dash_plot.html

https://altair-viz.github.io/gallery/multiline_tooltip_standard.html

https://altair-viz.github.io/gallery/scatter_point_paths_hover.html

.otherwise(...)

If we returned a new object - I think that would be more consistent with the rest of the Compliant* API

2. function_name

Is this solely to align with the CompliantExpr._function_name requirement?

Every backend defines this, but I haven't found anywhere that utilizes the distinction between the two:

class SomeWhen: def then(self, value): return SomeThen(..., function_name="whenthen", ...) class SomeThen: def otherwise(self, value): ... self._function_name = "whenotherwise" return self

3. otherwise() -> CompliantExpr

I'm assuming the intent is to "erase" the fact that this returns CompliantThen?

At this stage we have a Callable[[Frame], Sequence[Series]].
So, we could go a step further and actually wrap it in the parent class?

Note
All of this will need to be rewritten either way, if anyone picks up chaining again (#669)

chore: Add CompliantExpr._with_metadata #2262

so we can get rid of tracking function_name at the compliant level

@MarcoGorelli we're getting a bit too in sync lol

sure, sounds good to returns a new object and avoid mutating

I've held off allowing chaining of when-then statements until the current implementation is better, so definitely happy to make steps in the right direction

sure, sounds good to returns a new object and avoid mutating

I've held off allowing chaining of when-then statements until the current implementation is better, so definitely happy to make steps in the right direction

@MarcoGorelli lovely

Should be quite easy to tweak after this PR 😎

- Skipping `__init__` to internally handle a generic `CompliantExpr` - Should mean we only need: - `CompliantWhen.__call__(...)` - Stub `CompliantThen` as a mixin with the main `*Expr` class - Return that type from `CompliantWhen._then`

dangotbanned · 2025-03-21T22:04:24Z

narwhals/_arrow/namespace.py

-        self._call._otherwise_value = value
-        self._function_name = "whenotherwise"
-        return self
+class ArrowThen(CompliantThen[ArrowDataFrame, ArrowSeries, ArrowExpr], ArrowExpr): ...


Before anyone asks - yes this is a fully working implementation 😄

https://github.com/narwhals-dev/narwhals/actions/runs/14001596439/job/39208960165?pr=2261

narwhals/_compliant/when_then.py

dangotbanned · 2025-03-22T10:44:23Z

narwhals/_compliant/when_then.py

+    def __call__(
+        self, compliant_frame: CompliantFrameT, /
+    ) -> Sequence[CompliantSeriesOrNativeExprT]: ...


TODO

Try to find common parts and split up their implementation

Ideally be able to define more as private methods here

Less work for subclasses

This will allow most of (Arrow|Pandas)Then.__call__ to be the same

refactor: Add EagerDataFrame._extract_comparand

Wasn't expecting I'd be able to reduce it down this much ...

feat: Implement EagerWhen 😎

Resolved from (cc5cce1) onwards

Note
Mentioned in (b0b1a05) that I'm tentative on LazyWhen

Happy to merge as-is, but its just the least-bad idea I had so far

narwhals/_compliant/when_then.py

Resolves #2261 (comment)

- Reserve `_when`, `_then` for types - Always use `_*_value` for expression fragments stored on an instance - Positional-only, non-prefixed args for the method/function they correspond to #2261 (comment)

- The annotations in `when_then` will be much easier to read by just defining the `TypeVar`(s) internally - Also means one less place to keep the type params in sync

This **always** gets me 🤦

First step towards (#2261 (comment))

#2261 (comment)

Using `Incomplete` for now, instead of adding more `TypeVar`s

- Need to sleep on this one - It works, but seems like quite a battle to unify them

https://github.com/narwhals-dev/narwhals/actions/runs/14012533923/job/39234108456?pr=2261

MarcoGorelli

wow, this is awesome!

thanks a lot @dangotbanned , well done here!

MarcoGorelli · 2025-03-23T16:06:07Z

narwhals/_arrow/namespace.py

-        self._call._otherwise_value = value
-        self._function_name = "whenotherwise"
-        return self
+class ArrowThen(CompliantThen[ArrowDataFrame, ArrowSeries, ArrowExpr], ArrowExpr): ...


dangotbanned · 2025-03-23T16:09:45Z

wow, this is awesome!

thanks a lot @dangotbanned , well done here!

Thanks @MarcoGorelli 😍

Wasn't sure if I'd gone off the deep end - but I'm really happy with how clean each of the backends are looking now

…hals into compliant-when

dangotbanned · 2025-03-23T16:12:03Z

narwhals/_duckdb/utils.py

-def maybe_evaluate_expr(
-    df: DuckDBLazyFrame, obj: DuckDBExpr | object
-) -> duckdb.Expression:
-    from narwhals._duckdb.expr import DuckDBExpr
-
-    if isinstance(obj, DuckDBExpr):
-        column_results = obj._call(df)
-        assert len(column_results) == 1  # debug assertion  # noqa: S101
-        return column_results[0]
-    return lit(obj)


Beat me to it @MarcoGorelli 😉

(a339a5e)

dangotbanned · 2025-03-23T16:18:10Z

narwhals/_compliant/when_then.py

+    @classmethod
+    def from_when(
+        cls,
+        when: CompliantWhen[FrameT, SeriesT, ExprT],
+        then: IntoExpr[SeriesT, ExprT],
+        /,
+    ) -> Self:
+        when._then_value = then
+        obj = cls.__new__(cls)
+        obj._call = when
+        obj._when_value = when
+        obj._depth = 0
+        obj._function_name = "whenthen"
+        obj._evaluate_output_names = getattr(
+            then, "_evaluate_output_names", lambda _df: ["literal"]
+        )
+        obj._alias_output_names = getattr(then, "_alias_output_names", None)
+        obj._implementation = when._implementation
+        obj._backend_version = when._backend_version
+        obj._version = when._version
+        obj._call_kwargs = {}
+        return obj


Important
Probably going to try something like this for _compliant.selectors.py

Will wait for (#2266) to land first though

feat(DRAFT): Adds Compliant(When|Then)

f6deab5

- Very barebones to start with - Mostly just defining the current impl in a generic way - There are some quirks I'd like to iron out before merging

dangotbanned mentioned this pull request Mar 20, 2025

Formalise Compliant* protocols #2230

Closed

15 tasks

dangotbanned added api design internal typing labels Mar 21, 2025

dangotbanned commented Mar 21, 2025

View reviewed changes

dangotbanned added 5 commits March 21, 2025 19:01

Merge remote-tracking branch 'upstream/main' into compliant-when

db5c25f

chore(DRAFT): Try out more complete CompliantThen

3486f55

- Skipping `__init__` to internally handle a generic `CompliantExpr` - Should mean we only need: - `CompliantWhen.__call__(...)` - Stub `CompliantThen` as a mixin with the main `*Expr` class - Return that type from `CompliantWhen._then`

fix(typing): make SeriesT invariant

f1689e7

fix(typing): Sneaky _when to avoid _call overloading

de3ccad

feat: Implement Arrow(Then|When)

8d671b2

dangotbanned commented Mar 21, 2025

View reviewed changes

dangotbanned added 4 commits March 21, 2025 22:08

fix: Older Protocol compat?

1ac7daa

https://github.com/narwhals-dev/narwhals/actions/runs/14001596439/job/39208960165?pr=2261

chore: imports/exports

9e7d848

feat: Implement Dask(Then|When)

ebbefa8

feat: Implement Pandas(When|Then)

6767c38

dangotbanned commented Mar 22, 2025

View reviewed changes

narwhals/_compliant/when_then.py Outdated Show resolved Hide resolved

dangotbanned commented Mar 22, 2025

View reviewed changes

narwhals/_compliant/when_then.py Outdated Show resolved Hide resolved

dangotbanned added 11 commits March 22, 2025 11:54

Merge remote-tracking branch 'upstream/main' into compliant-when

396605e

refactor: CompliantWhen.__init__ -> .from_expr

5a22284

Resolves #2261 (comment)

refactor: Naming conventions

c002fcf

- Reserve `_when`, `_then` for types - Always use `_*_value` for expression fragments stored on an instance - Positional-only, non-prefixed args for the method/function they correspond to #2261 (comment)

refactor(typing): Clean up, expose bound aliases

3b81fbe

- The annotations in `when_then` will be much easier to read by just defining the `TypeVar`(s) internally - Also means one less place to keep the type params in sync

Make annotations readable again

f2149e6

typo

32d2741

This **always** gets me 🤦

feat: Implement DuckDB(When|Then)

541f532

Merge remote-tracking branch 'upstream/main' into compliant-when

6183be7

feat: Implement SparkLike(When|Then)

724ff96

refactor: Add EagerDataFrame._extract_comparand

a9a6ccd

First step towards (#2261 (comment))

feat: Implement EagerWhen 😎

5adda67

#2261 (comment)

dangotbanned added 3 commits March 22, 2025 18:56

chore(typing): Partially annotate .when

9c9ea48

Using `Incomplete` for now, instead of adding more `TypeVar`s

refactor: Initial LazyWhen idea

cc5cce1

refactor(DRAFT): LazyWhen

b0b1a05

- Need to sleep on this one - It works, but seems like quite a battle to unify them

dangotbanned mentioned this pull request Mar 22, 2025

chore: Track ExpansionKind in ExprMetadata #2266

Merged

10 tasks

dangotbanned added 3 commits March 22, 2025 21:43

remove dead code

9c680dc

https://github.com/narwhals-dev/narwhals/actions/runs/14012533923/job/39234108456?pr=2261

refactor: remove from spark_like as well

ededf10

Merge remote-tracking branch 'upstream/main' into compliant-when

2e38646

dangotbanned changed the title ~~feat(DRAFT): Adds Compliant(When|Then)~~ feat: Adds Compliant(When|Then) Mar 22, 2025

dangotbanned marked this pull request as ready for review March 22, 2025 21:53

dangotbanned and others added 2 commits March 22, 2025 23:19

Merge branch 'main' into compliant-when

9c85a3f

Merge branch 'main' into compliant-when

ed2e7e8

MarcoGorelli approved these changes Mar 23, 2025

View reviewed changes

MarcoGorelli and others added 3 commits March 23, 2025 16:10

merge fixup

07f51da

revert: redo merge conflict from (9c680dc)

a339a5e

Merge branch 'compliant-when' of https://github.com/narwhals-dev/narw…

9632f3a

…hals into compliant-when

dangotbanned commented Mar 23, 2025

View reviewed changes

Update imports/exports

328bb43

dangotbanned commented Mar 23, 2025

View reviewed changes

dangotbanned merged commit 36b0e9a into main Mar 23, 2025
27 of 28 checks passed

dangotbanned deleted the compliant-when branch March 23, 2025 16:21

This was referenced Mar 23, 2025

chore(typing): Move DepthTracking* typing #2272

Merged

refactor(DRAFT): (Dask|DepthTracking)When #2274

Closed

chore: Add EagerSeries*Namespace protocols #2294

Merged

aivanoved mentioned this pull request Mar 27, 2025

feat: add when chaining #669

Closed

12 tasks

dangotbanned mentioned this pull request Apr 1, 2025

chore: Transition selectors away from Protocol38 #2326

Merged

10 tasks

	def otherwise(self: Self, value: ArrowExpr \| ArrowSeries \| _Scalar) -> ArrowExpr:
	self._call._otherwise_value = value
	self._function_name = "whenotherwise"
	return self

	def otherwise(self: Self, value: DuckDBExpr \| Any) -> DuckDBExpr:
	# type ignore because we are setting the `_call` attribute to a
	# callable object of type `DuckDBWhen`, base class has the attribute as
	# only a `Callable`
	self._call._otherwise_value = value # type: ignore[attr-defined]
	self._function_name = "whenotherwise"
	return self

	class When:
	def __init__(self: Self, *predicates: IntoExpr \| Iterable[IntoExpr]) -> None:
	self._predicate = all_horizontal(*flatten(predicates))
	check_expressions_preserve_length(self._predicate, function_name="when")

	def then(self: Self, value: IntoExpr \| Any) -> Then:
	return Then(
	lambda plx: apply_n_ary_operation(
	plx,
	lambda *args: plx.when(args[0]).then(args[1]),
	self._predicate,
	value,
	str_as_lit=False,
	),
	combine_metadata(self._predicate, value, str_as_lit=False),
	)


	class Then(Expr):
	def otherwise(self: Self, value: IntoExpr \| Any) -> Expr:
	kind = infer_kind(value, str_as_lit=False)

	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):
	compliant_value = compliant_value.broadcast(kind)
	return compliant_expr.otherwise(compliant_value) # type: ignore[attr-defined, no-any-return]

	return Expr(
	func,
	combine_metadata(self, value, str_as_lit=False),
	)

feat: Adds Compliant(When|Then) #2261

feat: Adds Compliant(When|Then) #2261

Uh oh!

Conversation

dangotbanned commented Mar 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What type of PR is this? (check all applicable)

Related issues

Checklist

Tasks

Uh oh!

Choose a reason for hiding this comment

1. Mutability

Suggestions

If we change nothing

.otherwise(...)

2. function_name

3. otherwise() -> CompliantExpr

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dangotbanned Mar 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dangotbanned Mar 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

TODO

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

MarcoGorelli left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dangotbanned commented Mar 23, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

feat: Adds `Compliant(When|Then)` #2261

feat: Adds `Compliant(When|Then)` #2261

dangotbanned commented Mar 21, 2025 •

edited

Loading

`.otherwise(...)`

2. `function_name`

3. `otherwise() -> CompliantExpr`

dangotbanned Mar 22, 2025 •

edited

Loading

dangotbanned Mar 21, 2025 •

edited

Loading