Merge branch 'main' into sycai_fix

Author: sycai

bigframes/dataframe.py

@@ -2926,9 +2926,23 @@ def nunique(self) -> bigframes.series.Series:
         return bigframes.series.Series(block)
 
     def agg(
-        self, func: str | typing.Sequence[str]
+        self,
+        func: str
+        | typing.Sequence[str]
+        | typing.Mapping[blocks.Label, typing.Sequence[str] | str],
     ) -> DataFrame | bigframes.series.Series:
-        if utils.is_list_like(func):
+        if utils.is_dict_like(func):
+            # Must check dict-like first because dictionaries are list-like
+            # according to Pandas.
+            agg_cols = []
+            for col_label, agg_func in func.items():
+                agg_cols.append(self[col_label].agg(agg_func))
+
+            from bigframes.core.reshape import api as reshape
+
+            return reshape.concat(agg_cols, axis=1)
+
+        elif utils.is_list_like(func):
             aggregations = [agg_ops.lookup_agg_func(f) for f in func]
 
             for dtype, agg in itertools.product(self.dtypes, aggregations):
@@ -2942,6 +2956,7 @@ def agg(
                     aggregations,
                 )
             )
+
         else:
             return bigframes.series.Series(
                 self._block.aggregate_all_and_stack(

tests/system/small/test_dataframe.py

@@ -5652,3 +5652,29 @@ def test_astype_invalid_type_fail(scalars_dfs):
 
     with pytest.raises(TypeError, match=r".*Share your usecase with.*"):
         bf_df.astype(123)
+
+
+def test_agg_with_dict(scalars_dfs):
+    bf_df, pd_df = scalars_dfs
+    agg_funcs = {
+        "int64_too": ["min", "max"],
+        "int64_col": ["min", "count"],
+    }
+
+    bf_result = bf_df.agg(agg_funcs).to_pandas()
+    pd_result = pd_df.agg(agg_funcs)
+
+    pd.testing.assert_frame_equal(
+        bf_result, pd_result, check_dtype=False, check_index_type=False
+    )
+
+
+def test_agg_with_dict_containing_non_existing_col_raise_key_error(scalars_dfs):
+    bf_df, _ = scalars_dfs
+    agg_funcs = {
+        "int64_too": ["min", "max"],
+        "nonexisting_col": ["count"],
+    }
+
+    with pytest.raises(KeyError):
+        bf_df.agg(agg_funcs)