feat(expr-ir): Acero order_by, hashjoin , DataFrame.{filter,join}, Expr.is_{first,last}_distinct (#3173)

Author: dangotbanned

narwhals/_plan/_expr_ir.py

@@ -304,3 +304,7 @@ def is_column(self, *, allow_aliasing: bool = False) -> bool:
 
         ir = self.expr
         return isinstance(ir, Column) and ((self.name == ir.name) or allow_aliasing)
+
+
+def named_ir(name: str, expr: ExprIRT, /) -> NamedIR[ExprIRT]:
+    return NamedIR(expr=expr, name=name)

narwhals/_plan/_guards.py

@@ -15,7 +15,7 @@
     from narwhals._plan.compliant.series import CompliantSeries
     from narwhals._plan.expr import Expr
     from narwhals._plan.series import Series
-    from narwhals._plan.typing import NativeSeriesT, Seq
+    from narwhals._plan.typing import IntoExprColumn, NativeSeriesT, Seq
     from narwhals.typing import NonNestedLiteral
 
     T = TypeVar("T")
@@ -67,6 +67,10 @@ def is_series(obj: Series[NativeSeriesT] | Any) -> TypeIs[Series[NativeSeriesT]]
     return isinstance(obj, _series().Series)
 
 
+def is_into_expr_column(obj: Any) -> TypeIs[IntoExprColumn]:
+    return isinstance(obj, (str, _expr().Expr, _series().Series))
+
+
 def is_compliant_series(
     obj: CompliantSeries[NativeSeriesT] | Any,
 ) -> TypeIs[CompliantSeries[NativeSeriesT]]:

narwhals/_plan/_parse.py

@@ -6,13 +6,14 @@
 from itertools import chain
 from typing import TYPE_CHECKING
 
-from narwhals._plan._guards import is_expr, is_iterable_reject
+from narwhals._plan._guards import is_expr, is_into_expr_column, is_iterable_reject
 from narwhals._plan.exceptions import (
     invalid_into_expr_error,
     is_iterable_pandas_error,
     is_iterable_polars_error,
 )
 from narwhals.dependencies import get_polars, is_pandas_dataframe, is_pandas_series
+from narwhals.exceptions import InvalidOperationError
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
@@ -22,7 +23,13 @@
     from typing_extensions import TypeAlias, TypeIs
 
     from narwhals._plan.expressions import ExprIR
-    from narwhals._plan.typing import IntoExpr, IntoExprColumn, OneOrIterable, Seq
+    from narwhals._plan.typing import (
+        IntoExpr,
+        IntoExprColumn,
+        OneOrIterable,
+        PartialSeries,
+        Seq,
+    )
     from narwhals.typing import IntoDType
 
     T = TypeVar("T")
@@ -85,15 +92,33 @@
 
 
 def parse_into_expr_ir(
-    input: IntoExpr, *, str_as_lit: bool = False, dtype: IntoDType | None = None
+    input: IntoExpr | list[Any],
+    *,
+    str_as_lit: bool = False,
+    list_as_series: PartialSeries | None = None,
+    dtype: IntoDType | None = None,
 ) -> ExprIR:
-    """Parse a single input into an `ExprIR` node."""
+    """Parse a single input into an `ExprIR` node.
+
+    Arguments:
+        input: The input to be parsed as an expression.
+        str_as_lit: Interpret string input as a string literal. If set to `False` (default),
+            strings are parsed as column names.
+        list_as_series: Interpret list input as a Series literal, using the provided constructor.
+            If set to `None` (default), lists will raise when passed to `lit`.
+        dtype: If the input is expected to resolve to a literal with a known dtype, pass
+            this to the `lit` constructor.
+    """
     from narwhals._plan import col, lit
 
     if is_expr(input):
         expr = input
     elif isinstance(input, str) and not str_as_lit:
         expr = col(input)
+    elif isinstance(input, list):
+        if list_as_series is None:
+            raise TypeError(input)
+        expr = lit(list_as_series(input))
     else:
         expr = lit(input, dtype=dtype)
     return expr._ir
@@ -105,50 +130,90 @@ def parse_into_seq_of_expr_ir(
     **named_inputs: IntoExpr,
 ) -> Seq[ExprIR]:
     """Parse variadic inputs into a flat sequence of `ExprIR` nodes."""
-    return tuple(_parse_into_iter_expr_ir(first_input, *more_inputs, **named_inputs))
+    return tuple(
+        _parse_into_iter_expr_ir(
+            first_input, *more_inputs, _list_as_series=None, **named_inputs
+        )
+    )
 
 
 def parse_predicates_constraints_into_expr_ir(
-    first_predicate: OneOrIterable[IntoExprColumn] = (),
-    *more_predicates: IntoExprColumn | _RaisesInvalidIntoExprError,
+    first_predicate: OneOrIterable[IntoExprColumn] | list[bool] = (),
+    *more_predicates: IntoExprColumn | list[bool] | _RaisesInvalidIntoExprError,
+    _list_as_series: PartialSeries | None = None,
     **constraints: IntoExpr,
 ) -> ExprIR:
     """Parse variadic predicates and constraints into an `ExprIR` node.
 
     The result is an AND-reduction of all inputs.
     """
-    all_predicates = _parse_into_iter_expr_ir(first_predicate, *more_predicates)
+    all_predicates = _parse_into_iter_expr_ir(
+        first_predicate, *more_predicates, _list_as_series=_list_as_series
+    )
     if constraints:
         chained = chain(all_predicates, _parse_constraints(constraints))
         return _combine_predicates(chained)
     return _combine_predicates(all_predicates)
 
 
+def parse_sort_by_into_seq_of_expr_ir(
+    by: OneOrIterable[IntoExprColumn] = (), *more_by: IntoExprColumn
+) -> Seq[ExprIR]:
+    """Parse `DataFrame.sort` and `Expr.sort_by` keys into a flat sequence of `ExprIR` nodes."""
+    return tuple(_parse_sort_by_into_iter_expr_ir(by, more_by))
+
+
+# TODO @dangotbanned: Review the rejection predicate
+# It doesn't cover all length-changing expressions, only aggregations/literals
+def _parse_sort_by_into_iter_expr_ir(
+    by: OneOrIterable[IntoExprColumn], more_by: Iterable[IntoExprColumn]
+) -> Iterator[ExprIR]:
+    for e in _parse_into_iter_expr_ir(by, *more_by):
+        if e.is_scalar:
+            msg = f"All expressions sort keys must preserve length, but got:\n{e!r}"
+            raise InvalidOperationError(msg)
+        yield e
+
+
 def _parse_into_iter_expr_ir(
-    first_input: OneOrIterable[IntoExpr], *more_inputs: IntoExpr, **named_inputs: IntoExpr
+    first_input: OneOrIterable[IntoExpr],
+    *more_inputs: IntoExpr | list[Any],
+    _list_as_series: PartialSeries | None = None,
+    **named_inputs: IntoExpr,
 ) -> Iterator[ExprIR]:
     if not _is_empty_sequence(first_input):
         # NOTE: These need to be separated to introduce an intersection type
         # Otherwise, `str | bytes` always passes through typing
         if _is_iterable(first_input) and not is_iterable_reject(first_input):
-            if more_inputs:
+            if more_inputs and (
+                _list_as_series is None or not isinstance(first_input, list)
+            ):
                 raise invalid_into_expr_error(first_input, more_inputs, named_inputs)
+            # NOTE: Ensures `first_input = [False, True, True] -> lit(Series([False, True, True]))`
+            elif (
+                _list_as_series is not None
+                and isinstance(first_input, list)
+                and not is_into_expr_column(first_input[0])
+            ):
+                yield parse_into_expr_ir(first_input, list_as_series=_list_as_series)
             else:
-                yield from _parse_positional_inputs(first_input)
+                yield from _parse_positional_inputs(first_input, _list_as_series)
         else:
-            yield parse_into_expr_ir(first_input)
+            yield parse_into_expr_ir(first_input, list_as_series=_list_as_series)
     else:
         # NOTE: Passthrough case for no inputs - but gets skipped when calling next
         yield from ()
     if more_inputs:
-        yield from _parse_positional_inputs(more_inputs)
+        yield from _parse_positional_inputs(more_inputs, _list_as_series)
     if named_inputs:
         yield from _parse_named_inputs(named_inputs)
 
 
-def _parse_positional_inputs(inputs: Iterable[IntoExpr], /) -> Iterator[ExprIR]:
+def _parse_positional_inputs(
+    inputs: Iterable[IntoExpr | list[Any]], /, list_as_series: PartialSeries | None = None
+) -> Iterator[ExprIR]:
     for into in inputs:
-        yield parse_into_expr_ir(into)
+        yield parse_into_expr_ir(into, list_as_series=list_as_series)
 
 
 def _parse_named_inputs(named_inputs: dict[str, IntoExpr], /) -> Iterator[ExprIR]: