Skip to content

Commit c27929d

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 c27929d

File tree

2 files changed

+42
-9
lines changed

2 files changed

+42
-9
lines changed

pandas/core/reshape/reshape.py

+27-9
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from __future__ import annotations
22

3+
import collections
34
import itertools
45
from typing import (
56
TYPE_CHECKING,
@@ -162,8 +163,17 @@ def _indexer_and_to_sort(
162163
]:
163164
v = self.level
164165

165-
codes = list(self.index.codes)
166166
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+
167177
to_sort = codes[:v] + codes[v + 1 :] + [codes[v]]
168178
sizes = tuple(len(x) for x in levs[:v] + levs[v + 1 :] + [levs[v]])
169179

@@ -174,25 +184,33 @@ def _indexer_and_to_sort(
174184
return indexer, to_sort
175185

176186
@cache_readonly
177-
def sorted_labels(self) -> list[np.ndarray]:
187+
def labels(self) -> list[np.ndarray]:
178188
indexer, to_sort = self._indexer_and_to_sort
179189
if self.sort:
180190
return [line.take(indexer) for line in to_sort]
181191
return to_sort
182192

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

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
190207

191208
def _make_selectors(self):
192209
new_levels = self.new_index_levels
193210

194211
# make the mask
195-
remaining_labels = self.sorted_labels[:-1]
212+
remaining_labels = self.labels[:-1]
213+
choosen_labels = self.labels[-1]
196214
level_sizes = tuple(len(x) for x in new_levels)
197215

198216
comp_index, obs_ids = get_compressed_ids(remaining_labels, level_sizes)
@@ -202,7 +220,7 @@ def _make_selectors(self):
202220
stride = self.index.levshape[self.level] + self.lift
203221
self.full_shape = ngroups, stride
204222

205-
selector = self.sorted_labels[-1] + stride * comp_index + self.lift
223+
selector = choosen_labels + stride * comp_index + self.lift
206224
mask = np.zeros(np.prod(self.full_shape), dtype=bool)
207225
mask.put(selector, True)
208226

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)