Skip to content

Commit

Permalink
CoW: Set refs properly for non CoW mode (#54118)
Browse files Browse the repository at this point in the history
  • Loading branch information
phofl authored Jul 17, 2023
1 parent c539d1f commit 6dbc455
Show file tree
Hide file tree
Showing 3 changed files with 9 additions and 13 deletions.
12 changes: 4 additions & 8 deletions pandas/core/generic.py
Original file line number Diff line number Diff line change
Expand Up @@ -802,11 +802,7 @@ def swapaxes(self, axis1: Axis, axis2: Axis, copy: bool_t | None = None) -> Self

new_axes = [self._get_axis(mapping.get(k, k)) for k in range(self._AXIS_LEN)]
new_values = self._values.swapaxes(i, j) # type: ignore[union-attr]
if (
using_copy_on_write()
and self._mgr.is_single_block
and isinstance(self._mgr, BlockManager)
):
if self._mgr.is_single_block and isinstance(self._mgr, BlockManager):
# This should only get hit in case of having a single block, otherwise a
# copy is made, we don't have to set up references.
new_mgr = ndarray_to_mgr(
Expand All @@ -823,10 +819,10 @@ def swapaxes(self, axis1: Axis, axis2: Axis, copy: bool_t | None = None) -> Self
new_mgr.blocks[0].refs.add_reference(
new_mgr.blocks[0] # type: ignore[arg-type]
)
return self._constructor(new_mgr).__finalize__(self, method="swapaxes")
if not using_copy_on_write() and copy is not False:
new_mgr = new_mgr.copy(deep=True)

elif (copy or copy is None) and self._mgr.is_single_block:
new_values = new_values.copy()
return self._constructor(new_mgr).__finalize__(self, method="swapaxes")

return self._constructor(
new_values,
Expand Down
8 changes: 4 additions & 4 deletions pandas/core/internals/blocks.py
Original file line number Diff line number Diff line change
Expand Up @@ -486,7 +486,7 @@ def _downcast_2d(self, dtype, using_cow: bool = False) -> list[Block]:
"""
new_values = maybe_downcast_to_dtype(self.values, dtype=dtype)
new_values = maybe_coerce_values(new_values)
refs = self.refs if using_cow and new_values is self.values else None
refs = self.refs if new_values is self.values else None
return [self.make_block(new_values, refs=refs)]

@final
Expand Down Expand Up @@ -529,7 +529,7 @@ def convert(
refs = None
if copy and res_values is values:
res_values = values.copy()
elif res_values is values and using_cow:
elif res_values is values:
refs = self.refs

res_values = ensure_block_shape(res_values, self.ndim)
Expand Down Expand Up @@ -577,7 +577,7 @@ def astype(
new_values = maybe_coerce_values(new_values)

refs = None
if using_cow and astype_is_view(values.dtype, new_values.dtype):
if (using_cow or not copy) and astype_is_view(values.dtype, new_values.dtype):
refs = self.refs

newb = self.make_block(new_values, refs=refs)
Expand Down Expand Up @@ -914,7 +914,7 @@ def _replace_coerce(
nb = self.astype(np.dtype(object), copy=False, using_cow=using_cow)
if (nb is self or using_cow) and not inplace:
nb = nb.copy()
elif inplace and has_ref and nb.refs.has_reference():
elif inplace and has_ref and nb.refs.has_reference() and using_cow:
# no copy in astype and we had refs before
nb = nb.copy()
putmask_inplace(nb.values, mask, value)
Expand Down
2 changes: 1 addition & 1 deletion pandas/core/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -893,7 +893,7 @@ def view(self, dtype: Dtype | None = None) -> Series:
# implementation
res_values = self.array.view(dtype)
res_ser = self._constructor(res_values, index=self.index, copy=False)
if isinstance(res_ser._mgr, SingleBlockManager) and using_copy_on_write():
if isinstance(res_ser._mgr, SingleBlockManager):
blk = res_ser._mgr._block
blk.refs = cast("BlockValuesRefs", self._references)
blk.refs.add_reference(blk) # type: ignore[arg-type]
Expand Down

0 comments on commit 6dbc455

Please sign in to comment.