BUG: MultiIndex.unique losing ea dtype (pandas-dev#48335)

Author: phofl

asv_bench/benchmarks/multiindex_object.py

@@ -3,9 +3,11 @@
 import numpy as np
 
 from pandas import (
+    NA,
     DataFrame,
     MultiIndex,
     RangeIndex,
+    Series,
     date_range,
 )
 
@@ -255,4 +257,31 @@ def time_operation(self, index_structure, dtype, method):
         getattr(self.left, method)(self.right)
 
 
+class Unique:
+    params = [
+        (("Int64", NA), ("int64", 0)),
+    ]
+    param_names = ["dtype_val"]
+
+    def setup(self, dtype_val):
+        level = Series(
+            [1, 2, dtype_val[1], dtype_val[1]] + list(range(1_000_000)),
+            dtype=dtype_val[0],
+        )
+        self.midx = MultiIndex.from_arrays([level, level])
+
+        level_dups = Series(
+            [1, 2, dtype_val[1], dtype_val[1]] + list(range(500_000)) * 2,
+            dtype=dtype_val[0],
+        )
+
+        self.midx_dups = MultiIndex.from_arrays([level_dups, level_dups])
+
+    def time_unique(self, dtype_val):
+        self.midx.unique()
+
+    def time_unique_dups(self, dtype_val):
+        self.midx_dups.unique()
+
+
 from .pandas_vb_common import setup  # noqa: F401 isort:skip

doc/source/whatsnew/v1.6.0.rst

@@ -100,7 +100,7 @@ Deprecations
 
 Performance improvements
 ~~~~~~~~~~~~~~~~~~~~~~~~
--
+- Performance improvement for :meth:`MultiIndex.unique` (:issue:`48335`)
 -
 
 .. ---------------------------------------------------------------------------
@@ -161,6 +161,7 @@ Missing
 
 MultiIndex
 ^^^^^^^^^^
+- Bug in :meth:`MultiIndex.unique` losing extension array dtype (:issue:`48335`)
 - Bug in :meth:`MultiIndex.append` not checking names for equality (:issue:`48288`)
 -
 

pandas/core/indexes/multi.py

@@ -1728,7 +1728,7 @@ def get_level_values(self, level):
     def unique(self, level=None):
 
         if level is None:
-            return super().unique()
+            return self.drop_duplicates()
         else:
             level = self._get_level_number(level)
             return self._get_level_values(level=level, unique=True)

pandas/tests/indexes/multi/test_duplicates.py

@@ -6,6 +6,7 @@
 from pandas._libs import hashtable
 
 from pandas import (
+    NA,
     DatetimeIndex,
     MultiIndex,
     Series,
@@ -337,3 +338,16 @@ def test_multi_drop_duplicates_pos_args_deprecation():
         result = idx.drop_duplicates("last")
     expected = MultiIndex.from_arrays([[2, 3, 1], [2, 3, 1]])
     tm.assert_index_equal(expected, result)
+
+
+def test_midx_unique_ea_dtype():
+    # GH#48335
+    vals_a = Series([1, 2, NA, NA], dtype="Int64")
+    vals_b = np.array([1, 2, 3, 3])
+    midx = MultiIndex.from_arrays([vals_a, vals_b], names=["a", "b"])
+    result = midx.unique()
+
+    exp_vals_a = Series([1, 2, NA], dtype="Int64")
+    exp_vals_b = np.array([1, 2, 3])
+    expected = MultiIndex.from_arrays([exp_vals_a, exp_vals_b], names=["a", "b"])
+    tm.assert_index_equal(result, expected)