feat: Adds {Expr,Series}.{first,last}

docs/api-reference/expr.md

@@ -21,6 +21,8 @@
         - exp
         - fill_null
         - filter
+        - first
+        - last
         - clip
         - is_between
         - is_duplicated

docs/api-reference/series.md

@@ -29,7 +29,9 @@
         - exp
         - fill_null
         - filter
+        - first
         - from_numpy
+        - last
         - gather_every
         - head
         - hist

narwhals/_arrow/expr.py

@@ -2,11 +2,12 @@
 
 from typing import TYPE_CHECKING, Any
 
+import pyarrow as pa
 import pyarrow.compute as pc
 
 from narwhals._arrow.series import ArrowSeries
 from narwhals._compliant import EagerExpr
-from narwhals._expression_parsing import evaluate_output_names_and_aliases
+from narwhals._expression_parsing import ExprKind, evaluate_output_names_and_aliases
 from narwhals._utils import (
     Implementation,
     generate_temporary_column_name,
@@ -113,11 +114,8 @@ def _reuse_series_extra_kwargs(
         return {"_return_py_scalar": False} if returns_scalar else {}
 
     def over(self, partition_by: Sequence[str], order_by: Sequence[str]) -> Self:
-        if (
-            partition_by
-            and self._metadata is not None
-            and not self._metadata.is_scalar_like
-        ):
+        meta = self._metadata
+        if partition_by and meta is not None and not meta.is_scalar_like:
             msg = "Only aggregation or literal operations are supported in grouped `over` context for PyArrow."
             raise NotImplementedError(msg)
 
@@ -131,12 +129,20 @@ def func(df: ArrowDataFrame) -> Sequence[ArrowSeries]:
                 df = df.with_row_index(token, order_by=None).sort(
                     *order_by, descending=False, nulls_last=False
                 )
-                result = self(df.drop([token], strict=True))
+                results = self(df.drop([token], strict=True))
+                if meta is not None and meta.last_node is ExprKind.ORDERABLE_AGGREGATION:
+                    # Orderable aggregations require `order_by` columns and results in a
+                    # scalar output (well actually in a length 1 series).
+                    # Therefore we need to broadcast the results to the original size, since
+                    # `over` is not a length changing operation.
+                    size = len(df)
+                    return [s._with_native(pa.repeat(s.item(), size)) for s in results]
+
                 # TODO(marco): is there a way to do this efficiently without
                 # doing 2 sorts? Here we're sorting the dataframe and then
                 # again calling `sort_indices`. `ArrowSeries.scatter` would also sort.
                 sorting_indices = pc.sort_indices(df.get_column(token).native)
-                return [s._with_native(s.native.take(sorting_indices)) for s in result]
+                return [s._with_native(s.native.take(sorting_indices)) for s in results]
         else:
 
             def func(df: ArrowDataFrame) -> Sequence[ArrowSeries]:

narwhals/_arrow/group_by.py

@@ -9,7 +9,7 @@
 from narwhals._arrow.utils import cast_to_comparable_string_types, extract_py_scalar
 from narwhals._compliant import EagerGroupBy
 from narwhals._expression_parsing import evaluate_output_names_and_aliases
-from narwhals._utils import generate_temporary_column_name
+from narwhals._utils import generate_temporary_column_name, requires
 
 if TYPE_CHECKING:
     from collections.abc import Iterator, Mapping, Sequence
@@ -39,12 +39,23 @@ class ArrowGroupBy(EagerGroupBy["ArrowDataFrame", "ArrowExpr", "Aggregation"]):
         "count": "count",
         "all": "all",
         "any": "any",
+        "first": "first",
+        "last": "last",
     }
     _REMAP_UNIQUE: ClassVar[Mapping[UniqueKeepStrategy, Aggregation]] = {
         "any": "min",
         "first": "min",
         "last": "max",
     }
+    _OPTION_COUNT_ALL: ClassVar[frozenset[NarwhalsAggregation]] = frozenset(
+        ("len", "n_unique")
+    )
+    _OPTION_COUNT_VALID: ClassVar[frozenset[NarwhalsAggregation]] = frozenset(("count",))
+    _OPTION_ORDERED: ClassVar[frozenset[NarwhalsAggregation]] = frozenset(
+        ("first", "last")
+    )
+    _OPTION_VARIANCE: ClassVar[frozenset[NarwhalsAggregation]] = frozenset(("std", "var"))
+    _OPTION_SCALAR: ClassVar[frozenset[NarwhalsAggregation]] = frozenset(("any", "all"))
 
     def __init__(
         self,
@@ -60,12 +71,58 @@ def __init__(
         self._grouped = pa.TableGroupBy(self.compliant.native, self._keys)
         self._drop_null_keys = drop_null_keys
 
+    def _configure_agg(
+        self, grouped: pa.TableGroupBy, expr: ArrowExpr, /
+    ) -> tuple[pa.TableGroupBy, Aggregation, AggregateOptions | None]:
+        option: AggregateOptions | None = None
+        function_name = self._leaf_name(expr)
+        if function_name in self._OPTION_VARIANCE:
+            ddof = expr._scalar_kwargs.get("ddof", 1)
+            option = pc.VarianceOptions(ddof=ddof)
+        elif function_name in self._OPTION_COUNT_ALL:
+            option = pc.CountOptions(mode="all")
+        elif function_name in self._OPTION_COUNT_VALID:
+            option = pc.CountOptions(mode="only_valid")
+        elif function_name in self._OPTION_SCALAR:
+            option = pc.ScalarAggregateOptions(min_count=0)
+        elif function_name in self._OPTION_ORDERED:
+            grouped, option = self._ordered_agg(grouped, function_name)
+        return grouped, self._remap_expr_name(function_name), option
+
+    def _ordered_agg(
+        self, grouped: pa.TableGroupBy, name: NarwhalsAggregation, /
+    ) -> tuple[pa.TableGroupBy, AggregateOptions]:
+        """The default behavior of `pyarrow` raises when `first` or `last` are used.
+
+        You'd see an error like:
+
+            ArrowNotImplementedError: Using ordered aggregator in multiple threaded execution is not supported
+
+        We need to **disable** multi-threading to use them, but the ability to do so
+        wasn't possible before `14.0.0` ([pyarrow-36709])
+
+        [pyarrow-36709]: https://github.com/apache/arrow/issues/36709
+        """
+        backend_version = self.compliant._backend_version
+        if backend_version >= (14, 0) and grouped._use_threads:
+            native = self.compliant.native
+            grouped = pa.TableGroupBy(native, grouped.keys, use_threads=False)
+        elif backend_version < (14, 0):  # pragma: no cover
+            msg = (
+                f"Using `{name}()` in a `group_by().agg(...)` context is only available in 'pyarrow>=14.0.0', "
+                f"found version {requires._unparse_version(backend_version)!r}.\n\n"
+                f"See https://github.com/apache/arrow/issues/36709"
+            )
+            raise NotImplementedError(msg)
+        return grouped, pc.ScalarAggregateOptions(skip_nulls=False)
+
     def agg(self, *exprs: ArrowExpr) -> ArrowDataFrame:
         self._ensure_all_simple(exprs)
         aggs: list[tuple[str, Aggregation, AggregateOptions | None]] = []
         expected_pyarrow_column_names: list[str] = self._keys.copy()
         new_column_names: list[str] = self._keys.copy()
         exclude = (*self._keys, *self._output_key_names)
+        grouped = self._grouped
 
         for expr in exprs:
             output_names, aliases = evaluate_output_names_and_aliases(
@@ -83,20 +140,7 @@ def agg(self, *exprs: ArrowExpr) -> ArrowDataFrame:
                 aggs.append((self._keys[0], "count", pc.CountOptions(mode="all")))
                 continue
 
-            function_name = self._leaf_name(expr)
-            if function_name in {"std", "var"}:
-                assert "ddof" in expr._scalar_kwargs  # noqa: S101
-                option: Any = pc.VarianceOptions(ddof=expr._scalar_kwargs["ddof"])
-            elif function_name in {"len", "n_unique"}:
-                option = pc.CountOptions(mode="all")
-            elif function_name == "count":
-                option = pc.CountOptions(mode="only_valid")
-            elif function_name in {"all", "any"}:
-                option = pc.ScalarAggregateOptions(min_count=0)
-            else:
-                option = None
-
-            function_name = self._remap_expr_name(function_name)
+            grouped, function_name, option = self._configure_agg(grouped, expr)
             new_column_names.extend(aliases)
             expected_pyarrow_column_names.extend(
                 [f"{output_name}_{function_name}" for output_name in output_names]
@@ -105,7 +149,7 @@ def agg(self, *exprs: ArrowExpr) -> ArrowDataFrame:
                 [(output_name, function_name, option) for output_name in output_names]
             )
 
-        result_simple = self._grouped.aggregate(aggs)
+        result_simple = grouped.aggregate(aggs)
 
         # Rename columns, being very careful
         expected_old_names_indices: dict[str, list[int]] = collections.defaultdict(list)

narwhals/_arrow/series.py

@@ -323,6 +323,15 @@ def filter(self, predicate: ArrowSeries | list[bool | None]) -> Self:
             other_native = predicate
         return self._with_native(self.native.filter(other_native))
 
+    def first(self, *, _return_py_scalar: bool = True) -> PythonLiteral:
+        result = self.native[0] if len(self.native) else None
+        return maybe_extract_py_scalar(result, _return_py_scalar)
+
+    def last(self, *, _return_py_scalar: bool = True) -> PythonLiteral:
+        ca = self.native
+        result = ca[height - 1] if (height := len(ca)) else None
+        return maybe_extract_py_scalar(result, _return_py_scalar)
+
     def mean(self, *, _return_py_scalar: bool = True) -> float:
         return maybe_extract_py_scalar(pc.mean(self.native), _return_py_scalar)
 

narwhals/_compliant/expr.py

@@ -27,7 +27,7 @@
     LazyExprT,
     NativeExprT,
 )
-from narwhals._utils import _StoresCompliant
+from narwhals._utils import _StoresCompliant, not_implemented
 from narwhals.dependencies import get_numpy, is_numpy_array
 
 if TYPE_CHECKING:
@@ -142,6 +142,8 @@ def cum_min(self, *, reverse: bool) -> Self: ...
     def cum_max(self, *, reverse: bool) -> Self: ...
     def cum_prod(self, *, reverse: bool) -> Self: ...
     def is_in(self, other: Any) -> Self: ...
+    def first(self) -> Self: ...
+    def last(self) -> Self: ...
     def rank(self, method: RankMethod, *, descending: bool) -> Self: ...
     def replace_strict(
         self,
@@ -892,6 +894,12 @@ def exp(self) -> Self:
     def sqrt(self) -> Self:
         return self._reuse_series("sqrt")
 
+    def first(self) -> Self:
+        return self._reuse_series("first", returns_scalar=True)
+
+    def last(self) -> Self:
+        return self._reuse_series("last", returns_scalar=True)
+
     @property
     def cat(self) -> EagerExprCatNamespace[Self]:
         return EagerExprCatNamespace(self)
@@ -922,6 +930,10 @@ class LazyExpr(  # type: ignore[misc]
     CompliantExpr[CompliantLazyFrameT, NativeExprT],
     Protocol[CompliantLazyFrameT, NativeExprT],
 ):
+    # NOTE: See https://github.com/narwhals-dev/narwhals/issues/2526#issuecomment-3019303816
+    first: not_implemented = not_implemented()
+    last: not_implemented = not_implemented()
+
     def _with_alias_output_names(self, func: AliasNames | None, /) -> Self: ...
     def alias(self, name: str) -> Self:
         def fn(names: Sequence[str]) -> Sequence[str]:

narwhals/_compliant/series.py

@@ -49,6 +49,7 @@
         MultiIndexSelector,
         NonNestedLiteral,
         NumericLiteral,
+        PythonLiteral,
         RankMethod,
         RollingInterpolationMethod,
         SizedMultiIndexSelector,
@@ -193,6 +194,8 @@ def fill_null(
         limit: int | None,
     ) -> Self: ...
     def filter(self, predicate: Any) -> Self: ...
+    def first(self) -> PythonLiteral: ...
+    def last(self) -> PythonLiteral: ...
     def gather_every(self, n: int, offset: int) -> Self: ...
     def head(self, n: int) -> Self: ...
     def is_between(

narwhals/_compliant/typing.py

@@ -184,6 +184,8 @@ class ScalarKwargs(TypedDict, total=False):
     "quantile",
     "all",
     "any",
+    "first",
+    "last",
 ]
 """`Expr` methods we aim to support in `DepthTrackingGroupBy`.
 

narwhals/_pandas_like/expr.py

@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING
 
 from narwhals._compliant import EagerExpr
-from narwhals._expression_parsing import evaluate_output_names_and_aliases
+from narwhals._expression_parsing import ExprKind, evaluate_output_names_and_aliases
 from narwhals._pandas_like.group_by import PandasLikeGroupBy
 from narwhals._pandas_like.series import PandasLikeSeries
 from narwhals._utils import generate_temporary_column_name
@@ -227,6 +227,20 @@ def func(df: PandasLikeDataFrame) -> Sequence[PandasLikeSeries]:
                     *order_by, descending=False, nulls_last=False
                 )
                 results = self(df.drop([token], strict=True))
+                if (
+                    meta := self._metadata
+                ) is not None and meta.last_node is ExprKind.ORDERABLE_AGGREGATION:
+                    # Orderable aggregations require `order_by` columns and result in a
+                    # scalar output (well actually in a length 1 series).
+                    # Therefore we need to broadcast the result to the original size, since
+                    # `over` is not a length changing operation.
+                    index = df.native.index
+                    ns = self._implementation.to_native_namespace()
+                    return [
+                        s._with_native(ns.Series(s.item(), index=index, name=s.name))
+                        for s in results
+                    ]
+
                 sorting_indices = df.get_column(token)
                 for s in results:
                     s._scatter_in_place(sorting_indices, s)