ENH: Add sort parameter to set operations for some Indexes and adjust tests

Author: reidy-p

doc/source/whatsnew/v0.24.0.rst

@@ -405,7 +405,7 @@ Other Enhancements
 - :func:`read_fwf` now accepts keyword ``infer_nrows`` (:issue:`15138`).
 - :func:`~DataFrame.to_parquet` now supports writing a ``DataFrame`` as a directory of parquet files partitioned by a subset of the columns when ``engine = 'pyarrow'`` (:issue:`23283`)
 - :meth:`Timestamp.tz_localize`, :meth:`DatetimeIndex.tz_localize`, and :meth:`Series.tz_localize` have gained the ``nonexistent`` argument for alternative handling of nonexistent times. See :ref:`timeseries.timezone_nonexistent` (:issue:`8917`)
-- :meth:`Index.difference` now has an optional ``sort`` parameter to specify whether the results should be sorted if possible (:issue:`17839`)
+- :meth:`Index.difference`, :meth:`Index.intersection`, :meth:`Index.union`, and :meth:`Index.symmetric_difference` now have an optional ``sort`` parameter to control whether the results should be sorted if possible (:issue:`17839`, :issue:`24471`)
 - :meth:`read_excel()` now accepts ``usecols`` as a list of column names or callable (:issue:`18273`)
 - :meth:`MultiIndex.to_flat_index` has been added to flatten multiple levels into a single-level :class:`Index` object.
 - :meth:`DataFrame.to_stata` and :class:`pandas.io.stata.StataWriter117` can write mixed sting columns to Stata strl format (:issue:`23633`)

pandas/_libs/lib.pyx

@@ -199,7 +199,7 @@ def item_from_zerodim(val: object) -> object:
 
 @cython.wraparound(False)
 @cython.boundscheck(False)
-def fast_unique_multiple(list arrays):
+def fast_unique_multiple(list arrays, sort: bool=True):
     cdef:
         ndarray[object] buf
         Py_ssize_t k = len(arrays)
@@ -216,10 +216,11 @@ def fast_unique_multiple(list arrays):
             if val not in table:
                 table[val] = stub
                 uniques.append(val)
-    try:
-        uniques.sort()
-    except Exception:
-        pass
+    if sort:
+        try:
+            uniques.sort()
+        except Exception:
+            pass
 
     return uniques
 

pandas/core/indexes/api.py

@@ -112,7 +112,7 @@ def _get_combined_index(indexes, intersect=False, sort=False):
     elif intersect:
         index = indexes[0]
         for other in indexes[1:]:
-            index = index.intersection(other)
+            index = index.intersection(other, sort=sort)
     else:
         index = _union_indexes(indexes, sort=sort)
         index = ensure_index(index)

pandas/core/indexes/base.py

@@ -2235,13 +2235,17 @@ def _get_reconciled_name_object(self, other):
             return self._shallow_copy(name=name)
         return self
 
-    def union(self, other):
+    def union(self, other, sort=True):
         """
-        Form the union of two Index objects and sorts if possible.
+        Form the union of two Index objects.
 
         Parameters
         ----------
         other : Index or array-like
+        sort : bool, default True
+            Sort the resulting index if possible
+
+            .. versionadded:: 0.24.0
 
         Returns
         -------
@@ -2271,7 +2275,7 @@ def union(self, other):
         if not is_dtype_union_equal(self.dtype, other.dtype):
             this = self.astype('O')
             other = other.astype('O')
-            return this.union(other)
+            return this.union(other, sort=sort)
 
         # TODO(EA): setops-refactor, clean all this up
         if is_period_dtype(self) or is_datetime64tz_dtype(self):
@@ -2302,44 +2306,51 @@ def union(self, other):
                                            allow_fill=False)
                 result = _concat._concat_compat((lvals, other_diff))
 
-                try:
-                    lvals[0] < other_diff[0]
-                except TypeError as e:
-                    warnings.warn("%s, sort order is undefined for "
-                                  "incomparable objects" % e, RuntimeWarning,
-                                  stacklevel=3)
-                else:
-                    types = frozenset((self.inferred_type,
-                                       other.inferred_type))
-                    if not types & _unsortable_types:
-                        result.sort()
+                if sort:
+                    try:
+                        lvals[0] < other_diff[0]
+                    except TypeError as e:
+                        warnings.warn("%s, sort order is undefined for "
+                                      "incomparable objects" % e,
+                                      RuntimeWarning,
+                                      stacklevel=3)
+                    else:
+                        types = frozenset((self.inferred_type,
+                                           other.inferred_type))
+                        if not types & _unsortable_types:
+                            result.sort()
 
             else:
                 result = lvals
 
-                try:
-                    result = np.sort(result)
-                except TypeError as e:
-                    warnings.warn("%s, sort order is undefined for "
-                                  "incomparable objects" % e, RuntimeWarning,
-                                  stacklevel=3)
+                if sort:
+                    try:
+                        result = np.sort(result)
+                    except TypeError as e:
+                        warnings.warn("%s, sort order is undefined for "
+                                      "incomparable objects" % e,
+                                      RuntimeWarning,
+                                      stacklevel=3)
 
         # for subclasses
         return self._wrap_setop_result(other, result)
 
     def _wrap_setop_result(self, other, result):
         return self._constructor(result, name=get_op_result_name(self, other))
 
-    def intersection(self, other):
+    def intersection(self, other, sort=True):
         """
         Form the intersection of two Index objects.
 
-        This returns a new Index with elements common to the index and `other`,
-        preserving the order of the calling index.
+        This returns a new Index with elements common to the index and `other`.
 
         Parameters
         ----------
         other : Index or array-like
+        sort : bool, default True
+            Sort the resulting index if possible
+
+            .. versionadded:: 0.24.0
 
         Returns
         -------
@@ -2362,7 +2373,7 @@ def intersection(self, other):
         if not is_dtype_equal(self.dtype, other.dtype):
             this = self.astype('O')
             other = other.astype('O')
-            return this.intersection(other)
+            return this.intersection(other, sort=sort)
 
         # TODO(EA): setops-refactor, clean all this up
         if is_period_dtype(self):
@@ -2393,6 +2404,13 @@ def intersection(self, other):
         taken = other.take(indexer)
         if self.name != other.name:
             taken.name = None
+        if sort:
+            try:
+                taken = taken.sort_values()
+            except TypeError as e:
+                warnings.warn("%s, sort order is undefined for "
+                              "incomparable objects" % e, RuntimeWarning,
+                              stacklevel=3)
         return taken
 
     def difference(self, other, sort=True):
@@ -2448,16 +2466,18 @@ def difference(self, other, sort=True):
 
         return this._shallow_copy(the_diff, name=result_name, freq=None)
 
-    def symmetric_difference(self, other, result_name=None):
+    def symmetric_difference(self, other, result_name=None, sort=True):
         """
         Compute the symmetric difference of two Index objects.
 
-        It's sorted if sorting is possible.
-
         Parameters
         ----------
         other : Index or array-like
         result_name : str
+        sort : bool, default True
+            Sort the resulting index if possible
+
+            .. versionadded:: 0.24.0
 
         Returns
         -------
@@ -2502,10 +2522,11 @@ def symmetric_difference(self, other, result_name=None):
         right_diff = other.values.take(right_indexer)
 
         the_diff = _concat._concat_compat([left_diff, right_diff])
-        try:
-            the_diff = sorting.safe_sort(the_diff)
-        except TypeError:
-            pass
+        if sort:
+            try:
+                the_diff = sorting.safe_sort(the_diff)
+            except TypeError:
+                pass
 
         attribs = self._get_attributes_dict()
         attribs['name'] = result_name
@@ -3233,7 +3254,7 @@ def join(self, other, how='left', level=None, return_indexers=False,
         elif how == 'right':
             join_index = other
         elif how == 'inner':
-            join_index = self.intersection(other)
+            join_index = self.intersection(other, sort=sort)
         elif how == 'outer':
             join_index = self.union(other)
 

pandas/core/indexes/datetimes.py

@@ -603,7 +603,7 @@ def _wrap_setop_result(self, other, result):
             raise ValueError('Passed item and index have different timezone')
         return self._shallow_copy(result, name=name, freq=None, tz=self.tz)
 
-    def intersection(self, other):
+    def intersection(self, other, sort=True):
         """
         Specialized intersection for DatetimeIndex objects. May be much faster
         than Index.intersection
@@ -626,7 +626,7 @@ def intersection(self, other):
                 other = DatetimeIndex(other)
             except (TypeError, ValueError):
                 pass
-            result = Index.intersection(self, other)
+            result = Index.intersection(self, other, sort=sort)
             if isinstance(result, DatetimeIndex):
                 if result.freq is None:
                     result.freq = to_offset(result.inferred_freq)
@@ -636,7 +636,7 @@ def intersection(self, other):
               other.freq != self.freq or
               not other.freq.isAnchored() or
               (not self.is_monotonic or not other.is_monotonic)):
-            result = Index.intersection(self, other)
+            result = Index.intersection(self, other, sort=sort)
             result = self._shallow_copy(result._values, name=result.name,
                                         tz=result.tz, freq=None)
             if result.freq is None:

pandas/core/indexes/interval.py

@@ -1102,11 +1102,8 @@ def func(self, other, sort=True):
                        'objects that have compatible dtypes')
                 raise TypeError(msg.format(op=op_name))
 
-            if op_name == 'difference':
-                result = getattr(self._multiindex, op_name)(other._multiindex,
-                                                            sort)
-            else:
-                result = getattr(self._multiindex, op_name)(other._multiindex)
+            result = getattr(self._multiindex, op_name)(other._multiindex,
+                                                        sort=sort)
             result_name = get_op_result_name(self, other)
 
             # GH 19101: ensure empty results have correct dtype

pandas/core/indexes/multi.py

@@ -2878,13 +2878,17 @@ def equal_levels(self, other):
                 return False
         return True
 
-    def union(self, other):
+    def union(self, other, sort=True):
         """
-        Form the union of two MultiIndex objects, sorting if possible
+        Form the union of two MultiIndex objects
 
         Parameters
         ----------
         other : MultiIndex or array / Index of tuples
+        sort : bool, default True
+            Sort the resulting MultiIndex if possible
+
+            .. versionadded:: 0.24.0
 
         Returns
         -------
@@ -2899,17 +2903,23 @@ def union(self, other):
             return self
 
         uniq_tuples = lib.fast_unique_multiple([self._ndarray_values,
-                                                other._ndarray_values])
+                                                other._ndarray_values],
+                                               sort=sort)
+
         return MultiIndex.from_arrays(lzip(*uniq_tuples), sortorder=0,
                                       names=result_names)
 
-    def intersection(self, other):
+    def intersection(self, other, sort=True):
         """
-        Form the intersection of two MultiIndex objects, sorting if possible
+        Form the intersection of two MultiIndex objects.
 
         Parameters
         ----------
         other : MultiIndex or array / Index of tuples
+        sort : bool, default True
+            Sort the resulting MultiIndex if possible
+
+            .. versionadded:: 0.24.0
 
         Returns
         -------
@@ -2923,7 +2933,11 @@ def intersection(self, other):
 
         self_tuples = self._ndarray_values
         other_tuples = other._ndarray_values
-        uniq_tuples = sorted(set(self_tuples) & set(other_tuples))
+        uniq_tuples = set(self_tuples) & set(other_tuples)
+
+        if sort:
+            uniq_tuples = sorted(uniq_tuples)
+
         if len(uniq_tuples) == 0:
             return MultiIndex(levels=self.levels,
                               codes=[[]] * self.nlevels,
@@ -2934,7 +2948,7 @@ def intersection(self, other):
 
     def difference(self, other, sort=True):
         """
-        Compute sorted set difference of two MultiIndex objects
+        Compute set difference of two MultiIndex objects
 
         Parameters
         ----------

pandas/core/indexes/range.py

@@ -343,14 +343,17 @@ def equals(self, other):
 
         return super(RangeIndex, self).equals(other)
 
-    def intersection(self, other):
+    def intersection(self, other, sort=True):
         """
-        Form the intersection of two Index objects. Sortedness of the result is
-        not guaranteed
+        Form the intersection of two Index objects.
 
         Parameters
         ----------
         other : Index or array-like
+        sort : bool, default True
+            Sort the resulting index if possible
+
+            .. versionadded:: 0.24.0
 
         Returns
         -------
@@ -361,7 +364,7 @@ def intersection(self, other):
             return self._get_reconciled_name_object(other)
 
         if not isinstance(other, RangeIndex):
-            return super(RangeIndex, self).intersection(other)
+            return super(RangeIndex, self).intersection(other, sort=sort)
 
         if not len(self) or not len(other):
             return RangeIndex._simple_new(None)
@@ -398,6 +401,8 @@ def intersection(self, other):
 
         if (self._step < 0 and other._step < 0) is not (new_index._step < 0):
             new_index = new_index[::-1]
+        if sort:
+            new_index = new_index.sort_values()
         return new_index
 
     def _min_fitting_element(self, lower_limit):

pandas/tests/indexes/datetimes/test_setops.py

@@ -138,7 +138,8 @@ def test_intersection2(self):
 
     @pytest.mark.parametrize("tz", [None, 'Asia/Tokyo', 'US/Eastern',
                                     'dateutil/US/Pacific'])
-    def test_intersection(self, tz):
+    @pytest.mark.parametrize("sort", [True, False])
+    def test_intersection(self, tz, sort):
         # GH 4690 (with tz)
         base = date_range('6/1/2000', '6/30/2000', freq='D', name='idx')
 
@@ -185,7 +186,9 @@ def test_intersection(self, tz):
 
         for (rng, expected) in [(rng2, expected2), (rng3, expected3),
                                 (rng4, expected4)]:
-            result = base.intersection(rng)
+            result = base.intersection(rng, sort=sort)
+            if sort:
+                expected = expected.sort_values()
             tm.assert_index_equal(result, expected)
             assert result.name == expected.name
             assert result.freq is None