add tests

Author: chelsea-lin

bigframes/core/groupby/__init__.py

@@ -340,19 +340,29 @@ def _agg_list(self, func: typing.Sequence) -> df.DataFrame:
             for f in func
         ]
 
-        aggregated_columns = pd.MultiIndex.from_tuples(
-            [
-                self._block.col_id_to_label[col_id]
+        if self._block.column_labels.nlevels > 1:
+            # Restructure MultiIndex for proper format: (idx1, idx2, func)
+            # rather than ((idx1, idx2), func).
+            aggregated_columns = pd.MultiIndex.from_tuples(
+                [
+                    self._block.col_id_to_label[col_id]
+                    for col_id in self._aggregated_columns()
+                ],
+                names=[*self._block.column_labels.names],
+            ).to_frame(index=False)
+
+            column_labels = [
+                tuple(col_id) + (f,)
+                for col_id in aggregated_columns.to_numpy()
+                for f in func
+            ]
+        else:
+            column_labels = [
+                (self._block.col_id_to_label[col_id], f)
                 for col_id in self._aggregated_columns()
-            ],
-            names=[*self._block.column_labels.names],
-        ).to_frame(index=False)
+                for f in func
+            ]
 
-        column_labels = [
-            tuple(col_id) + (f,)
-            for col_id in aggregated_columns.to_numpy()
-            for f in func
-        ]
         agg_block, _ = self._block.aggregate(
             by_column_ids=self._by_col_ids,
             aggregations=aggregations,

tests/system/small/test_groupby.py

@@ -144,6 +144,23 @@ def test_dataframe_groupby_agg_list(scalars_df_index, scalars_pandas_df_index):
     pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False)
 
 
+def test_dataframe_groupby_agg_list_w_column_multi_index(
+    scalars_df_index, scalars_pandas_df_index
+):
+    columns = ["int64_too", "string_col", "bool_col"]
+    multi_columns = pd.MultiIndex.from_tuples(zip(["a", "b", "a"], columns))
+    bf_df = scalars_df_index[columns].copy()
+    bf_df.columns = multi_columns
+    pd_df = scalars_pandas_df_index[columns].copy()
+    pd_df.columns = multi_columns
+
+    bf_result = bf_df.groupby(level=0).agg(["count", "min"])
+    pd_result = pd_df.groupby(level=0).agg(["count", "min"])
+
+    bf_result_computed = bf_result.to_pandas()
+    pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False)
+
+
 @pytest.mark.parametrize(
     ("as_index"),
     [