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, --ignore-words-list=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/_namespace.py

@@ -203,7 +203,9 @@ def from_backend(cls, backend: EagerAllowed, /) -> EagerAllowedNamespace: ...
 
     @overload
     @classmethod
-    def from_backend(cls, backend: ModuleType, /) -> Namespace[CompliantNamespaceAny]: ...
+    def from_backend(
+        cls, backend: IntoBackend, /
+    ) -> Namespace[CompliantNamespaceAny]: ...
 
     @classmethod
     def from_backend(

narwhals/_plan/__init__.py

@@ -0,0 +1 @@
+from __future__ import annotations

narwhals/_plan/aggregation.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from narwhals._plan.common import ExprIR
+from narwhals.exceptions import InvalidOperationError
+
+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()
+
+    def __init__(self, *, expr: ExprIR, **kwds: Any) -> None:
+        if expr.is_scalar:
+            msg = "Can't apply aggregations to scalar-like expressions."
+            raise InvalidOperationError(msg)
+        super().__init__(expr=expr, **kwds)  # pyright: ignore[reportCallIssue]
+
+
+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
+
+
+class Sum(Agg): ...
+
+
+class Var(Agg):
+    __slots__ = (*Agg.__slots__, "ddof")
+
+    ddof: int
+
+
+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,39 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from narwhals._plan.common import ExprNamespace, Function, IRNamespace
+from narwhals._plan.options import FunctionOptions
+
+if TYPE_CHECKING:
+    from narwhals._plan.dummy import DummyExpr
+
+
+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:
+        return FunctionOptions.groupwise()
+
+    def __repr__(self) -> str:
+        return "cat.get_categories"
+
+
+class IRCatNamespace(IRNamespace):
+    def get_categories(self) -> GetCategories:
+        return GetCategories()
+
+
+class ExprCatNamespace(ExprNamespace[IRCatNamespace]):
+    @property
+    def _ir_namespace(self) -> type[IRCatNamespace]:
+        return IRCatNamespace
+
+    def get_categories(self) -> DummyExpr:
+        return self._to_narwhals(
+            self._ir.get_categories().to_function_expr(self._expr._ir)
+        )