Fix MultiIndex melt when col_level is used (#34329)

Author: ashtou

doc/source/whatsnew/v1.1.0.rst

@@ -885,6 +885,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
 ^^^^^^

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()
 

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):