PERF: improve efficiency of BaseMaskedArray.__setitem__

This somewhat deals with pandas-dev#44172, though that won't be fully resolved until 2D `ExtensionArray`s are supported (per the comments there).

Author: alexreg

pandas/core/apply.py

@@ -85,7 +85,7 @@ def frame_apply(
     args=None,
     kwargs=None,
 ) -> FrameApply:
-    """construct and return a row or column based frame apply object"""
+    """Construct and return a row- or column-based frame apply object."""
     axis = obj._get_axis_number(axis)
     klass: type[FrameApply]
     if axis == 0:
@@ -693,7 +693,7 @@ def dtypes(self) -> Series:
         return self.obj.dtypes
 
     def apply(self) -> DataFrame | Series:
-        """compute the results"""
+        """Compute the results."""
         # dispatch to agg
         if is_list_like(self.f):
             return self.apply_multiple()
@@ -1011,7 +1011,7 @@ def result_columns(self) -> Index:
     def wrap_results_for_axis(
         self, results: ResType, res_index: Index
     ) -> DataFrame | Series:
-        """return the results for the columns"""
+        """Return the results for the columns."""
         result: DataFrame | Series
 
         # we have requested to expand

pandas/core/arrays/boolean.py

@@ -44,6 +44,7 @@
     BaseMaskedArray,
     BaseMaskedDtype,
 )
+from pandas.core.indexers import check_array_indexer
 
 if TYPE_CHECKING:
     import pyarrow
@@ -364,6 +365,17 @@ def map_string(s):
 
     _HANDLED_TYPES = (np.ndarray, numbers.Number, bool, np.bool_)
 
+    def __setitem__(self, key, value):
+        if lib.is_bool(value):
+            key = check_array_indexer(self, key)
+            self._data[key] = value
+            self._mask[key] = False
+        elif value is libmissing.NA:
+            key = check_array_indexer(self, key)
+            self._mask[key] = True
+        else:
+            super().__setitem__(key, value)
+
     def _coerce_to_array(self, value) -> tuple[np.ndarray, np.ndarray]:
         return coerce_to_array(value)
 

pandas/core/arrays/floating.py

@@ -39,6 +39,7 @@
     NumericArray,
     NumericDtype,
 )
+from pandas.core.indexers import check_array_indexer
 from pandas.core.ops import invalid_comparison
 from pandas.core.tools.numeric import to_numeric
 
@@ -275,6 +276,17 @@ def _from_sequence_of_strings(
         scalars = to_numeric(strings, errors="raise")
         return cls._from_sequence(scalars, dtype=dtype, copy=copy)
 
+    def __setitem__(self, key, value):
+        if lib.is_float(value):
+            key = check_array_indexer(self, key)
+            self._data[key] = value
+            self._mask[key] = False
+        elif value is libmissing.NA:
+            key = check_array_indexer(self, key)
+            self._mask[key] = True
+        else:
+            super().__setitem__(key, value)
+
     def _coerce_to_array(self, value) -> tuple[np.ndarray, np.ndarray]:
         return coerce_to_array(value, dtype=self.dtype)
 

pandas/core/arrays/integer.py

@@ -46,6 +46,7 @@
     NumericArray,
     NumericDtype,
 )
+from pandas.core.indexers import check_array_indexer
 from pandas.core.ops import invalid_comparison
 from pandas.core.tools.numeric import to_numeric
 
@@ -342,6 +343,17 @@ def _from_sequence_of_strings(
         scalars = to_numeric(strings, errors="raise")
         return cls._from_sequence(scalars, dtype=dtype, copy=copy)
 
+    def __setitem__(self, key, value):
+        if lib.is_integer(value):
+            key = check_array_indexer(self, key)
+            self._data[key] = value
+            self._mask[key] = False
+        elif value is libmissing.NA:
+            key = check_array_indexer(self, key)
+            self._mask[key] = True
+        else:
+            super().__setitem__(key, value)
+
     def _coerce_to_array(self, value) -> tuple[np.ndarray, np.ndarray]:
         return coerce_to_array(value, dtype=self.dtype)
 

pandas/core/arrays/masked.py

@@ -77,7 +77,7 @@
 
 class BaseMaskedDtype(ExtensionDtype):
     """
-    Base class for dtypes for BasedMaskedArray subclasses.
+    Base class for dtypes for BaseMaskedArray subclasses.
     """
 
     name: str