Skip to content

PERF: Use generator expression for Blocks.replace_list #50778

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
Feb 28, 2023
30 changes: 30 additions & 0 deletions asv_bench/benchmarks/series_methods.py
Original file line number Diff line number Diff line change
Expand Up @@ -401,4 +401,34 @@ def time_to_numpy_copy(self):
self.ser.to_numpy(copy=True)


class Replace:

param_names = ["num_to_replace"]
params = [100, 1000]

def setup(self, num_to_replace):
N = 1_000_000
self.arr = np.random.randn(N)
self.arr1 = self.arr.copy()
np.random.shuffle(self.arr1)
self.ser = Series(self.arr)

self.to_replace_list = np.random.choice(self.arr, num_to_replace)
self.values_list = np.random.choice(self.arr1, num_to_replace)

self.replace_dict = dict(zip(self.to_replace_list, self.values_list))

def time_replace_dict(self, num_to_replace):
self.ser.replace(self.replace_dict)

def peakmem_replace_dict(self, num_to_replace):
self.ser.replace(self.replace_dict)

def time_replace_list(self, num_to_replace):
self.ser.replace(self.to_replace_list, self.values_list)

def peakmem_replace_list(self, num_to_replace):
self.ser.replace(self.to_replace_list, self.values_list)


from .pandas_vb_common import setup # noqa: F401 isort:skip
34 changes: 20 additions & 14 deletions pandas/core/internals/blocks.py
Original file line number Diff line number Diff line change
Expand Up @@ -668,22 +668,28 @@ def replace_list(
if is_string_dtype(values.dtype):
# Calculate the mask once, prior to the call of comp
# in order to avoid repeating the same computations
mask = ~isna(values)
masks = [
compare_or_regex_search(values, s[0], regex=regex, mask=mask)
na_mask = ~isna(values)
masks: Iterable[npt.NDArray[np.bool_]] = (
extract_bool_array(
cast(
ArrayLike,
compare_or_regex_search(
values, s[0], regex=regex, mask=na_mask
),
)
)
for s in pairs
]
)
else:
# GH#38086 faster if we know we dont need to check for regex
masks = [missing.mask_missing(values, s[0]) for s in pairs]

# error: Argument 1 to "extract_bool_array" has incompatible type
# "Union[ExtensionArray, ndarray, bool]"; expected "Union[ExtensionArray,
# ndarray]"
masks = [extract_bool_array(x) for x in masks] # type: ignore[arg-type]
masks = (missing.mask_missing(values, s[0]) for s in pairs)
# Materialize if inplace = True, since the masks can change
# as we replace
if inplace:
masks = list(masks)

rb = [self if inplace else self.copy()]
for i, (src, dest) in enumerate(pairs):
for i, ((src, dest), mask) in enumerate(zip(pairs, masks)):
convert = i == src_len # only convert once at the end
new_rb: list[Block] = []

Expand All @@ -692,9 +698,9 @@ def replace_list(
# where to index into the mask
for blk_num, blk in enumerate(rb):
if len(rb) == 1:
m = masks[i]
m = mask
else:
mib = masks[i]
mib = mask
assert not isinstance(mib, bool)
m = mib[blk_num : blk_num + 1]

Expand All @@ -704,7 +710,7 @@ def replace_list(
result = blk._replace_coerce(
to_replace=src,
value=dest,
mask=m, # type: ignore[arg-type]
mask=m,
inplace=inplace,
regex=regex,
)
Expand Down
12 changes: 12 additions & 0 deletions pandas/tests/series/methods/test_replace.py
Original file line number Diff line number Diff line change
Expand Up @@ -298,6 +298,18 @@ def test_replace2(self):
assert (ser[6:10] == -1).all()
assert (ser[20:30] == -1).all()

@pytest.mark.parametrize("inplace", [True, False])
def test_replace_cascade(self, inplace):
# Test that replaced values are not replaced again
ser = pd.Series([1, 2, 3])
expected = pd.Series([2, 3, 4])

res = ser.replace([1, 2, 3], [2, 3, 4], inplace=inplace)
if inplace:
tm.assert_series_equal(ser, expected)
else:
tm.assert_series_equal(res, expected)

def test_replace_with_dictlike_and_string_dtype(self, nullable_string_dtype):
# GH 32621, GH#44940
ser = pd.Series(["one", "two", np.nan], dtype=nullable_string_dtype)
Expand Down