11from __future__ import annotations
22
3+ import collections
34import itertools
45from typing import (
56 TYPE_CHECKING ,
@@ -162,8 +163,13 @@ def _indexer_and_to_sort(
162163 ]:
163164 v = self .level
164165
165- codes = list (self .index .codes )
166166 levs = list (self .index .levels )
167+ codes = list (self .index .codes )
168+
169+ if not self .sort :
170+ # Create new codes considering that labels are already sorted
171+ codes = [np .array (factorize (code )[0 ], dtype = code .dtype ) for code in codes ]
172+
167173 to_sort = codes [:v ] + codes [v + 1 :] + [codes [v ]]
168174 sizes = tuple (len (x ) for x in levs [:v ] + levs [v + 1 :] + [levs [v ]])
169175
@@ -174,25 +180,33 @@ def _indexer_and_to_sort(
174180 return indexer , to_sort
175181
176182 @cache_readonly
177- def sorted_labels (self ) -> list [np .ndarray ]:
183+ def labels (self ) -> list [np .ndarray ]:
178184 indexer , to_sort = self ._indexer_and_to_sort
179185 if self .sort :
180186 return [line .take (indexer ) for line in to_sort ]
181187 return to_sort
182188
183- def _make_sorted_values (self , values : np .ndarray ) -> np .ndarray :
189+ @cache_readonly
190+ def sorted_labels (self ) -> list [np .ndarray ]:
184191 if self .sort :
185- indexer , _ = self ._indexer_and_to_sort
192+ return self .labels
186193
187- sorted_values = algos .take_nd (values , indexer , axis = 0 )
188- return sorted_values
189- return values
194+ v = self .level
195+ codes = list (self .index .codes )
196+ to_sort = codes [:v ] + codes [v + 1 :] + [codes [v ]]
197+ return to_sort
198+
199+ def _make_sorted_values (self , values : np .ndarray ) -> np .ndarray :
200+ indexer , _ = self ._indexer_and_to_sort
201+ sorted_values = algos .take_nd (values , indexer , axis = 0 )
202+ return sorted_values
190203
191204 def _make_selectors (self ):
192205 new_levels = self .new_index_levels
193206
194207 # make the mask
195- remaining_labels = self .sorted_labels [:- 1 ]
208+ remaining_labels = self .labels [:- 1 ]
209+ choosen_labels = self .labels [- 1 ]
196210 level_sizes = tuple (len (x ) for x in new_levels )
197211
198212 comp_index , obs_ids = get_compressed_ids (remaining_labels , level_sizes )
@@ -202,7 +216,7 @@ def _make_selectors(self):
202216 stride = self .index .levshape [self .level ] + self .lift
203217 self .full_shape = ngroups , stride
204218
205- selector = self . sorted_labels [ - 1 ] + stride * comp_index + self .lift
219+ selector = choosen_labels + stride * comp_index + self .lift
206220 mask = np .zeros (np .prod (self .full_shape ), dtype = bool )
207221 mask .put (selector , True )
208222
0 commit comments