ENH: MultiIndex.from_frame

pandas/core/indexes/multi.py

@@ -2,6 +2,7 @@
 # pylint: disable=E1101,E1103,W0232
 import datetime
 import warnings
+from collections import OrderedDict
 from sys import getsizeof
 
 import numpy as np
@@ -1189,11 +1190,15 @@ def to_frame(self, index=True, name=None):
         else:
             idx_names = self.names
 
-        result = DataFrame({(name or level):
-                            self._get_level_values(level)
-                            for name, level in
-                            zip(idx_names, range(len(self.levels)))},
-                           copy=False)
+        result = DataFrame(
+            OrderedDict([
+                ((name or level), self._get_level_values(level))
+                for name, level in zip(idx_names, range(len(self.levels)))
+            ]),
+            copy=False
+        )
+
+
         if index:
             result.index = self
         return result
@@ -1294,6 +1299,7 @@ def from_arrays(cls, arrays, sortorder=None, names=None):
         MultiIndex.from_tuples : Convert list of tuples to MultiIndex
         MultiIndex.from_product : Make a MultiIndex from cartesian product
                                   of iterables
+        MultiIndex.from_frame : Make a MultiIndex from a DataFrame.
         """
         if not is_list_like(arrays):
             raise TypeError("Input must be a list / sequence of array-likes.")
@@ -1343,6 +1349,7 @@ def from_tuples(cls, tuples, sortorder=None, names=None):
         MultiIndex.from_arrays : Convert list of arrays to MultiIndex
         MultiIndex.from_product : Make a MultiIndex from cartesian product
                                   of iterables
+        MultiIndex.from_frame : Make a MultiIndex from a DataFrame.
         """
         if not is_list_like(tuples):
             raise TypeError('Input must be a list / sequence of tuple-likes.')
@@ -1399,6 +1406,7 @@ def from_product(cls, iterables, sortorder=None, names=None):
         --------
         MultiIndex.from_arrays : Convert list of arrays to MultiIndex
         MultiIndex.from_tuples : Convert list of tuples to MultiIndex
+        MultiIndex.from_frame : Make a MultiIndex from a DataFrame.
         """
         from pandas.core.arrays.categorical import _factorize_from_iterables
         from pandas.core.reshape.util import cartesian_product
@@ -1412,6 +1420,77 @@ def from_product(cls, iterables, sortorder=None, names=None):
         labels = cartesian_product(labels)
         return MultiIndex(levels, labels, sortorder=sortorder, names=names)
 
+    @classmethod
+    def from_frame(cls, df, squeeze=True, names=None):
+        """
+        Make a MultiIndex from a DataFrame.
+
+        Parameters
+        ----------
+        df : pd.DataFrame
+            DataFrame to be converted to MultiIndex.
+        squeeze : bool, default True
+            If df is a single column, squeeze MultiIndex to be a regular Index.
+        names : list / sequence / callable, optonal
+            If no names provided, use column names, or tuple of column names if
+            the columns is a MultiIndex. If sequence, overwrite names with the
+            given sequence. If callable, pass each column name or tuples of
+            names to the callable.
+
+        Returns
+        -------
+        MultiIndex or Index
+            The MultiIndex representation of the given DataFrame. Returns an
+            Index if the DataFrame is single column and squeeze is True.
+
+        Examples
+        --------
+        >>> df = pd.DataFrame([[0, 'happy'], [0, 'jolly'], [1, 'happy'],
+        ...                    [1, 'jolly'], [2, 'joy'], [2, 'joy']],
+        ...                   columns=['number', 'mood'])
+        >>> df
+           number   mood
+        0       0  happy
+        1       0  jolly
+        2       1  happy
+        3       1  jolly
+        4       2    joy
+        5       2    joy
+        >>> pd.MultiIndex.from_frame(df)
+        MultiIndex(levels=[[0, 1, 2], ['happy', 'jolly', 'joy']],
+                   labels=[[0, 0, 1, 1, 2, 2], [0, 1, 0, 1, 2, 2]],
+                   names=['number', 'mood'])
+
+        See Also
+        --------
+        MultiIndex.from_arrays : Convert list of arrays to MultiIndex
+        MultiIndex.from_tuples : Convert list of tuples to MultiIndex
+        MultiIndex.from_product : Make a MultiIndex from cartesian product
+                                  of iterables
+        """
+        from pandas import DataFrame
+        if not isinstance(df, DataFrame):
+            raise TypeError("Input must be a DataFrame")
+
+        # Get MultiIndex names
+        if names is None:
+            names = list(df)
+        else:
+            if callable(names):
+                names = [names(x) for x in list(df)]
+            else:
+                if not is_list_like(names):
+                    raise TypeError("'names' must be a list / sequence "
+                                    "of column names, or a callable.")
+
+                if len(names) != len(list(df)):
+                    raise ValueError("'names' should have same length as "
+                                     "number of columns in df.")
+
+        # This way will preserve dtype of columns
+        mi = cls.from_arrays([df[x] for x in df], names=names)
+        return mi.squeeze() if squeeze else mi
+
     def _sort_levels_monotonic(self):
         """
         .. versionadded:: 0.20.0
@@ -1474,6 +1553,30 @@ def _sort_levels_monotonic(self):
                           names=self.names, sortorder=self.sortorder,
                           verify_integrity=False)
 
+    def squeeze(self):
+        """
+        Squeeze a single level MultiIndex to be a regular Index instance.
+
+        Returns
+        -------
+        Index or MultiIndex
+            Returns Index equivalent of single level MultiIndex. Returns
+            copy of MultiIndex if multilevel.
+
+        Examples
+        --------
+        >>> mi = pd.MultiIndex.from_tuples([('a',), ('b',), ('c',)])
+        >>> mi
+        MultiIndex(levels=[['a', 'b', 'c']],
+                   labels=[[0, 1, 2]])
+        >>> mi.squeeze()
+        Index(['a', 'b', 'c'], dtype='object')
+        """
+        if len(self.levels) == 1:
+            return self.levels[0][self.labels[0]]
+        else:
+            return self.copy()
+
     def remove_unused_levels(self):
         """
         Create a new MultiIndex from the current that removes

pandas/tests/indexes/multi/test_constructor.py

@@ -472,3 +472,77 @@ def test_from_tuples_with_tuple_label():
     idx = pd.MultiIndex.from_tuples([(2, 1), (4, (1, 2))], names=('a', 'b'))
     result = pd.DataFrame([2, 3], columns=['c'], index=idx)
     tm.assert_frame_equal(expected, result)
+
+
+def test_from_frame():
+    df = pd.DataFrame([['a', 'a'], ['a', 'b'], ['b', 'a'], ['b', 'b']],
+                      columns=['L1', 'L2'])
+    expected = pd.MultiIndex.from_tuples([('a', 'a'), ('a', 'b'),
+                                          ('b', 'a'), ('b', 'b')],
+                                         names=['L1', 'L2'])
+    result = pd.MultiIndex.from_frame(df)
+    tm.assert_index_equal(expected, result)
+
+
+@pytest.mark.parametrize('squeeze,input_type,expected', [
+    (True, 'multi', pd.MultiIndex.from_tuples([('a', 'a'), ('a', 'b'),
+                                               ('b', 'a'), ('b', 'b')],
+                                              names=['L1', 'L2'])),
+    (True, 'single', pd.Index(['a', 'a', 'b', 'b'], name='L1')),
+    (False, 'multi', pd.MultiIndex.from_tuples([('a', 'a'), ('a', 'b'),
+                                                ('b', 'a'), ('b', 'b')],
+                                               names=['L1', 'L2'])),
+    (False, 'single', pd.MultiIndex.from_tuples([('a',), ('a',),
+                                                 ('b',), ('b',)],
+                                                names=['L1']))
+])
+def test_from_frame_squeeze(squeeze, input_type, expected):
+    if input_type == 'multi':
+        df = pd.DataFrame([['a', 'a'], ['a', 'b'], ['b', 'a'], ['b', 'b']],
+                          columns=['L1', 'L2'])
+    elif input_type == 'single':
+        df = pd.DataFrame([['a'], ['a'], ['b'], ['b']], columns=['L1'])
+
+    result = pd.MultiIndex.from_frame(df, squeeze=squeeze)
+    tm.assert_index_equal(expected, result)
+
+
+def test_from_frame_non_frame():
+    with tm.assert_raises_regex(TypeError, 'Input must be a DataFrame'):
+        pd.MultiIndex.from_frame([1, 2, 3, 4])
+
+
+def test_from_frame_dtype_fidelity():
+    df = pd.DataFrame({
+        'dates': pd.date_range('19910905', periods=6),
+        'a': [1, 1, 1, 2, 2, 2],
+        'b': pd.Categorical(['a', 'a', 'b', 'b', 'c', 'c'], ordered=True),
+        'c': ['x', 'x', 'y', 'z', 'x', 'y']
+    })
+    original_dtypes = df.dtypes.to_dict()
+    mi = pd.MultiIndex.from_frame(df)
+    mi_dtypes = {name: mi.levels[i].dtype for i, name in enumerate(mi.names)}
+    assert original_dtypes == mi_dtypes
+
+
+def test_from_frame_names_as_list():
+    df = pd.DataFrame([['a', 'a'], ['a', 'b'], ['b', 'a'], ['b', 'b']],
+                      columns=['L1', 'L2'])
+    mi = pd.MultiIndex.from_frame(df, names=['a', 'b'])
+    assert mi.names == ['a', 'b']
+
+
+def test_from_frame_names_as_callable():
+    df = pd.DataFrame([['a', 'a'], ['a', 'b'], ['b', 'a'], ['b', 'b']],
+                      columns=pd.MultiIndex.from_tuples([('L1', 'x'),
+                                                         ('L2', 'y')]))
+    mi = pd.MultiIndex.from_frame(df, names=lambda x: '_'.join(x))
+    assert mi.names == ['L1_x', 'L2_y']
+
+
+def test_from_frame_names_bad_input():
+    df = pd.DataFrame([['a', 'a'], ['a', 'b'], ['b', 'a'], ['b', 'b']],
+                      columns=['L1', 'L2'])
+    with tm.assert_raises_regex(TypeError, "names' must be a list / sequence "
+                                           "of column names, or a callable."):
+        pd.MultiIndex.from_frame(df, names='bad')

pandas/tests/indexes/multi/test_conversion.py

@@ -82,6 +82,20 @@ def test_to_frame():
     tm.assert_frame_equal(result, expected)
 
 
+def test_to_frame_dtype_fidelity():
+    mi = pd.MultiIndex.from_arrays([
+        pd.date_range('19910905', periods=6),
+        [1, 1, 1, 2, 2, 2],
+        pd.Categorical(['a', 'a', 'b', 'b', 'c', 'c'], ordered=True),
+        ['x', 'x', 'y', 'z', 'x', 'y']
+    ], names=['dates', 'a', 'b', 'c'])
+    original_dtypes = {name: mi.levels[i].dtype 
+                       for i, name in enumerate(mi.names)}
+    df = mi.to_frame()
+    df_dtypes = df.dtypes.to_dict()
+    assert original_dtypes == df_dtypes
+
+
 def test_to_hierarchical():
     index = MultiIndex.from_tuples([(1, 'one'), (1, 'two'), (2, 'one'), (
         2, 'two')])
@@ -169,3 +183,20 @@ def test_to_series_with_arguments(idx):
     assert s.values is not idx.values
     assert s.index is not idx
     assert s.name != idx.name
+
+
+def test_squeeze_single_level():
+    mi = pd.MultiIndex.from_tuples([('a',), ('a',), ('b',), ('b',)],
+                                   names=['L1'])
+    expected = pd.Index(['a', 'a', 'b', 'b'], name='L1')
+    result = mi.squeeze()
+    tm.assert_index_equal(expected, result)
+
+
+def test_squeeze_multi_level():
+    mi = pd.MultiIndex.from_tuples([('a', 'a'), ('a', 'b'), ('b', 'a'),
+                                   ('b', 'b')],
+                                   names=['L1', 'L2'])
+    expected = mi.copy()
+    result = mi.squeeze()
+    tm.assert_index_equal(expected, result)