ENH: per-level sort order in sort/sort_index. close pandas-dev#928

Author: wesm

RELEASE.rst

@@ -27,6 +27,10 @@ pandas 0.9.1
 
 **Release date:** NOT YET RELEASED
 
+**New features**
+
+  - Can specify multiple sort orders in DataFrame/Series.sort/sort_index (#928)
+
 **Improvements to existing features**
 
   - Improve performance of datetime + business day offset with large number of

pandas/core/frame.py

@@ -2796,13 +2796,18 @@ def sort(self, columns=None, column=None, axis=0, ascending=True,
         columns : object
             Column name(s) in frame. Accepts a column name or a list or tuple
             for a nested sort.
-        ascending : boolean, default True
-            Sort ascending vs. descending
+        ascending : boolean or list, default True
+            Sort ascending vs. descending. Specify list for multiple sort
+            orders
         axis : {0, 1}
             Sort index/rows versus columns
         inplace : boolean, default False
             Sort the DataFrame without creating a new instance
 
+        Examples
+        --------
+        >>> result = df.sort(['A', 'B'], ascending=[1, 0])
+
         Returns
         -------
         sorted : DataFrame
@@ -2826,11 +2831,16 @@ def sort_index(self, axis=0, by=None, ascending=True, inplace=False):
         by : object
             Column name(s) in frame. Accepts a column name or a list or tuple
             for a nested sort.
-        ascending : boolean, default True
-            Sort ascending vs. descending
+        ascending : boolean or list, default True
+            Sort ascending vs. descending. Specify list for multiple sort
+            orders
         inplace : boolean, default False
             Sort the DataFrame without creating a new instance
 
+        Examples
+        --------
+        >>> result = df.sort_index(by=['A', 'B'], ascending=[1, 0])
+
         Returns
         -------
         sorted : DataFrame
@@ -2846,14 +2856,17 @@ def sort_index(self, axis=0, by=None, ascending=True, inplace=False):
             assert(axis == 0)
             if isinstance(by, (tuple, list)):
                 keys = [self[x].values for x in by]
-                indexer = _lexsort_indexer(keys)
+                indexer = _lexsort_indexer(keys, orders=ascending)
             else:
                 indexer = self[by].values.argsort()
+                if not ascending:
+                    indexer = indexer[::-1]
+        elif isinstance(labels, MultiIndex):
+            indexer = _lexsort_indexer(labels.labels, orders=ascending)
         else:
             indexer = labels.argsort()
-
-        if not ascending:
-            indexer = indexer[::-1]
+            if not ascending:
+                indexer = indexer[::-1]
 
         if inplace:
             if axis == 1:

pandas/core/groupby.py

@@ -2103,18 +2103,31 @@ def _indexer_from_factorized(labels, shape, compress=True):
     return indexer
 
 
-def _lexsort_indexer(keys):
+def _lexsort_indexer(keys, orders=None):
     labels = []
     shape = []
-    for key in keys:
+
+    if isinstance(orders, bool):
+        orders = [orders] * len(keys)
+    elif orders is None:
+        orders = [True] * len(keys)
+
+    for key, order in zip(keys, orders):
         rizer = lib.Factorizer(len(key))
 
         if not key.dtype == np.object_:
             key = key.astype('O')
 
         ids, _ = rizer.factorize(key, sort=True)
+
+        n = len(rizer.uniques)
+        shape.append(n)
+        if not order:
+            mask = ids == -1
+            ids = np.where(mask, -1, n - ids)
+
         labels.append(ids)
-        shape.append(len(rizer.uniques))
+
     return _indexer_from_factorized(labels, shape)
 
 

pandas/core/series.py

@@ -1749,15 +1749,28 @@ def sort_index(self, ascending=True):
 
         Parameters
         ----------
-        ascending : boolean, default True
-            Sort ascending vs. descending
+        ascending : boolean or list, default True
+            Sort ascending vs. descending. Specify list for multiple sort
+            orders
+
+        Examples
+        --------
+        >>> result1 = s.sort_index(ascending=False)
+        >>> result2 = s.sort_index(ascending=[1, 0])
 
         Returns
         -------
         sorted_obj : Series
         """
-        new_labels, indexer = self.index.order(return_indexer=True,
-                                               ascending=ascending)
+        index = self.index
+        if isinstance(index, MultiIndex):
+            from pandas.core.groupby import _lexsort_indexer
+            indexer = _lexsort_indexer(index.labels, orders=ascending)
+            new_labels = index.take(indexer)
+        else:
+            new_labels, indexer = index.order(return_indexer=True,
+                                              ascending=ascending)
+
         new_values = self.values.take(indexer)
         return Series(new_values, new_labels, name=self.name)
 

pandas/tests/test_frame.py

@@ -5657,7 +5657,9 @@ def test_sort_index_multicolumn(self):
         assert_frame_equal(result, expected)
 
         result = frame.sort_index(by=['A', 'B'], ascending=False)
-        expected = frame.take(indexer[::-1])
+        indexer = np.lexsort((frame['B'].rank(ascending=False),
+                              frame['A'].rank(ascending=False)))
+        expected = frame.take(indexer)
         assert_frame_equal(result, expected)
 
         result = frame.sort_index(by=['B', 'A'])
@@ -5695,6 +5697,35 @@ def test_sort_index_inplace(self):
         expected = frame.ix[:, ::-1]
         assert_frame_equal(df, expected)
 
+    def test_sort_index_different_sortorder(self):
+        import random
+        A = np.arange(20).repeat(5)
+        B = np.tile(np.arange(5), 20)
+
+        indexer = np.random.permutation(100)
+        A = A.take(indexer)
+        B = B.take(indexer)
+
+        df = DataFrame({'A' : A, 'B' : B,
+                        'C' : np.random.randn(100)})
+
+        result = df.sort_index(by=['A', 'B'], ascending=[1, 0])
+
+        ex_indexer = np.lexsort((df.B.max() - df.B, df.A))
+        expected = df.take(ex_indexer)
+        assert_frame_equal(result, expected)
+
+        # test with multiindex, too
+        idf = df.set_index(['A', 'B'])
+
+        result = idf.sort_index(ascending=[1, 0])
+        expected = idf.take(ex_indexer)
+        assert_frame_equal(result, expected)
+
+        # also, Series!
+        result = idf['C'].sort_index(ascending=[1, 0])
+        assert_series_equal(result, expected['C'])
+
     def test_sort_inplace(self):
         frame = DataFrame(np.random.randn(4, 4), index=[1, 2, 3, 4],
                           columns=['A', 'B', 'C', 'D'])