BUG: MultiIndex.putmask losing ea dtype (#49847)

* BUG: MultiIndex.putmask losing ea dtype

* Fix typing

* Add asv

* Simplify and add whatsnew

Author: phofl

asv_bench/benchmarks/multiindex_object.py

@@ -379,4 +379,26 @@ def time_isin_large(self, dtype):
         self.midx.isin(self.values_large)
 
 
+class Putmask:
+    def setup(self):
+        N = 10**5
+        level1 = range(1_000)
+
+        level2 = date_range(start="1/1/2000", periods=N // 1000)
+        self.midx = MultiIndex.from_product([level1, level2])
+
+        level1 = range(1_000, 2_000)
+        self.midx_values = MultiIndex.from_product([level1, level2])
+
+        level2 = date_range(start="1/1/2010", periods=N // 1000)
+        self.midx_values_different = MultiIndex.from_product([level1, level2])
+        self.mask = np.array([True, False] * (N // 2))
+
+    def time_putmask(self):
+        self.midx.putmask(self.mask, self.midx_values)
+
+    def time_putmask_all_different(self):
+        self.midx.putmask(self.mask, self.midx_values_different)
+
+
 from .pandas_vb_common import setup  # noqa: F401 isort:skip

doc/source/whatsnew/v2.0.0.rst

@@ -588,6 +588,7 @@ Performance improvements
 - Performance improvement in :class:`MultiIndex` set operations with sort=None (:issue:`49010`)
 - Performance improvement in :meth:`.DataFrameGroupBy.mean`, :meth:`.SeriesGroupBy.mean`, :meth:`.DataFrameGroupBy.var`, and :meth:`.SeriesGroupBy.var` for extension array dtypes (:issue:`37493`)
 - Performance improvement in :meth:`MultiIndex.isin` when ``level=None`` (:issue:`48622`, :issue:`49577`)
+- Performance improvement in :meth:`MultiIndex.putmask` (:issue:`49830`)
 - Performance improvement in :meth:`Index.union` and :meth:`MultiIndex.union` when index contains duplicates (:issue:`48900`)
 - Performance improvement in :meth:`Series.fillna` for pyarrow-backed dtypes (:issue:`49722`)
 - Performance improvement for :meth:`Series.value_counts` with nullable dtype (:issue:`48338`)
@@ -703,6 +704,7 @@ MultiIndex
 - Bug in :meth:`MultiIndex.union` not sorting when sort=None and index contains missing values (:issue:`49010`)
 - Bug in :meth:`MultiIndex.append` not checking names for equality (:issue:`48288`)
 - Bug in :meth:`MultiIndex.symmetric_difference` losing extension array (:issue:`48607`)
+- Bug in :meth:`MultiIndex.putmask` losing extension array (:issue:`49830`)
 - Bug in :meth:`MultiIndex.value_counts` returning a :class:`Series` indexed by flat index of tuples instead of a :class:`MultiIndex` (:issue:`49558`)
 -
 

pandas/core/array_algos/putmask.py

@@ -3,7 +3,10 @@
 """
 from __future__ import annotations
 
-from typing import Any
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
 
 import numpy as np
 
@@ -19,6 +22,9 @@
 
 from pandas.core.arrays import ExtensionArray
 
+if TYPE_CHECKING:
+    from pandas import MultiIndex
+
 
 def putmask_inplace(values: ArrayLike, mask: npt.NDArray[np.bool_], value: Any) -> None:
     """
@@ -96,7 +102,7 @@ def putmask_without_repeat(
 
 
 def validate_putmask(
-    values: ArrayLike, mask: np.ndarray
+    values: ArrayLike | MultiIndex, mask: np.ndarray
 ) -> tuple[npt.NDArray[np.bool_], bool]:
     """
     Validate mask and check if this putmask operation is a no-op.

pandas/core/indexes/base.py

@@ -5144,7 +5144,6 @@ def _concat(self, to_concat: list[Index], name: Hashable) -> Index:
 
         return Index._with_infer(result, name=name)
 
-    @final
     def putmask(self, mask, value) -> Index:
         """
         Return a new Index of the values set with the mask.

pandas/core/indexes/multi.py

@@ -79,6 +79,7 @@
 )
 
 import pandas.core.algorithms as algos
+from pandas.core.array_algos.putmask import validate_putmask
 from pandas.core.arrays import Categorical
 from pandas.core.arrays.categorical import factorize_from_iterables
 import pandas.core.common as com
@@ -3660,6 +3661,46 @@ def _validate_fill_value(self, item):
             raise ValueError("Item must have length equal to number of levels.")
         return item
 
+    def putmask(self, mask, value: MultiIndex) -> MultiIndex:
+        """
+        Return a new MultiIndex of the values set with the mask.
+
+        Parameters
+        ----------
+        mask : array like
+        value : MultiIndex
+            Must either be the same length as self or length one
+
+        Returns
+        -------
+        MultiIndex
+        """
+        mask, noop = validate_putmask(self, mask)
+        if noop:
+            return self.copy()
+
+        if len(mask) == len(value):
+            subset = value[mask].remove_unused_levels()
+        else:
+            subset = value.remove_unused_levels()
+
+        new_levels = []
+        new_codes = []
+
+        for i, (value_level, level, level_codes) in enumerate(
+            zip(subset.levels, self.levels, self.codes)
+        ):
+            new_level = level.union(value_level, sort=False)
+            value_codes = new_level.get_indexer_for(subset.get_level_values(i))
+            new_code = ensure_int64(level_codes)
+            new_code[mask] = value_codes
+            new_levels.append(new_level)
+            new_codes.append(new_code)
+
+        return MultiIndex(
+            levels=new_levels, codes=new_codes, names=self.names, verify_integrity=False
+        )
+
     def insert(self, loc: int, item) -> MultiIndex:
         """
         Make new MultiIndex inserting new item at location

pandas/tests/indexes/multi/test_indexing.py

@@ -162,6 +162,34 @@ def test_putmask_multiindex_other(self):
         expected = MultiIndex.from_tuples([right[0], right[1], left[2]])
         tm.assert_index_equal(result, expected)
 
+    def test_putmask_keep_dtype(self, any_numeric_ea_dtype):
+        # GH#49830
+        midx = MultiIndex.from_arrays(
+            [pd.Series([1, 2, 3], dtype=any_numeric_ea_dtype), [10, 11, 12]]
+        )
+        midx2 = MultiIndex.from_arrays(
+            [pd.Series([5, 6, 7], dtype=any_numeric_ea_dtype), [-1, -2, -3]]
+        )
+        result = midx.putmask([True, False, False], midx2)
+        expected = MultiIndex.from_arrays(
+            [pd.Series([5, 2, 3], dtype=any_numeric_ea_dtype), [-1, 11, 12]]
+        )
+        tm.assert_index_equal(result, expected)
+
+    def test_putmask_keep_dtype_shorter_value(self, any_numeric_ea_dtype):
+        # GH#49830
+        midx = MultiIndex.from_arrays(
+            [pd.Series([1, 2, 3], dtype=any_numeric_ea_dtype), [10, 11, 12]]
+        )
+        midx2 = MultiIndex.from_arrays(
+            [pd.Series([5], dtype=any_numeric_ea_dtype), [-1]]
+        )
+        result = midx.putmask([True, False, False], midx2)
+        expected = MultiIndex.from_arrays(
+            [pd.Series([5, 2, 3], dtype=any_numeric_ea_dtype), [-1, 11, 12]]
+        )
+        tm.assert_index_equal(result, expected)
+
 
 class TestGetIndexer:
     def test_get_indexer(self):