@@ -162,8 +162,17 @@ def _indexer_and_to_sort(
162
162
]:
163
163
v = self .level
164
164
165
- codes = list (self .index .codes )
166
165
levs = list (self .index .levels )
166
+ codes = list (self .index .codes )
167
+
168
+ if not self .sort :
169
+ codes = [list (code ) for code in codes ]
170
+ ids_code = [
171
+ ({y : x for x , y in enumerate (sorted (set (code ), key = code .index ))}, code )
172
+ for code in codes
173
+ ]
174
+ codes = [np .array ([d [x ] for x in code ]) for d , code in ids_code ]
175
+
167
176
to_sort = codes [:v ] + codes [v + 1 :] + [codes [v ]]
168
177
sizes = tuple (len (x ) for x in levs [:v ] + levs [v + 1 :] + [levs [v ]])
169
178
@@ -174,25 +183,33 @@ def _indexer_and_to_sort(
174
183
return indexer , to_sort
175
184
176
185
@cache_readonly
177
- def sorted_labels (self ) -> list [np .ndarray ]:
186
+ def labels (self ) -> list [np .ndarray ]:
178
187
indexer , to_sort = self ._indexer_and_to_sort
179
188
if self .sort :
180
189
return [line .take (indexer ) for line in to_sort ]
181
190
return to_sort
182
191
183
- def _make_sorted_values (self , values : np .ndarray ) -> np .ndarray :
192
+ @cache_readonly
193
+ def sorted_labels (self ) -> list [np .ndarray ]:
184
194
if self .sort :
185
- indexer , _ = self ._indexer_and_to_sort
195
+ return self .labels
196
+
197
+ v = self .level
198
+ codes = list (self .index .codes )
199
+ to_sort = codes [:v ] + codes [v + 1 :] + [codes [v ]]
200
+ return to_sort
186
201
187
- sorted_values = algos .take_nd (values , indexer , axis = 0 )
188
- return sorted_values
189
- return values
202
+ def _make_sorted_values (self , values : np .ndarray ) -> np .ndarray :
203
+ indexer , _ = self ._indexer_and_to_sort
204
+ sorted_values = algos .take_nd (values , indexer , axis = 0 )
205
+ return sorted_values
190
206
191
207
def _make_selectors (self ):
192
208
new_levels = self .new_index_levels
193
209
194
210
# make the mask
195
- remaining_labels = self .sorted_labels [:- 1 ]
211
+ remaining_labels = self .labels [:- 1 ]
212
+ choosen_labels = self .labels [- 1 ]
196
213
level_sizes = tuple (len (x ) for x in new_levels )
197
214
198
215
comp_index , obs_ids = get_compressed_ids (remaining_labels , level_sizes )
@@ -202,7 +219,7 @@ def _make_selectors(self):
202
219
stride = self .index .levshape [self .level ] + self .lift
203
220
self .full_shape = ngroups , stride
204
221
205
- selector = self . sorted_labels [ - 1 ] + stride * comp_index + self .lift
222
+ selector = choosen_labels + stride * comp_index + self .lift
206
223
mask = np .zeros (np .prod (self .full_shape ), dtype = bool )
207
224
mask .put (selector , True )
208
225
0 commit comments