From 5331bc54b4fb853d2d5f544a06be320660e7482c Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Mon, 18 Mar 2024 18:20:18 +0000 Subject: [PATCH] fix: fix grouping series on multiple other series --- bigframes/series.py | 2 +- tests/system/small/test_series.py | 10 ++++++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/bigframes/series.py b/bigframes/series.py index ef2feb4f92..72b3c1e48b 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -1195,7 +1195,7 @@ def _groupby_values( get_column_right, ) = block.join(key._block, how="inner" if dropna else "left") - value_col = get_column_left[self._value_column] + value_col = get_column_left[value_col] grouping_cols = [ *[get_column_left[value] for value in grouping_cols], get_column_right[key._value_column], diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index e22037a1ce..4ffbefad59 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -1313,9 +1313,15 @@ def test_any(scalars_dfs): def test_groupby_sum(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs col_name = "int64_too" - bf_series = scalars_df[col_name].groupby(scalars_df["string_col"]).sum() + bf_series = ( + scalars_df[col_name] + .groupby([scalars_df["bool_col"], ~scalars_df["bool_col"]]) + .sum() + ) pd_series = ( - scalars_pandas_df[col_name].groupby(scalars_pandas_df["string_col"]).sum() + scalars_pandas_df[col_name] + .groupby([scalars_pandas_df["bool_col"], ~scalars_pandas_df["bool_col"]]) + .sum() ) # TODO(swast): Update groupby to use index based on group by key(s). bf_result = bf_series.to_pandas()