From 3971bd27c96b68b859399564dbb6abdb93de5f14 Mon Sep 17 00:00:00 2001 From: TrevorBergeron Date: Tue, 19 Mar 2024 10:31:34 -0700 Subject: [PATCH] fix: fix grouping series on multiple other series (#455) --- 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 8c3f1352f2..d01ee88cf5 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 4ce3dcfe2c..f63ea977ff 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()