feat(expr-ir): Finish* implementing ArrowExpr

narwhals/_plan/__init__.py

@@ -7,12 +7,15 @@
     all,
     all_horizontal,
     any_horizontal,
+    coalesce,
     col,
     concat_str,
     date_range,
     exclude,
+    format,
     int_range,
     len,
+    linear_space,
     lit,
     max,
     max_horizontal,
@@ -37,12 +40,15 @@
     "all",
     "all_horizontal",
     "any_horizontal",
+    "coalesce",
     "col",
     "concat_str",
     "date_range",
     "exclude",
+    "format",
     "int_range",
     "len",
+    "linear_space",
     "lit",
     "max",
     "max_horizontal",

narwhals/_plan/_dispatch.py

@@ -182,7 +182,20 @@ def _method_name(tp: type[ExprIRT | FunctionT]) -> str:
 
 
 def get_dispatch_name(expr: ExprIR | type[Function], /) -> str:
-    """Return the synthesized method name for `expr`."""
-    return (
-        repr(expr.function) if is_function_expr(expr) else expr.__expr_ir_dispatch__.name
-    )
+    """Return the synthesized method name for `expr`.
+
+    Note:
+        Refers to the `Compliant*` method name, which may be *either* more general
+        or more specialized than what the user called.
+    """
+    dispatch: Dispatcher[Any]
+    if is_function_expr(expr):
+        from narwhals._plan import expressions as ir
+
+        if isinstance(expr, (ir.RollingExpr, ir.AnonymousExpr)):
+            dispatch = expr.__expr_ir_dispatch__
+        else:
+            dispatch = expr.function.__expr_ir_dispatch__
+    else:
+        dispatch = expr.__expr_ir_dispatch__
+    return dispatch.name

narwhals/_plan/_function.py

@@ -52,8 +52,8 @@ def __init_subclass__(
         **kwds: Any,
     ) -> None:
         super().__init_subclass__(*args, **kwds)
-        if accessor:
-            config = replace(config or FEOptions.default(), accessor_name=accessor)
+        if accessor_name := accessor or cls.__expr_ir_config__.accessor_name:
+            config = replace(config or FEOptions.default(), accessor_name=accessor_name)
         if options:
             cls._function_options = staticmethod(options)
         if config:

narwhals/_plan/_guards.py

@@ -24,7 +24,7 @@
         NativeSeriesT,
         Seq,
     )
-    from narwhals.typing import NonNestedLiteral
+    from narwhals.typing import NonNestedLiteral, PythonLiteral
 
     T = TypeVar("T")
 
@@ -38,6 +38,7 @@
     bytes,
     Decimal,
 )
+_PYTHON_LITERAL_TPS = (*_NON_NESTED_LITERAL_TPS, list, tuple, type(None))
 
 
 def _ir(*_: Any):  # type: ignore[no-untyped-def]  # noqa: ANN202
@@ -68,6 +69,10 @@ def is_non_nested_literal(obj: Any) -> TypeIs[NonNestedLiteral]:
     return obj is None or isinstance(obj, _NON_NESTED_LITERAL_TPS)
 
 
+def is_python_literal(obj: Any) -> TypeIs[PythonLiteral]:
+    return isinstance(obj, _PYTHON_LITERAL_TPS)
+
+
 def is_expr(obj: Any) -> TypeIs[Expr]:
     return isinstance(obj, _expr().Expr)
 

narwhals/_plan/_parse.py

@@ -18,7 +18,11 @@
     is_selector,
 )
 from narwhals._plan.common import flatten_hash_safe
-from narwhals._plan.exceptions import invalid_into_expr_error, is_iterable_error
+from narwhals._plan.exceptions import (
+    invalid_into_expr_error,
+    is_iterable_error,
+    list_literal_error,
+)
 from narwhals._utils import qualified_type_name
 from narwhals.dependencies import get_polars
 from narwhals.exceptions import InvalidOperationError
@@ -127,7 +131,7 @@ def parse_into_expr_ir(
         expr = col(input)
     elif isinstance(input, list):
         if list_as_series is None:
-            raise TypeError(input)  # pragma: no cover
+            raise list_literal_error(input)
         expr = lit(list_as_series(input))
     else:
         expr = lit(input, dtype=dtype)
@@ -331,7 +335,7 @@ def _combine_predicates(predicates: Iterator[ExprIR], /) -> ExprIR:
         inputs = (first,)
     else:
         return first
-    return AllHorizontal().to_function_expr(*inputs)
+    return AllHorizontal(ignore_nulls=False).to_function_expr(*inputs)
 
 
 def _is_iterable(obj: Iterable[T] | Any) -> TypeIs[Iterable[T]]:

narwhals/_plan/arrow/acero.py

@@ -18,7 +18,7 @@
 import operator
 from functools import reduce
 from itertools import chain
-from typing import TYPE_CHECKING, Any, Final, Union, cast
+from typing import TYPE_CHECKING, Any, Final, Literal, Union, cast
 
 import pyarrow as pa  # ignore-banned-import
 import pyarrow.acero as pac
@@ -61,7 +61,9 @@
 """
 
 Target: TypeAlias = OneOrSeq[Field]
-Aggregation: TypeAlias = "_Aggregation"
+Aggregation: TypeAlias = Union[
+    "_Aggregation", Literal["hash_kurtosis", "hash_skew", "kurtosis", "skew"]
+]
 AggregateOptions: TypeAlias = "_AggregateOptions"
 Opts: TypeAlias = "AggregateOptions | None"
 OutputName: TypeAlias = str

narwhals/_plan/arrow/common.py

@@ -4,7 +4,7 @@
 
 from typing import TYPE_CHECKING, Any, ClassVar, Generic
 
-from narwhals._plan.arrow.functions import BACKEND_VERSION
+from narwhals._plan.arrow.functions import BACKEND_VERSION, random_indices
 from narwhals._typing_compat import TypeVar
 from narwhals._utils import Implementation, Version, _StoresNative
 
@@ -43,6 +43,10 @@ def _with_native(self, native: NativeT) -> Self:
         msg = f"{type(self).__name__}._with_native"
         raise NotImplementedError(msg)
 
+    def __len__(self) -> int:
+        msg = f"{type(self).__name__}.__len__"
+        raise NotImplementedError(msg)
+
     if BACKEND_VERSION >= (18,):
 
         def _gather(self, indices: Indices) -> NativeT:
@@ -57,5 +61,14 @@ def gather(self, indices: Indices | _StoresNative[ChunkedArrayAny]) -> Self:
         ca = self._gather(indices.native if is_series(indices) else indices)
         return self._with_native(ca)
 
+    def gather_every(self, n: int, offset: int = 0) -> Self:
+        return self._with_native(self.native[offset::n])
+
     def slice(self, offset: int, length: int | None = None) -> Self:
         return self._with_native(self.native.slice(offset=offset, length=length))
+
+    def sample_n(
+        self, n: int = 1, *, with_replacement: bool = False, seed: int | None = None
+    ) -> Self:
+        mask = random_indices(len(self), n, with_replacement=with_replacement, seed=seed)
+        return self.gather(mask)

narwhals/_plan/arrow/dataframe.py

@@ -16,6 +16,7 @@
 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.exceptions import shape_error
 from narwhals._plan.expressions import NamedIR
 from narwhals._utils import Version, generate_repr
 from narwhals.schema import Schema
@@ -24,16 +25,18 @@
     from collections.abc import Iterable, Iterator, Mapping, Sequence
 
     import polars as pl
-    from typing_extensions import Self
+    from typing_extensions import Self, TypeAlias
 
-    from narwhals._plan.arrow.typing import ChunkedArrayAny
+    from narwhals._plan.arrow.typing import ChunkedArrayAny, ChunkedOrArrayAny
     from narwhals._plan.compliant.group_by import GroupByResolver
     from narwhals._plan.expressions import ExprIR, NamedIR
-    from narwhals._plan.options import SortMultipleOptions
+    from narwhals._plan.options import ExplodeOptions, SortMultipleOptions
     from narwhals._plan.typing import NonCrossJoinStrategy
     from narwhals.dtypes import DType
     from narwhals.typing import IntoSchema
 
+Incomplete: TypeAlias = Any
+
 
 class ArrowDataFrame(
     FrameSeries["pa.Table"], EagerDataFrame[Series, "pa.Table", "ChunkedArrayAny"]
@@ -48,6 +51,10 @@ def _with_native(self, native: pa.Table) -> Self:
     def _group_by(self) -> type[GroupBy]:
         return GroupBy
 
+    @property
+    def shape(self) -> tuple[int, int]:
+        return self.native.shape
+
     def group_by_resolver(self, resolver: GroupByResolver, /) -> GroupBy:
         return self._group_by.from_resolver(self, resolver)
 
@@ -68,11 +75,16 @@ def __len__(self) -> int:
 
     @classmethod
     def from_dict(
-        cls, data: Mapping[str, Any], /, *, schema: IntoSchema | None = None
+        cls,
+        data: Mapping[str, Any],
+        /,
+        *,
+        schema: IntoSchema | None = None,
+        version: Version = Version.MAIN,
     ) -> Self:
         pa_schema = Schema(schema).to_arrow() if schema is not None else schema
         native = pa.Table.from_pydict(data, schema=pa_schema)
-        return cls.from_native(native, version=Version.MAIN)
+        return cls.from_native(native, version=version)
 
     def iter_columns(self) -> Iterator[Series]:
         for name, series in zip(self.columns, self.native.itercolumns()):
@@ -96,10 +108,12 @@ def to_polars(self) -> pl.DataFrame:
 
         return pl.DataFrame(self.native)
 
-    def _evaluate_irs(self, nodes: Iterable[NamedIR[ExprIR]], /) -> Iterator[Series]:
-        ns = namespace(self)
-        from_named_ir = ns._expr.from_named_ir
-        yield from ns._expr.align(from_named_ir(e, self) for e in nodes)
+    def _evaluate_irs(
+        self, nodes: Iterable[NamedIR[ExprIR]], /, *, length: int | None = None
+    ) -> Iterator[Series]:
+        expr = namespace(self)._expr
+        from_named_ir = expr.from_named_ir
+        yield from expr.align((from_named_ir(e, self) for e in nodes), default=length)
 
     def sort(self, by: Sequence[str], options: SortMultipleOptions | None = None) -> Self:
         return self.gather(fn.sort_indices(self.native, *by, options=options))
@@ -121,6 +135,19 @@ def with_row_index_by(
         column = fn.unsort_indices(indices)
         return self._with_native(self.native.add_column(0, name, column))
 
+    def to_struct(self, name: str = "") -> Series:
+        native = self.native
+        if fn.TO_STRUCT_ARRAY_ACCEPTS_EMPTY:
+            struct = native.to_struct_array()
+        elif fn.HAS_FROM_TO_STRUCT_ARRAY:
+            if len(native):
+                struct = native.to_struct_array()
+            else:
+                struct = fn.chunked_array([], pa.struct(native.schema))
+        else:
+            struct = fn.struct(native.column_names, native.columns)
+        return Series.from_native(struct, name, version=self.version)
+
     def get_column(self, name: str) -> Series:
         chunked = self.native.column(name)
         return Series.from_native(chunked, name, version=self.version)
@@ -136,6 +163,12 @@ def drop_nulls(self, subset: Sequence[str] | None) -> Self:
             native = self.native.filter(~to_drop)
         return self._with_native(native)
 
+    def explode(self, subset: Sequence[str], options: ExplodeOptions) -> Self:
+        builder = fn.ExplodeBuilder.from_options(options)
+        if len(subset) == 1:
+            return self._with_native(builder.explode_column(self.native, subset[0]))
+        return self._with_native(builder.explode_columns(self.native, subset))
+
     def rename(self, mapping: Mapping[str, str]) -> Self:
         names: dict[str, str] | list[str]
         if fn.BACKEND_VERSION >= (17,):
@@ -144,20 +177,26 @@ def rename(self, mapping: Mapping[str, str]) -> Self:
             names = [mapping.get(c, c) for c in self.columns]
         return self._with_native(self.native.rename_columns(names))
 
-    # NOTE: Use instead of `with_columns` for trivial cases
+    def with_series(self, series: Series) -> Self:
+        """Add a new column or replace an existing one.
+
+        Uses similar semantics as `with_columns`, but:
+        - for a single named `Series`
+        - no broadcasting (use `Scalar.broadcast` instead)
+        - no length checking (use `with_series_checked` instead)
+        """
+        return self._with_native(with_array(self.native, series.name, series.native))
+
+    def with_series_checked(self, series: Series) -> Self:
+        expected, actual = len(self), len(series)
+        if len(series) != len(self):
+            raise shape_error(expected, actual)
+        return self.with_series(series)
+
     def _with_columns(self, exprs: Iterable[Expr | Scalar], /) -> Self:
-        native = self.native
-        columns = self.columns
         height = len(self)
-        for into_series in exprs:
-            name = into_series.name
-            chunked = into_series.broadcast(height).native
-            if name in columns:
-                i = columns.index(name)
-                native = native.set_column(i, name, chunked)
-            else:
-                native = native.append_column(name, chunked)
-        return self._with_native(native)
+        names_and_columns = ((e.name, e.broadcast(height).native) for e in exprs)
+        return self._with_native(with_arrays(self.native, names_and_columns))
 
     def select_names(self, *column_names: str) -> Self:
         return self._with_native(self.native.select(list(column_names)))
@@ -200,3 +239,22 @@ def partition_by(self, by: Sequence[str], *, include_key: bool = True) -> list[S
         from_native = self._with_native
         partitions = partition_by(self.native, by, include_key=include_key)
         return [from_native(df) for df in partitions]
+
+
+def with_array(table: pa.Table, name: str, column: ChunkedOrArrayAny) -> pa.Table:
+    column_names = table.column_names
+    if name in column_names:
+        return table.set_column(column_names.index(name), name, column)
+    return table.append_column(name, column)
+
+
+def with_arrays(
+    table: pa.Table, names_and_columns: Iterable[tuple[str, ChunkedOrArrayAny]], /
+) -> pa.Table:
+    column_names = table.column_names
+    for name, column in names_and_columns:
+        if name in column_names:
+            table = table.set_column(column_names.index(name), name, column)
+        else:
+            table = table.append_column(name, column)
+    return table