feat(RFC): A richer Expr IR

[dangotbanned]

Will close #2571

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

• ✨ Feature

Related issues

• Closes [Enh]: A richer Expr internal representation #2571
  • See Related section there for more links
• Search PRs: (sort:updated-desc "(expr-ir)/" in:title)

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

Show 2025 May-July

• Investigate the sources for pl.Expr.meta
  • Reproducing how some of those methods work would be huge
  • See (exploring pl.Expr.meta)
• Implement iteration methods for ExprIR
  • Needed for meta methods
  • feat: Add Expr graph iteration methods
• chore: Add _typing_compat module #2578
  • Merge another PR with (perf: Avoid module-level importlib.util.find_spec #2391 (comment)) first
  • Wanting to use TypeVar defaults more
  • A lot of places should have TypeVar("T", bound=Thing, default=Thing), instead of an opaque ExprIR
• Selector(s)
  

  narwhals/narwhals/_plan/expr.py

  Lines 336 to 337 in 0bada48

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

  • Add a variant of BinaryExpr that describes the restricted set of operators that are allowed
    • Some WIP-rogress (87b8402)
  • Very close (b9eb25e)
  • (feat: Add selectors reprs)
    • Deviating a bit from polars, since they wrap pl.col internally
• Accept IntoExpr in more places (including and beyond what narwhals allows now)
  • demo.py
    • *_horizontal
    • (feat: Expressify concat_str)
  • dummy.py
    • (feat: Expressify over, sort_by)
    • Binary operators (af274d3)
• Flesh out (and properly understand) FunctionOptions + friends see comment
  • Top-level (0982b3a)
  • Where does the {flags: ...} -> {collect_groups: ..., flags: ...} expansion happen?
  • polars>=1.3.0 fixed the issue (see comment)
• Ternary
  • feat: Implement chained when-then-otherwise 🥳)
  • Related (enh: improve when-then-otherwise to include chaining #668 (comment))
• Namespaces
  • Meta
    • Support is_*, has_*, output_name meta methods
      • (882eff0)
      • (6d9b9e1)
      • (4f1ea19)
    • root_names
    • undo_aliases, pop
      • Too complex
  • Name
    • Redo name.py (Expr::KeepName, Expr::RenameAlias)
      • Changing to how they're represented in polars will help with the meta methods
      • See comments
      • Resolved in (4a8c2d5)
    • (737c83a)
  • Cat, Struct, List (a3e29d1)
  • String (72c33ce)
  • DateTime (aee0a7e)
• Port over various _expression_parsing.py rules
• Output name expansion
  • Review how the rust version works
  • ExpansionFlags
  • Pending upstream refactor
  • expand_function_inputs
    • rewrite_projections
      • replace_selector
        • expand_selector
        • replace_selector_inner
      • replace_and_add_to_results
        • replace_nth
        • prepare_excluded
        • expand_columns
        • expand_dtypes
          • replace_dtype_or_index_with_column
          • dtypes_match (probably can solve w/ existing narwhals)
        • expand_indices
          • replace_index_with_column
        • replace_wildcard
          • rewrite_special_aliases
          • replace_wildcard_with_column
        • replace_regex
          • expand_regex
• Vistor pattern/rewrites/ExprIR.map_ir
  • See (feat(RFC): A richer Expr IR #2572 (comment))
  • feat: Impl ExprIR.map_ir for most nodes
  • feat: Impl WindowExpr.map_ir
  • feat: Impl FunctionExpr.map_ir
    • RollingExpr, AnonymousExpr inherit
• Add tests for selectors
  • See (https://github.com/narwhals-dev/narwhals/blob/6414142765b8645f57888db71c7773145162e034/tests/plan/expr_expansion_test.py)

• refactor(expr-ir): Shrinking ExprIR (main) #3066
  • Leftover 2025 Aug-Sep refactor ideas (refactor(expr-ir): Shrinking ExprIR (main) #3066 (comment))
• refactor(expr-ir): Organize _plan package #3122
• feat(expr-ir): Support group_by, utilize pyarrow.acero #3143
  • Had left this pending (feat: Adds {Expr,Series}.{first,last} #2528)
• refactor(expr-ir): Split up and refine protocols.py #3166
• feat(expr-ir): Acero order_by, hashjoin , DataFrame.{filter,join}, Expr.is_{first,last}_distinct #3173
• Fix __dict__ appearing on Immutable subclasses (thread)
  • fix(expr-ir): Ensure only __slots__, and not __dict__ too #3201
• perf(expr-ir): Optimize ExpansionFlags.from_ir #3206
• refactor(expr-ir): Improve function dispatch #3215
• Review where ColumnNameOrSelector can be added
• Support Expr.meta.serialize, Expr.deserialize
  • Briefly played around in (feat(expr-ir): Start sketching out (de)serde)
  • Definitely made easier by using such a restricted set of types
• (low prio) Add optimistic-lazy predicate parser into pc.Expression
  • Mentioned in (thread)
  • Similar to https://github.com/apache/datafusion-python/blob/e75addfa64a91c0d91ef059917a451e04495b83a/src/pyarrow_filter_expression.rs#L96-L181
• feat(expr-ir): Support over(*partition_by) #3224
  • #3173 needed to build out more supporting features first
• feat(expr-ir): Impl rank, with_row_index_by, over(*partition_by, order_by=...) #3295
• ✨ feat(expr-ir): The big Selector overhaul #3233
• chore(expr-ir): Improve coverage #3265
• feat(expr-ir): Add date_range, support {date,int}_range(eager=...) #3294
• ✨ feat(expr-ir): Finish* implementing ArrowExpr #3325
• Add namespaced methods
  • Excluding dt.*, str.to_date(time)
  • dt.*
  • str.to_date(time)
• (ongoing) Keep ExprIR with in sync with narwhals
  • feat(expr-ir): Add list.* aggregate methods #3353
  • feat(expr-ir): Add list.sort #3359
  • feat(expr-ir): Add top-level struct function #3666
  • feat: Add {Expr,Series}.any_value #3315
  • feat: Add {Expr,Series}.sin #3365
  • feat: Add {Expr,Series}.cos #3392
  • feat: Add str.pad_{start,end} #3395
  • feat: Add nw.corr #3460
  • feat: Add {Expr,Series}.str.to_time #3538
  • [Enh]: Add the negation unary operator for expressions and series #3625
• chore(expr-ir): Re-sync over with current polars #3352
• Review sugar for reducing ArrowExpr boilerplate
  • See What's next?
• refactor(expr-ir): Make arrow.functions a package #3384
• ArrowListNamespace -> (ArrowExprListNamespace, ArrowScalarListNamespace)
  • They can still share a base for len, get
  • Everything else has needed a different impl for ListScalar
• Move things up from Eager{DataFrame,Namespace} when the narwhals-level depends on them
  • See concat comment
• fix(expr-ir): Reject Selectors that do more than select #3376
• (high prio) Add LogicalPlan 🚀
  • feat(expr-ir): Long-overdue sync #3611
  • Related: 1, 2, 3
  • First target will be sharing over(*partition_by) windows between expressions in the same context
    • main can do this for multi-output (by not expanding)
    • but this could be done for expressions with a different root, inside a LogicalPlan
      d3fc48a1a0b9291ae6bac21af27e1039e63ba3cf/python/cudf_polars/cudf_polars/dsl/ir.py)
  • Review DataType propagation and supertyping rules
    • 🙏 [Enh]: Add support for concat(..., how={"vertical_relaxed", "diagonal_relaxed"}) #3386
    • It's looking like some version of this will be required to correctly rewrite LogicalPlans
    • Otherwise, each sub-plan has to be eager and request a new schema from the backend after processing the previous step
      • That's current narwhals behavior, but likely more of an issue when Selectors can be used in more places
  • A little write-up comparing oh-nodes to polars internals
• (medium prio) Implement missing (non-Expr) features for LogicalPlan
  • Functions
    • feat(expr-ir): Add concat #3371
  • I/O
    • feat(expr-ir): Add read_{csv,parquet} #3370
    • feat(expr-ir): Add DataFrame.write_{csv,parquet} #3369
  • BaseFrame
    • feat(expr-ir): Support DataFrame.join_asof #3378
    • feat(expr-ir): Add BaseFrame.unpivot #3368
    • feat(expr-ir): Add BaseFrame.unnest #3414
  • DataFrame
    • feat(expr-ir): Support DataFrame.pivot #3373
    • feat(expr-ir): Impl DataFrame.unique #3364
• fix(expr-ir): Add compat for old polars #3617
• feat(expr-ir): Add top-level select #3647
• feat(expr-ir): Almost finish PolarsExpr #3659
  • fix(expr-ir): Ensure over(order_by=...) always broadcasts #3662

[dangotbanned]

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)

[MarcoGorelli]

@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

[dangotbanned]

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:

• https://github.com/pola-rs/polars/blob/b9dd8cdbd6e6ec8373110536955ed5940b9460ec/crates/polars-plan/src/dsl/mod.rs
• https://github.com/pola-rs/polars/blob/b9dd8cdbd6e6ec8373110536955ed5940b9460ec/crates/polars-plan/src/dsl/expr/mod.rs

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",
)

[dangotbanned]

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