Backport PR #54962 on branch 2.1.x (BUG: DataFrame.stack with future_…

…stack=True failing when columns are tuples) (#54971) Backport PR #54962: BUG: DataFrame.stack with future_stack=True failing when columns are tuples Co-authored-by: Richard Shadrach <[email protected]>
pandas-dev · Sep 2, 2023 · 0c03254 · 0c03254
1 parent c22f2a2
commit 0c03254
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 2 deletions.
diff --git a/doc/source/whatsnew/v2.1.1.rst b/doc/source/whatsnew/v2.1.1.rst
@@ -28,7 +28,7 @@ Fixed regressions
 
 Bug fixes
 ~~~~~~~~~
--
+- Fixed bug in :meth:`DataFrame.stack` with ``future_stack=True`` and columns a non-:class:`MultiIndex` consisting of tuples (:issue:`54948`)
 
 .. ---------------------------------------------------------------------------
 .. _whatsnew_211.other:

diff --git a/pandas/core/reshape/reshape.py b/pandas/core/reshape/reshape.py
@@ -908,7 +908,7 @@ def stack_v3(frame: DataFrame, level: list[int]) -> Series | DataFrame:
             data = frame.copy()
         else:
             # Take the data from frame corresponding to this idx value
-            if not isinstance(idx, tuple):
+            if len(level) == 1:
                 idx = (idx,)
             gen = iter(idx)
             column_indexer = tuple(

diff --git a/pandas/tests/frame/test_stack_unstack.py b/pandas/tests/frame/test_stack_unstack.py
@@ -2508,3 +2508,19 @@ def test_unstack_mixed_level_names(self):
             index=MultiIndex.from_tuples([(1, "red"), (2, "blue")], names=[0, "y"]),
         )
         tm.assert_frame_equal(result, expected)
+
+
+def test_stack_tuple_columns(future_stack):
+    # GH#54948 - test stack when the input has a non-MultiIndex with tuples
+    df = DataFrame(
+        [[1, 2, 3], [4, 5, 6], [7, 8, 9]], columns=[("a", 1), ("a", 2), ("b", 1)]
+    )
+    result = df.stack(future_stack=future_stack)
+    expected = Series(
+        [1, 2, 3, 4, 5, 6, 7, 8, 9],
+        index=MultiIndex(
+            levels=[[0, 1, 2], [("a", 1), ("a", 2), ("b", 1)]],
+            codes=[[0, 0, 0, 1, 1, 1, 2, 2, 2], [0, 1, 2, 0, 1, 2, 0, 1, 2]],
+        ),
+    )
+    tm.assert_series_equal(result, expected)