fetch correct grouping keys

rapidsai · Aug 12, 2021 · ad7f2b2 · ad7f2b2
1 parent e3f35af
commit ad7f2b2
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/python/dask_cudf/dask_cudf/groupby.py b/python/dask_cudf/dask_cudf/groupby.py
@@ -16,6 +16,8 @@
 from dask.dataframe.groupby import DataFrameGroupBy, SeriesGroupBy
 from dask.highlevelgraph import HighLevelGraph
 
+import cudf
+
 
 class CudfDataFrameGroupBy(DataFrameGroupBy):
     def __init__(self, *args, **kwargs):
@@ -79,9 +81,14 @@ def aggregate(self, arg, split_every=None, split_out=1):
             and isinstance(self.index, (str, list))
             and _is_supported(arg, _supported)
         ):
+            if isinstance(self._meta.grouping.keys, cudf.MultiIndex):
+                keys = self._meta.grouping.keys.names
+            else:
+                keys = self._meta.grouping.keys.name
+
             return groupby_agg(
                 self.obj,
-                self.index,
+                keys,
                 arg,
                 split_every=split_every,
                 split_out=split_out,