diff --git a/doc/source/whatsnew/v1.1.0.rst b/doc/source/whatsnew/v1.1.0.rst index 41d519e0765dc..bab8b5cd99255 100644 --- a/doc/source/whatsnew/v1.1.0.rst +++ b/doc/source/whatsnew/v1.1.0.rst @@ -853,6 +853,7 @@ Reshaping - Bug in :func:`cut` raised an error when non-unique labels (:issue:`33141`) - Bug in :meth:`DataFrame.replace` casts columns to ``object`` dtype if items in ``to_replace`` not in values (:issue:`32988`) - Ensure only named functions can be used in :func:`eval()` (:issue:`32460`) +- Fixed bug in :func:`melt` where melting MultiIndex columns with ``col_level`` > 0 would raise a ``KeyError`` on ``id_vars`` (:issue:`34129`) Sparse ^^^^^^ diff --git a/pandas/core/reshape/melt.py b/pandas/core/reshape/melt.py index 1a315ff0ead52..7d22b86c5c07c 100644 --- a/pandas/core/reshape/melt.py +++ b/pandas/core/reshape/melt.py @@ -72,7 +72,13 @@ def melt( "The following 'value_vars' are not present in " f"the DataFrame: {list(missing)}" ) - frame = frame.loc[:, id_vars + value_vars] + if col_level is not None: + idx = frame.columns.get_level_values(col_level).get_indexer( + id_vars + value_vars + ) + else: + idx = frame.columns.get_indexer(id_vars + value_vars) + frame = frame.iloc[:, idx] else: frame = frame.copy() diff --git a/pandas/tests/reshape/test_melt.py b/pandas/tests/reshape/test_melt.py index 6a670e6c729e9..000a6354277ab 100644 --- a/pandas/tests/reshape/test_melt.py +++ b/pandas/tests/reshape/test_melt.py @@ -100,15 +100,39 @@ def test_vars_work_with_multiindex(self): result = self.df1.melt(id_vars=[("A", "a")], value_vars=[("B", "b")]) tm.assert_frame_equal(result, expected) - def test_single_vars_work_with_multiindex(self): - expected = DataFrame( - { - "A": {0: 1.067683, 1: -1.321405, 2: -0.807333}, - "CAP": {0: "B", 1: "B", 2: "B"}, - "value": {0: -1.110463, 1: 0.368915, 2: 0.08298}, - } - ) - result = self.df1.melt(["A"], ["B"], col_level=0) + @pytest.mark.parametrize( + "id_vars, value_vars, col_level, expected", + [ + ( + ["A"], + ["B"], + 0, + DataFrame( + { + "A": {0: 1.067683, 1: -1.321405, 2: -0.807333}, + "CAP": {0: "B", 1: "B", 2: "B"}, + "value": {0: -1.110463, 1: 0.368915, 2: 0.08298}, + } + ), + ), + ( + ["a"], + ["b"], + 1, + DataFrame( + { + "a": {0: 1.067683, 1: -1.321405, 2: -0.807333}, + "low": {0: "b", 1: "b", 2: "b"}, + "value": {0: -1.110463, 1: 0.368915, 2: 0.08298}, + } + ), + ), + ], + ) + def test_single_vars_work_with_multiindex( + self, id_vars, value_vars, col_level, expected + ): + result = self.df1.melt(id_vars, value_vars, col_level=col_level) tm.assert_frame_equal(result, expected) def test_tuple_vars_fail_with_multiindex(self):