refactor: Finish sort_by_ir removal

Author: dangotbanned

narwhals/_plan/arrow/dataframe.py

@@ -16,8 +16,6 @@
 from narwhals._plan.compliant.dataframe import EagerDataFrame
 from narwhals._plan.compliant.typing import namespace
 from narwhals._plan.expressions import NamedIR
-from narwhals._plan.typing import Seq
-from narwhals._typing_compat import deprecated
 from narwhals._utils import Implementation, Version
 from narwhals.schema import Schema
 
@@ -30,7 +28,7 @@
     from narwhals._plan.arrow.namespace import ArrowNamespace
     from narwhals._plan.expressions import ExprIR, NamedIR
     from narwhals._plan.options import SortMultipleOptions
-    from narwhals._plan.typing import NonCrossJoinStrategy, Seq
+    from narwhals._plan.typing import NonCrossJoinStrategy
     from narwhals.dtypes import DType
     from narwhals.typing import IntoSchema
 
@@ -100,13 +98,6 @@ def sort(self, by: Sequence[str], options: SortMultipleOptions) -> Self:
         indices = pc.sort_indices(native.select(list(by)), options=options.to_arrow(by))
         return self._with_native(native.take(indices))
 
-    # TODO @dangotbanned: Finish scoping out what breaks if removed
-    @deprecated("Use `DataFrame.sort` instead", category=None)
-    def sort_by_ir(self, by: Seq[NamedIR], options: SortMultipleOptions) -> Self:
-        df_by = self.select(by)
-        indices = pc.sort_indices(df_by.native, options=options.to_arrow(df_by.columns))
-        return self._with_native(self.native.take(indices))
-
     def with_row_index(self, name: str) -> Self:
         return self._with_native(self.native.add_column(0, name, fn.int_range(len(self))))
 

narwhals/_plan/arrow/expr.py

@@ -15,7 +15,6 @@
 from narwhals._plan.compliant.expr import EagerExpr
 from narwhals._plan.compliant.scalar import EagerScalar
 from narwhals._plan.compliant.typing import namespace
-from narwhals._plan.expressions import NamedIR
 from narwhals._utils import Implementation, Version, _StoresNative, not_implemented
 from narwhals.exceptions import InvalidOperationError, ShapeError
 
@@ -363,10 +362,10 @@ def over_ordered(
             raise NotImplementedError(msg)
 
         # NOTE: Converting `over(order_by=..., options=...)` into the right shape for `DataFrame.sort`
-        sort_by = tuple(NamedIR.from_ir(e) for e in node.order_by)
-        options = node.sort_options.to_multiple(len(node.order_by))
+        sort_by = tuple(node.order_by_names())
+        options = node.sort_options.to_multiple(len(sort_by))
         idx_name = temp.column_name(frame)
-        sorted_context = frame.with_row_index(idx_name).sort_by_ir(sort_by, options)
+        sorted_context = frame.with_row_index(idx_name).sort(sort_by, options)
         evaluated = node.expr.dispatch(self, sorted_context.drop([idx_name]), name)
         if isinstance(evaluated, ArrowScalar):
             # NOTE: We're already sorted, defer broadcasting to the outer context

narwhals/_plan/exceptions.py

@@ -139,6 +139,24 @@ def over_row_separable_error(
     return InvalidOperationError(msg)
 
 
+def over_order_by_names_error(
+    expr: ir.OrderedWindowExpr, by: ir.ExprIR
+) -> InvalidOperationError:
+    if by.meta.is_column_selection(allow_aliasing=True):
+        # narwhals dev error
+        msg = (
+            f"Cannot use `{type(expr).__name__}.order_by_names()` before expression expansion.\n"
+            f"Found unresolved selection {by!r}, in:\n{expr!r}"
+        )
+    else:
+        # user error
+        msg = (
+            f"Only column selection expressions are supported in `over(order_by=...)`.\n"
+            f"Found {by!r}, in:\n{expr!r}"
+        )
+    return InvalidOperationError(msg)
+
+
 def invalid_into_expr_error(
     first_input: Iterable[IntoExpr],
     more_inputs: tuple[Any, ...],

narwhals/_plan/expressions/expr.py

@@ -7,7 +7,10 @@
 
 from narwhals._plan._expr_ir import ExprIR, SelectorIR
 from narwhals._plan.common import replace
-from narwhals._plan.exceptions import function_expr_invalid_operation_error
+from narwhals._plan.exceptions import (
+    function_expr_invalid_operation_error,
+    over_order_by_names_error,
+)
 from narwhals._plan.expressions import selectors as cs
 from narwhals._plan.options import ExprIROptions
 from narwhals._plan.typing import (
@@ -383,6 +386,19 @@ def iter_root_names(self) -> t.Iterator[ExprIR]:
             yield from e.iter_left()
         yield self
 
+    def order_by_names(self) -> Iterator[str]:
+        """Yield the names resolved from expanding `order_by`.
+
+        Raises:
+            InvalidOperationError: If used *before* expansion, or
+                `order_by` contains expressions that do more than select.
+        """
+        for by in self.order_by:
+            if isinstance(by, Column):
+                yield by.name
+            else:
+                raise over_order_by_names_error(self, by)
+
 
 class Len(ExprIR, config=ExprIROptions.namespaced()):
     @property

narwhals/_plan/options.py

@@ -167,7 +167,7 @@ def to_multiple(self, n_repeat: int = 1, /) -> SortMultipleOptions:
         if n_repeat == 1:
             desc: Seq[bool] = (self.descending,)
             nulls: Seq[bool] = (self.nulls_last,)
-        else:  # pragma: no cover
+        else:
             desc = tuple(repeat(self.descending, n_repeat))
             nulls = tuple(repeat(self.nulls_last, n_repeat))
         return SortMultipleOptions(descending=desc, nulls_last=nulls)

tests/plan/expr_expansion_test.py

@@ -726,3 +726,16 @@ def test_expand_binary_expr_combination_invalid(df_1: Frame) -> None:
     )
     with pytest.raises(MultiOutputExpressionError, match=pattern):
         df_1.project(ten_to_nine)
+
+
+def test_over_order_by_names() -> None:
+    expr = nwp.col("a").first().over(order_by=ncs.string())
+    e_ir = expr._ir
+    assert isinstance(e_ir, ir.OrderedWindowExpr)
+    with pytest.raises(
+        InvalidOperationError,
+        match=re_compile(
+            r"cannot use.+order_by_names.+before.+expansion.+ncs.string\(\)"
+        ),
+    ):
+        list(e_ir.order_by_names())

tests/plan/over_test.py

@@ -10,7 +10,7 @@
 import narwhals._plan as nwp
 from narwhals._plan import selectors as ncs
 from narwhals.exceptions import InvalidOperationError
-from tests.plan.utils import assert_equal_data, dataframe
+from tests.plan.utils import assert_equal_data, dataframe, re_compile
 
 if TYPE_CHECKING:
     from narwhals._plan.typing import IntoExprColumn, OneOrIterable
@@ -225,3 +225,22 @@ def test_shift_kitchen_sink(data_alt: Data) -> None:
     )
     expected = {"b": [0, 5, 10, 15, 500], "c": [5, 5, 10, 45, 500]}
     assert_equal_data(result, expected)
+
+
+def test_over_order_by_expr(data_alt: Data) -> None:
+    df = dataframe(data_alt)
+    result = df.select(
+        nwp.all()
+        + nwp.all().last().over(order_by=[nwp.nth(1), ncs.first()], descending=True)
+    )
+    expected = {"a": [6, 8, 4, 5, None], "b": [0, 1, 3, 2, 1], "c": [18, 10, 11, 10, 10]}
+    assert_equal_data(result, expected)
+
+
+def test_over_order_by_expr_invalid(data_alt: Data) -> None:
+    df = dataframe(data_alt)
+    with pytest.raises(
+        InvalidOperationError,
+        match=re_compile(r"only.+column.+selection.+in.+order_by.+found.+sort"),
+    ):
+        df.select(nwp.col("a").first().over(order_by=nwp.col("b").sort()))