Enable groupby list for strings

rapidsai · Dec 4, 2020 · 9c2741e · 9c2741e
1 parent 358d582
commit 9c2741e
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/python/cudf/cudf/_lib/groupby.pyx b/python/cudf/cudf/_lib/groupby.pyx
@@ -51,6 +51,7 @@ _STRING_AGGS = {
     "min",
     "nunique",
     "nth",
+    "collect"
 }
 
 _LIST_AGGS = {

diff --git a/python/cudf/cudf/tests/test_groupby.py b/python/cudf/cudf/tests/test_groupby.py
@@ -1268,6 +1268,18 @@ def test_groupby_list_single_element(list_agg):
     )
 
 
+@pytest.mark.parametrize("list_agg", [list, "collect"])
+def test_groupby_list_strings(list_agg):
+    pdf = pd.DataFrame({"a": [1, 1, 1, 2, 2], "b": ["b", "a", "c", "e", "d"]})
+    gdf = cudf.from_pandas(pdf)
+
+    assert_eq(
+        pdf.groupby("a").agg({"b": list}),
+        gdf.groupby("a").agg({"b": list_agg}),
+        check_dtype=False,
+    )
+
+
 def test_groupby_list_columns_excluded():
     pdf = pd.DataFrame(
         {
-Original file line number
+Diff line change
@@ Expand Up / @@ -51,6 +51,7 @@ _STRING_AGGS = { @@
         "min",
         "nunique",
         "nth",
+        "collect"
     }
     _LIST_AGGS = {
@@ Expand Down @@