Repository: narwhals-dev/narwhals
Pull request #1951 (2025-02-06 10:26:09)

feat: restore support of `is_duplicated` for pyarrow
<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [ ] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues

- Related issue #\<issue number\>
- Closes #\<issue number\>

## Checklist

- [ ] Code follows style guide (ruff)
- [ ] Tests added
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below



Patch:
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index 71f15c1b..e9783954 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -693,46 +693,22 @@ class ArrowDataFrame(CompliantDataFrame, CompliantLazyFrame):
     def is_duplicated(self: Self) -> ArrowSeries:
         from narwhals._arrow.series import ArrowSeries
 
-        columns = self.columns
-        if (
-            len(columns) > 1
-            and pc.any(
-                [pc.any(pc.is_null(self._native_frame[col])).as_py() for col in columns]
-            ).as_py()
-        ):
-            # blocked on https://github.com/apache/arrow/issues/32204
-            msg = "DataFrame.is_duplicated with more than one column when nulls are present is not yet supported for PyArrow."
-            raise NotImplementedError(msg)
-
-        index_token = generate_temporary_column_name(n_bytes=8, columns=columns)
-        col_token = generate_temporary_column_name(
-            n_bytes=8, columns=[*columns, index_token]
+        col_token = generate_temporary_column_name(n_bytes=8, columns=self.columns)
+        row_index = pa.array(range(len(self)))
+        keep_idx = (
+            self._native_frame.append_column(col_token, row_index)
+            .group_by(self.columns)
+            .aggregate([(col_token, "min"), (col_token, "max")])
         )
-        df = self.with_row_index(index_token)._native_frame
-        row_count = (
-            df.append_column(col_token, pa.repeat(pa.scalar(1), len(self)))
-            .group_by(columns)
-            .aggregate([(col_token, "sum")])
-        )
-        is_duplicated = pc.greater(
-            df.join(
-                row_count,
-                keys=columns,
-                right_keys=columns,
-                join_type="left outer",
-                use_threads=False,
-            )
-            .sort_by(index_token)
-            .column(f"{col_token}_sum"),
-            1,
-        )
-        res = ArrowSeries(
-            is_duplicated,
+        return ~ArrowSeries(
+            pc.and_(
+                pc.is_in(row_index, keep_idx[f"{col_token}_min"]),
+                pc.is_in(row_index, keep_idx[f"{col_token}_max"]),
+            ),
             name="",
             backend_version=self._backend_version,
             version=self._version,
         )
-        return res.fill_null(res.null_count() > 1, strategy=None, limit=None)
 
     def is_unique(self: Self) -> ArrowSeries:
         from narwhals._arrow.series import ArrowSeries


Test patch:
diff --git a/tests/frame/is_duplicated_test.py b/tests/frame/is_duplicated_test.py
index cb0ee4ea..f3caa19c 100644
--- a/tests/frame/is_duplicated_test.py
+++ b/tests/frame/is_duplicated_test.py
@@ -1,9 +1,5 @@
 from __future__ import annotations
 
-from contextlib import nullcontext as does_not_raise
-
-import pytest
-
 import narwhals.stable.v1 as nw
 from tests.utils import ConstructorEager
 from tests.utils import assert_equal_data
@@ -19,15 +15,9 @@ def test_is_duplicated(constructor_eager: ConstructorEager) -> None:
 
 
 def test_is_duplicated_with_nulls(constructor_eager: ConstructorEager) -> None:
-    context = (
-        pytest.raises(NotImplementedError)
-        if "pyarrow_table" in str(constructor_eager)
-        else does_not_raise()
-    )
     data = {"col1": [1, 2, 3], "col2": ["one", None, None]}
     df_raw = constructor_eager(data)
     df = nw.from_native(df_raw, eager_only=True)
-    with context:
-        result = df.is_duplicated()
-        expected = {"is_duplicated": [False, False, False]}
-        assert_equal_data({"is_duplicated": result}, expected)
+    result = df.is_duplicated()
+    expected = {"is_duplicated": [False, False, False]}
+    assert_equal_data({"is_duplicated": result}, expected)
diff --git a/tests/frame/is_unique_test.py b/tests/frame/is_unique_test.py
index 88dacd01..3d329414 100644
--- a/tests/frame/is_unique_test.py
+++ b/tests/frame/is_unique_test.py
@@ -12,3 +12,12 @@ def test_is_unique(constructor_eager: ConstructorEager) -> None:
     result = nw.concat([df, df.head(1)]).is_unique()
     expected = {"is_unique": [False, True, True, False]}
     assert_equal_data({"is_unique": result}, expected)
+
+
+def test_is_unique_with_nulls(constructor_eager: ConstructorEager) -> None:
+    data = {"col1": [1, 2, 3, 3], "col2": ["one", None, None, None]}
+    df_raw = constructor_eager(data)
+    df = nw.from_native(df_raw, eager_only=True)
+    result = df.is_unique()
+    expected = {"is_unique": [True, True, False, False]}
+    assert_equal_data({"is_unique": result}, expected)