Revert "BUG: DataFrame.stack sometimes sorting the resulting index" (#54068)

Revert "BUG: DataFrame.stack sometimes sorting the resulting index (#53825)"

This reverts commit 5307062.

Author: rhshadrach

doc/source/whatsnew/v2.1.0.rst

@@ -154,6 +154,7 @@ Other enhancements
 - Let :meth:`DataFrame.to_feather` accept a non-default :class:`Index` and non-string column names (:issue:`51787`)
 - Performance improvement in :func:`read_csv` (:issue:`52632`) with ``engine="c"``
 - :meth:`Categorical.from_codes` has gotten a ``validate`` parameter (:issue:`50975`)
+- :meth:`DataFrame.stack` gained the ``sort`` keyword to dictate whether the resulting :class:`MultiIndex` levels are sorted (:issue:`15105`)
 - :meth:`DataFrame.unstack` gained the ``sort`` keyword to dictate whether the resulting :class:`MultiIndex` levels are sorted (:issue:`15105`)
 - :meth:`DataFrameGroupby.agg` and :meth:`DataFrameGroupby.transform` now support grouping by multiple keys when the index is not a :class:`MultiIndex` for ``engine="numba"`` (:issue:`53486`)
 - :meth:`Series.explode` now supports pyarrow-backed list types (:issue:`53602`)
@@ -570,8 +571,7 @@ Reshaping
 - Bug in :meth:`DataFrame.idxmin` and :meth:`DataFrame.idxmax`, where the axis dtype would be lost for empty frames (:issue:`53265`)
 - Bug in :meth:`DataFrame.merge` not merging correctly when having ``MultiIndex`` with single level (:issue:`52331`)
 - Bug in :meth:`DataFrame.stack` losing extension dtypes when columns is a :class:`MultiIndex` and frame contains mixed dtypes (:issue:`45740`)
-- Bug in :meth:`DataFrame.stack` sorting columns lexicographically in rare cases (:issue:`53786`)
-- Bug in :meth:`DataFrame.stack` sorting index lexicographically in rare cases (:issue:`53824`)
+- Bug in :meth:`DataFrame.stack` sorting columns lexicographically (:issue:`53786`)
 - Bug in :meth:`DataFrame.transpose` inferring dtype for object column (:issue:`51546`)
 - Bug in :meth:`Series.combine_first` converting ``int64`` dtype to ``float64`` and losing precision on very large integers (:issue:`51764`)
 - Bug when joining empty :class:`DataFrame` objects, where the joined index would be a :class:`RangeIndex` instead of the joined index type (:issue:`52777`)

pandas/core/frame.py

@@ -9031,7 +9031,7 @@ def pivot_table(
             sort=sort,
         )
 
-    def stack(self, level: IndexLabel = -1, dropna: bool = True):
+    def stack(self, level: IndexLabel = -1, dropna: bool = True, sort: bool = True):
         """
         Stack the prescribed level(s) from columns to index.
 
@@ -9057,6 +9057,8 @@ def stack(self, level: IndexLabel = -1, dropna: bool = True):
             axis can create combinations of index and column values
             that are missing from the original dataframe. See Examples
             section.
+        sort : bool, default True
+            Whether to sort the levels of the resulting MultiIndex.
 
         Returns
         -------
@@ -9156,15 +9158,15 @@ def stack(self, level: IndexLabel = -1, dropna: bool = True):
 
         >>> df_multi_level_cols2.stack(0)
                      kg    m
-        cat weight  1.0  NaN
-            height  NaN  2.0
-        dog weight  3.0  NaN
-            height  NaN  4.0
+        cat height  NaN  2.0
+            weight  1.0  NaN
+        dog height  NaN  4.0
+            weight  3.0  NaN
         >>> df_multi_level_cols2.stack([0, 1])
-        cat  weight  kg    1.0
-             height  m     2.0
-        dog  weight  kg    3.0
-             height  m     4.0
+        cat  height  m     2.0
+             weight  kg    1.0
+        dog  height  m     4.0
+             weight  kg    3.0
         dtype: float64
 
         **Dropping missing values**
@@ -9200,9 +9202,9 @@ def stack(self, level: IndexLabel = -1, dropna: bool = True):
         )
 
         if isinstance(level, (tuple, list)):
-            result = stack_multiple(self, level, dropna=dropna)
+            result = stack_multiple(self, level, dropna=dropna, sort=sort)
         else:
-            result = stack(self, level, dropna=dropna)
+            result = stack(self, level, dropna=dropna, sort=sort)
 
         return result.__finalize__(self, method="stack")
 

pandas/core/reshape/reshape.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import itertools
 from typing import (
     TYPE_CHECKING,
     cast,
@@ -498,7 +499,7 @@ def unstack(obj: Series | DataFrame, level, fill_value=None, sort: bool = True):
         if isinstance(obj.index, MultiIndex):
             return _unstack_frame(obj, level, fill_value=fill_value, sort=sort)
         else:
-            return obj.T.stack(dropna=False)
+            return obj.T.stack(dropna=False, sort=sort)
     elif not isinstance(obj.index, MultiIndex):
         # GH 36113
         # Give nicer error messages when unstack a Series whose
@@ -571,7 +572,7 @@ def _unstack_extension_series(
     return result
 
 
-def stack(frame: DataFrame, level=-1, dropna: bool = True):
+def stack(frame: DataFrame, level=-1, dropna: bool = True, sort: bool = True):
     """
     Convert DataFrame to Series with multi-level Index. Columns become the
     second level of the resulting hierarchical index
@@ -593,7 +594,9 @@ def factorize(index):
     level_num = frame.columns._get_level_number(level)
 
     if isinstance(frame.columns, MultiIndex):
-        return _stack_multi_columns(frame, level_num=level_num, dropna=dropna)
+        return _stack_multi_columns(
+            frame, level_num=level_num, dropna=dropna, sort=sort
+        )
     elif isinstance(frame.index, MultiIndex):
         new_levels = list(frame.index.levels)
         new_codes = [lab.repeat(K) for lab in frame.index.codes]
@@ -646,13 +649,13 @@ def factorize(index):
     return frame._constructor_sliced(new_values, index=new_index)
 
 
-def stack_multiple(frame: DataFrame, level, dropna: bool = True):
+def stack_multiple(frame: DataFrame, level, dropna: bool = True, sort: bool = True):
     # If all passed levels match up to column names, no
     # ambiguity about what to do
     if all(lev in frame.columns.names for lev in level):
         result = frame
         for lev in level:
-            result = stack(result, lev, dropna=dropna)
+            result = stack(result, lev, dropna=dropna, sort=sort)
 
     # Otherwise, level numbers may change as each successive level is stacked
     elif all(isinstance(lev, int) for lev in level):
@@ -665,7 +668,7 @@ def stack_multiple(frame: DataFrame, level, dropna: bool = True):
 
         while level:
             lev = level.pop(0)
-            result = stack(result, lev, dropna=dropna)
+            result = stack(result, lev, dropna=dropna, sort=sort)
             # Decrement all level numbers greater than current, as these
             # have now shifted down by one
             level = [v if v <= lev else v - 1 for v in level]
@@ -691,14 +694,7 @@ def _stack_multi_column_index(columns: MultiIndex) -> MultiIndex:
 
     # Remove duplicate tuples in the MultiIndex.
     tuples = zip(*levs)
-    seen = set()
-    # mypy doesn't like our trickery to get `set.add` to work in a comprehension
-    # error: "add" of "set" does not return a value
-    unique_tuples = (
-        key
-        for key in tuples
-        if not (key in seen or seen.add(key))  # type: ignore[func-returns-value]
-    )
+    unique_tuples = (key for key, _ in itertools.groupby(tuples))
     new_levs = zip(*unique_tuples)
 
     # The dtype of each level must be explicitly set to avoid inferring the wrong type.
@@ -714,7 +710,7 @@ def _stack_multi_column_index(columns: MultiIndex) -> MultiIndex:
 
 
 def _stack_multi_columns(
-    frame: DataFrame, level_num: int = -1, dropna: bool = True
+    frame: DataFrame, level_num: int = -1, dropna: bool = True, sort: bool = True
 ) -> DataFrame:
     def _convert_level_number(level_num: int, columns: Index):
         """
@@ -744,22 +740,31 @@ def _convert_level_number(level_num: int, columns: Index):
             roll_columns = roll_columns.swaplevel(lev1, lev2)
         this.columns = mi_cols = roll_columns
 
+    if not mi_cols._is_lexsorted() and sort:
+        # Workaround the edge case where 0 is one of the column names,
+        # which interferes with trying to sort based on the first
+        # level
+        level_to_sort = _convert_level_number(0, mi_cols)
+        this = this.sort_index(level=level_to_sort, axis=1)
+        mi_cols = this.columns
+
+    mi_cols = cast(MultiIndex, mi_cols)
     new_columns = _stack_multi_column_index(mi_cols)
 
     # time to ravel the values
     new_data = {}
     level_vals = mi_cols.levels[-1]
     level_codes = unique(mi_cols.codes[-1])
+    if sort:
+        level_codes = np.sort(level_codes)
     level_vals_nan = level_vals.insert(len(level_vals), None)
 
     level_vals_used = np.take(level_vals_nan, level_codes)
     levsize = len(level_codes)
     drop_cols = []
     for key in new_columns:
         try:
-            with warnings.catch_warnings():
-                warnings.simplefilter("ignore", PerformanceWarning)
-                loc = this.columns.get_loc(key)
+            loc = this.columns.get_loc(key)
         except KeyError:
             drop_cols.append(key)
             continue
@@ -769,12 +774,9 @@ def _convert_level_number(level_num: int, columns: Index):
         # but if unsorted can get a boolean
         # indexer
         if not isinstance(loc, slice):
-            slice_len = loc.sum()
+            slice_len = len(loc)
         else:
             slice_len = loc.stop - loc.start
-            if loc.step is not None:
-                # Integer division using ceiling instead of floor
-                slice_len = -(slice_len // -loc.step)
 
         if slice_len != levsize:
             chunk = this.loc[:, this.columns[loc]]

pandas/tests/frame/test_stack_unstack.py

@@ -1099,18 +1099,18 @@ def test_stack_preserve_categorical_dtype(self, ordered, labels):
         "labels,data",
         [
             (list("xyz"), [10, 11, 12, 13, 14, 15]),
-            (list("zyx"), [10, 11, 12, 13, 14, 15]),
+            (list("zyx"), [14, 15, 12, 13, 10, 11]),
         ],
     )
     def test_stack_multi_preserve_categorical_dtype(self, ordered, labels, data):
         # GH-36991
         cidx = pd.CategoricalIndex(labels, categories=sorted(labels), ordered=ordered)
         cidx2 = pd.CategoricalIndex(["u", "v"], ordered=ordered)
         midx = MultiIndex.from_product([cidx, cidx2])
-        df = DataFrame([data], columns=midx)
+        df = DataFrame([sorted(data)], columns=midx)
         result = df.stack([0, 1])
 
-        s_cidx = pd.CategoricalIndex(labels, ordered=ordered)
+        s_cidx = pd.CategoricalIndex(sorted(labels), ordered=ordered)
         expected = Series(data, index=MultiIndex.from_product([[0], s_cidx, cidx2]))
 
         tm.assert_series_equal(result, expected)
@@ -1400,16 +1400,16 @@ def test_unstack_non_slice_like_blocks(using_array_manager):
     tm.assert_frame_equal(res, expected)
 
 
-def test_stack_nosort():
-    # GH 15105, GH 53825
+def test_stack_sort_false():
+    # GH 15105
     data = [[1, 2, 3.0, 4.0], [2, 3, 4.0, 5.0], [3, 4, np.nan, np.nan]]
     df = DataFrame(
         data,
         columns=MultiIndex(
             levels=[["B", "A"], ["x", "y"]], codes=[[0, 0, 1, 1], [0, 1, 0, 1]]
         ),
     )
-    result = df.stack(level=0)
+    result = df.stack(level=0, sort=False)
     expected = DataFrame(
         {"x": [1.0, 3.0, 2.0, 4.0, 3.0], "y": [2.0, 4.0, 3.0, 5.0, 4.0]},
         index=MultiIndex.from_arrays([[0, 0, 1, 1, 2], ["B", "A", "B", "A", "B"]]),
@@ -1421,15 +1421,15 @@ def test_stack_nosort():
         data,
         columns=MultiIndex.from_arrays([["B", "B", "A", "A"], ["x", "y", "x", "y"]]),
     )
-    result = df.stack(level=0)
+    result = df.stack(level=0, sort=False)
     tm.assert_frame_equal(result, expected)
 
 
-def test_stack_nosort_multi_level():
-    # GH 15105, GH 53825
+def test_stack_sort_false_multi_level():
+    # GH 15105
     idx = MultiIndex.from_tuples([("weight", "kg"), ("height", "m")])
     df = DataFrame([[1.0, 2.0], [3.0, 4.0]], index=["cat", "dog"], columns=idx)
-    result = df.stack([0, 1])
+    result = df.stack([0, 1], sort=False)
     expected_index = MultiIndex.from_tuples(
         [
             ("cat", "weight", "kg"),
@@ -1999,12 +1999,13 @@ def __init__(self, *args, **kwargs) -> None:
         ),
     )
     @pytest.mark.parametrize("stack_lev", range(2))
-    def test_stack_order_with_unsorted_levels(self, levels, stack_lev):
+    @pytest.mark.parametrize("sort", [True, False])
+    def test_stack_order_with_unsorted_levels(self, levels, stack_lev, sort):
         # GH#16323
         # deep check for 1-row case
         columns = MultiIndex(levels=levels, codes=[[0, 0, 1, 1], [0, 1, 0, 1]])
         df = DataFrame(columns=columns, data=[range(4)])
-        df_stacked = df.stack(stack_lev)
+        df_stacked = df.stack(stack_lev, sort=sort)
         for row in df.index:
             for col in df.columns:
                 expected = df.loc[row, col]
@@ -2036,7 +2037,7 @@ def test_stack_order_with_unsorted_levels_multi_row_2(self):
         stack_lev = 1
         columns = MultiIndex(levels=levels, codes=[[0, 0, 1, 1], [0, 1, 0, 1]])
         df = DataFrame(columns=columns, data=[range(4)], index=[1, 0, 2, 3])
-        result = df.stack(stack_lev)
+        result = df.stack(stack_lev, sort=True)
         expected_index = MultiIndex(
             levels=[[0, 1, 2, 3], [0, 1]],
             codes=[[1, 1, 0, 0, 2, 2, 3, 3], [1, 0, 1, 0, 1, 0, 1, 0]],