ENH: DataFrame.pivot now returns MultiIndex columns with multiple fields. Return view when indexing homogeneous hierarhical columns. various related unit tests

address GH #124

Author: wesm

pandas/core/frame.py

@@ -795,10 +795,10 @@ def _getitem_multilevel(self, key):
         loc = self.columns.get_loc(key)
         if isinstance(loc, slice):
             new_columns = self.columns[loc]
-            result = self.reindex(columns=new_columns)
-
-            # HACK: need a more general way of addressing this problem
-            result.columns = _maybe_droplevels(new_columns, key)
+            new_columns = _maybe_droplevels(new_columns, key)
+            new_values = self.values[:, loc]
+            result = DataFrame(new_values, index=self.index,
+                               columns=new_columns)
             return result
         else:
             return self._getitem_single(key)
@@ -1631,11 +1631,9 @@ def pivot(self, index=None, columns=None, values=None):
         pivoted : DataFrame (value column specified) or WidePanel (no value
         column specified)
         """
-        from pandas.core.panel import _make_long_index, LongPanel
-
         index_vals = self[index]
         column_vals = self[columns]
-        long_index = _make_long_index(index_vals, column_vals)
+        mindex = MultiIndex.from_arrays([index_vals, column_vals])
 
         if values is None:
             items = self.columns - [index, columns]
@@ -1644,14 +1642,15 @@ def pivot(self, index=None, columns=None, values=None):
             items = [values]
             mat = np.atleast_2d(self[values].values).T
 
-        lp = LongPanel(mat, index=long_index, columns=items)
-        lp = lp.sortlevel(level=0)
+        stacked = DataFrame(mat, index=mindex, columns=items)
 
-        wp = lp.to_wide()
+        if not mindex.is_lexsorted:
+            stacked = stacked.sortlevel(level=0)
+
+        unstacked = stacked.unstack()
         if values is not None:
-            return wp[values]
-        else:
-            return wp
+            unstacked.columns = unstacked.columns.droplevel(0)
+        return unstacked
 
     def stack(self):
         """

pandas/core/index.py

@@ -652,18 +652,20 @@ def levshape(self):
     def __reduce__(self):
         """Necessary for making this object picklable"""
         object_state = list(np.ndarray.__reduce__(self))
-        subclass_state = (self.levels, self.labels)
+        subclass_state = (self.levels, self.labels, self.sortorder, self.names)
         object_state[2] = (object_state[2], subclass_state)
         return tuple(object_state)
 
     def __setstate__(self, state):
         """Necessary for making this object picklable"""
         nd_state, own_state = state
         np.ndarray.__setstate__(self, nd_state)
-        levels, labels, = own_state
+        levels, labels, sortorder, names = own_state
 
         self.levels = [Index(x) for x in levels]
         self.labels = labels
+        self.names = names
+        self.sortorder = sortorder
 
     def __getitem__(self, key):
         arr_idx = self.view(np.ndarray)
@@ -959,12 +961,12 @@ def get_loc(self, key):
                     raise KeyError(key)
                 return result
         else:
-            # assert(self.sortorder == 0)
-            # slice level 0
             level = self.levels[0]
             labels = self.labels[0]
-
             loc = level.get_loc(key)
+
+            assert(self.lexsort_depth >= 1)
+
             i = labels.searchsorted(loc, side='left')
             j = labels.searchsorted(loc, side='right')
             return slice(i, j)

pandas/tests/test_frame.py

@@ -1721,7 +1721,7 @@ def test_pivot(self):
         wp = tm.makeWidePanel()
         lp = wp.to_long()
         df = DataFrame.from_records(lp.toRecords())
-        tm.assert_panel_equal(df.pivot('major', 'minor'), wp)
+        assert_frame_equal(df.pivot('major', 'minor'), lp.unstack())
 
     def test_reindex(self):
         newFrame = self.frame.reindex(self.ts1.index)

pandas/tests/test_multilevel.py

@@ -262,6 +262,22 @@ def test_is_lexsorted(self):
         self.assert_(not index.is_lexsorted())
         self.assert_(index.lexsort_depth == 0)
 
+    def test_frame_getitem_view(self):
+        df = self.frame.T
+        df['foo'].values[:] = 0
+        self.assert_((df['foo'].values == 0).all())
+
+        # but not if it's mixed-type
+        df['foo', 'four'] = 'foo'
+        df = df.sortlevel(0, axis=1)
+        df['foo']['one'] = 2
+        self.assert_((df['foo', 'one'] == 0).all())
+
+    def test_frame_getitem_not_sorted(self):
+        df = self.frame.T
+        df['foo', 'four'] = 'foo'
+        self.assertRaises(Exception, df.__getitem__, 'foo')
+
 if __name__ == '__main__':
 
     # unittest.main()