Skip to content

Commit

Permalink
PERF: improve efficiency of BaseMaskedArray.__setitem__
Browse files Browse the repository at this point in the history
This somewhat deals with #44172, though that won't be fully resolved until 2D `ExtensionArray`s are supported (per the comments there).
  • Loading branch information
alexreg committed Oct 26, 2021
1 parent 7c00e0c commit e0a5472
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 15 deletions.
6 changes: 3 additions & 3 deletions pandas/core/apply.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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
Expand Down
23 changes: 11 additions & 12 deletions pandas/core/arrays/masked.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@

class BaseMaskedDtype(ExtensionDtype):
"""
Base class for dtypes for BasedMaskedArray subclasses.
Base class for dtypes for BaseMaskedArray subclasses.
"""

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

def __setitem__(self, key, value) -> None:
_is_scalar = is_scalar(value)
if _is_scalar:
value = [value]
value, mask = self._coerce_to_array(value)

if _is_scalar:
value = value[0]
mask = mask[0]

key = check_array_indexer(self, key)
self._data[key] = value
self._mask[key] = mask
if is_scalar(value):
if isna(value):
self._mask[key] = True
else:
self._data[key] = value
self._mask[key] = False
else:
value, mask = self._coerce_to_array(value)
self._data[key] = value
self._mask[key] = mask

def __iter__(self):
if self.ndim == 1:
Expand Down

0 comments on commit e0a5472

Please sign in to comment.