@@ -153,6 +153,16 @@ def __init__(
153
153
154
154
self ._make_selectors ()
155
155
156
+ @cache_readonly
157
+ def sorted_labels (self ) -> list [np .ndarray ]:
158
+ if self .sort :
159
+ return self .labels
160
+
161
+ v = self .level
162
+ codes = list (self .index .codes )
163
+ to_sort = codes [:v ] + codes [v + 1 :] + [codes [v ]]
164
+ return to_sort
165
+
156
166
@cache_readonly
157
167
def _indexer_and_to_sort (
158
168
self ,
@@ -162,8 +172,14 @@ def _indexer_and_to_sort(
162
172
]:
163
173
v = self .level
164
174
165
- codes = list (self .index .codes )
166
175
levs = list (self .index .levels )
176
+ codes = list (self .index .codes )
177
+
178
+ if not self .sort :
179
+ codes = [list (l ) for l in codes ]
180
+ ids_code = [(dict ([(y , x ) for x , y in enumerate (sorted (set (l ), key = l .index ))]), l ) for l in codes ]
181
+ codes = [np .array ([d [x ] for x in code ]) for d , code in ids_code ]
182
+
167
183
to_sort = codes [:v ] + codes [v + 1 :] + [codes [v ]]
168
184
sizes = tuple (len (x ) for x in levs [:v ] + levs [v + 1 :] + [levs [v ]])
169
185
@@ -172,27 +188,24 @@ def _indexer_and_to_sort(
172
188
173
189
indexer = get_group_index_sorter (comp_index , ngroups )
174
190
return indexer , to_sort
175
-
191
+
176
192
@cache_readonly
177
- def sorted_labels (self ) -> list [np .ndarray ]:
193
+ def labels (self ) -> list [np .ndarray ]:
178
194
indexer , to_sort = self ._indexer_and_to_sort
179
195
if self .sort :
180
196
return [line .take (indexer ) for line in to_sort ]
181
197
return to_sort
182
198
183
199
def _make_sorted_values (self , values : np .ndarray ) -> np .ndarray :
184
- if self .sort :
185
- indexer , _ = self ._indexer_and_to_sort
186
-
187
- sorted_values = algos .take_nd (values , indexer , axis = 0 )
188
- return sorted_values
189
- return values
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
- # make the mask
195
- remaining_labels = self .sorted_labels [: - 1 ]
207
+ remaining_labels = self . labels [: - 1 ]
208
+ choosen_labels = self .labels [ - 1 ]
196
209
level_sizes = tuple (len (x ) for x in new_levels )
197
210
198
211
comp_index , obs_ids = get_compressed_ids (remaining_labels , level_sizes )
@@ -202,7 +215,7 @@ def _make_selectors(self):
202
215
stride = self .index .levshape [self .level ] + self .lift
203
216
self .full_shape = ngroups , stride
204
217
205
- selector = self . sorted_labels [ - 1 ] + stride * comp_index + self .lift
218
+ selector = choosen_labels + stride * comp_index + self .lift
206
219
mask = np .zeros (np .prod (self .full_shape ), dtype = bool )
207
220
mask .put (selector , True )
208
221
0 commit comments