Update indices property from groupby

Author: T. Koskamp

pandas/core/groupby/groupby.py

@@ -636,7 +636,7 @@ def indices(self) -> dict[Hashable, npt.NDArray[np.intp]]:
         return self._grouper.indices
 
     @final
-    def _get_indices(self, names):
+    def _get_indices(self, name):
         """
         Safe get multiple indices, translate keys for
         datelike to underlying repr.
@@ -649,28 +649,27 @@ def get_converter(s):
                 return lambda key: Timestamp(key)
             elif isinstance(s, np.datetime64):
                 return lambda key: Timestamp(key).asm8
-            elif isna(s):
-                return lambda key: np.nan
             else:
                 return lambda key: key
 
-        if len(names) == 0:
-            return []
+        if isna(name):
+            return self.indices.get(np.nan, [])
+        if isinstance(name, tuple):
+            name = tuple(np.nan if isna(comp) else comp for comp in name)
 
         if len(self.indices) > 0:
             index_sample = next(iter(self.indices))
         else:
             index_sample = None  # Dummy sample
 
-        name_sample = names[0]
         if isinstance(index_sample, tuple):
-            if not isinstance(name_sample, tuple):
+            if not isinstance(name, tuple):
                 msg = "must supply a tuple to get_group with multiple grouping keys"
                 raise ValueError(msg)
-            if not len(name_sample) == len(index_sample):
+            if not len(name) == len(index_sample):
                 try:
                     # If the original grouper was a tuple
-                    return [self.indices[name] for name in names]
+                    return self.indices[name]
                 except KeyError as err:
                     # turns out it wasn't a tuple
                     msg = (
@@ -679,41 +678,20 @@ def get_converter(s):
                     )
                     raise ValueError(msg) from err
 
-            has_nan = any(isna(n) for n in name_sample)
-
-            sample = name_sample if has_nan else index_sample
-            converters = (get_converter(s) for s in sample)
-
-            names = (
-                tuple(f(n) for f, n in zip(converters, name, strict=True))
-                for name in names
-            )
-
-            indices = self.indices
-            if not self.dropna and has_nan:
-                indices = {}
-                for k, v in self.indices.items():
-                    k = tuple(np.nan if isna(e) else e for e in k)
-                    indices[k] = v
+            converters = (get_converter(s) for s in index_sample)
+            name = tuple(f(n) for f, n in zip(converters, name, strict=True))
         else:
-            has_nan = isna(name_sample)
-
-            convert_sample = name_sample if has_nan else index_sample
-            converter = get_converter(convert_sample)
-            names = (converter(name) for name in names)
-
-            indices = self.indices
-            if not self.dropna and has_nan:
-                indices = {np.nan if isna(k) else k: v for k, v in indices.items()}
+            converter = get_converter(index_sample)
+            name = converter(name)
 
-        return [indices.get(name, []) for name in names]
+        return self.indices.get(name, [])
 
     @final
     def _get_index(self, name):
         """
         Safe get index, translate keys for datelike to underlying repr.
         """
-        return self._get_indices([name])[0]
+        return self._get_indices(name)
 
     @final
     @cache_readonly

pandas/core/groupby/ops.py

@@ -652,9 +652,24 @@ def indices(self) -> dict[Hashable, npt.NDArray[np.intp]]:
         """dict {group name -> group indices}"""
         if len(self.groupings) == 1 and isinstance(self.result_index, CategoricalIndex):
             # This shows unused categories in indices GH#38642
-            return self.groupings[0].indices
-        codes_list = [ping.codes for ping in self.groupings]
-        return get_indexer_dict(codes_list, self.levels)
+            result = self.groupings[0].indices
+        else:
+            codes_list = [ping.codes for ping in self.groupings]
+            result = get_indexer_dict(codes_list, self.levels)
+        if not self.dropna:
+            has_mi = isinstance(self.result_index, MultiIndex)
+            if not has_mi and self.result_index.hasnans:
+                result = {
+                    np.nan if isna(key) else key: value for key, value in result.items()
+                }
+            elif has_mi:
+                # MultiIndex has no efficient way to tell if there are NAs
+                result = {
+                    tuple(np.nan if isna(comp) else comp for comp in key): value
+                    for key, value in result.items()
+                }
+
+        return result
 
     @final
     @cache_readonly