REF: Compute complete result_index upfront in groupby

pandas/core/groupby/grouper.py

@@ -678,7 +678,7 @@ def _ilevel(self) -> int | None:
 
  @property
  def ngroups(self) -> int:
- return len(self._codes_and_uniques[1])
+ return len(self.uniques)
 
  @cache_readonly
  def indices(self) -> dict[Hashable, npt.NDArray[np.intp]]:
@@ -715,32 +715,31 @@ def _codes_and_uniques(self) -> tuple[npt.NDArray[np.signedinteger], ArrayLike]:
  else:
  ucodes = np.arange(len(categories))
 
+ has_dropped_na = False
  if not self._dropna:
  na_mask = cat.codes < 0
  if np.any(na_mask):
+ has_dropped_na = True
  if self._sort:
- # Replace NA codes with `largest code + 1`
+ # NA goes at the end, gets `largest non-NA code + 1`
  na_code = len(categories)
  else:
- # Insert NA code into the codes based on first appearance
- # A negative code must exist, no need to check codes[na_idx] < 0
+ # Insert NA in result based on first appearance, need
+ # the number of unique codes prior
  na_idx = na_mask.argmax()
- # count number of unique codes that comes before the nan value
  na_code = algorithms.nunique_ints(cat.codes[:na_idx])
  ucodes = np.insert(ucodes, na_code, -1)
 
  uniques = Categorical.from_codes(
  codes=ucodes, categories=categories, ordered=cat.ordered, validate=False
  )
  codes = cat.codes
- if not self._dropna:
- na_mask = codes < 0
- if np.any(na_mask):
- if self._sort:
- codes = np.where(na_mask, na_code, codes)
- else:
- codes = np.where(codes >= na_code, codes + 1, codes)
- codes = np.where(na_mask, na_code, codes)
+
+ if has_dropped_na:
+ if not self._sort:
+ # NA code is based on first appearance, increment higher codes
+ codes = np.where(codes >= na_code, codes + 1, codes)
+ codes = np.where(na_mask, na_code, codes)
 
  return codes, uniques
 

pandas/core/groupby/ops.py

@@ -832,12 +832,15 @@ def result_index_and_ids(self) -> tuple[Index, np.ndarray]:
  ids = len(unob_index) * ob_ids + unob_ids
 
  if self._sort:
+ # Sort result_index and recode ids using the new order
  sorter = result_index.argsort()
  result_index = result_index.take(sorter)
  _, index = np.unique(sorter, return_index=True)
  ids = ensure_platform_int(ids)
  ids = index.take(ids)
  else:
+ # Recode ids and reorder result_index with observed groups up front,
+ # unobserved at the end
  ids, uniques = compress_group_index(ids, sort=False)
  ids = ensure_platform_int(ids)
  taker = np.concatenate(

pandas/tests/groupby/methods/test_value_counts.py

@@ -617,7 +617,7 @@ def test_categorical_single_grouper_with_only_observed_categories(
  )
 
  gp = education_df.astype("category").groupby(
- "country", as_index=as_index, observed=observed, sort=True
+ "country", as_index=as_index, observed=observed
  )
  result = gp.value_counts(normalize=normalize)