This is an automated email from the ASF dual-hosted git repository.
damccorm 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 3486966c4a1 [Python] Keep count/size labels in DataFrame agg with a
list (#40506)
3486966c4a1 is described below
commit 3486966c4a102eb6705a52fb1e74a615ff73293f
Author: Anish Mehta <[email protected]>
AuthorDate: Mon Oct 12 20:20:16 2026 +0530
[Python] Keep count/size labels in DataFrame agg with a list (#40506)
---
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 6b01d163823..487d4e67d6a 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -100,6 +100,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 452f519c58d..ea634beadee 100644
--- a/sdks/python/apache_beam/dataframe/frames.py
+++ b/sdks/python/apache_beam/dataframe/frames.py
@@ -2081,8 +2081,12 @@ class DeferredSeries(DeferredDataFrameOrSeries):
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 290d13adc40..41f3a913fb7 100644
--- a/sdks/python/apache_beam/dataframe/frames_test.py
+++ b/sdks/python/apache_beam/dataframe/frames_test.py
@@ -2190,6 +2190,14 @@ class AggregationTest(_AbstractFrameTest):
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(
