REF: Move Categorical logic from Grouping (#46203)

Author: NickCrews

pandas/core/groupby/categorical.py

@@ -1,14 +1,72 @@
 from __future__ import annotations
 
+import dataclasses
+from typing import final
+
 import numpy as np
 
+from pandas._typing import npt
+
 from pandas.core.algorithms import unique1d
 from pandas.core.arrays.categorical import (
  Categorical,
  CategoricalDtype,
  recode_for_categories,
 )
-from pandas.core.indexes.api import CategoricalIndex
+from pandas.core.indexes.api import (
+ CategoricalIndex,
+ Index,
+)
+
+
+@final
+@dataclasses.dataclass
+class CategoricalGrouper:
+ original_grouping_vector: Categorical | None
+ new_grouping_vector: Categorical
+ observed: bool
+ sort: bool
+
+ def result_index(self, group_index: Index) -> Index:
+ if self.original_grouping_vector is None:
+ return group_index
+ assert isinstance(group_index, CategoricalIndex)
+ return recode_from_groupby(
+ self.original_grouping_vector, self.sort, group_index
+ )
+
+ @classmethod
+ def make(cls, grouping_vector, sort: bool, observed: bool) -> CategoricalGrouper:
+ new_grouping_vector, original_grouping_vector = recode_for_groupby(
+ grouping_vector, sort, observed
+ )
+ return cls(
+ original_grouping_vector=original_grouping_vector,
+ new_grouping_vector=new_grouping_vector,
+ observed=observed,
+ sort=sort,
+ )
+
+ def codes_and_uniques(
+ self, cat: Categorical
+ ) -> tuple[npt.NDArray[np.signedinteger], Categorical]:
+ """
+ This does a version of `algorithms.factorize()` that works on categoricals.
+ """
+ categories = cat.categories
+
+ if self.observed:
+ ucodes = unique1d(cat.codes)
+ ucodes = ucodes[ucodes != -1]
+ if self.sort or cat.ordered:
+ ucodes = np.sort(ucodes)
+ else:
+ ucodes = np.arange(len(categories))
+
+ uniques = Categorical.from_codes(
+ codes=ucodes, categories=categories, ordered=cat.ordered
+ )
+ return cat.codes, uniques
 
 
 def recode_for_groupby(

pandas/core/groupby/grouper.py

@@ -8,13 +8,15 @@
  TYPE_CHECKING,
  Any,
  Hashable,
+ cast,
  final,
 )
 import warnings
 
 import numpy as np
 
 from pandas._typing import (
+ AnyArrayLike,
  ArrayLike,
  NDFrameT,
  npt,
@@ -38,12 +40,8 @@
 import pandas.core.common as com
 from pandas.core.frame import DataFrame
 from pandas.core.groupby import ops
-from pandas.core.groupby.categorical import (
- recode_for_groupby,
- recode_from_groupby,
-)
+from pandas.core.groupby.categorical import CategoricalGrouper
 from pandas.core.indexes.api import (
- CategoricalIndex,
  Index,
  MultiIndex,
 )
@@ -461,8 +459,7 @@ class Grouping:
 
  _codes: npt.NDArray[np.signedinteger] | None = None
  _group_index: Index | None = None
- _passed_categorical: bool
- _all_grouper: Categorical | None
+ _cat_grouper: CategoricalGrouper | None = None
  _index: Index
 
  def __init__(
@@ -479,16 +476,12 @@ def __init__(
  self.level = level
  self._orig_grouper = grouper
  self.grouping_vector = _convert_grouper(index, grouper)
- self._all_grouper = None
  self._index = index
  self._sort = sort
  self.obj = obj
- self._observed = observed
  self.in_axis = in_axis
  self._dropna = dropna
 
- self._passed_categorical = False
-
  # we have a single grouper which may be a myriad of things,
  # some of which are dependent on the passing in level
 
@@ -527,13 +520,10 @@ def __init__(
  self.grouping_vector = Index(ng, name=newgrouper.result_index.name)
 
  elif is_categorical_dtype(self.grouping_vector):
- # a passed Categorical
- self._passed_categorical = True
-
- self.grouping_vector, self._all_grouper = recode_for_groupby(
+ self._cat_grouper = CategoricalGrouper.make(
  self.grouping_vector, sort, observed
  )
-
+ self.grouping_vector = self._cat_grouper.new_grouping_vector
  elif not isinstance(
  self.grouping_vector, (Series, Index, ExtensionArray, np.ndarray)
  ):
@@ -631,20 +621,23 @@ def group_arraylike(self) -> ArrayLike:
  # _group_index is set in __init__ for MultiIndex cases
  return self._group_index._values
 
- elif self._all_grouper is not None:
+ elif (
+ self._cat_grouper is not None
+ and self._cat_grouper.original_grouping_vector is not None
+ ):
  # retain dtype for categories, including unobserved ones
  return self.result_index._values
 
- return self._codes_and_uniques[1]
+ return cast(ArrayLike, self._codes_and_uniques[1])
 
  @cache_readonly
  def result_index(self) -> Index:
- # result_index retains dtype for categories, including unobserved ones,
- # which group_index does not
- if self._all_grouper is not None:
-  group_idx = self.group_index
-  assert isinstance(group_idx, CategoricalIndex)
- return recode_from_groupby(self._all_grouper, self._sort, group_idx)
+ """
+ result_index retains dtype for categories, including unobserved ones,
+ which group_index does not
+ """
+ if self._cat_grouper is not None:
+ return self._cat_grouper.result_index(self.group_index)
  return self.group_index
 
  @cache_readonly
@@ -657,26 +650,10 @@ def group_index(self) -> Index:
  return Index._with_infer(uniques, name=self.name)
 
  @cache_readonly
- def _codes_and_uniques(self) -> tuple[npt.NDArray[np.signedinteger], ArrayLike]:
- if self._passed_categorical:
- # we make a CategoricalIndex out of the cat grouper
- # preserving the categories / ordered attributes
- cat = self.grouping_vector
- categories = cat.categories
-
- if self._observed:
- ucodes = algorithms.unique1d(cat.codes)
- ucodes = ucodes[ucodes != -1]
- if self._sort or cat.ordered:
- ucodes = np.sort(ucodes)
- else:
- ucodes = np.arange(len(categories))
-
- uniques = Categorical.from_codes(
- codes=ucodes, categories=categories, ordered=cat.ordered
- )
- return cat.codes, uniques
-
+ def _codes_and_uniques(self) -> tuple[npt.NDArray[np.signedinteger], AnyArrayLike]:
+ uniques: AnyArrayLike
+ if self._cat_grouper is not None:
+ return self._cat_grouper.codes_and_uniques(self.grouping_vector)
  elif isinstance(self.grouping_vector, ops.BaseGrouper):
  # we have a list of groupers
  codes = self.grouping_vector.codes_info