Skip to content

Commit e0a5472

Browse files
committed
PERF: improve efficiency of BaseMaskedArray.__setitem__
This somewhat deals with #44172, though that won't be fully resolved until 2D `ExtensionArray`s are supported (per the comments there).
1 parent 7c00e0c commit e0a5472

File tree

2 files changed

+14
-15
lines changed

2 files changed

+14
-15
lines changed

pandas/core/apply.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ def frame_apply(
8585
args=None,
8686
kwargs=None,
8787
) -> FrameApply:
88-
"""construct and return a row or column based frame apply object"""
88+
"""Construct and return a row- or column-based frame apply object."""
8989
axis = obj._get_axis_number(axis)
9090
klass: type[FrameApply]
9191
if axis == 0:
@@ -693,7 +693,7 @@ def dtypes(self) -> Series:
693693
return self.obj.dtypes
694694

695695
def apply(self) -> DataFrame | Series:
696-
"""compute the results"""
696+
"""Compute the results."""
697697
# dispatch to agg
698698
if is_list_like(self.f):
699699
return self.apply_multiple()
@@ -1011,7 +1011,7 @@ def result_columns(self) -> Index:
10111011
def wrap_results_for_axis(
10121012
self, results: ResType, res_index: Index
10131013
) -> DataFrame | Series:
1014-
"""return the results for the columns"""
1014+
"""Return the results for the columns."""
10151015
result: DataFrame | Series
10161016

10171017
# we have requested to expand

pandas/core/arrays/masked.py

+11-12
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@
7777

7878
class BaseMaskedDtype(ExtensionDtype):
7979
"""
80-
Base class for dtypes for BasedMaskedArray subclasses.
80+
Base class for dtypes for BaseMaskedArray subclasses.
8181
"""
8282

8383
name: str
@@ -209,18 +209,17 @@ def _coerce_to_array(self, values) -> tuple[np.ndarray, np.ndarray]:
209209
raise AbstractMethodError(self)
210210

211211
def __setitem__(self, key, value) -> None:
212-
_is_scalar = is_scalar(value)
213-
if _is_scalar:
214-
value = [value]
215-
value, mask = self._coerce_to_array(value)
216-
217-
if _is_scalar:
218-
value = value[0]
219-
mask = mask[0]
220-
221212
key = check_array_indexer(self, key)
222-
self._data[key] = value
223-
self._mask[key] = mask
213+
if is_scalar(value):
214+
if isna(value):
215+
self._mask[key] = True
216+
else:
217+
self._data[key] = value
218+
self._mask[key] = False
219+
else:
220+
value, mask = self._coerce_to_array(value)
221+
self._data[key] = value
222+
self._mask[key] = mask
224223

225224
def __iter__(self):
226225
if self.ndim == 1:

0 commit comments

Comments
 (0)