[Python] Keep count/size labels in DataFrame agg with a list

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

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

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(