Skip to content

ENH: Optimize putmask implementation for CoW #51268

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.

Sign up for GitHub

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 9 commits into from
Feb 10, 2023
Merged

ENH: Optimize putmask implementation for CoW #51268

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
fdedea4
ENH: Optimize putmask implementation for CoW
phofl Feb 9, 2023
2ce28c6
Merge remote-tracking branch 'upstream/main' into cow_putmask
phofl Feb 9, 2023
94cd34a
Fix
phofl Feb 9, 2023
1411b6c
remove comment
phofl Feb 9, 2023
c4bd3a7
Refactor
phofl Feb 9, 2023
beb70b0
Fix test
phofl Feb 9, 2023
ce084ad
Merge remote-tracking branch 'upstream/main' into cow_putmask
phofl Feb 10, 2023
c5c3a2e
Address review
phofl Feb 10, 2023
01bf0b2
Fix mypy
phofl Feb 10, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions pandas/core/generic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9612,7 +9612,8 @@ def _where(
# align the cond to same shape as myself
cond = common.apply_if_callable(cond, self)
if isinstance(cond, NDFrame):
cond, _ = cond.align(self, join="right", broadcast_axis=1, copy=False)
# CoW: Make sure reference is not kept alive
cond = cond.align(self, join="right", broadcast_axis=1, copy=False)[0]
else:
if not hasattr(cond, "shape"):
cond = np.asanyarray(cond)
Expand Down Expand Up @@ -9648,14 +9649,15 @@ def _where(
# align with me
if other.ndim <= self.ndim:

_, other = self.align(
# CoW: Make sure reference is not kept alive
other = self.align(
other,
join="left",
axis=axis,
level=level,
fill_value=None,
copy=False,
)
)[1]

# if we are NOT aligned, raise as we cannot where index
if axis is None and not other._indexed_same(self):
Expand Down
41 changes: 35 additions & 6 deletions pandas/core/internals/blocks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -946,7 +946,7 @@ def _unstack(

# ---------------------------------------------------------------------

def setitem(self, indexer, value) -> Block:
def setitem(self, indexer, value, using_cow: bool = False) -> Block:
"""
Attempt self.values[indexer] = value, possibly creating a new array.

Expand All @@ -956,6 +956,8 @@ def setitem(self, indexer, value) -> Block:
The subset of self.values to set
value : object
The value being set
using_cow: bool, default False
Signaling if CoW is used.

Returns
-------
Expand Down Expand Up @@ -991,10 +993,17 @@ def setitem(self, indexer, value) -> Block:
# checking lib.is_scalar here fails on
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can't comment on it, but a few lines above we have:

        except LossySetitemError:
            # current dtype cannot store value, coerce to common dtype
            nb = self.coerce_to_target_dtype(value)
            return nb.setitem(indexer, value)

Is it theoretically possible that the dtype cannot store the value, but that the astype operation can still be done with a view? And should we in that case pass through using_cow=using_cow to the recursive setitem call?

I can't directly think of a case (maybe an EA that is backed by object array but can't hold the object you are setting?), but maybe the safest option is to just pass it.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hm I was thinking about it this as well, but I could not think of a case where we would get a LossySetitemError and then would not make a copy in the astype. Can add it if you prefer anyways

# test_iloc_setitem_custom_object
casted = setitem_datetimelike_compat(values, len(vi), casted)

if using_cow and self.refs.has_reference():
values = values.copy()
self = self.make_block_same_class(
values.T if values.ndim == 2 else values
)

values[indexer] = casted
return self

def putmask(self, mask, new) -> list[Block]:
def putmask(self, mask, new, using_cow: bool = False) -> list[Block]:
"""
putmask the data to the block; it is possible that we may create a
new dtype of block
Expand Down Expand Up @@ -1022,11 +1031,21 @@ def putmask(self, mask, new) -> list[Block]:
new = extract_array(new, extract_numpy=True)

if noop:
if using_cow:
return [self.copy(deep=False)]
return [self]
Comment on lines +1034 to 1036
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It also doesn't necessarily harm to always do return [self.copy(deep=False)] to avoid the extra if branch?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think so, but it could also mess up something I am not thinking about right now. I‘d like to keep as is if you don’t feel strongly about it. Could our caching mechanism be a problem (only theoretically, because we should clear the cache when getting here I think), but there might be edge cases we are missing right now

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's fine to leave as is for now. It will also be an easy clean-up once CoW is enforced.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yep exactly my thinking as well


try:
casted = np_can_hold_element(values.dtype, new)

if using_cow and self.refs.has_reference():
# Do this here to avoid copying twice
values = values.copy()
self = self.make_block_same_class(values)

putmask_without_repeat(values.T, mask, casted)
if using_cow:
return [self.copy(deep=False)]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't need to shallow copy if we already did a hard copy above?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, but the intermediate block will go out of scope, so making a shallow copy here will still only have a single ref in that case?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Exactly, it does not hurt, that’s why I wanted to avoid making the check more complex

return [self]
except LossySetitemError:

Expand All @@ -1038,7 +1057,7 @@ def putmask(self, mask, new) -> list[Block]:
return self.coerce_to_target_dtype(new).putmask(mask, new)
else:
indexer = mask.nonzero()[0]
nb = self.setitem(indexer, new[indexer])
nb = self.setitem(indexer, new[indexer], using_cow=using_cow)
return [nb]

else:
Expand All @@ -1053,7 +1072,7 @@ def putmask(self, mask, new) -> list[Block]:
n = new[:, i : i + 1]

submask = orig_mask[:, i : i + 1]
rbs = nb.putmask(submask, n)
rbs = nb.putmask(submask, n, using_cow=using_cow)
res_blocks.extend(rbs)
return res_blocks

Expand Down Expand Up @@ -1448,7 +1467,7 @@ class EABackedBlock(Block):

values: ExtensionArray

def setitem(self, indexer, value):
def setitem(self, indexer, value, using_cow: bool = False):
"""
Attempt self.values[indexer] = value, possibly creating a new array.

Expand All @@ -1461,6 +1480,8 @@ def setitem(self, indexer, value):
The subset of self.values to set
value : object
The value being set
using_cow: bool, default False
Signaling if CoW is used.

Returns
-------
Expand Down Expand Up @@ -1567,7 +1588,7 @@ def where(self, other, cond, _downcast: str | bool = "infer") -> list[Block]:
nb = self.make_block_same_class(res_values)
return [nb]

def putmask(self, mask, new) -> list[Block]:
def putmask(self, mask, new, using_cow: bool = False) -> list[Block]:
"""
See Block.putmask.__doc__
"""
Expand All @@ -1585,8 +1606,16 @@ def putmask(self, mask, new) -> list[Block]:
mask = self._maybe_squeeze_arg(mask)

if not mask.any():
if using_cow:
return [self.copy(deep=False)]
return [self]

if using_cow and self.refs.has_reference():
values = values.copy()
self = self.make_block_same_class( # type: ignore[assignment]
values.T if values.ndim == 2 else values
)

try:
# Caller is responsible for ensuring matching lengths
values._putmask(mask, new)
Expand Down
9 changes: 1 addition & 8 deletions pandas/core/internals/managers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -362,14 +362,6 @@ def setitem(self: T, indexer, value) -> T:
return self.apply("setitem", indexer=indexer, value=value)

def putmask(self, mask, new, align: bool = True):
if using_copy_on_write() and any(
not self._has_no_reference_block(i) for i in range(len(self.blocks))
):
# some reference -> copy full dataframe
# TODO(CoW) this could be optimized to only copy the blocks that would
# get modified
self = self.copy()

if align:
align_keys = ["new", "mask"]
else:
Expand All @@ -381,6 +373,7 @@ def putmask(self, mask, new, align: bool = True):
align_keys=align_keys,
mask=mask,
new=new,
using_cow=using_copy_on_write(),
)

def diff(self: T, n: int, axis: AxisInt) -> T:
Expand Down
50 changes: 48 additions & 2 deletions pandas/tests/copy_view/test_methods.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1237,8 +1237,9 @@ def test_replace(using_copy_on_write, replace_kwargs):
tm.assert_frame_equal(df, df_orig)


def test_putmask(using_copy_on_write):
df = DataFrame({"a": [1, 2], "b": 1, "c": 2})
@pytest.mark.parametrize("dtype", ["int64", "Int64"])
def test_putmask(using_copy_on_write, dtype):
df = DataFrame({"a": [1, 2], "b": 1, "c": 2}, dtype=dtype)
view = df[:]
df_orig = df.copy()
df[df == df] = 5
Expand All @@ -1252,6 +1253,51 @@ def test_putmask(using_copy_on_write):
assert view.iloc[0, 0] == 5


@pytest.mark.parametrize("dtype", ["int64", "Int64"])
def test_putmask_no_reference(using_copy_on_write, dtype):
df = DataFrame({"a": [1, 2], "b": 1, "c": 2}, dtype=dtype)
arr_a = get_array(df, "a")
df[df == df] = 5

if using_copy_on_write:
assert np.shares_memory(arr_a, get_array(df, "a"))


@pytest.mark.parametrize("dtype", ["float64", "Float64"])
def test_putmask_aligns_rhs_no_reference(using_copy_on_write, dtype):
df = DataFrame({"a": [1.5, 2], "b": 1.5}, dtype=dtype)
arr_a = get_array(df, "a")
df[df == df] = DataFrame({"a": [5.5, 5]})

if using_copy_on_write:
assert np.shares_memory(arr_a, get_array(df, "a"))


@pytest.mark.parametrize("val, exp", [(5.5, True), (5, False)])
def test_putmask_dont_copy_some_blocks(using_copy_on_write, val, exp):
df = DataFrame({"a": [1, 2], "b": 1, "c": 1.5})
view = df[:]
df_orig = df.copy()
indexer = DataFrame(
[[True, False, False], [True, False, False]], columns=list("abc")
)
df[indexer] = val

if using_copy_on_write:
assert not np.shares_memory(get_array(view, "a"), get_array(df, "a"))
# TODO(CoW): Could split blocks to avoid copying the whole block
assert np.shares_memory(get_array(view, "b"), get_array(df, "b")) is exp
assert np.shares_memory(get_array(view, "c"), get_array(df, "c"))
assert df._mgr._has_no_reference(1) is not exp
assert not df._mgr._has_no_reference(2)
tm.assert_frame_equal(view, df_orig)
elif val == 5:
# Without CoW the original will be modified, the other case upcasts, e.g. copy
assert np.shares_memory(get_array(view, "a"), get_array(df, "a"))
assert np.shares_memory(get_array(view, "c"), get_array(df, "c"))
assert view.iloc[0, 0] == 5


def test_asfreq_noop(using_copy_on_write):
df = DataFrame(
{"a": [0.0, None, 2.0, 3.0]},
Expand Down