Skip to content

Commit 6c75f4f

Browse files
authored
REF: helpers to de-duplicate CoW checks (#53882)
* REF: refactor out Block helpers for CoW * REF: de-duplicate helpers * De-duplicate helper, revert _replace_coerce edits
1 parent 5788685 commit 6c75f4f

File tree

1 file changed

+43
-41
lines changed

1 file changed

+43
-41
lines changed

pandas/core/internals/blocks.py

+43-41
Original file line numberDiff line numberDiff line change
@@ -571,6 +571,43 @@ def copy(self, deep: bool = True) -> Self:
571571
refs = self.refs
572572
return type(self)(values, placement=self._mgr_locs, ndim=self.ndim, refs=refs)
573573

574+
# ---------------------------------------------------------------------
575+
# Copy-on-Write Helpers
576+
577+
@final
578+
def _maybe_copy(self, using_cow: bool, inplace: bool) -> Self:
579+
if using_cow and inplace:
580+
deep = self.refs.has_reference()
581+
blk = self.copy(deep=deep)
582+
else:
583+
blk = self if inplace else self.copy()
584+
return blk
585+
586+
@final
587+
def _get_values_and_refs(self, using_cow, inplace):
588+
if using_cow:
589+
if inplace and not self.refs.has_reference():
590+
refs = self.refs
591+
new_values = self.values
592+
else:
593+
refs = None
594+
new_values = self.values.copy()
595+
else:
596+
refs = None
597+
new_values = self.values if inplace else self.values.copy()
598+
return new_values, refs
599+
600+
@final
601+
def _get_refs_and_copy(self, using_cow: bool, inplace: bool):
602+
refs = None
603+
arr_inplace = inplace
604+
if inplace:
605+
if using_cow and self.refs.has_reference():
606+
arr_inplace = False
607+
else:
608+
refs = self.refs
609+
return arr_inplace, refs
610+
574611
# ---------------------------------------------------------------------
575612
# Replace
576613

@@ -597,12 +634,7 @@ def replace(
597634
if isinstance(values, Categorical):
598635
# TODO: avoid special-casing
599636
# GH49404
600-
if using_cow and (self.refs.has_reference() or not inplace):
601-
blk = self.copy()
602-
elif using_cow:
603-
blk = self.copy(deep=False)
604-
else:
605-
blk = self if inplace else self.copy()
637+
blk = self._maybe_copy(using_cow, inplace)
606638
values = cast(Categorical, blk.values)
607639
values._replace(to_replace=to_replace, value=value, inplace=True)
608640
return [blk]
@@ -630,13 +662,7 @@ def replace(
630662
elif self._can_hold_element(value):
631663
# TODO(CoW): Maybe split here as well into columns where mask has True
632664
# and rest?
633-
if using_cow:
634-
if inplace:
635-
blk = self.copy(deep=self.refs.has_reference())
636-
else:
637-
blk = self.copy()
638-
else:
639-
blk = self if inplace else self.copy()
665+
blk = self._maybe_copy(using_cow, inplace)
640666
putmask_inplace(blk.values, mask, value)
641667
if not (self.is_object and value is None):
642668
# if the user *explicitly* gave None, we keep None, otherwise
@@ -712,16 +738,7 @@ def _replace_regex(
712738

713739
rx = re.compile(to_replace)
714740

715-
if using_cow:
716-
if inplace and not self.refs.has_reference():
717-
refs = self.refs
718-
new_values = self.values
719-
else:
720-
refs = None
721-
new_values = self.values.copy()
722-
else:
723-
refs = None
724-
new_values = self.values if inplace else self.values.copy()
741+
new_values, refs = self._get_values_and_refs(using_cow, inplace)
725742

726743
replace_regex(new_values, rx, value, mask)
727744

@@ -745,10 +762,7 @@ def replace_list(
745762
if isinstance(values, Categorical):
746763
# TODO: avoid special-casing
747764
# GH49404
748-
if using_cow and inplace:
749-
blk = self.copy(deep=self.refs.has_reference())
750-
else:
751-
blk = self if inplace else self.copy()
765+
blk = self._maybe_copy(using_cow, inplace)
752766
values = cast(Categorical, blk.values)
753767
values._replace(to_replace=src_list, value=dest_list, inplace=True)
754768
return [blk]
@@ -1421,13 +1435,7 @@ def interpolate(
14211435
**kwargs,
14221436
)
14231437

1424-
refs = None
1425-
arr_inplace = inplace
1426-
if inplace:
1427-
if using_cow and self.refs.has_reference():
1428-
arr_inplace = False
1429-
else:
1430-
refs = self.refs
1438+
arr_inplace, refs = self._get_refs_and_copy(using_cow, inplace)
14311439

14321440
# Dispatch to the PandasArray method.
14331441
# We know self.array_values is a PandasArray bc EABlock overrides
@@ -2268,13 +2276,7 @@ def interpolate(
22682276
# "Literal['linear']") [comparison-overlap]
22692277
if method == "linear": # type: ignore[comparison-overlap]
22702278
# TODO: GH#50950 implement for arbitrary EAs
2271-
refs = None
2272-
arr_inplace = inplace
2273-
if using_cow:
2274-
if inplace and not self.refs.has_reference():
2275-
refs = self.refs
2276-
else:
2277-
arr_inplace = False
2279+
arr_inplace, refs = self._get_refs_and_copy(using_cow, inplace)
22782280

22792281
new_values = self.values.interpolate(
22802282
method=method,

0 commit comments

Comments
 (0)