ENH: MultiIndex.intersection now keeping EA dtypes (#48604)

* ENH: MultiIndex.intersection now keeping EA dtypes

* Improve performance

* Improve performance

* Improve performance

* Add test and whatsnew

* Add gh ref

* Fix typing

* Fix typing

Author: phofl

asv_bench/benchmarks/index_object.py

@@ -19,28 +19,38 @@
 class SetOperations:
 
     params = (
-        ["datetime", "date_string", "int", "strings"],
+        ["monotonic", "non_monotonic"],
+        ["datetime", "date_string", "int", "strings", "ea_int"],
         ["intersection", "union", "symmetric_difference"],
     )
-    param_names = ["dtype", "method"]
+    param_names = ["index_structure", "dtype", "method"]
 
-    def setup(self, dtype, method):
+    def setup(self, index_structure, dtype, method):
         N = 10**5
         dates_left = date_range("1/1/2000", periods=N, freq="T")
         fmt = "%Y-%m-%d %H:%M:%S"
         date_str_left = Index(dates_left.strftime(fmt))
         int_left = Index(np.arange(N))
+        ea_int_left = Index(np.arange(N), dtype="Int64")
         str_left = tm.makeStringIndex(N)
+
         data = {
-            "datetime": {"left": dates_left, "right": dates_left[:-1]},
-            "date_string": {"left": date_str_left, "right": date_str_left[:-1]},
-            "int": {"left": int_left, "right": int_left[:-1]},
-            "strings": {"left": str_left, "right": str_left[:-1]},
+            "datetime": dates_left,
+            "date_string": date_str_left,
+            "int": int_left,
+            "strings": str_left,
+            "ea_int": ea_int_left,
         }
+
+        if index_structure == "non_monotonic":
+            data = {k: mi[::-1] for k, mi in data.items()}
+
+        data = {k: {"left": idx, "right": idx[:-1]} for k, idx in data.items()}
+
         self.left = data[dtype]["left"]
         self.right = data[dtype]["right"]
 
-    def time_operation(self, dtype, method):
+    def time_operation(self, index_structure, dtype, method):
         getattr(self.left, method)(self.right)
 
 

asv_bench/benchmarks/multiindex_object.py

@@ -237,7 +237,7 @@ class SetOperations:
 
     params = [
         ("monotonic", "non_monotonic"),
-        ("datetime", "int", "string"),
+        ("datetime", "int", "string", "ea_int"),
         ("intersection", "union", "symmetric_difference"),
     ]
     param_names = ["index_structure", "dtype", "method"]
@@ -255,10 +255,14 @@ def setup(self, index_structure, dtype, method):
         level2 = tm.makeStringIndex(N // 1000).values
         str_left = MultiIndex.from_product([level1, level2])
 
+        level2 = range(N // 1000)
+        ea_int_left = MultiIndex.from_product([level1, Series(level2, dtype="Int64")])
+
         data = {
             "datetime": dates_left,
             "int": int_left,
             "string": str_left,
+            "ea_int": ea_int_left,
         }
 
         if index_structure == "non_monotonic":

doc/source/whatsnew/v1.6.0.rst

@@ -112,6 +112,7 @@ Performance improvements
 - Performance improvement for :class:`Series` constructor passing integer numpy array with nullable dtype (:issue:`48338`)
 - Performance improvement in :meth:`DataFrame.loc` and :meth:`Series.loc` for tuple-based indexing of a :class:`MultiIndex` (:issue:`48384`)
 - Performance improvement for :meth:`MultiIndex.unique` (:issue:`48335`)
+- Performance improvement for :meth:`MultiIndex.intersection` (:issue:`48604`)
 - Performance improvement in ``var`` for nullable dtypes (:issue:`48379`).
 - Performance improvement to :func:`read_sas` with ``blank_missing=True`` (:issue:`48502`)
 -
@@ -179,6 +180,7 @@ MultiIndex
 ^^^^^^^^^^
 - Bug in :class:`MultiIndex.set_levels` raising ``IndexError`` when setting empty level (:issue:`48636`)
 - Bug in :meth:`MultiIndex.unique` losing extension array dtype (:issue:`48335`)
+- Bug in :meth:`MultiIndex.intersection` losing extension array (:issue:`48604`)
 - Bug in :meth:`MultiIndex.union` losing extension array (:issue:`48498`, :issue:`48505`)
 - Bug in :meth:`MultiIndex.append` not checking names for equality (:issue:`48288`)
 - Bug in :meth:`MultiIndex.symmetric_difference` losing extension array (:issue:`48607`)

pandas/core/indexes/base.py

@@ -3559,13 +3559,19 @@ def _intersection(self, other: Index, sort: bool = False):
             and self._can_use_libjoin
         ):
             try:
-                result = self._inner_indexer(other)[0]
+                res_indexer, indexer, _ = self._inner_indexer(other)
             except TypeError:
                 # non-comparable; should only be for object dtype
                 pass
             else:
                 # TODO: algos.unique1d should preserve DTA/TDA
-                res = algos.unique1d(result)
+                if self.is_numeric():
+                    # This is faster, because Index.unique() checks for uniqueness
+                    # before calculating the unique values.
+                    res = algos.unique1d(res_indexer)
+                else:
+                    result = self.take(indexer)
+                    res = result.drop_duplicates()
                 return ensure_wrapped_if_datetimelike(res)
 
         res_values = self._intersection_via_get_indexer(other, sort=sort)
@@ -3577,7 +3583,9 @@ def _wrap_intersection_result(self, other, result):
         return self._wrap_setop_result(other, result)
 
     @final
-    def _intersection_via_get_indexer(self, other: Index, sort) -> ArrayLike:
+    def _intersection_via_get_indexer(
+        self, other: Index | MultiIndex, sort
+    ) -> ArrayLike | MultiIndex:
         """
         Find the intersection of two Indexes using get_indexer.
 
@@ -3600,7 +3608,10 @@ def _intersection_via_get_indexer(self, other: Index, sort) -> ArrayLike:
             # unnecessary in the case with sort=None bc we will sort later
             taker = np.sort(taker)
 
-        result = left_unique.take(taker)._values
+        if isinstance(left_unique, ABCMultiIndex):
+            result = left_unique.take(taker)
+        else:
+            result = left_unique.take(taker)._values
         return result
 
     @final

pandas/core/indexes/multi.py

@@ -3725,16 +3725,7 @@ def _maybe_match_names(self, other):
 
     def _wrap_intersection_result(self, other, result) -> MultiIndex:
         _, result_names = self._convert_can_do_setop(other)
-
-        if len(result) == 0:
-            return MultiIndex(
-                levels=self.levels,
-                codes=[[]] * self.nlevels,
-                names=result_names,
-                verify_integrity=False,
-            )
-        else:
-            return MultiIndex.from_arrays(zip(*result), sortorder=0, names=result_names)
+        return result.set_names(result_names)
 
     def _wrap_difference_result(self, other, result) -> MultiIndex:
         _, result_names = self._convert_can_do_setop(other)

pandas/tests/indexes/multi/test_setops.py

@@ -525,7 +525,7 @@ def test_intersection_with_missing_values_on_both_sides(nulls_fixture):
     mi1 = MultiIndex.from_arrays([[3, nulls_fixture, 4, nulls_fixture], [1, 2, 4, 2]])
     mi2 = MultiIndex.from_arrays([[3, nulls_fixture, 3], [1, 2, 4]])
     result = mi1.intersection(mi2)
-    expected = MultiIndex.from_arrays([[3.0, nulls_fixture], [1, 2]])
+    expected = MultiIndex.from_arrays([[3, nulls_fixture], [1, 2]])
     tm.assert_index_equal(result, expected)
 
 
@@ -631,4 +631,18 @@ def test_intersection_lexsort_depth(levels1, levels2, codes1, codes2, names):
     mi_int = mi1.intersection(mi2)
 
     with tm.assert_produces_warning(FutureWarning, match="MultiIndex.lexsort_depth"):
-        assert mi_int.lexsort_depth == 0
+        assert mi_int.lexsort_depth == 2
+
+
+@pytest.mark.parametrize("val", [pd.NA, 100])
+def test_intersection_keep_ea_dtypes(val, any_numeric_ea_dtype):
+    # GH#48604
+    midx = MultiIndex.from_arrays(
+        [Series([1, 2], dtype=any_numeric_ea_dtype), [2, 1]], names=["a", None]
+    )
+    midx2 = MultiIndex.from_arrays(
+        [Series([1, 2, val], dtype=any_numeric_ea_dtype), [1, 1, 3]]
+    )
+    result = midx.intersection(midx2)
+    expected = MultiIndex.from_arrays([Series([2], dtype=any_numeric_ea_dtype), [1]])
+    tm.assert_index_equal(result, expected)