Repository: narwhals-dev/narwhals
Pull request #565 (2024-07-21 12:51:24)

feat: arrow `unique`, `is_duplicated`, `is_unique`
## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue #372 

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added 
- [x] Documented the changes

## If you have comments or can explain your changes, please do so below.

[Relevant arrow issue/conversation](https://github.com/apache/arrow/issues/30950)


Patch:
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index 7526e528..efc825d1 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -12,6 +12,7 @@ from narwhals._arrow.utils import validate_dataframe_comparand
 from narwhals._expression_parsing import evaluate_into_exprs
 from narwhals.dependencies import get_numpy
 from narwhals.dependencies import get_pyarrow
+from narwhals.dependencies import get_pyarrow_compute
 from narwhals.dependencies import get_pyarrow_parquet
 from narwhals.utils import flatten
 from narwhals.utils import generate_unique_token
@@ -387,3 +388,76 @@ class ArrowDataFrame:
     def write_parquet(self, file: Any) -> Any:
         pp = get_pyarrow_parquet()
         pp.write_table(self._native_dataframe, file)
+
+    def is_duplicated(self: Self) -> ArrowSeries:
+        from narwhals._arrow.series import ArrowSeries
+
+        np = get_numpy()
+        pa = get_pyarrow()
+        pc = get_pyarrow_compute()
+        df = self._native_dataframe
+
+        columns = self.columns
+        col_token = generate_unique_token(n_bytes=8, columns=columns)
+        row_count = (
+            df.append_column(col_token, pa.array(np.arange(len(self))))
+            .group_by(columns)
+            .aggregate([(col_token, "count")])
+        )
+        is_duplicated = pc.greater(
+            df.join(
+                row_count, keys=columns, right_keys=columns, join_type="inner"
+            ).column(f"{col_token}_count"),
+            1,
+        )
+        return ArrowSeries(is_duplicated, name="", backend_version=self._backend_version)
+
+    def is_unique(self: Self) -> ArrowSeries:
+        from narwhals._arrow.series import ArrowSeries
+
+        pc = get_pyarrow_compute()
+        is_duplicated = self.is_duplicated()._native_series
+
+        return ArrowSeries(
+            pc.invert(is_duplicated), name="", backend_version=self._backend_version
+        )
+
+    def unique(
+        self: Self,
+        subset: str | list[str] | None,
+        *,
+        keep: Literal["any", "first", "last", "none"] = "any",
+        maintain_order: bool = False,  # noqa: ARG002
+    ) -> Self:
+        """
+        NOTE:
+            The param `maintain_order` is only here for compatibility with the polars API
+            and has no effect on the output.
+        """
+
+        np = get_numpy()
+        pa = get_pyarrow()
+        pc = get_pyarrow_compute()
+
+        df = self._native_dataframe
+
+        if isinstance(subset, str):
+            subset = [subset]
+        subset = subset or self.columns
+
+        if keep in {"any", "first", "last"}:
+            agg_func_map = {"any": "min", "first": "min", "last": "max"}
+
+            agg_func = agg_func_map[keep]
+            col_token = generate_unique_token(n_bytes=8, columns=self.columns)
+            keep_idx = (
+                df.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._from_native_dataframe(pc.take(df, keep_idx))
+
+        keep_idx = self.select(*subset).is_unique()
+        return self.filter(keep_idx)
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
index 6606e8f1..70670287 100644
--- a/narwhals/_pandas_like/dataframe.py
+++ b/narwhals/_pandas_like/dataframe.py
@@ -446,10 +446,23 @@ class PandasLikeDataFrame:
     def tail(self, n: int) -> Self:
         return self._from_native_dataframe(self._native_dataframe.tail(n))
 
-    def unique(self, subset: str | list[str]) -> Self:
-        subset = flatten(subset)
+    def unique(
+        self: Self,
+        subset: str | list[str] | None,
+        *,
+        keep: Literal["any", "first", "last", "none"] = "any",
+        maintain_order: bool = False,  # noqa: ARG002
+    ) -> Self:
+        """
+        NOTE:
+            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)
+        subset = flatten(subset) if subset else None
         return self._from_native_dataframe(
-            self._native_dataframe.drop_duplicates(subset=subset)
+            self._native_dataframe.drop_duplicates(subset=subset, keep=mapped_keep)
         )
 
     # --- lazy-only ---
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index f4aa363d..ba3fa3de 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -140,8 +140,18 @@ class BaseFrame(Generic[FrameT]):
     def drop(self, *columns: str | Iterable[str]) -> Self:
         return self._from_compliant_dataframe(self._compliant_frame.drop(*columns))
 
-    def unique(self, subset: str | list[str]) -> Self:
-        return self._from_compliant_dataframe(self._compliant_frame.unique(subset=subset))
+    def unique(
+        self,
+        subset: str | list[str] | None = None,
+        *,
+        keep: Literal["any", "first", "last", "none"] = "any",
+        maintain_order: bool = False,
+    ) -> Self:
+        return self._from_compliant_dataframe(
+            self._compliant_frame.unique(
+                subset=subset, keep=keep, maintain_order=maintain_order
+            )
+        )
 
     def filter(self, *predicates: IntoExpr | Iterable[IntoExpr]) -> Self:
         predicates, _ = self._flatten_and_extract(*predicates)
@@ -1298,12 +1308,29 @@ class DataFrame(BaseFrame[FrameT]):
         """
         return super().drop(*columns)
 
-    def unique(self, subset: str | list[str]) -> Self:
+    def unique(
+        self,
+        subset: str | list[str] | None = None,
+        *,
+        keep: Literal["any", "first", "last", "none"] = "any",
+        maintain_order: bool = False,
+    ) -> Self:
         """
         Drop duplicate rows from this dataframe.
 
         Arguments:
             subset: Column name(s) to consider when identifying duplicate rows.
+            keep: {'first', 'last', 'any', 'none'}
+                Which of the duplicate rows to keep.
+
+                * 'any': Does not give any guarantee of which row is kept.
+                        This allows more optimizations.
+                * 'none': Don't keep duplicate rows.
+                * 'first': Keep first unique row.
+                * 'last': Keep last unique row.
+            maintain_order: Keep the same order as the original DataFrame. This is more
+                expensive to compute. Settings this to `True` blocks the possibility
+                to run on the streaming engine for polars.
 
         Examples:
             >>> import pandas as pd
@@ -1338,7 +1365,7 @@ class DataFrame(BaseFrame[FrameT]):
             │ 1   ┆ a   ┆ b   │
             └─────┴─────┴─────┘
         """
-        return super().unique(subset)
+        return super().unique(subset, keep=keep, maintain_order=maintain_order)
 
     def filter(self, *predicates: IntoExpr | Iterable[IntoExpr]) -> Self:
         r"""
@@ -2658,13 +2685,30 @@ class LazyFrame(BaseFrame[FrameT]):
         """
         return super().drop(*columns)
 
-    def unique(self, subset: str | list[str]) -> Self:
+    def unique(
+        self,
+        subset: str | list[str] | None = None,
+        *,
+        keep: Literal["any", "first", "last", "none"] = "any",
+        maintain_order: bool = False,
+    ) -> Self:
         """
         Drop duplicate rows from this LazyFrame.
 
         Arguments:
             subset: Column name(s) to consider when identifying duplicate rows.
                      If set to `None`, use all columns.
+            keep: {'first', 'last', 'any', 'none'}
+                Which of the duplicate rows to keep.
+
+                * 'any': Does not give any guarantee of which row is kept.
+                        This allows more optimizations.
+                * 'none': Don't keep duplicate rows.
+                * 'first': Keep first unique row.
+                * 'last': Keep last unique row.
+            maintain_order: Keep the same order as the original DataFrame. This is more
+                expensive to compute. Settings this to `True` blocks the possibility
+                to run on the streaming engine for polars.
 
         Returns:
             LazyFrame: LazyFrame with unique rows.
@@ -2702,7 +2746,7 @@ class LazyFrame(BaseFrame[FrameT]):
             │ 1   ┆ a   ┆ b   │
             └─────┴─────┴─────┘
         """
-        return super().unique(subset)
+        return super().unique(subset, keep=keep, maintain_order=maintain_order)
 
     def filter(self, *predicates: IntoExpr | Iterable[IntoExpr]) -> Self:
         r"""
diff --git a/utils/check_backend_completeness.py b/utils/check_backend_completeness.py
index 1c74f68a..dacd7ead 100644
--- a/utils/check_backend_completeness.py
+++ b/utils/check_backend_completeness.py
@@ -12,11 +12,8 @@ import narwhals as nw
 from narwhals._arrow.dataframe import ArrowDataFrame
 
 MISSING = [
-    "DataFrame.is_duplicated",
-    "DataFrame.is_unique",
     "DataFrame.iter_rows",
     "DataFrame.pipe",
-    "DataFrame.unique",
     "Series.drop_nulls",
     "Series.from_iterable",
     "Series.is_between",


Test patch:
diff --git a/tests/frame/is_duplicated_test.py b/tests/frame/is_duplicated_test.py
index 55fa70a4..041d6841 100644
--- a/tests/frame/is_duplicated_test.py
+++ b/tests/frame/is_duplicated_test.py
@@ -3,15 +3,11 @@ from __future__ import annotations
 from typing import Any
 
 import numpy as np
-import pytest
 
 import narwhals.stable.v1 as nw
 
 
-def test_is_duplicated(request: Any, constructor: Any) -> None:
-    if "pyarrow_table" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
-
+def test_is_duplicated(constructor: Any) -> None:
     data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
     df_raw = constructor(data)
     df = nw.from_native(df_raw, eager_only=True)
diff --git a/tests/frame/is_unique_test.py b/tests/frame/is_unique_test.py
index 9a7aa475..2f41b672 100644
--- a/tests/frame/is_unique_test.py
+++ b/tests/frame/is_unique_test.py
@@ -3,14 +3,11 @@ from __future__ import annotations
 from typing import Any
 
 import numpy as np
-import pytest
 
 import narwhals.stable.v1 as nw
 
 
-def test_is_unique(request: Any, constructor: Any) -> None:
-    if "pyarrow_table" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
+def test_is_unique(constructor: Any) -> None:
     data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
     df_raw = constructor(data)
     df = nw.from_native(df_raw, eager_only=True)
diff --git a/tests/frame/unique_test.py b/tests/frame/unique_test.py
index 3f6ff4c9..2556358a 100644
--- a/tests/frame/unique_test.py
+++ b/tests/frame/unique_test.py
@@ -7,16 +7,35 @@ import pytest
 import narwhals.stable.v1 as nw
 from tests.utils import compare_dicts
 
+data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
 
-def test_unique(request: Any, constructor_with_lazy: Any) -> None:
-    if "pyarrow_table" in str(constructor_with_lazy):
-        request.applymarker(pytest.mark.xfail)
-    data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
+
+@pytest.mark.parametrize("subset", ["b", ["b"]])
+@pytest.mark.parametrize(
+    ("keep", "expected"),
+    [
+        ("first", {"a": [1, 2], "b": [4, 6], "z": [7.0, 9.0]}),
+        ("last", {"a": [3, 2], "b": [4, 6], "z": [8.0, 9.0]}),
+        ("any", {"a": [1, 2], "b": [4, 6], "z": [7.0, 9.0]}),
+        ("none", {"a": [2], "b": [6], "z": [9]}),
+    ],
+)
+def test_unique(
+    constructor_with_lazy: Any,
+    subset: str | list[str] | None,
+    keep: str,
+    expected: dict[str, list[float]],
+) -> None:
     df_raw = constructor_with_lazy(data)
-    df = nw.from_native(df_raw).lazy()
-    result = nw.to_native(df.unique("b").sort("b"))
-    expected = {"a": [1, 2], "b": [4, 6], "z": [7.0, 9.0]}
-    compare_dicts(result, expected)
-    result = nw.to_native(df.collect().unique("b").sort("b"))
-    expected = {"a": [1, 2], "b": [4, 6], "z": [7.0, 9.0]}
+    df = nw.from_native(df_raw)
+
+    result = df.unique(subset, keep=keep, maintain_order=True)  # type: ignore[arg-type]
     compare_dicts(result, expected)
+
+
+def test_unique_none(constructor_with_lazy: Any) -> None:
+    df_raw = constructor_with_lazy(data)
+    df = nw.from_native(df_raw)
+
+    result = df.unique(maintain_order=True)
+    compare_dicts(result, data)