Skip to content

Commit e057c9f

Browse files
BUG: unstack with sort=False fails when used with the level parameter (#54987)
Assign new codes to labels when sort=False. This is done so that the data appears to be already sorted, fixing the bug.
1 parent 46c8da3 commit e057c9f

File tree

2 files changed

+41
-9
lines changed

2 files changed

+41
-9
lines changed

pandas/core/reshape/reshape.py

+26-9
Original file line numberDiff line numberDiff line change
@@ -162,8 +162,17 @@ def _indexer_and_to_sort(
162162
]:
163163
v = self.level
164164

165-
codes = list(self.index.codes)
166165
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+
167176
to_sort = codes[:v] + codes[v + 1 :] + [codes[v]]
168177
sizes = tuple(len(x) for x in levs[:v] + levs[v + 1 :] + [levs[v]])
169178

@@ -174,25 +183,33 @@ def _indexer_and_to_sort(
174183
return indexer, to_sort
175184

176185
@cache_readonly
177-
def sorted_labels(self) -> list[np.ndarray]:
186+
def labels(self) -> list[np.ndarray]:
178187
indexer, to_sort = self._indexer_and_to_sort
179188
if self.sort:
180189
return [line.take(indexer) for line in to_sort]
181190
return to_sort
182191

183-
def _make_sorted_values(self, values: np.ndarray) -> np.ndarray:
192+
@cache_readonly
193+
def sorted_labels(self) -> list[np.ndarray]:
184194
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
186201

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
190206

191207
def _make_selectors(self):
192208
new_levels = self.new_index_levels
193209

194210
# make the mask
195-
remaining_labels = self.sorted_labels[:-1]
211+
remaining_labels = self.labels[:-1]
212+
choosen_labels = self.labels[-1]
196213
level_sizes = tuple(len(x) for x in new_levels)
197214

198215
comp_index, obs_ids = get_compressed_ids(remaining_labels, level_sizes)
@@ -202,7 +219,7 @@ def _make_selectors(self):
202219
stride = self.index.levshape[self.level] + self.lift
203220
self.full_shape = ngroups, stride
204221

205-
selector = self.sorted_labels[-1] + stride * comp_index + self.lift
222+
selector = choosen_labels + stride * comp_index + self.lift
206223
mask = np.zeros(np.prod(self.full_shape), dtype=bool)
207224
mask.put(selector, True)
208225

pandas/tests/frame/test_stack_unstack.py

+15
Original file line numberDiff line numberDiff line change
@@ -1318,6 +1318,21 @@ def test_unstack_sort_false(frame_or_series, dtype):
13181318
[("two", "z", "b"), ("two", "y", "a"), ("one", "z", "b"), ("one", "y", "a")]
13191319
)
13201320
obj = frame_or_series(np.arange(1.0, 5.0), index=index, dtype=dtype)
1321+
1322+
result = obj.unstack(level=0, sort=False)
1323+
1324+
if frame_or_series is DataFrame:
1325+
expected_columns = MultiIndex.from_tuples([(0, "two"), (0, "one")])
1326+
else:
1327+
expected_columns = ["two", "one"]
1328+
expected = DataFrame(
1329+
[[1.0, 3.0], [2.0, 4.0]],
1330+
index=MultiIndex.from_tuples([('z', 'b'), ('y', 'a')]),
1331+
columns=expected_columns,
1332+
dtype=dtype,
1333+
)
1334+
tm.assert_frame_equal(result, expected)
1335+
13211336
result = obj.unstack(level=-1, sort=False)
13221337

13231338
if frame_or_series is DataFrame:

0 commit comments

Comments
 (0)