From 93827f08eafc4c998be9f63bcbf9271d3e0be653 Mon Sep 17 00:00:00 2001 From: Anish Mehta Date: Sun, 11 Oct 2026 06:56:29 +0530 Subject: [PATCH] [Python] Keep count/size labels in DataFrame agg with a list --- CHANGES.md | 1 + sdks/python/apache_beam/dataframe/frames.py | 6 +++++- sdks/python/apache_beam/dataframe/frames_test.py | 8 ++++++++ 3 files changed, 14 insertions(+), 1 deletion(-) diff --git a/CHANGES.md b/CHANGES.md index cb5eced42f7c..ce1c373cec22 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -99,6 +99,7 @@ * (Python) Fixed stateful DoFns with side inputs sometimes taking the timer key coder from a side input instead of the main input, which could make the worker fail to decode timer keys with `Unknown type tag` ([#40374](https://github.com/apache/beam/issues/40374)). * (Python) `Duration` built from float seconds now rounds to the nearest microsecond instead of truncating, which could lose a microsecond ([#40263](https://github.com/apache/beam/issues/40263)). * (Python) `ReadFromText` with `escapechar` no longer skips a delimiter that directly follows an escaped delimiter, which merged two records into one ([#40459](https://github.com/apache/beam/issues/40459)). +* (Python) DataFrame API: `agg` with a list that includes `'count'` or `'size'` now labels those rows `count`/`size` instead of `sum` ([#40504](https://github.com/apache/beam/issues/40504)). * 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 452f519c58db..ea634beadeec 100644 --- a/sdks/python/apache_beam/dataframe/frames.py +++ b/sdks/python/apache_beam/dataframe/frames.py @@ -2081,8 +2081,12 @@ def aggregate(self, func, axis, *args, **kwargs): allow_nonparallel_final = True if _is_associative(base_func): agg_func = func + elif isinstance(func, list): + # Post-aggregate with sum, but keep the original label (e.g. 'count') + # in the result's index. + agg_func = {base_func: 'sum'} else: - agg_func = ['sum'] if isinstance(func, list) else 'sum' + agg_func = 'sum' else: intermediate = self._expr allow_nonparallel_final = None # i.e. don't change the value diff --git a/sdks/python/apache_beam/dataframe/frames_test.py b/sdks/python/apache_beam/dataframe/frames_test.py index 290d13adc405..41f3a913fb70 100644 --- a/sdks/python/apache_beam/dataframe/frames_test.py +++ b/sdks/python/apache_beam/dataframe/frames_test.py @@ -2190,6 +2190,14 @@ def test_dataframe_agg_modes(self): self._run_test(lambda df: df.agg({'A': ['sum', 'mean']}), df) self._run_test(lambda df: df.agg({'A': ['sum', 'mean'], 'B': 'min'}), df) + def test_agg_list_with_count_and_size(self): + s = pd.Series([1.0, None, 3.0, 4.0]) + self._run_test(lambda s: s.agg(['count']), s) + self._run_test(lambda s: s.agg(['count', 'size', 'mean']), s) + + df = pd.DataFrame({'A': [1.0, None, 3.0, 4.0], 'B': [2, 3, 5, 7]}) + self._run_test(lambda df: df.agg(['count', 'mean']), df) + @unittest.skipIf(PD_VERSION >= (2, 0), "level argument removed in Pandas 2") def test_series_agg_level(self): self._run_test(