Add sort argument to Index.join

Issue pandas-dev#15582: fix the sort=True bug in DataFrame.join

Author: Albert Villanova del Moral

pandas/indexes/base.py

@@ -2831,7 +2831,8 @@ def _reindex_non_unique(self, target):
         new_index = self._shallow_copy_with_infer(new_labels, freq=None)
         return new_index, indexer, new_indexer
 
-    def join(self, other, how='left', level=None, return_indexers=False):
+    def join(self, other, how='left', level=None, return_indexers=False,
+             sort=False):
         """
         *this is an internal non-public method*
 
@@ -2844,6 +2845,7 @@ def join(self, other, how='left', level=None, return_indexers=False):
         how : {'left', 'right', 'inner', 'outer'}
         level : int or level name, default None
         return_indexers : boolean, default False
+        sort : boolean, default False
 
         Returns
         -------
@@ -2929,6 +2931,9 @@ def join(self, other, how='left', level=None, return_indexers=False):
         elif how == 'outer':
             join_index = self.union(other)
 
+        if sort:
+            join_index = join_index.sort_values()
+
         if return_indexers:
             if join_index is self:
                 lindexer = None

pandas/indexes/range.py

@@ -431,7 +431,8 @@ def union(self, other):
 
         return self._int64index.union(other)
 
-    def join(self, other, how='left', level=None, return_indexers=False):
+    def join(self, other, how='left', level=None, return_indexers=False,
+             sort=False):
         """
         *this is an internal non-public method*
 
@@ -444,16 +445,19 @@ def join(self, other, how='left', level=None, return_indexers=False):
         how : {'left', 'right', 'inner', 'outer'}
         level : int or level name, default None
         return_indexers : boolean, default False
+        sort : boolean, default False
 
         Returns
         -------
         join_index, (left_indexer, right_indexer)
         """
         if how == 'outer' and self is not other:
             # note: could return RangeIndex in more circumstances
-            return self._int64index.join(other, how, level, return_indexers)
+            return self._int64index.join(other, how, level, return_indexers,
+                                         sort)
 
-        return super(RangeIndex, self).join(other, how, level, return_indexers)
+        return super(RangeIndex, self).join(other, how, level, return_indexers,
+                                            sort)
 
     def __len__(self):
         """

pandas/tests/frame/test_join.py

@@ -41,6 +41,41 @@ def test_join(self):
                                 index=[0, 1, 2, 3])
         tm.assert_frame_equal(result, expected)
 
+    def test_join_sort(self):
+        df1 = pd.DataFrame({'a': [20, 10, 0]}, index=[2, 1, 0])
+        df2 = pd.DataFrame({'b': [100, 200, 300]}, index=[1, 2, 3])
+
+        result = df1.join(df2, sort=True)
+        expected = pd.DataFrame({'a': [0, 10, 20], 'b': [None, 100, 200]},
+                                index=[0, 1, 2])
+        tm.assert_frame_equal(result, expected)
+
+        result = df1.join(df2, how='left', sort=True)
+        expected = pd.DataFrame({'a': [0, 10, 20], 'b': [None, 100, 200]},
+                                index=[0, 1, 2])
+        tm.assert_frame_equal(result, expected)
+
+        result = df1.join(df2, how='right', sort=True)
+        expected = pd.DataFrame({'a': [10, 20, None], 'b': [100, 200, 300]},
+                                index=[1, 2, 3])
+        tm.assert_frame_equal(result, expected)
+
+        result = df2.join(df1, how='right', sort=True)
+        expected = pd.DataFrame([[None, 0], [100, 10], [200, 20]],
+                                columns=['b', 'a'], index=[0, 1, 2])
+        tm.assert_frame_equal(result, expected)
+
+        result = df1.join(df2, how='inner', sort=True)
+        expected = pd.DataFrame({'a': [10, 20], 'b': [100, 200]},
+                                index=[1, 2])
+        tm.assert_frame_equal(result, expected)
+
+        result = df1.join(df2, how='outer', sort=True)
+        expected = pd.DataFrame({'a': [0, 10, 20, None],
+                                 'b': [None, 100, 200, 300]},
+                                index=[0, 1, 2, 3])
+        tm.assert_frame_equal(result, expected)
+
     def test_join_index(self):
         # left / right
 

pandas/tools/merge.py

@@ -733,7 +733,8 @@ def _get_join_info(self):
 
         if self.left_index and self.right_index and self.how != 'asof':
             join_index, left_indexer, right_indexer = \
-                left_ax.join(right_ax, how=self.how, return_indexers=True)
+                left_ax.join(right_ax, how=self.how, return_indexers=True,
+                             sort=self.sort)
         elif self.right_index and self.how == 'left':
             join_index, left_indexer, right_indexer = \
                 _left_join_on_index(left_ax, right_ax, self.left_join_keys,

pandas/tseries/index.py

@@ -1037,7 +1037,8 @@ def union_many(self, others):
             this.offset = to_offset(this.inferred_freq)
         return this
 
-    def join(self, other, how='left', level=None, return_indexers=False):
+    def join(self, other, how='left', level=None, return_indexers=False,
+             sort=False):
         """
         See Index.join
         """
@@ -1051,7 +1052,7 @@ def join(self, other, how='left', level=None, return_indexers=False):
 
         this, other = self._maybe_utc_convert(other)
         return Index.join(this, other, how=how, level=level,
-                          return_indexers=return_indexers)
+                          return_indexers=return_indexers, sort=sort)
 
     def _maybe_utc_convert(self, other):
         this = self