handle multiindex columns

Author: michaelsilverstein

pandas/core/reshape/melt.py

@@ -25,6 +25,10 @@
 def melt(frame, id_vars=None, value_vars=None, var_name=None,
          value_name='value', col_level=None):
     # TODO: what about the existing index?
+    if isinstance(frame.columns, ABCMultiIndex):
+        cols = [x for c in frame.columns for x in c]
+    else:
+        cols = list(frame.columns)
     if id_vars is not None:
         if not is_list_like(id_vars):
             id_vars = [id_vars]
@@ -35,7 +39,7 @@ def melt(frame, id_vars=None, value_vars=None, var_name=None,
         else:
             # Check that `id_vars` are in frame
             id_vars = list(id_vars)
-            missing = Index(id_vars).difference(frame.columns)
+            missing = Index(np.ravel(id_vars)).difference(cols)
             if not missing.empty:
                 raise KeyError("The following 'id_vars' are not present"
                                " in the DataFrame: {missing}"
@@ -53,7 +57,7 @@ def melt(frame, id_vars=None, value_vars=None, var_name=None,
         else:
             value_vars = list(value_vars)
             # Check that `value_vars` are in frame
-            missing = Index(value_vars).difference(frame.columns)
+            missing = Index(np.ravel(value_vars)).difference(cols)
             if not missing.empty:
                 raise KeyError("The following 'value_vars' are not present in"
                                " the DataFrame: {missing}"