feat: Get 3x over() rules passing

Aaaaand added comments on rule origin

Author: dangotbanned

narwhals/_plan/expr.py

@@ -22,6 +22,7 @@
     RollingT,
     SelectorOperatorT,
 )
+from narwhals.exceptions import InvalidOperationError
 
 if t.TYPE_CHECKING:
     from typing_extensions import Self
@@ -403,6 +404,34 @@ def iter_right(self) -> t.Iterator[ExprIR]:
             yield from e.iter_right()
         yield from self.expr.iter_right()
 
+    def __init__(
+        self,
+        *,
+        expr: ExprIR,
+        partition_by: Seq[ExprIR],
+        order_by: tuple[Seq[ExprIR], SortOptions] | None,
+        options: Window,
+    ) -> None:
+        if isinstance(expr, WindowExpr):
+            msg = "Cannot nest `over` statements."
+            raise InvalidOperationError(msg)
+
+        if isinstance(expr, FunctionExpr):
+            if expr.options.is_elementwise():
+                msg = f"Cannot use `over` on expressions which are elementwise.\n{expr!r}"
+                raise InvalidOperationError(msg)
+            if expr.options.is_row_separable():
+                msg = f"Cannot use `over` on expressions which change length.\n{expr!r}"
+                raise InvalidOperationError(msg)
+
+        kwds = {
+            "expr": expr,
+            "partition_by": partition_by,
+            "order_by": order_by,
+            "options": options,
+        }
+        super().__init__(**kwds)
+
 
 class Len(ExprIR):
     @property

narwhals/_plan/options.py

@@ -27,7 +27,7 @@ class FunctionFlags(enum.Flag):
     """Automatically explode on unit length if it ran as final aggregation."""
 
     ROW_SEPARABLE = 1 << 8
-    """Not sure lol.
+    """`drop_nulls` is the only one we've got that is *just* this.
 
     https://github.com/pola-rs/polars/pull/22573
     """
@@ -36,14 +36,17 @@ class FunctionFlags(enum.Flag):
     """mutually exclusive with `RETURNS_SCALAR`"""
 
     def is_elementwise(self) -> bool:
-        return self in (FunctionFlags.ROW_SEPARABLE | FunctionFlags.LENGTH_PRESERVING)
+        return (FunctionFlags.ROW_SEPARABLE | FunctionFlags.LENGTH_PRESERVING) in self
 
     def returns_scalar(self) -> bool:
         return FunctionFlags.RETURNS_SCALAR in self
 
     def is_length_preserving(self) -> bool:
         return FunctionFlags.LENGTH_PRESERVING in self
 
+    def is_row_separable(self) -> bool:
+        return FunctionFlags.ROW_SEPARABLE in self
+
     @staticmethod
     def default() -> FunctionFlags:
         return FunctionFlags.ALLOW_GROUP_AWARE
@@ -75,6 +78,9 @@ def returns_scalar(self) -> bool:
     def is_length_preserving(self) -> bool:
         return self.flags.is_length_preserving()
 
+    def is_row_separable(self) -> bool:
+        return self.flags.is_row_separable()
+
     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."

tests/plan/expr_parsing_test.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import re
 from typing import TYPE_CHECKING, Callable, Iterable
 
 import pytest
@@ -81,17 +82,10 @@ def test_function_expr_horizontal(
     assert sequence_node != unrelated_node
 
 
-# TODO @dangotbanned: Get partity with the existing tests
+# TODO @dangotbanned: Get parity with the existing tests
 # https://github.com/narwhals-dev/narwhals/blob/63c8e4771a1df4e0bfeea5559c303a4a447d5cc2/tests/expression_parsing_test.py#L48-L105
 
 
-def test_misleading_order_by() -> None:
-    with pytest.raises(InvalidOperationError):
-        nw.col("a").mean().over(order_by="b")
-    with pytest.raises(InvalidOperationError):
-        nw.col("a").rank().over(order_by="b")
-
-
 # `test_double_over` is already covered in the later `test_nested_over`
 
 
@@ -107,6 +101,9 @@ def test_invalid_repeat_agg() -> None:
         nwd.col("a").all().quantile(0.5, "linear")
 
 
+# TODO @dangotbanned: Weirdly, `polars` suggestion **does** resolve it
+# InvalidOperationError: Series idx, length 1 doesn't match the DataFrame height of 9
+# If you want expression: col("idx").mean().drop_nulls() to be broadcasted, ensure it is a scalar (for instance by adding '.first()')
 def test_filter_aggregation() -> None:
     with pytest.raises(InvalidOperationError):
         nwd.col("a").mean().drop_nulls()
@@ -118,32 +115,48 @@ def test_head_aggregation() -> None:
         nwd.col("a").mean().head()  # type: ignore[attr-defined]
 
 
+# TODO @dangotbanned: (Same as `test_filter_aggregation`)
 def test_rank_aggregation() -> None:
     with pytest.raises(InvalidOperationError):
         nwd.col("a").mean().rank()
 
 
+# TODO @dangotbanned: No error in `polars`, but results in all `null`s
 def test_diff_aggregation() -> None:
     with pytest.raises(InvalidOperationError):
         nwd.col("a").mean().diff()
 
 
-def test_invalid_over() -> None:
+# TODO @dangotbanned: Non-`polars`` rule
+def test_misleading_order_by() -> None:
+    with pytest.raises(InvalidOperationError):
+        nwd.col("a").mean().over(order_by="b")
     with pytest.raises(InvalidOperationError):
+        nwd.col("a").rank().over(order_by="b")
+
+
+# NOTE: Non-`polars`` rule
+def test_invalid_over() -> None:
+    pattern = re.compile(r"cannot use.+over.+elementwise", re.IGNORECASE)
+    with pytest.raises(InvalidOperationError, match=pattern):
         nwd.col("a").fill_null(3).over("b")
 
 
 def test_nested_over() -> None:
-    with pytest.raises(InvalidOperationError):
+    pattern = re.compile(r"cannot nest.+over", re.IGNORECASE)
+    with pytest.raises(InvalidOperationError, match=pattern):
         nwd.col("a").mean().over("b").over("c")
-    with pytest.raises(InvalidOperationError):
+    with pytest.raises(InvalidOperationError, match=pattern):
         nwd.col("a").mean().over("b").over("c", order_by="i")
 
 
+# NOTE: This *can* error in polars, but only if the length **actualy changes**
+# The rule then breaks down to needing the same length arrays in all parts of the over
 def test_filtration_over() -> None:
-    with pytest.raises(InvalidOperationError):
+    pattern = re.compile(r"cannot use.+over.+change length", re.IGNORECASE)
+    with pytest.raises(InvalidOperationError, match=pattern):
         nwd.col("a").drop_nulls().over("b")
-    with pytest.raises(InvalidOperationError):
+    with pytest.raises(InvalidOperationError, match=pattern):
         nwd.col("a").drop_nulls().over("b", order_by="i")
-    with pytest.raises(InvalidOperationError):
+    with pytest.raises(InvalidOperationError, match=pattern):
         nwd.col("a").diff().drop_nulls().over("b", order_by="i")