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,17 @@ 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
+ for i in range (len (codes )):
172
+ dd : collections .defaultdict = collections .defaultdict (
173
+ itertools .count ().__next__
174
+ )
175
+ codes [i ] = np .array ([dd [c ] for c in codes [i ]], dtype = codes [i ].dtype )
176
+
167
177
to_sort = codes [:v ] + codes [v + 1 :] + [codes [v ]]
168
178
sizes = tuple (len (x ) for x in levs [:v ] + levs [v + 1 :] + [levs [v ]])
169
179
@@ -174,25 +184,33 @@ def _indexer_and_to_sort(
174
184
return indexer , to_sort
175
185
176
186
@cache_readonly
177
- def sorted_labels (self ) -> list [np .ndarray ]:
187
+ def labels (self ) -> list [np .ndarray ]:
178
188
indexer , to_sort = self ._indexer_and_to_sort
179
189
if self .sort :
180
190
return [line .take (indexer ) for line in to_sort ]
181
191
return to_sort
182
192
183
- def _make_sorted_values (self , values : np .ndarray ) -> np .ndarray :
193
+ @cache_readonly
194
+ def sorted_labels (self ) -> list [np .ndarray ]:
184
195
if self .sort :
185
- indexer , _ = self ._indexer_and_to_sort
196
+ return self .labels
197
+
198
+ v = self .level
199
+ codes = list (self .index .codes )
200
+ to_sort = codes [:v ] + codes [v + 1 :] + [codes [v ]]
201
+ return to_sort
186
202
187
- sorted_values = algos .take_nd (values , indexer , axis = 0 )
188
- return sorted_values
189
- return values
203
+ def _make_sorted_values (self , values : np .ndarray ) -> np .ndarray :
204
+ indexer , _ = self ._indexer_and_to_sort
205
+ sorted_values = algos .take_nd (values , indexer , axis = 0 )
206
+ return sorted_values
190
207
191
208
def _make_selectors (self ):
192
209
new_levels = self .new_index_levels
193
210
194
211
# make the mask
195
- remaining_labels = self .sorted_labels [:- 1 ]
212
+ remaining_labels = self .labels [:- 1 ]
213
+ choosen_labels = self .labels [- 1 ]
196
214
level_sizes = tuple (len (x ) for x in new_levels )
197
215
198
216
comp_index , obs_ids = get_compressed_ids (remaining_labels , level_sizes )
@@ -202,7 +220,7 @@ def _make_selectors(self):
202
220
stride = self .index .levshape [self .level ] + self .lift
203
221
self .full_shape = ngroups , stride
204
222
205
- selector = self . sorted_labels [ - 1 ] + stride * comp_index + self .lift
223
+ selector = choosen_labels + stride * comp_index + self .lift
206
224
mask = np .zeros (np .prod (self .full_shape ), dtype = bool )
207
225
mask .put (selector , True )
208
226
0 commit comments