[Python] Handle NaN in DataFrame drop_duplicates, duplicated and nunique

CHANGES.md

@@ -135,6 +135,7 @@
 * (Python) Setting `max_batch_weight` on a RunInference `ModelHandler` no longer fails with a `TypeError`: `BatchElements` now supports `max_batch_weight`, which limits the total `element_size_fn` of a batch in addition to its number of elements ([#40644](https://github.com/apache/beam/issues/40644)).
 * (Python) `BatchElements` now honors `element_size_fn` when `max_batch_duration_secs` is set, and no longer emits an empty batch before an element that exceeds the batch size on its own ([#40644](https://github.com/apache/beam/issues/40644)).
 * (Python) `KeyedModelHandler` and `MaybeKeyedModelHandler` now apply the `element_size_fn` and `batch_length_fn` of the wrapped `ModelHandler` to examples, instead of to `(key, example)` pairs ([#40644](https://github.com/apache/beam/issues/40644)).
+* (Python) DataFrame API: `drop_duplicates`, `duplicated` and `nunique` no longer fail on rows with NaN values or on categorical columns with unused categories, and `nunique(dropna=False)` now counts NaN ([#40570](https://github.com/apache/beam/issues/40570)).
 * (Java) BigQuery Storage Write API `TableRow` conversion now honors `allowMissingRequiredFields` when building protocol-buffer messages, so callers that pass `true` with required descriptors get partial messages instead of a conversion-time `UninitializedMessageException` ([#40189](https://github.com/apache/beam/issues/40189)).
 * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
 

sdks/python/apache_beam/dataframe/frames.py

@@ -1589,7 +1589,9 @@ def dot(self, other):
   @frame_base.with_docs_from(pd.Series)
   @frame_base.args_to_kwargs(pd.Series)
   @frame_base.populate_defaults(pd.Series)
-  def nunique(self, **kwargs):
+  def nunique(self, dropna, **kwargs):
+    if dropna:
+      self = self.dropna()  # pylint: disable=self-cls-assignment
     return self.drop_duplicates(keep="any").size
 
   @frame_base.with_docs_from(pd.Series)
@@ -2843,11 +2845,12 @@ def duplicated(self, keep, subset):
           "keep=\"any\" are supported.",
           reason="order-sensitive")
 
-    by = subset or list(self.columns)
-
-    return self.groupby(by).apply(
-        lambda df: pd.DataFrame(df.duplicated(keep=keep, subset=subset),
-                                columns=[None]))[None].droplevel(by)
+    return self._apply_with_duplicates_colocated(
+        'duplicated',
+        # set_axis keeps the index when df is empty (e.g. for the proxy),
+        # pandas returns a default index in that case.
+        lambda df: df.duplicated(keep=keep, subset=subset).set_axis(df.index),
+        subset)
 
   @frame_base.with_docs_from(pd.DataFrame)
   @frame_base.args_to_kwargs(pd.DataFrame)
@@ -2876,10 +2879,42 @@ def drop_duplicates(self, keep, subset, ignore_index):
           "the data.",
           reason="order-sensitive")
 
+    return self._apply_with_duplicates_colocated(
+        'drop_duplicates',
+        lambda df: df.drop_duplicates(keep=keep, subset=subset),
+        subset)
+
+  def _apply_with_duplicates_colocated(self, name, func, subset):
+    """Applies ``func`` to partitions in which all duplicate rows (considering
+    only the columns in ``subset``) are colocated.
+
+    Rows are colocated by a hash of their values instead of with groupby,
+    which drops rows with NaN values and fails on unobserved categories."""
     by = subset or list(self.columns)
+    hash_level = '__apache_beam_duplicates_hash'
+
+    def add_hash_level(df):
+      key = df[by].copy()
+      for col in key.columns:
+        if key[col].dtype.kind in 'fc':
+          # -0.0 and 0.0 hash differently but are duplicates in pandas
+          key[col] = key[col] + 0
+      hashes = pd.util.hash_pandas_object(key, index=False)
+      return df.set_index(
+          pd.Index(hashes.to_numpy(), name=hash_level), append=True)
+
+    with_hash = expressions.ComputedExpression(
+        'add_hash_level',
+        add_hash_level, [self._expr],
+        requires_partition_by=partitionings.Arbitrary(),
+        preserves_partition_by=partitionings.Singleton())
 
-    return self.groupby(by).apply(
-        lambda df: df.drop_duplicates(keep=keep, subset=subset)).droplevel(by)
+    return frame_base.DeferredFrame.wrap(
+        expressions.ComputedExpression(
+            name,
+            lambda df: func(df.droplevel(hash_level)), [with_hash],
+            requires_partition_by=partitionings.Index([hash_level]),
+            preserves_partition_by=partitionings.Singleton()))
 
   @frame_base.with_docs_from(pd.DataFrame)
   @frame_base.args_to_kwargs(pd.DataFrame)

sdks/python/apache_beam/dataframe/frames_test.py

@@ -1199,6 +1199,25 @@ def test_drop_duplicates(self):
         lambda df: df.drop_duplicates(subset=['brand', 'style'], keep=False),
         df)
 
+  def test_duplicates_with_nan(self):
+    df = pd.DataFrame({
+        'a': [1.0, 1.0, np.nan, np.nan, 2.0, 3.0, 0.0, -0.0],
+        'b': ['x', 'x', 'y', 'y', None, 'z', 'w', 'w'],
+    })
+
+    self._run_test(lambda df: df.drop_duplicates(keep=False), df)
+    self._run_test(lambda df: df.duplicated(keep=False), df)
+    self._run_test(lambda df: df.drop_duplicates(subset=['b'], keep=False), df)
+
+    s = pd.Series([1.0, 1.0, np.nan, 2.0], name='a')
+    self._run_test(lambda s: s.drop_duplicates(keep=False), s)
+    self._run_test(lambda s: s.nunique(), s)
+    self._run_test(lambda s: s.nunique(dropna=False), s)
+
+  def test_nunique_categorical_with_unused_categories(self):
+    s = pd.Series(['a', 'b', 'a'], dtype=pd.CategoricalDtype(['a', 'b', 'c']))
+    self._run_test(lambda s: s.nunique(), s)
+
   @parameterized.expand([
       (
           lambda base: base.from_dict({