enh: Support keep={'first', 'last'} in {DataFrame,LazyFrame}.unique (#3118)

Author: MarcoGorelli

docs/api-reference/typing.md

@@ -35,7 +35,6 @@ Narwhals comes fully statically typed. In addition to `nw.DataFrame`, `nw.Expr`,
         - RankMethod
         - RollingInterpolationMethod
         - UniqueKeepStrategy
-        - LazyUniqueKeepStrategy
       show_source: false
       show_bases: false
 

narwhals/_arrow/dataframe.py

@@ -700,6 +700,7 @@ def unique(
         *,
         keep: UniqueKeepStrategy,
         maintain_order: bool | None = None,
+        order_by: Sequence[str] | None,
     ) -> Self:
         # The param `maintain_order` is only here for compatibility with the Polars API
         # and has no effect on the output.
@@ -714,14 +715,30 @@ def unique(
 
             agg_func = ArrowGroupBy._REMAP_UNIQUE[keep]
             col_token = generate_temporary_column_name(n_bytes=8, columns=self.columns)
+            if order_by and maintain_order:
+                idx_token = generate_temporary_column_name(
+                    n_bytes=8, columns=[*self.columns, col_token]
+                )
+                df = (
+                    self.with_row_index(idx_token, order_by=None)
+                    .sort(*order_by, nulls_last=False, descending=False)
+                    .unique(subset=subset, keep=keep, maintain_order=False, order_by=None)
+                )
+                return df.sort(idx_token, descending=False, nulls_last=False).drop(
+                    [idx_token], strict=False
+                )
+            if order_by:
+                native = self.sort(*order_by, nulls_last=False, descending=False).native
+            else:
+                native = self.native
             keep_idx_native = (
-                self.native.append_column(col_token, pa.array(np.arange(len(self))))
+                native.append_column(col_token, pa.array(np.arange(len(self))))
                 .group_by(subset)
                 .aggregate([(col_token, agg_func)])
                 .column(f"{col_token}_{agg_func}")
             )
             return self._with_native(
-                self.native.take(keep_idx_native), validate_column_names=False
+                native.take(keep_idx_native), validate_column_names=False
             )
 
         keep_idx = self.simple_select(*subset).is_unique()

narwhals/_compliant/dataframe.py

@@ -61,7 +61,6 @@
         AsofJoinStrategy,
         IntoSchema,
         JoinStrategy,
-        LazyUniqueKeepStrategy,
         MultiColSelector,
         MultiIndexSelector,
         PivotAgg,
@@ -155,7 +154,11 @@ def sort(
     ) -> Self: ...
     def tail(self, n: int) -> Self: ...
     def unique(
-        self, subset: Sequence[str] | None, *, keep: LazyUniqueKeepStrategy
+        self,
+        subset: Sequence[str] | None,
+        *,
+        keep: UniqueKeepStrategy,
+        order_by: Sequence[str] | None,
     ) -> Self: ...
     def unpivot(
         self,
@@ -265,6 +268,7 @@ def unique(
         *,
         keep: UniqueKeepStrategy,
         maintain_order: bool | None = None,
+        order_by: Sequence[str] | None,
     ) -> Self: ...
     def with_row_index(self, name: str, order_by: Sequence[str] | None) -> Self: ...
     @overload

narwhals/_dask/dataframe.py

@@ -39,7 +39,7 @@
     from narwhals.dataframe import LazyFrame
     from narwhals.dtypes import DType
     from narwhals.exceptions import ColumnNotFoundError
-    from narwhals.typing import AsofJoinStrategy, JoinStrategy, LazyUniqueKeepStrategy
+    from narwhals.typing import AsofJoinStrategy, JoinStrategy, UniqueKeepStrategy
 
 Incomplete: TypeAlias = "Any"
 """Using `_pandas_like` utils with `_dask`.
@@ -237,7 +237,11 @@ def head(self, n: int) -> Self:
         return self._with_native(self.native.head(n=n, compute=False, npartitions=-1))
 
     def unique(
-        self, subset: Sequence[str] | None, *, keep: LazyUniqueKeepStrategy
+        self,
+        subset: Sequence[str] | None,
+        *,
+        keep: UniqueKeepStrategy,
+        order_by: Sequence[str] | None,
     ) -> Self:
         if subset and (error := self._check_columns_exist(subset)):
             raise error
@@ -250,7 +254,11 @@ def unique(
             result = self.native.merge(unique, on=subset, how="inner")
         else:
             mapped_keep = {"any": "first"}.get(keep, keep)
-            result = self.native.drop_duplicates(subset=subset, keep=mapped_keep)
+            if order_by:
+                native = self.sort(*order_by, descending=False, nulls_last=False).native
+            else:
+                native = self.native
+            result = native.drop_duplicates(subset=subset, keep=mapped_keep)
         return self._with_native(result)
 
     def sort(self, *by: str, descending: bool | Sequence[bool], nulls_last: bool) -> Self:

narwhals/_dask/group_by.py

@@ -112,7 +112,7 @@ def agg(self, *exprs: DaskExpr) -> DaskLazyFrame:
             # No aggregation provided
             return (
                 self.compliant.simple_select(*self._keys)
-                .unique(self._keys, keep="any")
+                .unique(self._keys, keep="any", order_by=None)
                 .rename(dict(zip(self._keys, self._output_key_names)))
             )
 

narwhals/_duckdb/dataframe.py

@@ -55,7 +55,7 @@
     from narwhals.dataframe import LazyFrame
     from narwhals.dtypes import DType
     from narwhals.stable.v1 import DataFrame as DataFrameV1
-    from narwhals.typing import AsofJoinStrategy, JoinStrategy, LazyUniqueKeepStrategy
+    from narwhals.typing import AsofJoinStrategy, JoinStrategy, UniqueKeepStrategy
 
 
 class DuckDBLazyFrame(
@@ -188,7 +188,7 @@ def select(self, *exprs: DuckDBExpr) -> Self:
 
     def drop(self, columns: Sequence[str], *, strict: bool) -> Self:
         columns_to_drop = parse_columns_to_drop(self, columns, strict=strict)
-        selection = (name for name in self.columns if name not in columns_to_drop)
+        selection = [col(name) for name in self.columns if name not in columns_to_drop]
         return self._with_native(self.native.select(*selection))
 
     def lazy(self, backend: None = None, **_: None) -> Self:
@@ -380,25 +380,43 @@ def collect_schema(self) -> dict[str, DType]:
         return self.schema
 
     def unique(
-        self, subset: Sequence[str] | None, *, keep: LazyUniqueKeepStrategy
+        self,
+        subset: Sequence[str] | None,
+        *,
+        keep: UniqueKeepStrategy,
+        order_by: Sequence[str] | None,
     ) -> Self:
-        if subset_ := subset if keep == "any" else (subset or self.columns):
-            # Sanitise input
-            if error := self._check_columns_exist(subset_):
-                raise error
-            idx_name = generate_temporary_column_name(8, self.columns)
-            count_name = generate_temporary_column_name(8, [*self.columns, idx_name])
-            name = count_name if keep == "none" else idx_name
-            idx_expr = window_expression(F("row_number"), subset_).alias(idx_name)
-            count_expr = window_expression(
-                F("count", StarExpression()), subset_, ()
-            ).alias(count_name)
-            return self._with_native(
-                self.native.select(StarExpression(), idx_expr, count_expr)
-                .filter(col(name) == lit(1))
-                .select(StarExpression(exclude=[count_name, idx_name]))
+        subset_ = subset or self.columns
+        if error := self._check_columns_exist(subset_):
+            raise error
+        tmp_name = generate_temporary_column_name(8, self.columns)
+        if order_by and keep == "last":
+            descending = [True] * len(order_by)
+            nulls_last = [True] * len(order_by)
+        else:
+            descending = None
+            nulls_last = None
+        if keep == "none":
+            expr = window_expression(
+                F("count", StarExpression()),
+                subset_,
+                order_by or (),
+                descending=descending,
+                nulls_last=nulls_last,
+            )
+        else:
+            expr = window_expression(
+                F("row_number"),
+                subset_,
+                order_by or (),
+                descending=descending,
+                nulls_last=nulls_last,
+            )
+        return self._with_native(
+            self.native.select(StarExpression(), expr.alias(tmp_name)).filter(
+                col(tmp_name) == lit(1)
             )
-        return self._with_native(self.native.unique(join_column_names(*self.columns)))
+        ).drop([tmp_name], strict=False)
 
     def sort(self, *by: str, descending: bool | Sequence[bool], nulls_last: bool) -> Self:
         if isinstance(descending, bool):

narwhals/_ibis/dataframe.py

@@ -2,23 +2,25 @@
 
 import operator
 from io import BytesIO
-from typing import TYPE_CHECKING, Any, Literal, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import ibis
 import ibis.expr.types as ir
 
-from narwhals._ibis.utils import evaluate_exprs, native_to_narwhals_dtype
+from narwhals._ibis.expr import IbisExpr
+from narwhals._ibis.utils import evaluate_exprs, lit, native_to_narwhals_dtype
 from narwhals._sql.dataframe import SQLLazyFrame
 from narwhals._utils import (
     Implementation,
     ValidateBackendVersion,
     Version,
+    generate_temporary_column_name,
     not_implemented,
     parse_columns_to_drop,
     to_pyarrow_table,
     zip_strict,
 )
-from narwhals.exceptions import ColumnNotFoundError, InvalidOperationError
+from narwhals.exceptions import InvalidOperationError
 
 if TYPE_CHECKING:
     from collections.abc import Iterable, Iterator, Mapping, Sequence
@@ -31,7 +33,6 @@
     from typing_extensions import Self, TypeAlias, TypeIs
 
     from narwhals._compliant.typing import CompliantDataFrameAny
-    from narwhals._ibis.expr import IbisExpr
     from narwhals._ibis.group_by import IbisGroupBy
     from narwhals._ibis.namespace import IbisNamespace
     from narwhals._ibis.series import IbisInterchangeSeries
@@ -40,7 +41,7 @@
     from narwhals.dataframe import LazyFrame
     from narwhals.dtypes import DType
     from narwhals.stable.v1 import DataFrame as DataFrameV1
-    from narwhals.typing import AsofJoinStrategy, JoinStrategy, LazyUniqueKeepStrategy
+    from narwhals.typing import AsofJoinStrategy, JoinStrategy, UniqueKeepStrategy
 
     JoinPredicates: TypeAlias = "Sequence[ir.BooleanColumn] | Sequence[str]"
 
@@ -320,21 +321,33 @@ def collect_schema(self) -> dict[str, DType]:
         }
 
     def unique(
-        self, subset: Sequence[str] | None, *, keep: LazyUniqueKeepStrategy
+        self,
+        subset: Sequence[str] | None,
+        *,
+        keep: UniqueKeepStrategy,
+        order_by: Sequence[str] | None,
     ) -> Self:
-        if subset_ := subset if keep == "any" else (subset or self.columns):
-            # Sanitise input
-            if any(x not in self.columns for x in subset_):
-                msg = f"Columns {set(subset_).difference(self.columns)} not found in {self.columns}."
-                raise ColumnNotFoundError(msg)
-
-            mapped_keep: dict[str, Literal["first"] | None] = {
-                "any": "first",
-                "none": None,
-            }
-            to_keep = mapped_keep[keep]
-            return self._with_native(self.native.distinct(on=subset_, keep=to_keep))
-        return self._with_native(self.native.distinct(on=subset))
+        subset_ = subset or self.columns
+        if error := self._check_columns_exist(subset_):
+            raise error
+        tmp_name = generate_temporary_column_name(8, self.columns)
+        if order_by and keep == "last":
+            order_by_ = IbisExpr._sort(*order_by, descending=True, nulls_last=True)
+        elif order_by:
+            order_by_ = IbisExpr._sort(*order_by, descending=False, nulls_last=False)
+        else:
+            order_by_ = lit(1)
+        window = ibis.window(group_by=subset_, order_by=order_by_)
+        if keep == "none":
+            expr = self.native.count().over(window)
+        else:
+            expr = ibis.row_number().over(window) + lit(1)
+        df = (
+            self.native.mutate(**{tmp_name: expr})
+            .filter(ibis._[tmp_name] == lit(1))
+            .drop(tmp_name)
+        )
+        return self._with_native(df)
 
     def sort(self, *by: str, descending: bool | Sequence[bool], nulls_last: bool) -> Self:
         from narwhals._ibis.expr import IbisExpr

narwhals/_pandas_like/dataframe.py

@@ -761,16 +761,29 @@ def unique(
         *,
         keep: UniqueKeepStrategy,
         maintain_order: bool | None = None,
+        order_by: Sequence[str] | None,
     ) -> Self:
         # The param `maintain_order` is only here for compatibility with the Polars API
         # and has no effect on the output.
         mapped_keep = {"none": False, "any": "first"}.get(keep, keep)
         if subset and (error := self._check_columns_exist(subset)):
             raise error
-        return self._with_native(
-            self.native.drop_duplicates(subset=subset, keep=mapped_keep),
-            validate_column_names=False,
-        )
+        if order_by and maintain_order:
+            token = generate_temporary_column_name(8, self.columns)
+            res = (
+                self.with_row_index(token, order_by=None)
+                .sort(*order_by, nulls_last=False, descending=False)
+                .native.drop_duplicates(subset or self.columns, keep=mapped_keep)
+                .sort_values(token)
+            )
+            res.drop(columns=token, inplace=True)  # noqa: PD002
+        elif order_by:
+            res = self.sort(
+                *order_by, nulls_last=False, descending=False
+            ).native.drop_duplicates(subset or self.columns, keep=mapped_keep)
+        else:
+            res = self.native.drop_duplicates(subset or self.columns, keep=mapped_keep)
+        return self._with_native(res, validate_column_names=False)
 
     # --- lazy-only ---
     def lazy(

narwhals/_polars/dataframe.py

@@ -16,6 +16,7 @@
     Implementation,
     _into_arrow_table,
     convert_str_slice_to_int_slice,
+    generate_temporary_column_name,
     is_compliant_series,
     is_index_selector,
     is_range,
@@ -53,6 +54,7 @@
         MultiIndexSelector,
         PivotAgg,
         SingleIndexSelector,
+        UniqueKeepStrategy,
         _2DArray,
     )
 
@@ -89,7 +91,6 @@
         "tail",
         "to_arrow",
         "to_pandas",
-        "unique",
         "with_columns",
         "write_csv",
         "write_parquet",
@@ -110,7 +111,6 @@ class PolarsBaseFrame(Generic[NativePolarsFrame]):
     select: Method[Self]
     sort: Method[Self]
     tail: Method[Self]
-    unique: Method[Self]
     with_columns: Method[Self]
 
     _native_frame: NativePolarsFrame
@@ -175,6 +175,31 @@ def simple_select(self, *column_names: str) -> Self:
     def aggregate(self, *exprs: Any) -> Self:
         return self.select(*exprs)
 
+    def unique(
+        self,
+        subset: Sequence[str] | None,
+        *,
+        keep: UniqueKeepStrategy,
+        maintain_order: bool | None = None,
+        order_by: Sequence[str] | None = None,
+    ) -> Self:
+        if order_by and maintain_order:
+            token = generate_temporary_column_name(8, self.columns)
+            res = (
+                self.native.with_row_index(token)
+                .sort(order_by, nulls_last=False)
+                .unique(subset or self.columns, keep=keep)
+                .sort(token)
+                .drop(token)
+            )
+        elif order_by:
+            res = self.native.sort(order_by).unique(subset, keep=keep)
+        else:
+            res = self.native.unique(
+                subset, keep=keep, maintain_order=maintain_order or False
+            )
+        return self._with_native(res)
+
     @property
     def schema(self) -> dict[str, DType]:
         return self.collect_schema()