feat(RFC): A richer `Expr` IR #2572

dangotbanned · 2025-05-18T20:39:27Z

Will close #2571

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

Related issues

Closes [Enh]: A richer Expr internal representation #2571
See Related section there for more links

Checklist

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

If you have comments or can explain your changes, please do so below

Important

See (#2571) for detail!!!!!!!

Very open to feedback

Tasks

dangotbanned · 2025-05-19T16:12:54Z

`FunctionFlags`, `FunctionOptions`, `Function`, `FunctionExpr`

Feel like there's more I need to understand on the correct way to propagate the flags to the top-level.

@MarcoGorelli whenever you get to this - I've left loads of notes w/ references - hoping that you'd be able to demystify the rust magic 😄

All 4 classes

They're a bit trimmed down from the polars versions, where I couldn't see a need for us to mirror everything if we wouldn't use it

narwhals/narwhals/_plan/options.py

Lines 13 to 49 in 0bada48

    
           class FunctionFlags(enum.Flag): 
        
               ALLOW_GROUP_AWARE = 1 << 0 
        
               """> Raise if use in group by 
        
               Not sure where this is disabled. 
        
               """ 
        
               INPUT_WILDCARD_EXPANSION = 1 << 4 
        
               """Appears on all the horizontal aggs. 
        
               https://github.com/pola-rs/polars/blob/e8ad1059721410e65a3d5c1d84055fb22a4d6d43/crates/polars-plan/src/plans/options.rs#L49-L58 
        
               """ 
        
               RETURNS_SCALAR = 1 << 5 
        
               """Automatically explode on unit length if it ran as final aggregation.""" 
        
               ROW_SEPARABLE = 1 << 8 
        
               """Not sure lol. 
        
               https://github.com/pola-rs/polars/pull/22573 
        
               """ 
        
               LENGTH_PRESERVING = 1 << 9 
        
               """mutually exclusive with `RETURNS_SCALAR`""" 
        
               def is_elementwise(self) -> bool: 
        
                   return self in (FunctionFlags.ROW_SEPARABLE | FunctionFlags.LENGTH_PRESERVING) 
        
               def returns_scalar(self) -> bool: 
        
                   return self in FunctionFlags.RETURNS_SCALAR 
        
               def is_length_preserving(self) -> bool: 
        
                   return self in FunctionFlags.LENGTH_PRESERVING 
        
               @staticmethod 
        
               def default() -> FunctionFlags: 
        
                   return FunctionFlags.ALLOW_GROUP_AWARE

narwhals/narwhals/_plan/options.py

Lines 52 to 108 in 0bada48

    
           class FunctionOptions(Immutable): 
        
               """ExprMetadata` but less god object. 
        
               https://github.com/pola-rs/polars/blob/3fd7ecc5f9de95f62b70ea718e7e5dbf951b6d1c/crates/polars-plan/src/plans/options.rs 
        
               """ 
        
               __slots__ = ("flags",) 
        
               flags: FunctionFlags 
        
               def is_elementwise(self) -> bool: 
        
                   return self.flags.is_elementwise() 
        
               def returns_scalar(self) -> bool: 
        
                   return self.flags.returns_scalar() 
        
               def is_length_preserving(self) -> bool: 
        
                   return self.flags.is_length_preserving() 
        
               def with_flags(self, flags: FunctionFlags, /) -> FunctionOptions: 
        
                   if (FunctionFlags.RETURNS_SCALAR | FunctionFlags.LENGTH_PRESERVING) in flags: 
        
                       msg = "A function cannot both return a scalar and preserve length, they are mutually exclusive." 
        
                       raise TypeError(msg) 
        
                   obj = FunctionOptions.__new__(FunctionOptions) 
        
                   object.__setattr__(obj, "flags", self.flags | flags) 
        
                   return obj 
        
               def with_elementwise(self) -> FunctionOptions: 
        
                   return self.with_flags( 
        
                       FunctionFlags.ROW_SEPARABLE | FunctionFlags.LENGTH_PRESERVING 
        
                   ) 
        
               @staticmethod 
        
               def default() -> FunctionOptions: 
        
                   obj = FunctionOptions.__new__(FunctionOptions) 
        
                   object.__setattr__(obj, "flags", FunctionFlags.default()) 
        
                   return obj 
        
               @staticmethod 
        
               def elementwise() -> FunctionOptions: 
        
                   return FunctionOptions.default().with_elementwise() 
        
               @staticmethod 
        
               def row_separable() -> FunctionOptions: 
        
                   return FunctionOptions.groupwise().with_flags(FunctionFlags.ROW_SEPARABLE) 
        
               @staticmethod 
        
               def length_preserving() -> FunctionOptions: 
        
                   return FunctionOptions.default().with_flags(FunctionFlags.LENGTH_PRESERVING) 
        
               @staticmethod 
        
               def groupwise() -> FunctionOptions: 
        
                   return FunctionOptions.default() 
        
               @staticmethod 
        
               def aggregation() -> FunctionOptions: 
        
                   return FunctionOptions.groupwise().with_flags(FunctionFlags.RETURNS_SCALAR)

narwhals/narwhals/_plan/common.py

Lines 149 to 172 in 0bada48

    
           class Function(ExprIR): 
        
               """Shared by expr functions and namespace functions. 
        
               https://github.com/pola-rs/polars/blob/112cab39380d8bdb82c6b76b31aca9b58c98fd93/crates/polars-plan/src/dsl/expr.rs#L114 
        
               """ 
        
               @property 
        
               def function_options(self) -> FunctionOptions: 
        
                   from narwhals._plan.options import FunctionOptions 
        
                   return FunctionOptions.default() 
        
               @property 
        
               def is_scalar(self) -> bool: 
        
                   return self.function_options.returns_scalar() 
        
               def to_function_expr(self, *inputs: ExprIR) -> FunctionExpr[Self]: 
        
                   from narwhals._plan.expr import FunctionExpr 
        
                   from narwhals._plan.options import FunctionOptions 
        
                   # NOTE: Still need to figure out how these should be generated 
        
                   # Feel like it should be the union of `input` & `function` 
        
                   PLACEHOLDER = FunctionOptions.default()  # noqa: N806 
        
                   return FunctionExpr(input=inputs, function=self, options=PLACEHOLDER)

narwhals/narwhals/_plan/expr.py

Lines 157 to 185 in 0bada48

    
           class FunctionExpr(ExprIR, t.Generic[_FunctionT]): 
        
               """**Representing `Expr::Function`**. 
        
               https://github.com/pola-rs/polars/blob/dafd0a2d0e32b52bcfa4273bffdd6071a0d5977a/crates/polars-plan/src/dsl/expr.rs#L114-L120 
        
               https://github.com/pola-rs/polars/blob/112cab39380d8bdb82c6b76b31aca9b58c98fd93/crates/polars-plan/src/dsl/function_expr/mod.rs#L123 
        
               """ 
        
               __slots__ = ("function", "input", "options") 
        
               input: Seq[ExprIR] 
        
               function: _FunctionT 
        
               """Enum type is named `FunctionExpr` in `polars`. 
        
               Mirroring *exactly* doesn't make much sense in OOP. 
        
               https://github.com/pola-rs/polars/blob/112cab39380d8bdb82c6b76b31aca9b58c98fd93/crates/polars-plan/src/dsl/function_expr/mod.rs#L123 
        
               """ 
        
               options: FunctionOptions 
        
               """Assuming this is **either**: 
        
               1. `function.function_options` 
        
               2. The union of (1) and any `FunctionOptions` in `inputs` 
        
               """ 
        
               def with_options(self, options: FunctionOptions, /) -> Self: 
        
                   options = self.options.with_flags(options.flags) 
        
                   return type(self)(input=self.input, function=self.function, options=options)

- Mentioned in (#2391 (comment)) - Needed again for #2572

narwhals/_plan/common.py

MarcoGorelli · 2025-05-21T08:58:53Z

@MarcoGorelli whenever you get to this - I've left loads of notes w/ references - hoping that you'd be able to demystify the rust magic 😄

at the moment it looks like this adds a self-standing _plan, that's not integrated into the rest of Narwhals? if it's possible to integrate it with the rest to prove that this is feasible, i'd be very interested in taking a close look

* chore(typing): Add `_typing_compat.py` - Mentioned in (#2391 (comment)) - Needed again for #2572 * refactor: Reuse `TypeVar` import * refactor: Reuse `@deprecated` import * refactor: Reuse `Protocol38` import * docs: Add module-level docstring

Still need: - reprs - fix the hierarchy issue (#2572 (comment)) - Flag summing (#2572 (comment))

Resolves (#2572 (comment))

- 1 step closer to the understanding for (#2572 (comment)) - There's still some magic going on when `polars` serializes - Need to track down where `'collect_groups': 'ElementWise'` and `'collect_groups': 'GroupWise'` first appear - Seems like the flags get reduced

dangotbanned · 2025-05-21T21:55:31Z

narwhals/_plan/functions.py

+    @property
+    def function_options(self) -> FunctionOptions:
+        return FunctionOptions.length_preserving()


Seems like FunctionOptions.length_preserving is one we need to pay more attention to

dangotbanned · 2025-05-22T12:40:02Z

Thanks for peeking @MarcoGorelli

at the moment it looks like this adds a self-standing _plan, that's not integrated into the rest of Narwhals? if it's possible to integrate it with the rest to prove that this is feasible, i'd be very interested in taking a close look

That is definitely the eventual goal! 🤞

Despite how quickly things have progressed, I still feel I'm a few steps behind being ready for that just yet.

General overview

I'm trying to focus on modeling these structures and how they interact:

My thought was that narwhals currently solves similar problems, but in different ways to polars.
By getting this subset of polars that we need translated from rust -> python first, we're then in a good position to make decisions on how to bridge any gaps that narwhals is occupying at the moment, if that makes sense?

So like what I have in narwhals/_plan/dummy.py is all about creating an accurate expression graph.
Consuming the graph (evaluating an expression) is a very important step - but the main difference I'm anticipating is:

Current

Expressions are based on evaluating callables
with other callables used for output names & aliases
They also optionally store a ExprMetadata object
Some backends do other things with function names, kwargs, depth
Lazy backends use more callables for handling windows

`ExprIR`

Expressions are (likely still) based on evaluating callables
All of the remaining details are either:
- Encoded into a node on the graph
  - E.g. the nodes for Expr.var, Expr.std have the name and ddof already
  - Depth can be computed anywhere by traversing the graph
- Or, can be expressed in terms of operating on a node
  - meta namespace, examples
  - WindowExpr is just another albeit complex 😅 node

I am confident we'll end up with something that's easier to maintain - but trying to integrate the two mid-solve and maintaining that branch over time seems like it'd be a real challenge 😔

`FunctionOptions`, `FunctionFlags`

The parts mentioned in (#2572 (comment)) ~~are some~~ were one of the main hurdles left.
However I think it is really just a skill issue on my part in understanding the rust code. That was all I was hoping for some help with 🙏
In narwhals terms, it is closest to (https://github.com/narwhals-dev/narwhals/blob/1b93c0ed2dc2bf47d7b8e4b4cab4c9e9cad59800/narwhals/_expression_parsing.py) but only a subset of the rules - since it only concerns functions and how they compose.

Note

Right before I was about to send this comment, I managed to fix the issue I had by updating to 1.30.0 🤦‍♂️
See (diff), which removed the flags I was having trouble finding in (0982b3a)

Example

Now all the flags I've been using are propagated in the same way as in polars! 🎉

Repro code

import polars as pl

from narwhals._plan import demo as nwd  # noqa: F811
from narwhals._plan import meta  # noqa: F811

expr_pl = (
    pl.col("a")
    .sort()
    .fill_null(1)
    .shift(1)
    .abs()
    .drop_nulls()
    .skew()
    .alias("col->sort->fill_null->shift->abs->drop_nulls->skew->alias")
)
expr_nwd = (
    nwd.col("a")
    .sort()
    .fill_null(1)
    .shift(1)
    .abs()
    .drop_nulls()
    .skew()
    .alias("col->sort->fill_null->shift->abs->drop_nulls->skew->alias")
)
roundtrip_pl = meta.polars_expr_to_dict(expr_pl)
roundtrip_nw = str(expr_nwd._ir)

roundtrip_pl

I haven't added ALLOW_EMPTY_INPUTS - so that's an expected difference between the two

>>> roundtrip_pl

{'Alias': [{'Function': {'input': [{'Function': {'input': [{'Function': {'input': [{'Function': {'input': [{'Function': {'input': [{'Sort': {'expr': {'Column': 'a'},
                   'options': {'descending': False,
                    'nulls_last': False,
                    'multithreaded': True,
                    'maintain_order': False,
                    'limit': None}}},
                 {'Literal': {'Dyn': {'Int': 1}}}],
                'function': 'FillNull',
                'options': {'check_lengths': True,
                 'flags': 'ALLOW_GROUP_AWARE | ROW_SEPARABLE | LENGTH_PRESERVING'}}},
              {'Literal': {'Dyn': {'Int': 1}}}],
             'function': 'Shift',
             'options': {'check_lengths': True,
              'flags': 'ALLOW_GROUP_AWARE | LENGTH_PRESERVING'}}}],
          'function': 'Abs',
          'options': {'check_lengths': True,
           'flags': 'ALLOW_GROUP_AWARE | ROW_SEPARABLE | LENGTH_PRESERVING'}}}],
       'function': 'DropNulls',
       'options': {'check_lengths': True,
        'flags': 'ALLOW_GROUP_AWARE | ALLOW_EMPTY_INPUTS | ROW_SEPARABLE'}}}],
    'function': {'Skew': True},
    'options': {'check_lengths': True,
     'flags': 'ALLOW_GROUP_AWARE | RETURNS_SCALAR'}}},
  'col->sort->fill_null->shift->abs->drop_nulls->skew->alias']}

roundtrip_nw

To produce this I removed the outer ", " and pasted back to ruff to format

The overall shape is very similar and the deviations from polars have been documented

Alias(
    expr=FunctionExpr(
        function=Skew(),
        input=[
            FunctionExpr(
                function=DropNulls(),
                input=[
                    FunctionExpr(
                        function=Abs(),
                        input=[
                            FunctionExpr(
                                function=Shift(n=1),
                                input=[
                                    FunctionExpr(
                                        function=FillNull(),
                                        input=[
                                            Sort(
                                                expr=Column(name="a"),
                                                options=SortOptions(
                                                    descending=False, nulls_last=False
                                                ),
                                            ),
                                            Literal(
                                                value=ScalarLiteral(
                                                    dtype=Unknown, value=1
                                                )
                                            ),
                                        ],
                                        options=FunctionOptions(
                                            flags="ALLOW_GROUP_AWARE | ROW_SEPARABLE | LENGTH_PRESERVING"
                                        ),
                                    )
                                ],
                                options=FunctionOptions(
                                    flags="ALLOW_GROUP_AWARE | LENGTH_PRESERVING"
                                ),
                            )
                        ],
                        options=FunctionOptions(
                            flags="ALLOW_GROUP_AWARE | ROW_SEPARABLE | LENGTH_PRESERVING"
                        ),
                    )
                ],
                options=FunctionOptions(flags="ALLOW_GROUP_AWARE | ROW_SEPARABLE"),
            )
        ],
        options=FunctionOptions(flags="ALLOW_GROUP_AWARE | RETURNS_SCALAR"),
    ),
    name="col->sort->fill_null->shift->abs->drop_nulls->skew->alias",
)

Can't tell if this means `FirstT` will match the entry `firstt`, but preserve the `firstt` fix (https://github.com/codespell-project/codespell#ignoring-words) (#2572 (comment))

.pre-commit-config.yaml

#2572 (comment)

dangotbanned · 2025-05-23T14:23:40Z

I should've expected this, but it was a nice suprise to find we get hashable selectors for free 😄

from narwhals._plan import selectors as ndcs

>>> ndcs.matches("[^z]a")._ir == ndcs.matches("[^z]a")._ir
True

>>> ndcs.matches("[^z]a")._ir == ndcs.matches("abc")._ir
False

@MarcoGorelli regarding (#2291)

from narwhals._plan import selectors as ndcs

>>> ndcs.all()._ir == ndcs.all()._ir
True

lhs = ndcs.all()
rhs = ndcs.all().mean()

>>> lhs._ir == rhs._ir
False

>>> lhs._ir == rhs._ir.expr
True

And the same holds for the non-selectors all 🥳

from narwhals._plan import demo as nwd

lhs = nwd.all()
rhs = nwd.all().mean()
>>> lhs._ir == rhs._ir
False

>>> lhs._ir == rhs._ir.expr
True

>>> type(rhs._ir)
narwhals._plan.aggregation.Mean

pyproject.toml

Related #2572 (comment)

An experiment towards (#2572 (comment))

dangotbanned · 2025-05-28T14:59:22Z

tests/plan/expr_parsing_test.py

+def test_valid_windows() -> None:
+    """Was planning to test this matched, but we seem to allow elementwise horizontal?
+
+    https://github.com/narwhals-dev/narwhals/blob/63c8e4771a1df4e0bfeea5559c303a4a447d5cc2/tests/expression_parsing_test.py#L10-L45
+    """
+    ELEMENTWISE_ERR = re.compile(r"cannot use.+over.+elementwise", re.IGNORECASE)  # noqa: N806
+    a = nwd.col("a")
+    assert a.cum_sum()
+    assert a.cum_sum().over(order_by="id")
+    with pytest.raises(InvalidOperationError, match=ELEMENTWISE_ERR):
+        assert a.cum_sum().abs().over(order_by="id")
+
+    assert (a.cum_sum() + 1).over(order_by="id")
+    assert a.cum_sum().cum_sum().over(order_by="id")
+    assert a.cum_sum().cum_sum()
+    assert nwd.sum_horizontal(a, a.cum_sum())
+    with pytest.raises(InvalidOperationError, match=ELEMENTWISE_ERR):
+        assert nwd.sum_horizontal(a, a.cum_sum()).over(order_by="a")
+
+    assert nwd.sum_horizontal(a, a.cum_sum().over(order_by="i"))
+    assert nwd.sum_horizontal(a.diff(), a.cum_sum().over(order_by="i"))
+    with pytest.raises(InvalidOperationError, match=ELEMENTWISE_ERR):
+        assert nwd.sum_horizontal(a.diff(), a.cum_sum()).over(order_by="i")
+
+    with pytest.raises(InvalidOperationError, match=ELEMENTWISE_ERR):
+        assert nwd.sum_horizontal(a.diff().abs(), a.cum_sum()).over(order_by="i")


@MarcoGorelli quick question

This is adapted from an existing test:

tests.expression_parsing_test.test_window_kind

narwhals/tests/expression_parsing_test.py

Lines 10 to 45 in 63c8e47

@pytest.mark.parametrize(

("expr", "expected"),

[

(nw.col("a"), 0),

(nw.col("a").mean(), 0),

(nw.col("a").cum_sum(), 1),

(nw.col("a").cum_sum().over(order_by="id"), 0),

(nw.col("a").cum_sum().abs().over(order_by="id"), 1),

((nw.col("a").cum_sum() + 1).over(order_by="id"), 1),

(nw.col("a").cum_sum().cum_sum().over(order_by="id"), 1),

(nw.col("a").cum_sum().cum_sum(), 2),

(nw.sum_horizontal(nw.col("a"), nw.col("a").cum_sum()), 1),

(nw.sum_horizontal(nw.col("a"), nw.col("a").cum_sum()).over(order_by="a"), 1),

(nw.sum_horizontal(nw.col("a"), nw.col("a").cum_sum().over(order_by="i")), 0),

(

nw.sum_horizontal(

nw.col("a").diff(), nw.col("a").cum_sum().over(order_by="i")

),

1,

),

(

nw.sum_horizontal(nw.col("a").diff(), nw.col("a").cum_sum()).over(

order_by="i"

),

2,

),

(

nw.sum_horizontal(nw.col("a").diff().abs(), nw.col("a").cum_sum()).over(

order_by="i"

),

2,

),

],

)

def test_window_kind(expr: nw.Expr, expected: int) -> None:

assert expr._metadata.n_orderable_ops == expected

AFAICT, all of the expressions I've needed a InvalidOperationError for shouldn't be valid.

But they aren't raising in current narwhals 🤔

1

import narwhals as nw a = nw.col("a") a.cum_sum().abs().over(order_by="id")

This error explicitly mentions abs

narwhals/narwhals/_expression_parsing.py

Lines 357 to 362 in 9bd10ad

if self.is_elementwise or self.is_filtration:

msg = (

"Cannot use `over` on expressions which are elementwise\n"

"(e.g. `abs`) or which change length (e.g. `drop_nulls`)."

)

raise InvalidOperationError(msg)

2, 3, 4

These are all raising the same as (1), but the issue seems to be that horizontal functions aren't being treated as elementwise

import narwhals as nw a = nw.col("a") nw.sum_horizontal(a, a.cum_sum()).over(order_by="a") nw.sum_horizontal(a.diff(), a.cum_sum()).over(order_by="i") nw.sum_horizontal(a.diff().abs(), a.cum_sum()).over(order_by="i")

In polars, they all seem to be elementwise but with an additional flag

https://github.com/pola-rs/polars/blob/944bf553f1111c31259b55348a7dd0a512ae51a1/crates/polars-plan/src/dsl/function_expr/mod.rs#L1388-L1392

I've done the same in this PR, but I don't think that flag would factor into this?

narwhals/narwhals/_plan/functions.py

Lines 291 to 299 in 9bd10ad

class SumHorizontal(Function):

@property

def function_options(self) -> FunctionOptions:

return FunctionOptions.elementwise().with_flags(

FunctionFlags.INPUT_WILDCARD_EXPANSION

)

def __repr__(self) -> str:

return "sum_horizontal"

narwhals/_namespace.py

tests/plan/expr_parsing_test.py

Related #2528

#2944

Child of #2572

dangotbanned mentioned this pull request May 18, 2025

[Enh]: A richer Expr internal representation #2571

Open

7 tasks

dangotbanned added enhancement New feature or request api design needs discussion labels May 18, 2025

dangotbanned added a commit that referenced this pull request May 20, 2025

chore(typing): Add _typing_compat.py

603c4e9

- Mentioned in (#2391 (comment)) - Needed again for #2572

dangotbanned mentioned this pull request May 20, 2025

chore: Add _typing_compat module #2578

Merged

10 tasks

dangotbanned commented May 21, 2025

View reviewed changes

narwhals/_plan/common.py Outdated Show resolved Hide resolved

dangotbanned added a commit that referenced this pull request May 21, 2025

feat: Add methods for all of functions

be6ec3e

Still need: - reprs - fix the hierarchy issue (#2572 (comment)) - Flag summing (#2572 (comment))

dangotbanned added a commit that referenced this pull request May 21, 2025

fix: Remove ExprIR from Function bases

e6ca72b

Resolves (#2572 (comment))

dangotbanned commented May 21, 2025

View reviewed changes

dangotbanned mentioned this pull request May 22, 2025

Add Selector.__xor__ (symmetric difference) #2589

Open

This comment was marked as resolved.

Sign in to view

dangotbanned added a commit that referenced this pull request May 23, 2025

ci: Maybe codespell ignore FirstT

fb303f6

Can't tell if this means `FirstT` will match the entry `firstt`, but preserve the `firstt` fix (https://github.com/codespell-project/codespell#ignoring-words) (#2572 (comment))

MarcoGorelli reviewed May 23, 2025

View reviewed changes

.pre-commit-config.yaml Outdated Show resolved Hide resolved

dangotbanned added a commit that referenced this pull request May 23, 2025

ci: fix codespell

89c4587

#2572 (comment)

dangotbanned mentioned this pull request May 23, 2025

enh: improve when-then-otherwise to include chaining #668

Open

dangotbanned commented May 23, 2025

View reviewed changes

pyproject.toml Outdated Show resolved Hide resolved

dangotbanned added a commit that referenced this pull request May 24, 2025

chore(ruff): Add tool.ruff.lint.typing-modules

7c4c6ee

Related #2572 (comment)

dangotbanned mentioned this pull request May 24, 2025

chore(ruff): Add tool.ruff.lint.typing-modules #2610

Merged

10 tasks

dangotbanned added a commit that referenced this pull request May 26, 2025

feat(DRAFT): Mock up to_compliant as an adapter

92694ce

An experiment towards (#2572 (comment))

dangotbanned commented May 28, 2025

View reviewed changes

narwhals/_namespace.py Outdated Show resolved Hide resolved

dangotbanned mentioned this pull request May 31, 2025

Tracking: Drop support for python 3.8 #2084

Closed

5 tasks

dangotbanned commented May 31, 2025

View reviewed changes

tests/plan/expr_parsing_test.py Outdated Show resolved Hide resolved

dangotbanned added 8 commits July 30, 2025 14:53

add _with_columns

2c85701

Merge remote-tracking branch 'upstream/main' into oh-nodes

4de8579

perf: Add early return path for over_ordered

94d533d

Related #2528

Merge branch 'main' into oh-nodes

c893b2f

Merge remote-tracking branch 'upstream/main' into oh-nodes

b03d31f

style(ruff): re-run updated config

25549de

Merge remote-tracking branch 'upstream/main' into oh-nodes

3ec5f7b

chore(ruff): re-run updated config x2

a12a985

dangotbanned mentioned this pull request Aug 19, 2025

refactor(expr-ir): Rename Dummy* everything #3014

Merged

7 tasks

refactor(expr-ir): Rename Dummy* everything (#3014)

62030ae

dangotbanned mentioned this pull request Aug 20, 2025

refactor(expr-ir): Heavily sugar Function defs #3017

Merged

10 tasks

dangotbanned added 4 commits August 21, 2025 09:26

refactor(expr-ir): Heavily sugar Function defs (#3017)

25b744d

Merge remote-tracking branch 'upstream/main' into oh-nodes

fcfcabf

refactor(typing): Align NativeDataFrame

de884c5

#2944

Merge remote-tracking branch 'upstream/main' into oh-nodes

dbe2a45

dangotbanned mentioned this pull request Aug 25, 2025

refactor(expr-ir): Add more builder sugar #3040

Merged

10 tasks

dangotbanned added 2 commits August 25, 2025 22:15

refactor(expr-ir): Add more builder sugar (#3040)

f5ecc6e

Merge remote-tracking branch 'upstream/main' into oh-nodes

69eede5

dangotbanned mentioned this pull request Aug 28, 2025

refactor(expr-ir): Clearing out the cobwebs #3053

Merged

10 tasks

dangotbanned added 2 commits August 29, 2025 21:30

refactor(expr-ir): Clearing out the cobwebs (#3053)

c934d25

Merge branch 'main' into oh-nodes

dc3b9ea

dangotbanned mentioned this pull request Aug 30, 2025

Untangling CompliantLazyFrame._iter_columns #3058

Open

Merge remote-tracking branch 'upstream/main' into oh-nodes

607ea5a

dangotbanned added a commit that referenced this pull request Aug 30, 2025

refactor: copy.replace most with_* impls

1bf8bba

Child of #2572

dangotbanned mentioned this pull request Aug 30, 2025

refactor(expr-ir): copy.replace most with_* methods #3063

Merged

10 tasks

refactor(expr-ir): copy.replace most with_* methods (#3063)

631d3a3

This was referenced Aug 31, 2025

refactor(expr-ir): Deduplicate ExprIR.iter_{left,right} #3065

Merged

refactor(expr-ir): Shrinking ExprIR (main) #3066

Draft

dangotbanned added 2 commits September 1, 2025 10:41

Merge branch 'main' into oh-nodes

192de64

Merge remote-tracking branch 'upstream/main' into oh-nodes

215386f

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat(RFC): A richer `Expr` IR #2572

feat(RFC): A richer `Expr` IR #2572

dangotbanned commented May 18, 2025 •

edited

Loading

Uh oh!

dangotbanned commented May 19, 2025 •

edited

Loading

Uh oh!

Uh oh!

MarcoGorelli commented May 21, 2025

Uh oh!

dangotbanned May 21, 2025

Uh oh!

dangotbanned commented May 22, 2025

Uh oh!

This comment was marked as resolved.

Uh oh!

dangotbanned commented May 23, 2025

Uh oh!

Uh oh!

dangotbanned May 28, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

	class Selector(ExprIR):
	"""by_dtype, matches, numeric, boolean, string, categorical, datetime, all."""

	@pytest.mark.parametrize(
	("expr", "expected"),
	[
	(nw.col("a"), 0),
	(nw.col("a").mean(), 0),
	(nw.col("a").cum_sum(), 1),
	(nw.col("a").cum_sum().over(order_by="id"), 0),
	(nw.col("a").cum_sum().abs().over(order_by="id"), 1),
	((nw.col("a").cum_sum() + 1).over(order_by="id"), 1),
	(nw.col("a").cum_sum().cum_sum().over(order_by="id"), 1),
	(nw.col("a").cum_sum().cum_sum(), 2),
	(nw.sum_horizontal(nw.col("a"), nw.col("a").cum_sum()), 1),
	(nw.sum_horizontal(nw.col("a"), nw.col("a").cum_sum()).over(order_by="a"), 1),
	(nw.sum_horizontal(nw.col("a"), nw.col("a").cum_sum().over(order_by="i")), 0),
	(
	nw.sum_horizontal(
	nw.col("a").diff(), nw.col("a").cum_sum().over(order_by="i")
	),
	1,
	),
	(
	nw.sum_horizontal(nw.col("a").diff(), nw.col("a").cum_sum()).over(
	order_by="i"
	),
	2,
	),
	(
	nw.sum_horizontal(nw.col("a").diff().abs(), nw.col("a").cum_sum()).over(
	order_by="i"
	),
	2,
	),
	],
	)
	def test_window_kind(expr: nw.Expr, expected: int) -> None:
	assert expr._metadata.n_orderable_ops == expected

	if self.is_elementwise or self.is_filtration:
	msg = (
	"Cannot use `over` on expressions which are elementwise\n"
	"(e.g. `abs`) or which change length (e.g. `drop_nulls`)."
	)
	raise InvalidOperationError(msg)

	class SumHorizontal(Function):
	@property
	def function_options(self) -> FunctionOptions:
	return FunctionOptions.elementwise().with_flags(
	FunctionFlags.INPUT_WILDCARD_EXPANSION
	)

	def __repr__(self) -> str:
	return "sum_horizontal"

feat(RFC): A richer Expr IR #2572

Are you sure you want to change the base?

feat(RFC): A richer Expr IR #2572

Conversation

dangotbanned commented May 18, 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

If you have comments or can explain your changes, please do so below

Tasks

Uh oh!

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

Uh oh!

FunctionFlags, FunctionOptions, Function, FunctionExpr

Uh oh!

Uh oh!

MarcoGorelli commented May 21, 2025

Uh oh!

dangotbanned May 21, 2025

Choose a reason for hiding this comment

Uh oh!

dangotbanned commented May 22, 2025

General overview

Current

ExprIR

FunctionOptions, FunctionFlags

Example

Uh oh!

This comment was marked as resolved.

Uh oh!

dangotbanned commented May 23, 2025

Uh oh!

Uh oh!

dangotbanned May 28, 2025

Choose a reason for hiding this comment

1

2, 3, 4

Uh oh!

Uh oh!

Uh oh!

Uh oh!

feat(RFC): A richer `Expr` IR #2572

feat(RFC): A richer `Expr` IR #2572

dangotbanned commented May 18, 2025 •

edited

Loading

dangotbanned commented May 19, 2025 •

edited

Loading

`FunctionFlags`, `FunctionOptions`, `Function`, `FunctionExpr`

`ExprIR`

`FunctionOptions`, `FunctionFlags`