From 3d4db26cfa4ace0a0f2fbb602f422fe30670c35f Mon Sep 17 00:00:00 2001 From: Brian Hulette Date: Fri, 16 Apr 2021 14:02:52 -0700 Subject: [PATCH] [BEAM-9547] Add tests for groupby(dropna) (#14553) --- sdks/python/apache_beam/dataframe/frames_test.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/sdks/python/apache_beam/dataframe/frames_test.py b/sdks/python/apache_beam/dataframe/frames_test.py index 78031a4ee7d..f8197572632 100644 --- a/sdks/python/apache_beam/dataframe/frames_test.py +++ b/sdks/python/apache_beam/dataframe/frames_test.py @@ -686,6 +686,14 @@ def test_quantile_axis_columns(self): r"df\.quantile\(q=0\.1, axis='columns'\)"): self._run_test(lambda df: df.quantile([0.1, 0.5], axis='columns'), df) + @unittest.skipIf(PD_VERSION < (1, 1), "drop_na added in pandas 1.1.0") + def test_groupby_count_na(self): + # Verify we can do a groupby.count() that doesn't drop NaN values + self._run_test( + lambda df: df.groupby('foo', dropna=True).bar.count(), GROUPBY_DF) + self._run_test( + lambda df: df.groupby('foo', dropna=False).bar.count(), GROUPBY_DF) + class AllowNonParallelTest(unittest.TestCase): def _use_non_parallel_operation(self):