feat: Adds {Expr,Series}.{first,last} (#2528)

---------

Co-authored-by: Francesco Bruzzesi <[email protected]>
Co-authored-by: FBruzzesi <[email protected]>
Co-authored-by: Marco Gorelli <[email protected]>

Author: 3 people

docs/api-reference/expr.md

@@ -23,6 +23,7 @@
         - fill_nan
         - fill_null
         - filter
+        - first
         - is_between
         - is_close
         - is_duplicated
@@ -34,6 +35,7 @@
         - is_null
         - is_unique
         - kurtosis
+        - last
         - len
         - log
         - map_batches

docs/api-reference/series.md

@@ -30,6 +30,7 @@
         - fill_nan
         - fill_null
         - filter
+        - first
         - from_iterable
         - from_numpy
         - gather_every
@@ -50,6 +51,7 @@
         - is_unique
         - item
         - kurtosis
+        - last
         - len
         - log
         - max

narwhals/_arrow/expr.py

@@ -2,6 +2,7 @@
 
 from typing import TYPE_CHECKING, Any
 
+import pyarrow as pa
 import pyarrow.compute as pc
 
 from narwhals._arrow.series import ArrowSeries
@@ -111,11 +112,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)
 
@@ -129,15 +127,24 @@ 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.is_scalar_like:
+                    # We need to broadcast the results to the original size, since
+                    # `over` is a length-preserving 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]:
+                if order_by:
+                    df = df.sort(*order_by, descending=False, nulls_last=False)
+
                 output_names, aliases = evaluate_output_names_and_aliases(self, df, [])
                 if overlap := set(output_names).intersection(partition_by):
                     # E.g. `df.select(nw.all().sum().over('a'))`. This is well-defined,

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

@@ -330,6 +330,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

@@ -124,6 +124,8 @@ def max(self) -> Self: ...
     def mean(self) -> Self: ...
     def sum(self) -> Self: ...
     def median(self) -> Self: ...
+    def first(self) -> Self: ...
+    def last(self) -> Self: ...
     def skew(self) -> Self: ...
     def kurtosis(self) -> Self: ...
     def std(self, *, ddof: int) -> Self: ...
@@ -867,6 +869,12 @@ def is_close(
             nans_equal=nans_equal,
         )
 
+    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)

narwhals/_compliant/series.py

@@ -45,6 +45,7 @@
         Into1DArray,
         IntoDType,
         MultiIndexSelector,
+        PythonLiteral,
         RollingInterpolationMethod,
         SizedMultiIndexSelector,
         _1DArray,
@@ -131,6 +132,8 @@ def arg_min(self) -> int: ...
     def arg_true(self) -> Self: ...
     def count(self) -> int: ...
     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_empty(self) -> bool:

narwhals/_compliant/typing.py

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

narwhals/_dask/expr.py

@@ -729,6 +729,8 @@ def dt(self) -> DaskExprDateTimeNamespace:
         return DaskExprDateTimeNamespace(self)
 
     rank = not_implemented()
+    first = not_implemented()
+    last = not_implemented()
 
     # namespaces
     list: not_implemented = not_implemented()  # type: ignore[assignment]

narwhals/_duckdb/expr.py

@@ -13,8 +13,10 @@
     DeferredTimeZone,
     F,
     col,
+    generate_order_by_sql,
     lit,
     narwhals_to_native_dtype,
+    sql_expression,
     when,
     window_expression,
 )
@@ -93,6 +95,24 @@ def _window_expression(
             nulls_last=nulls_last,
         )
 
+    def _first(self, expr: Expression, *order_by: str) -> Expression:
+        # https://github.com/duckdb/duckdb/discussions/19252
+        order_by_sql = generate_order_by_sql(
+            *order_by,
+            descending=[False] * len(order_by),
+            nulls_last=[False] * len(order_by),
+        )
+        return sql_expression(f"first({expr} {order_by_sql})")
+
+    def _last(self, expr: Expression, *order_by: str) -> Expression:
+        # https://github.com/duckdb/duckdb/discussions/19252
+        order_by_sql = generate_order_by_sql(
+            *order_by,
+            descending=[False] * len(order_by),
+            nulls_last=[False] * len(order_by),
+        )
+        return sql_expression(f"last({expr} {order_by_sql})")
+
     def __narwhals_namespace__(self) -> DuckDBNamespace:  # pragma: no cover
         from narwhals._duckdb.namespace import DuckDBNamespace