API: improve dtype in df.where with EA other

Author: jbrockmendel

doc/source/whatsnew/v3.0.0.rst

@@ -95,7 +95,7 @@ Other enhancements
 - Support passing a :class:`Iterable[Hashable]` input to :meth:`DataFrame.drop_duplicates` (:issue:`59237`)
 - Support reading Stata 102-format (Stata 1) dta files (:issue:`58978`)
 - Support reading Stata 110-format (Stata 7) dta files (:issue:`47176`)
-
+- Improve the resulting dtypes in :meth:`DataFrame.where` and :meth:`DataFrame.mask` with :class:`ExtensionDtype` ``other`` (:issue:`??`)
 .. ---------------------------------------------------------------------------
 .. _whatsnew_300.notable_bug_fixes:
 

pandas/core/generic.py

@@ -9788,14 +9788,42 @@ def _where(
                     raise InvalidIndexError
 
                 if other.ndim < self.ndim:
-                    # TODO(EA2D): avoid object-dtype cast in EA case GH#38729
                     other = other._values
-                    if axis == 0:
-                        other = np.reshape(other, (-1, 1))
-                    elif axis == 1:
-                        other = np.reshape(other, (1, -1))
-
-                    other = np.broadcast_to(other, self.shape)
+                    if isinstance(other, np.ndarray):
+                        # TODO(EA2D): could also do this for NDArrayBackedEA cases?
+                        if axis == 0:
+                            other = np.reshape(other, (-1, 1))
+                        elif axis == 1:
+                            other = np.reshape(other, (1, -1))
+
+                        other = np.broadcast_to(other, self.shape)
+                    else:
+                        # GH#38729 avoid lossy casting or object-casting
+                        if axis == 0:
+                            res_cols = [
+                                self.iloc[:, i]._where(
+                                    cond.iloc[:, i],
+                                    other,
+                                )
+                                for i in range(self.shape[1])
+                            ]
+                        elif axis == 1:
+                            # TODO: can we use a zero-copy alternative to "repeat"?
+                            res_cols = [
+                                self.iloc[:, i]._where(
+                                    cond.iloc[:, i],
+                                    other[i : i + 1].repeat(len(self)),
+                                )
+                                for i in range(self.shape[1])
+                            ]
+                        res = self._constructor(
+                            {i: res_cols[i] for i in range(len(res_cols))}
+                        )
+                        res.index = self.index
+                        res.columns = self.columns
+                        if inplace:
+                            return self._update_inplace(res)
+                        return res.__finalize__(self)
 
             # slice me out of the other
             else:

pandas/tests/frame/indexing/test_where.py

@@ -705,15 +705,23 @@ def test_where_ea_other(self):
         mask = np.ones(df.shape, dtype=bool)
         mask[1, :] = False
 
-        # TODO: ideally we would get Int64 instead of object
-        result = df.where(mask, ser, axis=0)
-        expected = DataFrame({"A": [1, np.nan, 3], "B": [4, np.nan, 6]})
-        tm.assert_frame_equal(result, expected)
+        result1 = df.where(mask, ser, axis=0)
+        expected1 = DataFrame({"A": [1, pd.NA, 3], "B": [4, pd.NA, 6]}, dtype="Int64")
+        tm.assert_frame_equal(result1, expected1)
 
         ser2 = Series(arr[:2], index=["A", "B"])
-        expected = DataFrame({"A": [1, 7, 3], "B": [4, np.nan, 6]})
-        result = df.where(mask, ser2, axis=1)
-        tm.assert_frame_equal(result, expected)
+        expected2 = DataFrame({"A": [1, 7, 3], "B": [4, pd.NA, 6]})
+        expected2["B"] = expected2["B"].astype("Int64")
+        result2 = df.where(mask, ser2, axis=1)
+        tm.assert_frame_equal(result2, expected2)
+
+        result3 = df.copy()
+        result3.mask(mask, ser, axis=0, inplace=True)
+        tm.assert_frame_equal(result3, expected1)
+
+        result4 = df.copy()
+        result4.mask(mask, ser2, axis=1, inplace=True)
+        tm.assert_frame_equal(result4, expected2)
 
     def test_where_interval_noop(self):
         # GH#44181