ENH MultiIndex columns with melt

Author: hayd

doc/source/release.rst

@@ -77,6 +77,7 @@ pandas 0.12
     to specify custom column names of the returned DataFrame (:issue:`3649`),
     thanks @hoechenberger. If ``var_name`` is not specified and ``dataframe.columns.name`` 
     is not None, then this will be used as the ``var_name`` (:issue:`4144`).
+    Also support for MultiIndex columns.
   - clipboard functions use pyperclip (no dependencies on Windows, alternative
     dependencies offered for Linux) (:issue:`3837`).
   - Plotting functions now raise a ``TypeError`` before trying to plot anything

pandas/core/index.py

@@ -1653,7 +1653,9 @@ def get_level_values(self, level):
         num = self._get_level_number(level)
         unique_vals = self.levels[num]  # .values
         labels = self.labels[num]
-        return unique_vals.take(labels)
+        values = unique_vals.take(labels)
+        values.name = self.names[num]
+        return values
 
     def format(self, space=2, sparsify=None, adjoin=True, names=False,
                na_rep='NaN', formatter=None):

pandas/core/reshape.py

@@ -617,24 +617,42 @@ def melt(frame, id_vars=None, value_vars=None,
 
     Examples
     --------
+    >>> df = pd.DataFrame({'A': {0: 'a', 1: 'b', 2: 'c'},
+              'B': {0: 1, 1: 3, 2: 5},
+              'C': {0: 2, 1: 4, 2: 6}})
+
     >>> df
-    A B C
-    a 1 2
-    b 3 4
-    c 5 6
+       A  B  C
+    0  a  1  2
+    1  b  3  4
+    2  c  5  6
 
     >>> melt(df, id_vars=['A'], value_vars=['B'])
-    A variable value
-    a        B     1
-    b        B     3
-    c        B     5
+       A variable  value
+    0  a        B      1
+    1  b        B      3
+    2  c        B      5
     
     >>> melt(df, id_vars=['A'], value_vars=['B'],
     ... var_name='myVarname', value_name='myValname')
-    A myVarname  myValname
-    a         B          1
-    b         B          3
-    c         B          5
+       A myVarname  myValname
+    0  a         B          1
+    1  b         B          3
+    2  c         B          5
+
+    >>> df.columns = [list('ABC'), list('DEF')]
+
+    >>> melt(df, col_level=0, id_vars=['A'], value_vars=['B'])
+       A variable  value
+    0  a        B      1
+    1  b        B      3
+    2  c        B      5
+
+    >>> melt(df, id_vars=[('A', 'D')], value_vars=[('B', 'E')])
+      (A, D) variable_0 variable_1  value
+    0      a          B          E      1
+    1      b          B          E      3
+    2      c          B          E      5
 
     """
     # TODO: what about the existing index?
@@ -653,11 +671,17 @@ def melt(frame, id_vars=None, value_vars=None,
     else:
         frame = frame.copy()
 
-    if col_level:  # allow list?
+    if col_level is not None:  # allow list or other?
         frame.columns = frame.columns.get_level_values(col_level) #  frame is a copy
 
     if var_name is None:
-        var_name = frame.columns.name if frame.columns.name is not None else 'variable'
+        if isinstance(frame.columns, MultiIndex):
+            if len(frame.columns.names) == len(set(frame.columns.names)):
+                var_name = frame.columns.names
+            else:
+                var_name = ['variable_%s' % i for i in range(len(frame.columns.names))]
+        else:
+            var_name = frame.columns.name if frame.columns.name is not None else 'variable'
 
     N, K = frame.shape
     K -= len(id_vars)
@@ -666,11 +690,18 @@ def melt(frame, id_vars=None, value_vars=None,
     for col in id_vars:
         mdata[col] = np.tile(frame.pop(col).values, K)
 
-    mcolumns = id_vars + [var_name, value_name]
+    if isinstance(var_name, list):
+        mcolumns = id_vars + var_name + [value_name]
+    else:
+        mcolumns = id_vars + [var_name, value_name]
 
     mdata[value_name] = frame.values.ravel('F')
-    mdata[var_name] = np.asarray(frame.columns).repeat(N)
-    
+    if isinstance(frame.columns, MultiIndex):
+        for i, col in enumerate(var_name):
+            mdata[col] = np.asarray(frame.columns.get_level_values(i)).repeat(N)
+    else: # assume isinstance(frame.columns, Index):
+        mdata[var_name] = np.asarray(frame.columns).repeat(N)
+
     return DataFrame(mdata, columns=mcolumns)
 
 

pandas/tests/test_index.py

@@ -1029,6 +1029,8 @@ def test_get_level_values(self):
         expected = ['foo', 'foo', 'bar', 'baz', 'qux', 'qux']
         self.assert_(np.array_equal(result, expected))
 
+        self.assertEquals(result.name, 'first')
+
         result = self.index.get_level_values('first')
         expected = self.index.get_level_values(0)
         self.assert_(np.array_equal(result, expected))

pandas/tests/test_reshape.py

@@ -10,6 +10,7 @@
 import nose
 
 from pandas import DataFrame
+import pandas as pd
 
 from numpy import nan
 import numpy as np
@@ -30,6 +31,12 @@ def setUp(self):
         self.var_name = 'var'
         self.value_name = 'val'
 
+        self.df1 = pd.DataFrame([[ 1.067683, -1.110463,  0.20867 ],
+                                 [-1.321405,  0.368915, -1.055342],
+                                 [-0.807333,  0.08298 , -0.873361]])
+        self.df1.columns = [list('ABC'), list('abc')]
+        self.df1.columns.names = ['CAP', 'low']
+
     def test_default_col_names(self):
         result = melt(self.df)
         self.assertEqual(result.columns.tolist(), ['variable', 'value'])
@@ -128,6 +135,17 @@ def test_custom_var_and_value_name(self):
         result20 = melt(self.df)
         self.assertEqual(result20.columns.tolist(), ['foo', 'value'])
 
+    def test_col_level(self):
+        res1 = melt(self.df1, col_level=0)
+        res2 = melt(self.df1, col_level='CAP')
+        self.assertEqual(res1.columns.tolist(), ['CAP', 'value'])
+        self.assertEqual(res1.columns.tolist(), ['CAP', 'value'])
+
+    def test_multiindex(self):
+        res = pd.melt(self.df1)
+        self.assertEqual(res.columns.tolist(), ['CAP', 'low', 'value'])
+
+
 class TestConvertDummies(unittest.TestCase):
     def test_convert_dummies(self):
         df = DataFrame({'A': ['foo', 'bar', 'foo', 'bar',