feat(RFC): A richer Expr IR

.pre-commit-config.yaml

@@ -19,7 +19,7 @@ repos:
   hooks:
     - id: codespell
       files: \.(py|rst|md)$
-      args: [--ignore-words-list=ser]
+      args: [--ignore-words-list=ser,rightt]
       exclude: ^docs/api-completeness.md$
 -   repo: https://github.com/pycqa/flake8
     rev: '7.2.0'  # todo: remove once https://github.com/astral-sh/ruff/issues/458 is addressed

narwhals/_plan/__init__.py

@@ -0,0 +1,41 @@
+"""Brainstorming an `Expr` internal representation.
+
+Notes:
+- Each `Expr` method should be representable by a single node
+  - But the node does not need to be unique to the method
+- A chain of `Expr` methods should form a plan of operations
+- We must be able to enforce rules on what plans are permitted:
+  - Must be flexible to both eager/lazy and individual backends
+  - Must be flexible to a given context (select, with_columns, filter, group_by)
+- Nodes & plans are:
+  - Immutable, but
+    - Can be extended/re-written at both the Narwhals & Compliant levels
+  - Introspectable, but
+    - Store as little-as-needed for the common case
+    - Provide properties/methods for computing the less frequent metadata
+
+References:
+- https://github.com/pola-rs/polars/blob/dafd0a2d0e32b52bcfa4273bffdd6071a0d5977a/crates/polars-python/src/lazyframe/visitor/expr_nodes.rs
+- https://github.com/pola-rs/polars/blob/dafd0a2d0e32b52bcfa4273bffdd6071a0d5977a/crates/polars-plan/src/dsl/expr.rs
+- https://github.com/pola-rs/polars/blob/dafd0a2d0e32b52bcfa4273bffdd6071a0d5977a/crates/polars-plan/src/dsl/function_expr/mod.rs
+- https://github.com/pola-rs/polars/blob/dafd0a2d0e32b52bcfa4273bffdd6071a0d5977a/crates/polars-plan/src/dsl/options/mod.rs#L137-L172
+- https://github.com/pola-rs/polars/blob/3fd7ecc5f9de95f62b70ea718e7e5dbf951b6d1c/crates/polars-plan/src/plans/options.rs#L35-L106
+- https://github.com/pola-rs/polars/blob/3fd7ecc5f9de95f62b70ea718e7e5dbf951b6d1c/crates/polars-plan/src/plans/options.rs#L131-L236
+- https://github.com/pola-rs/polars/blob/3fd7ecc5f9de95f62b70ea718e7e5dbf951b6d1c/crates/polars-plan/src/plans/options.rs#L240-L267
+- https://github.com/pola-rs/polars/blob/6df23a09a81c640c21788607611e09d9f43b1abc/crates/polars-plan/src/plans/aexpr/mod.rs
+
+Related:
+- https://github.com/narwhals-dev/narwhals/pull/2483#issuecomment-2866902903
+- https://github.com/narwhals-dev/narwhals/pull/2483#issuecomment-2867331343
+- https://github.com/narwhals-dev/narwhals/pull/2483#issuecomment-2867446959
+- https://github.com/narwhals-dev/narwhals/pull/2483#issuecomment-2869070157
+- (https://github.com/narwhals-dev/narwhals/pull/2538/commits/a7eeb0d23e67cb70e7cfa73cec2c7b69a15c8bef#r2083562677)
+- https://github.com/narwhals-dev/narwhals/issues/2225
+- https://github.com/narwhals-dev/narwhals/issues/1848
+- https://github.com/narwhals-dev/narwhals/issues/2534#issuecomment-2875676729
+- https://github.com/narwhals-dev/narwhals/issues/2291
+- https://github.com/narwhals-dev/narwhals/issues/2522
+- https://github.com/narwhals-dev/narwhals/pull/2555
+"""
+
+from __future__ import annotations

narwhals/_plan/aggregation.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from narwhals._plan.common import ExprIR
+
+if TYPE_CHECKING:
+    from typing import Iterator
+
+    from narwhals.typing import RollingInterpolationMethod
+
+
+class Agg(ExprIR):
+    __slots__ = ("expr",)
+
+    expr: ExprIR
+
+    @property
+    def is_scalar(self) -> bool:
+        return True
+
+    def __repr__(self) -> str:
+        tp = type(self)
+        if tp in {Agg, OrderableAgg}:
+            return tp.__name__
+        m = {ArgMin: "arg_min", ArgMax: "arg_max", NUnique: "n_unique"}
+        name = m.get(tp, tp.__name__.lower())
+        return f"{self.expr!r}.{name}()"
+
+    def iter_left(self) -> Iterator[ExprIR]:
+        yield from self.expr.iter_left()
+        yield self
+
+    def iter_right(self) -> Iterator[ExprIR]:
+        yield self
+        yield from self.expr.iter_right()
+
+
+class Count(Agg): ...
+
+
+class Max(Agg): ...
+
+
+class Mean(Agg): ...
+
+
+class Median(Agg): ...
+
+
+class Min(Agg): ...
+
+
+class NUnique(Agg): ...
+
+
+class Quantile(Agg):
+    __slots__ = (*Agg.__slots__, "interpolation", "quantile")
+
+    quantile: float
+    interpolation: RollingInterpolationMethod
+
+
+class Std(Agg):
+    __slots__ = (*Agg.__slots__, "ddof")
+
+    ddof: int
+    """https://github.com/narwhals-dev/narwhals/pull/2555"""
+
+
+class Sum(Agg): ...
+
+
+class Var(Agg):
+    __slots__ = (*Agg.__slots__, "ddof")
+
+    ddof: int
+    """https://github.com/narwhals-dev/narwhals/pull/2555"""
+
+
+class OrderableAgg(Agg): ...
+
+
+class First(OrderableAgg):
+    """https://github.com/narwhals-dev/narwhals/issues/2526."""
+
+
+class Last(OrderableAgg):
+    """https://github.com/narwhals-dev/narwhals/issues/2526."""
+
+
+class ArgMin(OrderableAgg): ...
+
+
+class ArgMax(OrderableAgg): ...

narwhals/_plan/boolean.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+# NOTE: Needed to avoid naming collisions
+# - Any
+import typing as t
+
+from narwhals._plan.common import Function
+from narwhals._plan.options import FunctionFlags, FunctionOptions
+
+if t.TYPE_CHECKING:
+    from narwhals.typing import ClosedInterval
+
+
+class BooleanFunction(Function):
+    def __repr__(self) -> str:
+        tp = type(self)
+        if tp is BooleanFunction:
+            return tp.__name__
+        m: dict[type[BooleanFunction], str] = {
+            All: "all",
+            Any: "any",
+            AllHorizontal: "all_horizontal",
+            AnyHorizontal: "any_horizontal",
+            IsBetween: "is_between",
+            IsDuplicated: "is_duplicated",
+            IsFinite: "is_finite",
+            IsNan: "is_nan",
+            IsNull: "is_null",
+            IsFirstDistinct: "is_first_distinct",
+            IsLastDistinct: "is_last_distinct",
+            IsUnique: "is_unique",
+            IsIn: "is_in",
+            Not: "not",
+        }
+        return m[tp]
+
+
+class All(BooleanFunction):
+    @property
+    def function_options(self) -> FunctionOptions:
+        return FunctionOptions.aggregation()
+
+
+class AllHorizontal(BooleanFunction):
+    @property
+    def function_options(self) -> FunctionOptions:
+        return FunctionOptions.elementwise().with_flags(
+            FunctionFlags.INPUT_WILDCARD_EXPANSION
+        )
+
+
+class Any(BooleanFunction):
+    @property
+    def function_options(self) -> FunctionOptions:
+        return FunctionOptions.aggregation()
+
+
+class AnyHorizontal(BooleanFunction):
+    @property
+    def function_options(self) -> FunctionOptions:
+        return FunctionOptions.elementwise().with_flags(
+            FunctionFlags.INPUT_WILDCARD_EXPANSION
+        )
+
+
+class IsBetween(BooleanFunction):
+    """`lower_bound`, `upper_bound` aren't spec'd in the function enum.
+
+    Assuming the `FunctionExpr.input` becomes `s` in the impl
+
+    https://github.com/pola-rs/polars/blob/62257860a43ec44a638e8492ed2cf98a49c05f2e/crates/polars-plan/src/dsl/function_expr/boolean.rs#L225-L237
+    """
+
+    __slots__ = ("closed",)
+
+    closed: ClosedInterval
+
+    @property
+    def function_options(self) -> FunctionOptions:
+        return FunctionOptions.elementwise()
+
+
+class IsDuplicated(BooleanFunction):
+    @property
+    def function_options(self) -> FunctionOptions:
+        return FunctionOptions.length_preserving()
+
+
+class IsFinite(BooleanFunction):
+    @property
+    def function_options(self) -> FunctionOptions:
+        return FunctionOptions.elementwise()
+
+
+class IsFirstDistinct(BooleanFunction):
+    @property
+    def function_options(self) -> FunctionOptions:
+        return FunctionOptions.length_preserving()
+
+
+class IsIn(BooleanFunction):
+    """``other` isn't spec'd in the function enum.
+
+    See `IsBetween` comment.
+    """
+
+    @property
+    def function_options(self) -> FunctionOptions:
+        return FunctionOptions.elementwise()
+
+
+class IsLastDistinct(BooleanFunction):
+    @property
+    def function_options(self) -> FunctionOptions:
+        return FunctionOptions.length_preserving()
+
+
+class IsNan(BooleanFunction):
+    @property
+    def function_options(self) -> FunctionOptions:
+        return FunctionOptions.elementwise()
+
+
+class IsNull(BooleanFunction):
+    @property
+    def function_options(self) -> FunctionOptions:
+        return FunctionOptions.elementwise()
+
+
+class IsUnique(BooleanFunction):
+    @property
+    def function_options(self) -> FunctionOptions:
+        return FunctionOptions.length_preserving()
+
+
+class Not(BooleanFunction):
+    """`__invert__`."""
+
+    @property
+    def function_options(self) -> FunctionOptions:
+        return FunctionOptions.elementwise()

narwhals/_plan/categorical.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from narwhals._plan.common import Function
+from narwhals._plan.options import FunctionOptions
+
+
+class CategoricalFunction(Function): ...
+
+
+class GetCategories(CategoricalFunction):
+    """https://github.com/pola-rs/polars/blob/62257860a43ec44a638e8492ed2cf98a49c05f2e/crates/polars-plan/src/dsl/function_expr/cat.rs#L7."""
+
+    @property
+    def function_options(self) -> FunctionOptions:
+        """https://github.com/pola-rs/polars/blob/62257860a43ec44a638e8492ed2cf98a49c05f2e/crates/polars-plan/src/dsl/function_expr/cat.rs#L41."""
+        return FunctionOptions.groupwise()
+
+    def __repr__(self) -> str:
+        return "cat.get_categories"