Fix dtype usage

apache · Oct 3, 2023 · 1473eeb · 1473eeb
1 parent 94327fd
commit 1473eeb
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/sdks/python/apache_beam/dataframe/frames.py b/sdks/python/apache_beam/dataframe/frames.py
@@ -4594,8 +4594,9 @@ def wrapper(self, *args, **kwargs):
       return _unliftable_agg(meth)(self, *args, **kwargs)
 
     to_group = self._ungrouped.proxy().index
-    is_categorical_grouping = any(to_group.get_level_values(i).is_categorical()
-                                  for i in self._grouping_indexes)
+    is_categorical_grouping = any(
+        isinstance(to_group.get_level_values(i).dtype, pd.CategoricalDtype)
+        for i in self._grouping_indexes)
     groupby_kwargs = self._kwargs
     group_keys = self._group_keys
 
@@ -4647,8 +4648,9 @@ def wrapper(self, *args, **kwargs):
 
     to_group = self._ungrouped.proxy().index
     group_keys = self._group_keys
-    is_categorical_grouping = any(to_group.get_level_values(i).is_categorical()
-                                  for i in self._grouping_indexes)
+    is_categorical_grouping = any(
+        isinstance(to_group.get_level_values(i).dtype, pd.CategoricalDtype)
+        for i in self._grouping_indexes)
 
     groupby_kwargs = self._kwargs
     project = _maybe_project_func(self._projection)