Tests

Author: abeltavares

pandas/tests/groupby/transform/test_transform.py

@@ -19,6 +19,10 @@
     date_range,
 )
 import pandas._testing as tm
+from pandas.core.groupby import (
+    DataFrameGroupBy,
+    NamedAgg,
+)
 from pandas.tests.groupby import get_groupby_method_args
 
 
@@ -84,6 +88,42 @@ def demean(arr):
     tm.assert_frame_equal(result, expected)
 
 
+def test_transform_with_namedagg():
+    df = DataFrame({"A": list("aaabbbccc"), "B": range(9), "D": range(9, 18)})
+    result = df.groupby("A").transform(
+        b_min=NamedAgg(column="B", aggfunc="min"),
+        d_sum=NamedAgg(column="D", aggfunc="sum"),
+    )
+    expected = DataFrame(
+        {
+            "b_min": [0, 0, 0, 3, 3, 3, 6, 6, 6],
+            "d_sum": [30, 30, 30, 39, 39, 39, 48, 48, 48],
+        }
+    )
+    tm.assert_frame_equal(result, expected)
+
+
+def test_transform_with_list_like():
+    df = DataFrame({"col": list("aab"), "val": range(3)})
+    result = df.groupby("col").transform(["sum", "min"])
+    expected = DataFrame({"val_sum": [1, 1, 2], "val_min": [0, 0, 2]})
+    expected.columns = MultiIndex.from_tuples([("val", "sum"), ("val", "min")])
+    tm.assert_frame_equal(result, expected)
+
+
+def test_named_agg_to_dict():
+    named_agg_items = [
+        ("b_min", (NamedAgg(column="B", aggfunc="min"))),
+        ("d_sum", (NamedAgg(column="D", aggfunc="sum"))),
+    ]
+    result = DataFrameGroupBy._named_agg_to_dict(*named_agg_items)
+    expected = {
+        "b_min": NamedAgg(column="B", aggfunc="min"),
+        "d_sum": NamedAgg(column="D", aggfunc="sum"),
+    }
+    assert result == expected
+
+
 def test_transform_fast():
     df = DataFrame(
         {