From 4689e16e76774e8c45cca03cd9ae0bd26e9d0803 Mon Sep 17 00:00:00 2001 From: Anish Mehta Date: Wed, 14 Oct 2026 12:33:46 +0530 Subject: [PATCH] [Python] Handle NaN in DataFrame drop_duplicates, duplicated and nunique --- CHANGES.md | 1 + sdks/python/apache_beam/dataframe/frames.py | 51 ++++++++++++++++--- .../apache_beam/dataframe/frames_test.py | 19 +++++++ 3 files changed, 63 insertions(+), 8 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 74dcb41af94b..79938c9935c1 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -129,6 +129,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)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Security Fixes diff --git a/sdks/python/apache_beam/dataframe/frames.py b/sdks/python/apache_beam/dataframe/frames.py index 4b62fed72f0d..f99df35afe83 100644 --- a/sdks/python/apache_beam/dataframe/frames.py +++ b/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) diff --git a/sdks/python/apache_beam/dataframe/frames_test.py b/sdks/python/apache_beam/dataframe/frames_test.py index 004000452004..fbe3ef36d804 100644 --- a/sdks/python/apache_beam/dataframe/frames_test.py +++ b/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({