ENH: store tuples in MultiIndex, add legacy pickle testing and other speed tweaks

Author: wesm

pandas/core/index.py

@@ -8,6 +8,7 @@
 from pandas.core.common import (adjoin as _adjoin, _stringify,
                                 _is_bool_indexer, _asarray_tuplesafe)
 from pandas.util.decorators import cache_readonly
+import pandas.core.common as com
 import pandas._tseries as lib
 import pandas._engines as _engines
 
@@ -868,27 +869,32 @@ def __new__(cls, levels=None, labels=None, sortorder=None, names=None):
 
             return Index(levels[0], name=name).take(labels[0])
 
-        return np.arange(len(labels[0]), dtype=object).view(cls)
+        levels = [_ensure_index(lev) for lev in levels]
+        labels = [np.asarray(labs, dtype=np.int32) for labs in labels]
 
-    def __init__(self, levels, labels, sortorder=None, names=None,
-                 consistent=None):
-        self.levels = [_ensure_index(lev) for lev in levels]
-        self.labels = [np.asarray(labs, dtype=np.int32) for labs in labels]
+        values = [np.asarray(lev).take(lab)
+                  for lev, lab in zip(levels, labels)]
+        subarr = lib.fast_zip(values).view(cls)
+
+        subarr.levels = levels
+        subarr.labels = labels
 
         if names is None:
-            self.names = [None] * self.nlevels
+            subarr.names = [None] * subarr.nlevels
         else:
-            assert(len(names) == self.nlevels)
-            self.names = list(names)
+            assert(len(names) == subarr.nlevels)
+            subarr.names = list(names)
 
         # set the name
-        for i, name in enumerate(self.names):
-            self.levels[i].name = name
+        for i, name in enumerate(subarr.names):
+            subarr.levels[i].name = name
 
         if sortorder is not None:
-            self.sortorder = int(sortorder)
+            subarr.sortorder = int(sortorder)
         else:
-            self.sortorder = sortorder
+            subarr.sortorder = sortorder
+
+        return subarr
 
     @property
     def dtype(self):
@@ -908,9 +914,18 @@ def _get_level_number(self, level):
 
     @property
     def values(self):
-        values = [np.asarray(lev).take(lab)
-                  for lev, lab in zip(self.levels, self.labels)]
-        return lib.fast_zip(values)
+        if self._is_legacy_format:
+            # for legacy MultiIndex
+            values = [np.asarray(lev).take(lab)
+                      for lev, lab in zip(self.levels, self.labels)]
+            return lib.fast_zip(values)
+        else:
+            return self.view(np.ndarray)
+
+    @property
+    def _is_legacy_format(self):
+        contents = self.view(np.ndarray)
+        return len(contents) > 0 and not isinstance(contents[0], tuple)
 
     def get_level_values(self, level):
         """
@@ -1112,7 +1127,7 @@ def append(self, other):
         return MultiIndex.from_tuples(new_tuples, names=self.names)
 
     def argsort(self, *args, **kwargs):
-        return self.get_tuple_index().argsort()
+        return self.values.argsort()
 
     def drop(self, labels):
         """
@@ -1290,10 +1305,12 @@ def get_indexer(self, target, method=None):
         method = self._get_method(method)
 
         target_index = target
-        if isinstance(target, MultiIndex):
+        if isinstance(target, MultiIndex) and target._is_legacy_format:
             target_index = target.get_tuple_index()
 
-        self_index = self.get_tuple_index()
+        self_index = self
+        if self._is_legacy_format:
+            self_index = self.get_tuple_index()
 
         if method == 'pad':
             indexer = self._pad(self_index, target_index, self_index.indexMap,
@@ -1332,7 +1349,7 @@ def get_tuple_index(self):
         -------
         index : Index
         """
-        return Index(list(self))
+        return Index(self.values)
 
     def slice_locs(self, start=None, end=None, strict=False):
         """

pandas/src/reduce.pyx

@@ -89,18 +89,18 @@ cdef class Reducer:
             raise ValueError('function does not reduce')
         return result
 
-cdef class Grouper:
+cdef class SeriesGrouper:
     '''
     Performs generic grouping operation while avoiding ndarray construction
     overhead
     '''
     cdef:
-        Py_ssize_t nresults, ngroups
+        Py_ssize_t nresults, ngroup
         object arr, dummy, f, labels, counts
         bint passed_dummy
 
-    def __init__(self, object arr, object index, object f,
-                 object labels, ngroups, dummy=None):
+    def __init__(self, object arr, object f, object labels, ngroups,
+                 dummy=None):
         n = len(arr)
 
         assert(arr.ndim == 1)
@@ -111,6 +111,7 @@ cdef class Grouper:
         self.labels = labels
         self.f = f
         self.arr = arr
+
         self.dummy = self._check_dummy(dummy)
         self.passed_dummy = dummy is not None
 

pandas/tests/test_groupby.py

@@ -800,8 +800,8 @@ def test_groupby_level(self):
         result0 = frame.groupby(level=0).sum()
         result1 = frame.groupby(level=1).sum()
 
-        expected0 = frame.groupby(deleveled['first']).sum()
-        expected1 = frame.groupby(deleveled['second']).sum()
+        expected0 = frame.groupby(deleveled['first'].values).sum()
+        expected1 = frame.groupby(deleveled['second'].values).sum()
 
         self.assert_(result0.index.name == 'first')
         self.assert_(result1.index.name == 'second')

pandas/tests/test_index.py

@@ -680,6 +680,26 @@ def test_pickle(self):
         unpickled = pickle.loads(pickled)
         self.assert_(self.index.equals(unpickled))
 
+    def test_legacy_pickle(self):
+        import os
+        def curpath():
+            pth, _ = os.path.split(os.path.abspath(__file__))
+            return pth
+
+        ppath = os.path.join(curpath(), 'data/multiindex_v1.pickle')
+        obj = pickle.load(open(ppath, 'r'))
+
+        self.assert_(obj._is_legacy_format)
+
+        obj2 = MultiIndex.from_tuples(obj.values)
+        self.assert_(obj.equals(obj2))
+
+        res = obj.get_indexer(obj2[::-1])
+        exp = obj.get_indexer(obj[::-1])
+        exp2 = obj2.get_indexer(obj2[::-1])
+        assert_almost_equal(res, exp)
+        assert_almost_equal(exp, exp2)
+
     def test_contains(self):
         self.assert_(('foo', 'two') in self.index)
         self.assert_(('bar', 'two') not in self.index)