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

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/_parse.py

@@ -13,6 +13,7 @@
     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
@@ -124,6 +125,25 @@ def parse_predicates_constraints_into_expr_ir(
     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
 ) -> Iterator[ExprIR]:

narwhals/_plan/arrow/acero.py

@@ -25,26 +25,40 @@
 import pyarrow.compute as pc  # ignore-banned-import
 from pyarrow.acero import Declaration as Decl
 
+from narwhals._plan.common import ensure_list_str, flatten_hash_safe, temp
+from narwhals._plan.options import SortMultipleOptions
 from narwhals._plan.typing import OneOrSeq
-from narwhals.typing import SingleColSelector
+from narwhals.typing import JoinStrategy, SingleColSelector
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Collection, Iterable, Iterator
+    from collections.abc import (
+        Callable,
+        Collection,
+        Iterable,
+        Iterator,
+        Mapping,
+        Sequence,
+    )
 
-    from typing_extensions import TypeAlias
+    from typing_extensions import TypeAlias, TypeIs
 
     from narwhals._arrow.typing import (  # type: ignore[attr-defined]
         AggregateOptions as _AggregateOptions,
         Aggregation as _Aggregation,
     )
     from narwhals._plan.arrow.group_by import AggSpec
-    from narwhals._plan.arrow.typing import NullPlacement
+    from narwhals._plan.arrow.typing import (
+        ArrowAny,
+        JoinTypeSubset,
+        NullPlacement,
+        ScalarAny,
+    )
     from narwhals._plan.typing import OneOrIterable, Order, Seq
     from narwhals.typing import NonNestedLiteral
 
 Incomplete: TypeAlias = Any
 Expr: TypeAlias = pc.Expression
-IntoExpr: TypeAlias = "Expr | NonNestedLiteral"
+IntoExpr: TypeAlias = "Expr | NonNestedLiteral | ScalarAny"
 Field: TypeAlias = Union[Expr, SingleColSelector]
 """Anything that passes as a single item in [`_compute._ensure_field_ref`].
 
@@ -57,12 +71,28 @@
 Opts: TypeAlias = "AggregateOptions | None"
 OutputName: TypeAlias = str
 
+IntoDecl: TypeAlias = Union[pa.Table, Decl]
+"""An in-memory table, or a plan that began with one."""
+
 _THREAD_UNSAFE: Final = frozenset[Aggregation](
     ("hash_first", "hash_last", "first", "last")
 )
 col = pc.field
-lit = cast("Callable[[NonNestedLiteral], Expr]", pc.scalar)
-"""Alias for `pyarrow.compute.scalar`."""
+lit = cast("Callable[[NonNestedLiteral | ScalarAny], Expr]", pc.scalar)
+"""Alias for `pyarrow.compute.scalar`.
+
+Extends the signature from `bool | float | str`.
+
+See https://github.com/apache/arrow/pull/47609#discussion_r2392499842
+"""
+
+_HOW_JOIN: Mapping[JoinStrategy, JoinTypeSubset] = {
+    "inner": "inner",
+    "left": "left outer",
+    "full": "full outer",
+    "anti": "left anti",
+    "semi": "left semi",
+}
 
 
 # NOTE: ATOW there are 304 valid function names, 46 can be used for some kind of agg
@@ -72,8 +102,17 @@ def can_thread(function_name: str, /) -> bool:
     return function_name not in _THREAD_UNSAFE
 
 
+def cols_iter(names: Iterable[str], /) -> Iterator[Expr]:
+    for name in names:
+        yield col(name)
+
+
+def _is_expr(obj: Any) -> TypeIs[pc.Expression]:
+    return isinstance(obj, pc.Expression)
+
+
 def _parse_into_expr(into: IntoExpr, /, *, str_as_lit: bool = False) -> Expr:
-    if isinstance(into, pc.Expression):
+    if _is_expr(into):
         return into
     if isinstance(into, str) and not str_as_lit:
         return col(into)
@@ -177,6 +216,35 @@ def project(**named_exprs: IntoExpr) -> Decl:
     return _project(names=named_exprs.keys(), exprs=exprs)
 
 
+def _add_column(
+    native: pa.Table, index: int, name: str, values: IntoExpr | ArrowAny
+) -> pa.Table:
+    if isinstance(values, (pa.ChunkedArray, pa.Array)):
+        return native.add_column(index, name, values)
+    column = values if _is_expr(values) else lit(values)
+    schema = native.schema
+    schema_names = schema.names
+    if index == 0:
+        names: Sequence[str] = (name, *schema_names)
+        exprs = (column, *cols_iter(schema_names))
+    elif index == native.num_columns:
+        names = (*schema_names, name)
+        exprs = (*cols_iter(schema_names), column)
+    else:
+        schema_names.insert(index, name)
+        names = schema_names
+        exprs = tuple(_parse_into_iter_expr(nm if nm != name else column for nm in names))
+    return collect(table_source(native), _project(exprs, names))
+
+
+def append_column(native: pa.Table, name: str, values: IntoExpr | ArrowAny) -> pa.Table:
+    return _add_column(native, native.num_columns, name, values)
+
+
+def prepend_column(native: pa.Table, name: str, values: IntoExpr | ArrowAny) -> pa.Table:
+    return _add_column(native, 0, name, values)
+
+
 def _order_by(
     sort_keys: Iterable[tuple[str, Order]] = (),
     *,
@@ -189,10 +257,84 @@ def _order_by(
     return Decl("order_by", pac.OrderByNodeOptions(keys, null_placement=null_placement))
 
 
-# TODO @dangotbanned: Utilize `SortMultipleOptions.to_arrow_acero`
-def sort_by(*args: Any, **kwds: Any) -> Decl:
-    msg = "Should convert from polars args -> use `_order_by"
-    raise NotImplementedError(msg)
+def sort_by(
+    by: OneOrIterable[str],
+    *more_by: str,
+    descending: OneOrIterable[bool] = False,
+    nulls_last: bool = False,
+) -> Decl:
+    return SortMultipleOptions.parse(
+        descending=descending, nulls_last=nulls_last
+    ).to_arrow_acero(tuple(flatten_hash_safe((by, more_by))))
+
+
+# TODO @dangotbanned: Add a variant that doesn't depend on 2x tables
+# - `_join_options`: just needs iterables (currently sourced from `schema.names``)
+# -`_hashjoin`: can now accept `Declaration`s in either case
+def join(
+    left: pa.Table,
+    right: pa.Table,
+    how: JoinTypeSubset,
+    left_on: OneOrIterable[str],
+    right_on: OneOrIterable[str],
+    suffix: str = "_right",
+    *,
+    coalesce_keys: bool = True,
+) -> Decl:
+    """Heavily based on [`pyarrow.acero._perform_join`].
+
+    [`pyarrow.acero._perform_join`]: https://github.com/apache/arrow/blob/f7320c9a40082639f9e0cf8b3075286e3fc6c0b9/python/pyarrow/acero.py#L82-L260
+    """
+    opts = _join_options(
+        how,
+        left_on,
+        right_on,
+        suffix,
+        left.schema.names,
+        right.schema.names,
+        coalesce_keys=coalesce_keys,
+    )
+    return _hashjoin(left, right, opts)
+
+
+def _join_options(
+    how: JoinTypeSubset,
+    left_on: OneOrIterable[str],
+    right_on: OneOrIterable[str],
+    suffix: str = "_right",
+    left_names: Iterable[str] | None = None,
+    right_names: Iterable[str] = (),
+    *,
+    coalesce_keys: bool = True,
+) -> pac.HashJoinNodeOptions:
+    right_on = ensure_list_str(right_on)
+    rhs_names: Iterable[str] | None = None
+    # polars full join does not coalesce keys
+    if not (coalesce_keys and (how != "full outer")):
+        lhs_names = None
+    else:
+        lhs_names = left_names
+        if how in {"inner", "left outer"}:
+            rhs_names = (name for name in right_names if name not in right_on)
+    tp: Incomplete = pac.HashJoinNodeOptions
+    return tp(  # type: ignore[no-any-return]
+        join_type=how,
+        left_keys=ensure_list_str(left_on),
+        right_keys=right_on,
+        left_output=lhs_names,
+        right_output=rhs_names,
+        output_suffix_for_right=suffix,
+    )
+
+
+def _into_decl(source: IntoDecl, /) -> Decl:
+    return source if not isinstance(source, pa.Table) else table_source(source)
+
+
+def _hashjoin(
+    left: IntoDecl, right: IntoDecl, /, options: pac.HashJoinNodeOptions
+) -> Decl:
+    return Decl("hashjoin", options, [_into_decl(left), _into_decl(right)])
 
 
 def collect(*declarations: Decl, use_threads: bool = True) -> pa.Table:
@@ -251,3 +393,44 @@ def select_names_table(
     native: pa.Table, column_names: OneOrIterable[str], *more_names: str
 ) -> pa.Table:
     return collect(table_source(native), select_names(column_names, *more_names))
+
+
+def join_tables(
+    left: pa.Table,
+    right: pa.Table,
+    how: JoinStrategy,
+    left_on: OneOrIterable[str] | None = (),
+    right_on: OneOrIterable[str] | None = (),
+    suffix: str = "_right",
+    *,
+    coalesce_keys: bool = True,
+) -> pa.Table:
+    if how == "cross":
+        return _join_cross_tables(left, right, suffix)
+    join_type = _HOW_JOIN[how]
+    left_on = left_on or ()
+    right_on = right_on or left_on
+    decl = join(
+        left, right, join_type, left_on, right_on, suffix, coalesce_keys=coalesce_keys
+    )
+    return collect(decl)
+
+
+# TODO @dangotbanned: Very rough start to get tests passing
+# - Decouple from `pa.Table` & collecting 3 times
+# - Reuse the plan from `_add_column`
+# - Write some more specialized parsers for
+#   [x] column names
+#   [ ] indices?
+def _join_cross_tables(
+    left: pa.Table, right: pa.Table, suffix: str = "_right"
+) -> pa.Table:
+    key_token = temp.column_name(chain(left.column_names, right.column_names))
+    result = join_tables(
+        prepend_column(left, key_token, 0),
+        prepend_column(right, key_token, 0),
+        how="inner",
+        left_on=key_token,
+        suffix=suffix,
+    )
+    return result.remove_column(0)

narwhals/_plan/arrow/dataframe.py

@@ -9,32 +9,33 @@
 import pyarrow.compute as pc  # ignore-banned-import
 
 from narwhals._arrow.utils import native_to_narwhals_dtype
-from narwhals._plan.arrow import functions as fn
+from narwhals._plan.arrow import acero, functions as fn
+from narwhals._plan.arrow.expr import ArrowExpr as Expr, ArrowScalar as Scalar
 from narwhals._plan.arrow.group_by import ArrowGroupBy as GroupBy
 from narwhals._plan.arrow.series import ArrowSeries as Series
 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._utils import Version, parse_columns_to_drop
+from narwhals._utils import Implementation, Version, parse_columns_to_drop
 from narwhals.schema import Schema
 
 if TYPE_CHECKING:
     from collections.abc import Iterable, Iterator, Mapping, Sequence
 
     from typing_extensions import Self
 
-    from narwhals._arrow.typing import ChunkedArrayAny  # noqa: F401
-    from narwhals._plan.arrow.expr import ArrowExpr as Expr, ArrowScalar as Scalar
+    from narwhals._arrow.typing import ChunkedArrayAny
     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 Seq
     from narwhals.dtypes import DType
-    from narwhals.typing import IntoSchema
+    from narwhals.typing import IntoSchema, JoinStrategy
 
 
 class ArrowDataFrame(EagerDataFrame[Series, "pa.Table", "ChunkedArrayAny"]):
+    implementation = Implementation.PYARROW
     _native: pa.Table
     _version: Version
 
@@ -144,3 +145,28 @@ def select_names(self, *column_names: str) -> Self:
     def row(self, index: int) -> tuple[Any, ...]:
         row = self.native.slice(index, 1)
         return tuple(chain.from_iterable(row.to_pydict().values()))
+
+    def join(
+        self,
+        other: Self,
+        *,
+        how: JoinStrategy,
+        left_on: Sequence[str] | None,
+        right_on: Sequence[str] | None,
+        suffix: str = "_right",
+    ) -> Self:
+        left, right = self.native, other.native
+        result = acero.join_tables(left, right, how, left_on, right_on, suffix=suffix)
+        return self._with_native(result)
+
+    def filter(self, predicate: NamedIR | Series) -> Self:
+        mask: pc.Expression | ChunkedArrayAny
+        if not fn.is_series(predicate):
+            resolved = Expr.from_named_ir(predicate, self)
+            if isinstance(resolved, Expr):
+                mask = resolved.broadcast(len(self)).native
+            else:
+                mask = acero.lit(resolved.native)
+        else:
+            mask = predicate.native
+        return self._with_native(self.native.filter(mask))