chore: add has_open_windows to ExprMetadata instead of is_order_dependent (#2078)

Author: MarcoGorelli

narwhals/_arrow/expr.py

@@ -14,6 +14,7 @@
 from narwhals._arrow.series import ArrowSeries
 from narwhals._expression_parsing import ExprKind
 from narwhals._expression_parsing import evaluate_output_names_and_aliases
+from narwhals._expression_parsing import is_scalar_like
 from narwhals._expression_parsing import reuse_series_implementation
 from narwhals.dependencies import get_numpy
 from narwhals.dependencies import is_numpy_array
@@ -414,10 +415,8 @@ def clip(self: Self, lower_bound: Any | None, upper_bound: Any | None) -> Self:
         )
 
     def over(self: Self, keys: list[str], kind: ExprKind) -> Self:
-        if kind is ExprKind.TRANSFORM:
-            msg = (
-                "Elementwise operations in `over` context are not supported for PyArrow."
-            )
+        if not is_scalar_like(kind):
+            msg = "Only aggregation or literal operations are supported in `over` context for PyArrow."
             raise NotImplementedError(msg)
 
         def func(df: ArrowDataFrame) -> list[ArrowSeries]:

narwhals/_expression_parsing.py

@@ -9,6 +9,7 @@
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Callable
+from typing import Literal
 from typing import Sequence
 from typing import TypeVar
 from typing import overload
@@ -331,8 +332,8 @@ class ExprKind(Enum):
     Commutative composition rules are:
     - LITERAL vs LITERAL -> LITERAL
     - CHANGES_LENGTH vs (LITERAL | AGGREGATION) -> CHANGES_LENGTH
-    - CHANGES_LENGTH vs (CHANGES_LENGTH | TRANSFORM) -> raise
-    - TRANSFORM vs (LITERAL | AGGREGATION) -> TRANSFORM
+    - CHANGES_LENGTH vs (CHANGES_LENGTH | TRANSFORM | WINDOW) -> raise
+    - (TRANSFORM | WINDOW) vs (LITERAL | AGGREGATION) -> TRANSFORM
     - AGGREGATION vs (LITERAL | AGGREGATION) -> AGGREGATION
     """
 
@@ -343,18 +344,49 @@ class ExprKind(Enum):
     """e.g. `nw.col('a').mean()`"""
 
     TRANSFORM = auto()
-    """length-preserving, e.g. `nw.col('a').round()`"""
+    """preserves length, e.g. `nw.col('a').round()`"""
+
+    WINDOW = auto()
+    """transform in which last node is order-dependent
+
+    examples:
+    - `nw.col('a').cum_sum()`
+    - `(nw.col('a')+1).cum_sum()`
+
+    non-examples:
+    - `nw.col('a').cum_sum()+1`
+    - `nw.col('a').cum_sum().mean()`
+    """
 
     CHANGES_LENGTH = auto()
     """e.g. `nw.col('a').drop_nulls()`"""
 
+    def preserves_length(self) -> bool:
+        return self in {ExprKind.TRANSFORM, ExprKind.WINDOW}
+
+    def is_window(self) -> bool:
+        return self is ExprKind.WINDOW
+
+    def is_changes_length(self) -> bool:
+        return self is ExprKind.CHANGES_LENGTH
+
+    def is_scalar_like(self) -> bool:
+        return is_scalar_like(self)
+
+
+def is_scalar_like(
+    kind: ExprKind,
+) -> TypeIs[Literal[ExprKind.AGGREGATION, ExprKind.LITERAL]]:
+    # Like ExprKind.is_scalar_like, but uses TypeIs for better type checking.
+    return kind in {ExprKind.AGGREGATION, ExprKind.LITERAL}
+
 
 class ExprMetadata:
-    __slots__ = ("_kind", "_order_dependent")
+    __slots__ = ("_kind", "_n_open_windows")
 
-    def __init__(self, kind: ExprKind, /, *, order_dependent: bool) -> None:
+    def __init__(self, kind: ExprKind, /, *, n_open_windows: int) -> None:
         self._kind: ExprKind = kind
-        self._order_dependent: bool = order_dependent
+        self._n_open_windows = n_open_windows
 
     def __init_subclass__(cls, /, *args: Any, **kwds: Any) -> Never:  # pragma: no cover
         msg = f"Cannot subclass {cls.__name__!r}"
@@ -364,105 +396,98 @@ def __init_subclass__(cls, /, *args: Any, **kwds: Any) -> Never:  # pragma: no c
     def kind(self) -> ExprKind:
         return self._kind
 
-    def is_order_dependent(self) -> bool:
-        return self._order_dependent
-
-    def is_transform(self) -> bool:
-        return self.kind is ExprKind.TRANSFORM
-
-    def is_aggregation_or_literal(self) -> bool:
-        return self.kind in {ExprKind.AGGREGATION, ExprKind.LITERAL}
-
-    def is_changes_length(self) -> bool:
-        return self.kind is ExprKind.CHANGES_LENGTH
+    @property
+    def n_open_windows(self) -> int:
+        return self._n_open_windows
 
     def with_kind(self, kind: ExprKind, /) -> ExprMetadata:
         """Change metadata kind, leaving all other attributes the same."""
-        return ExprMetadata(kind, order_dependent=self.is_order_dependent())
+        return ExprMetadata(kind, n_open_windows=self._n_open_windows)
 
-    def with_order_dependence(self) -> ExprMetadata:
-        """Set `order_dependent` to True, leaving all other attributes the same."""
-        return ExprMetadata(self.kind, order_dependent=True)
+    def with_extra_open_window(self) -> ExprMetadata:
+        """Increment `n_open_windows` leaving other attributes the same."""
+        return ExprMetadata(self.kind, n_open_windows=self._n_open_windows + 1)
 
-    def with_kind_and_order_dependence(self, kind: ExprKind, /) -> ExprMetadata:
-        """Change kind and set `order_dependent` to True."""
-        return ExprMetadata(kind, order_dependent=True)
+    def with_kind_and_extra_open_window(self, kind: ExprKind, /) -> ExprMetadata:
+        """Change metadata kind and increment `n_open_windows`."""
+        return ExprMetadata(kind, n_open_windows=self._n_open_windows + 1)
 
     @staticmethod
     def selector() -> ExprMetadata:
-        return ExprMetadata(ExprKind.TRANSFORM, order_dependent=False)
+        return ExprMetadata(ExprKind.TRANSFORM, n_open_windows=0)
 
 
 def combine_metadata(*args: IntoExpr | object | None, str_as_lit: bool) -> ExprMetadata:
     # Combine metadata from `args`.
 
     n_changes_length = 0
-    has_transforms = False
+    has_transforms_or_windows = False
     has_aggregations = False
     has_literals = False
-    result_is_order_dependent = False
+    result_n_open_windows = 0
 
     for arg in args:
         if isinstance(arg, str) and not str_as_lit:
-            has_transforms = True
+            has_transforms_or_windows = True
         elif is_expr(arg):
-            if arg._metadata.is_order_dependent():
-                result_is_order_dependent = True
+            if arg._metadata.n_open_windows:
+                result_n_open_windows += 1
             kind = arg._metadata.kind
             if kind is ExprKind.AGGREGATION:
                 has_aggregations = True
             elif kind is ExprKind.LITERAL:
                 has_literals = True
             elif kind is ExprKind.CHANGES_LENGTH:
                 n_changes_length += 1
-            elif kind is ExprKind.TRANSFORM:
-                has_transforms = True
+            elif kind.preserves_length():
+                has_transforms_or_windows = True
             else:  # pragma: no cover
                 msg = "unreachable code"
                 raise AssertionError(msg)
     if (
         has_literals
         and not has_aggregations
-        and not has_transforms
+        and not has_transforms_or_windows
         and not n_changes_length
     ):
         result_kind = ExprKind.LITERAL
     elif n_changes_length > 1:
         msg = "Length-changing expressions can only be used in isolation, or followed by an aggregation"
         raise LengthChangingExprError(msg)
-    elif n_changes_length and has_transforms:
+    elif n_changes_length and has_transforms_or_windows:
         msg = "Cannot combine length-changing expressions with length-preserving ones or aggregations"
         raise ShapeError(msg)
     elif n_changes_length:
         result_kind = ExprKind.CHANGES_LENGTH
-    elif has_transforms:
+    elif has_transforms_or_windows:
         result_kind = ExprKind.TRANSFORM
     else:
         result_kind = ExprKind.AGGREGATION
 
-    return ExprMetadata(result_kind, order_dependent=result_is_order_dependent)
+    return ExprMetadata(result_kind, n_open_windows=result_n_open_windows)
 
 
-def check_expressions_transform(*args: IntoExpr, function_name: str) -> None:
+def check_expressions_preserve_length(*args: IntoExpr, function_name: str) -> None:
     # Raise if any argument in `args` isn't length-preserving.
     # For Series input, we don't raise (yet), we let such checks happen later,
     # as this function works lazily and so can't evaluate lengths.
     from narwhals.series import Series
 
     if not all(
-        (is_expr(x) and x._metadata.is_transform()) or isinstance(x, (str, Series))
+        (is_expr(x) and x._metadata.kind.preserves_length())
+        or isinstance(x, (str, Series))
         for x in args
     ):
         msg = f"Expressions which aggregate or change length cannot be passed to '{function_name}'."
         raise ShapeError(msg)
 
 
-def all_exprs_are_aggs_or_literals(*args: IntoExpr, **kwargs: IntoExpr) -> bool:
+def all_exprs_are_scalar_like(*args: IntoExpr, **kwargs: IntoExpr) -> bool:
     # Raise if any argument in `args` isn't an aggregation or literal.
     # For Series input, we don't raise (yet), we let such checks happen later,
     # as this function works lazily and so can't evaluate lengths.
     exprs = chain(args, kwargs.values())
-    return all(is_expr(x) and x._metadata.is_aggregation_or_literal() for x in exprs)
+    return all(is_expr(x) and x._metadata.kind.is_scalar_like() for x in exprs)
 
 
 def infer_kind(obj: IntoExpr | _1DArray | object, *, str_as_lit: bool) -> ExprKind:
@@ -489,7 +514,7 @@ def apply_n_ary_operation(
     )
     kinds = [infer_kind(comparand, str_as_lit=str_as_lit) for comparand in comparands]
 
-    broadcast = any(kind is ExprKind.TRANSFORM for kind in kinds)
+    broadcast = any(kind.preserves_length() for kind in kinds)
     compliant_exprs = (
         compliant_expr.broadcast(kind)
         if broadcast

narwhals/_pandas_like/expr.py

@@ -9,6 +9,7 @@
 
 from narwhals._expression_parsing import ExprKind
 from narwhals._expression_parsing import evaluate_output_names_and_aliases
+from narwhals._expression_parsing import is_scalar_like
 from narwhals._expression_parsing import is_simple_aggregation
 from narwhals._expression_parsing import reuse_series_implementation
 from narwhals._pandas_like.expr_cat import PandasLikeExprCatNamespace
@@ -475,9 +476,9 @@ def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
                     )
                 )
                 return [result_frame[name] for name in aliases]
-        elif kind is ExprKind.TRANSFORM:
+        elif not is_scalar_like(kind):
             msg = (
-                "Elementwise operations are only supported in `over` context "
+                "Length-preserving operations are only supported in `over` context "
                 "for pandas if they are elementary "
                 "(e.g. `nw.col('a').cum_sum().over('b'))`)."
             )

narwhals/dataframe.py

@@ -16,9 +16,10 @@
 from warnings import warn
 
 from narwhals._expression_parsing import ExprKind
-from narwhals._expression_parsing import all_exprs_are_aggs_or_literals
-from narwhals._expression_parsing import check_expressions_transform
+from narwhals._expression_parsing import all_exprs_are_scalar_like
+from narwhals._expression_parsing import check_expressions_preserve_length
 from narwhals._expression_parsing import infer_kind
+from narwhals._expression_parsing import is_scalar_like
 from narwhals.dependencies import get_polars
 from narwhals.dependencies import is_numpy_array
 from narwhals.dependencies import is_numpy_array_1d
@@ -137,9 +138,7 @@ def with_columns(
     ) -> Self:
         compliant_exprs, kinds = self._flatten_and_extract(*exprs, **named_exprs)
         compliant_exprs = [
-            compliant_expr.broadcast(kind)
-            if (kind is ExprKind.LITERAL or kind is ExprKind.AGGREGATION)
-            else compliant_expr
+            compliant_expr.broadcast(kind) if is_scalar_like(kind) else compliant_expr
             for compliant_expr, kind in zip(compliant_exprs, kinds)
         ]
         return self._from_compliant_dataframe(
@@ -166,14 +165,12 @@ def select(
                     missing_columns, available_columns
                 ) from e
         compliant_exprs, kinds = self._flatten_and_extract(*flat_exprs, **named_exprs)
-        if compliant_exprs and all_exprs_are_aggs_or_literals(*flat_exprs, **named_exprs):
+        if compliant_exprs and all_exprs_are_scalar_like(*flat_exprs, **named_exprs):
             return self._from_compliant_dataframe(
                 self._compliant_frame.aggregate(*compliant_exprs),
             )
         compliant_exprs = [
-            compliant_expr.broadcast(kind)
-            if (kind is ExprKind.LITERAL or kind is ExprKind.AGGREGATION)
-            else compliant_expr
+            compliant_expr.broadcast(kind) if is_scalar_like(kind) else compliant_expr
             for compliant_expr, kind in zip(compliant_exprs, kinds)
         ]
         return self._from_compliant_dataframe(
@@ -205,7 +202,7 @@ def filter(
             and all(isinstance(x, bool) for x in predicates[0])
         ):
             flat_predicates = flatten(predicates)
-            check_expressions_transform(*flat_predicates, function_name="filter")
+            check_expressions_preserve_length(*flat_predicates, function_name="filter")
             compliant_predicates, _kinds = self._flatten_and_extract(*flat_predicates)
             plx = self.__narwhals_namespace__()
             predicate = plx.all_horizontal(
@@ -2204,18 +2201,18 @@ def _extract_compliant(self: Self, arg: Any) -> Any:
             plx = self.__narwhals_namespace__()
             return plx.col(arg)
         if isinstance(arg, Expr):
-            if arg._metadata.is_order_dependent():
+            if arg._metadata.n_open_windows > 0:
                 msg = (
                     "Order-dependent expressions are not supported for use in LazyFrame.\n\n"
                     "Hints:\n"
                     "- Instead of `lf.select(nw.col('a').sort())`, use `lf.select('a').sort()\n"
                     "- Instead of `lf.select(nw.col('a').head())`, use `lf.select('a').head()\n"
                     "- `Expr.cum_sum`, and other such expressions, are not currently supported.\n"
-                    "  In a future version of Narwhals, a `order_by` argument will be added and \n"
-                    "  they will be supported."
+                    "  In a future version of Narwhals, a `order_by` argument will be added to\n"
+                    "  `over` and they will be supported."
                 )
                 raise OrderDependentExprError(msg)
-            if arg._metadata.is_changes_length():
+            if arg._metadata.kind.is_changes_length():
                 msg = (
                     "Length-changing expressions are not supported for use in LazyFrame, unless\n"
                     "followed by an aggregation.\n\n"