ENH: implemented lexsort_depth and is_lexsorted

Author: wesm

pandas/core/index.py

@@ -8,7 +8,7 @@
 from pandas.core.common import (_format, adjoin as _adjoin, _stringify,
                                 _ensure_index, _is_bool_indexer,
                                 _asarray_tuplesafe)
-from pandas.util.decorators import deprecate
+from pandas.util.decorators import deprecate, cache_readonly
 import pandas.core.common as common
 import pandas._tseries as _tseries
 
@@ -555,8 +555,25 @@ def format(self, space=2, sparsify=True, vertical=False, adjoin=True):
     def is_all_dates(self):
         return False
 
-    def is_sorted(self):
-        raise NotImplementedError
+    def is_lexsorted(self):
+        """
+        Return True if the labels are lexicographically sorted
+        """
+        return self.lexsort_depth == self.nlevels
+
+    @cache_readonly
+    def lexsort_depth(self):
+        if self.sortorder is not None:
+            if self.sortorder == 0:
+                return self.nlevels
+            else:
+                return 0
+
+        for k in range(self.nlevels, 0, -1):
+            if _tseries.is_lexsorted(self.labels[:k]):
+                return k
+
+        return 0
 
     @classmethod
     def from_arrays(cls, arrays, sortorder=None):
@@ -864,9 +881,6 @@ def slice_locs(self, start=None, end=None):
         -----
         This function assumes that the data is sorted by the first level
         """
-        # relax for now
-        # assert(self.sortorder == 0)
-
         if start is None:
             start_slice = 0
         else:
@@ -884,8 +898,9 @@ def slice_locs(self, start=None, end=None):
         return start_slice, end_slice
 
     def _partial_tup_index(self, tup, side='left'):
-        # relax for now
-        # assert(self.sortorder == 0)
+        if len(tup) > self.lexsort_depth:
+            raise Exception('MultiIndex lexsort depth %d, key was %d long' %
+                            (self.lexsort_depth, len(tup)))
 
         n = len(tup)
         start, end = 0, len(self)

pandas/src/groupby.pyx

@@ -133,6 +133,36 @@ def groupby_indices(ndarray values):
     return result
 
 
+@cython.wraparound(False)
+@cython.boundscheck(False)
+def is_lexsorted(list list_of_arrays):
+    cdef:
+        int i
+        Py_ssize_t n, nlevels
+        int32_t k, cur, pre
+
+    nlevels = len(list_of_arrays)
+    n = len(list_of_arrays[0])
+
+    cdef int32_t **vecs = <int32_t **> malloc(nlevels * sizeof(int32_t*))
+    for i from 0 <= i < nlevels:
+        vecs[i] = <int32_t *> (<ndarray> list_of_arrays[i]).data
+
+    # assume uniqueness??
+
+    for i from 1 <= i < n:
+        for k from 0 <= k < nlevels:
+            cur = vecs[k][i]
+            pre = vecs[k][i-1]
+            if cur == pre:
+                continue
+            elif cur > pre:
+                break
+            else:
+                return False
+    free(vecs)
+    return True
+
 @cython.wraparound(False)
 @cython.boundscheck(False)
 def group_labels(ndarray[object] values):

pandas/tests/test_multilevel.py

@@ -231,8 +231,26 @@ def test_insert_index(self):
     def test_alignment(self):
         pass
 
+    def test_is_lexsorted(self):
+        levels = [[0, 1], [0, 1, 2]]
+
+        index = MultiIndex(levels=levels,
+                           labels=[[0, 0, 0, 1, 1, 1],
+                                   [0, 1, 2, 0, 1, 2]])
+        self.assert_(index.is_lexsorted())
+
+        index = MultiIndex(levels=levels,
+                           labels=[[0, 0, 0, 1, 1, 1],
+                                   [0, 1, 2, 0, 2, 1]])
+        self.assert_(not index.is_lexsorted())
+
+        index = MultiIndex(levels=levels,
+                           labels=[[0, 0, 1, 0, 1, 1],
+                                   [0, 1, 0, 2, 2, 1]])
+        self.assert_(not index.is_lexsorted())
 
 if __name__ == '__main__':
+
     # unittest.main()
     import nose
     # nose.runmodule(argv=[__file__,'-vvs','-x', '--pdb-failure'],