multi-index column support

Author: jreback

pandas/core/frame.py

@@ -15,7 +15,7 @@
 import itertools
 import sys
 from textwrap import dedent
-from typing import FrozenSet, List, Optional, Set, Type, Union
+from typing import FrozenSet, List, Optional, Set, Tuple, Type, Union
 import warnings
 
 import numpy as np
@@ -6233,7 +6233,7 @@ def stack(self, level=-1, dropna=True):
         else:
             return stack(self, level, dropna=dropna)
 
-    def explode(self, column: str) -> "DataFrame":
+    def explode(self, column: Union[str, Tuple]) -> "DataFrame":
         """
         Transform each element of a list-like to a row, replicating the
         index values.
@@ -6242,7 +6242,7 @@ def explode(self, column: str) -> "DataFrame":
 
         Parameters
         ----------
-        column : str
+        column : str or tuple
 
         Returns
         -------
@@ -6290,7 +6290,7 @@ def explode(self, column: str) -> "DataFrame":
         3    4  1
         """
 
-        if not is_scalar(column):
+        if not (is_scalar(column) or isinstance(column, tuple)):
             raise ValueError("column must be a scalar")
         if not self.columns.is_unique:
             raise ValueError("columns must be unique")

pandas/tests/frame/test_explode.py

@@ -33,7 +33,7 @@ def test_basic():
     tm.assert_frame_equal(result, expected)
 
 
-def test_multi_index():
+def test_multi_index_rows():
     df = pd.DataFrame(
         {"A": np.array([[0, 1, 2], np.nan, [], (3, 4)], dtype=object), "B": 1},
         index=pd.MultiIndex.from_tuples([("a", 1), ("a", 2), ("b", 1), ("b", 2)]),
@@ -63,6 +63,25 @@ def test_multi_index():
     tm.assert_frame_equal(result, expected)
 
 
+def test_multi_index_columns():
+    df = pd.DataFrame(
+        {("A", 1): np.array([[0, 1, 2], np.nan, [], (3, 4)], dtype=object), ("A", 2): 1}
+    )
+
+    result = df.explode(("A", 1))
+    expected = pd.DataFrame(
+        {
+            ("A", 1): pd.Series(
+                [0, 1, 2, np.nan, np.nan, 3, 4],
+                index=pd.Index([0, 0, 0, 1, 2, 3, 3]),
+                dtype=object,
+            ),
+            ("A", 2): 1,
+        }
+    )
+    tm.assert_frame_equal(result, expected)
+
+
 def test_usecase():
     # explode a single column
     # gh-10511