This is an automated email from the ASF dual-hosted git repository.
shunping pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 5c153878385 [Python] Handle NaN in DataFrame drop_duplicates,
duplicated and nunique (#40571)
5c153878385 is described below
commit 5c153878385e27ad1a67580d99d07da9a38ef94f
Author: Anish Mehta <[email protected]>
AuthorDate: Fri Oct 23 22:33:10 2026 +0530
[Python] Handle NaN in DataFrame drop_duplicates, duplicated and nunique
(#40571)
Co-authored-by: Shunping Huang <[email protected]>
---
CHANGES.md | 1 +
sdks/python/apache_beam/dataframe/frames.py | 51 ++++++++++++++++++++----
sdks/python/apache_beam/dataframe/frames_test.py | 19 +++++++++
3 files changed, 63 insertions(+), 8 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index 9d3f6fbac24..33e66880343 100644
--- a/CHANGES.md
+++ b/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)).
diff --git a/sdks/python/apache_beam/dataframe/frames.py
b/sdks/python/apache_beam/dataframe/frames.py
index 4b62fed72f0..f99df35afe8 100644
--- a/sdks/python/apache_beam/dataframe/frames.py
+++ b/sdks/python/apache_beam/dataframe/frames.py
@@ -1589,7 +1589,9 @@ class DeferredSeries(DeferredDataFrameOrSeries):
@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 @@ class DeferredDataFrame(DeferredDataFrameOrSeries):
"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 @@ class DeferredDataFrame(DeferredDataFrameOrSeries):
"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 00400045200..fbe3ef36d80 100644
--- a/sdks/python/apache_beam/dataframe/frames_test.py
+++ b/sdks/python/apache_beam/dataframe/frames_test.py
@@ -1199,6 +1199,25 @@ class DeferredFrameTest(_AbstractFrameTest):
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({
