1
1
from __future__ import annotations
2
2
3
+ import collections
3
4
import itertools
4
5
from typing import (
5
6
TYPE_CHECKING ,
@@ -162,8 +163,13 @@ def _indexer_and_to_sort(
162
163
]:
163
164
v = self .level
164
165
165
- codes = list (self .index .codes )
166
166
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
+
167
173
to_sort = codes [:v ] + codes [v + 1 :] + [codes [v ]]
168
174
sizes = tuple (len (x ) for x in levs [:v ] + levs [v + 1 :] + [levs [v ]])
169
175
@@ -174,25 +180,33 @@ def _indexer_and_to_sort(
174
180
return indexer , to_sort
175
181
176
182
@cache_readonly
177
- def sorted_labels (self ) -> list [np .ndarray ]:
183
+ def labels (self ) -> list [np .ndarray ]:
178
184
indexer , to_sort = self ._indexer_and_to_sort
179
185
if self .sort :
180
186
return [line .take (indexer ) for line in to_sort ]
181
187
return to_sort
182
188
183
- def _make_sorted_values (self , values : np .ndarray ) -> np .ndarray :
189
+ @cache_readonly
190
+ def sorted_labels (self ) -> list [np .ndarray ]:
184
191
if self .sort :
185
- indexer , _ = self ._indexer_and_to_sort
192
+ return self .labels
186
193
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
190
203
191
204
def _make_selectors (self ):
192
205
new_levels = self .new_index_levels
193
206
194
207
# make the mask
195
- remaining_labels = self .sorted_labels [:- 1 ]
208
+ remaining_labels = self .labels [:- 1 ]
209
+ choosen_labels = self .labels [- 1 ]
196
210
level_sizes = tuple (len (x ) for x in new_levels )
197
211
198
212
comp_index , obs_ids = get_compressed_ids (remaining_labels , level_sizes )
@@ -202,7 +216,7 @@ def _make_selectors(self):
202
216
stride = self .index .levshape [self .level ] + self .lift
203
217
self .full_shape = ngroups , stride
204
218
205
- selector = self . sorted_labels [ - 1 ] + stride * comp_index + self .lift
219
+ selector = choosen_labels + stride * comp_index + self .lift
206
220
mask = np .zeros (np .prod (self .full_shape ), dtype = bool )
207
221
mask .put (selector , True )
208
222
0 commit comments