Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations
3import numpy as np
5from pandas.core.algorithms import unique1d
6from pandas.core.arrays.categorical import (
7 Categorical,
8 CategoricalDtype,
9 recode_for_categories,
10)
13def recode_for_groupby(c: Categorical, sort: bool, observed: bool) -> Categorical:
14 """
15 Code the categories to ensure we can groupby for categoricals.
17 If observed=True, we return a new Categorical with the observed
18 categories only.
20 If sort=False, return a copy of self, coded with categories as
21 returned by .unique(), followed by any categories not appearing in
22 the data. If sort=True, return self.
24 This method is needed solely to ensure the categorical index of the
25 GroupBy result has categories in the order of appearance in the data
26 (GH-8868).
28 Parameters
29 ----------
30 c : Categorical
31 sort : bool
32 The value of the sort parameter groupby was called with.
33 observed : bool
34 Account only for the observed values
36 Returns
37 -------
38 Categorical
39 If sort=False, the new categories are set to the order of
40 appearance in codes (unless ordered=True, in which case the
41 original order is preserved), followed by any unrepresented
42 categories in the original order.
43 """
44 # we only care about observed values
45 if observed:
46 # In cases with c.ordered, this is equivalent to
47 # return c.remove_unused_categories(), c
49 take_codes = unique1d(c.codes[c.codes != -1])
51 if sort:
52 take_codes = np.sort(take_codes)
54 # we recode according to the uniques
55 categories = c.categories.take(take_codes)
56 codes = recode_for_categories(c.codes, c.categories, categories, copy=False)
58 # return a new categorical that maps our new codes
59 # and categories
60 dtype = CategoricalDtype(categories, ordered=c.ordered)
61 return Categorical._simple_new(codes, dtype=dtype)
63 # Already sorted according to c.categories; all is fine
64 if sort:
65 return c
67 # sort=False should order groups in as-encountered order (GH-8868)
69 # GH:46909: Re-ordering codes faster than using (set|add|reorder)_categories
70 # GH 38140: exclude nan from indexer for categories
71 unique_notnan_codes = unique1d(c.codes[c.codes != -1])
72 if sort:
73 unique_notnan_codes = np.sort(unique_notnan_codes)
74 if (num_cat := len(c.categories)) > len(unique_notnan_codes):
75 # GH 13179: All categories need to be present, even if missing from the data
76 missing_codes = np.setdiff1d(
77 np.arange(num_cat), unique_notnan_codes, assume_unique=True
78 )
79 take_codes = np.concatenate((unique_notnan_codes, missing_codes))
80 else:
81 take_codes = unique_notnan_codes
83 return Categorical(c, c.categories.take(take_codes))