diff --git a/sdks/python/apache_beam/dataframe/frames.py b/sdks/python/apache_beam/dataframe/frames.py index 500152f44905..4b62fed72f0d 100644 --- a/sdks/python/apache_beam/dataframe/frames.py +++ b/sdks/python/apache_beam/dataframe/frames.py @@ -2057,11 +2057,13 @@ def aggregate(self, func, axis, *args, **kwargs): if base_func in HAND_IMPLEMENTED_GLOBAL_AGGREGATIONS: result = getattr(self, base_func)(*args, **kwargs) if isinstance(func, list): + name = self._expr.proxy().name with expressions.allow_non_parallel_operations(True): return frame_base.DeferredFrame.wrap( expressions.ComputedExpression( f'wrap_aggregate_{base_func}', - lambda x: pd.Series(x, index=[base_func]), [result._expr], + lambda x: pd.Series(x, index=[base_func], name=name), + [result._expr], requires_partition_by=partitionings.Singleton(), preserves_partition_by=partitionings.Singleton())) else: diff --git a/sdks/python/apache_beam/dataframe/frames_test.py b/sdks/python/apache_beam/dataframe/frames_test.py index 67f61aa68c6a..004000452004 100644 --- a/sdks/python/apache_beam/dataframe/frames_test.py +++ b/sdks/python/apache_beam/dataframe/frames_test.py @@ -2190,6 +2190,12 @@ def test_series_agg_modes(self): self._run_test(lambda s: s.agg(['mean']), s) self._run_test(lambda s: s.agg('mean'), s) + def test_series_agg_list_keeps_name(self): + s = pd.Series(list(range(16)), name='x') + self._run_test(lambda s: s.agg(['mean']), s) + self._run_test(lambda s: s.agg(['sum', 'mean']), s) + self._run_test(lambda s: s.agg(['std', 'var']), s) + def test_dataframe_agg_modes(self): df = pd.DataFrame({'A': [1, 2, 3, 4], 'B': [2, 3, 5, 7]}) self._run_test(lambda df: df.agg('sum'), df)