Skip to content

REF: share methods AM/BM #53859

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 1 commit into from
Jun 27, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
86 changes: 4 additions & 82 deletions pandas/core/internals/array_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
import itertools
from typing import (
TYPE_CHECKING,
Any,
Callable,
Hashable,
Literal,
Expand All @@ -16,10 +15,8 @@

from pandas._libs import (
NaT,
algos as libalgos,
lib,
)
from pandas.util._validators import validate_bool_kwarg

from pandas.core.dtypes.astype import (
astype_array,
Expand Down Expand Up @@ -261,10 +258,11 @@ def apply(
# expected "List[Union[ndarray, ExtensionArray]]"
return type(self)(result_arrays, new_axes) # type: ignore[arg-type]

def apply_with_block(
self, f, align_keys=None, swap_axis: bool = True, **kwargs
) -> Self:
def apply_with_block(self, f, align_keys=None, **kwargs) -> Self:
# switch axis to follow BlockManager logic
swap_axis = True
if f == "interpolate":
swap_axis = False
if swap_axis and "axis" in kwargs and self.ndim == 2:
kwargs["axis"] = 1 if kwargs["axis"] == 0 else 0

Expand Down Expand Up @@ -319,50 +317,13 @@ def apply_with_block(

return type(self)(result_arrays, self._axes)

def where(self, other, cond, align: bool) -> Self:
if align:
align_keys = ["other", "cond"]
else:
align_keys = ["cond"]
other = extract_array(other, extract_numpy=True)

return self.apply_with_block(
"where",
align_keys=align_keys,
other=other,
cond=cond,
)

def round(self, decimals: int, using_cow: bool = False) -> Self:
return self.apply_with_block("round", decimals=decimals, using_cow=using_cow)

def setitem(self, indexer, value) -> Self:
return self.apply_with_block("setitem", indexer=indexer, value=value)

def putmask(self, mask, new, align: bool = True) -> Self:
if align:
align_keys = ["new", "mask"]
else:
align_keys = ["mask"]
new = extract_array(new, extract_numpy=True)

return self.apply_with_block(
"putmask",
align_keys=align_keys,
mask=mask,
new=new,
)

def diff(self, n: int) -> Self:
assert self.ndim == 2 # caller ensures
return self.apply(algos.diff, n=n)

def pad_or_backfill(self, **kwargs) -> Self:
return self.apply_with_block("pad_or_backfill", swap_axis=False, **kwargs)

def interpolate(self, **kwargs) -> Self:
return self.apply_with_block("interpolate", swap_axis=False, **kwargs)

def shift(self, periods: int, axis: AxisInt, fill_value) -> Self:
if fill_value is lib.no_default:
fill_value = None
Expand All @@ -375,15 +336,6 @@ def shift(self, periods: int, axis: AxisInt, fill_value) -> Self:
"shift", periods=periods, axis=axis, fill_value=fill_value
)

def fillna(self, value, limit: int | None, inplace: bool, downcast) -> Self:
if limit is not None:
# Do this validation even if we go through one of the no-op paths
limit = libalgos.validate_limit(None, limit=limit)

return self.apply_with_block(
"fillna", value=value, limit=limit, inplace=inplace, downcast=downcast
)

def astype(self, dtype, copy: bool | None = False, errors: str = "raise") -> Self:
if copy is None:
copy = True
Expand All @@ -410,36 +362,6 @@ def _convert(arr):

return self.apply(_convert)

def replace_regex(self, **kwargs) -> Self:
return self.apply_with_block("_replace_regex", **kwargs)

def replace(self, to_replace, value, inplace: bool) -> Self:
inplace = validate_bool_kwarg(inplace, "inplace")
assert np.ndim(value) == 0, value
# TODO "replace" is right now implemented on the blocks, we should move
# it to general array algos so it can be reused here
return self.apply_with_block(
"replace", value=value, to_replace=to_replace, inplace=inplace
)

def replace_list(
self,
src_list: list[Any],
dest_list: list[Any],
inplace: bool = False,
regex: bool = False,
) -> Self:
"""do a list replace"""
inplace = validate_bool_kwarg(inplace, "inplace")

return self.apply_with_block(
"replace_list",
src_list=src_list,
dest_list=dest_list,
inplace=inplace,
regex=regex,
)

def to_native_types(self, **kwargs) -> Self:
return self.apply(to_native_types, **kwargs)

Expand Down
127 changes: 127 additions & 0 deletions pandas/core/internals/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,20 +6,29 @@

from typing import (
TYPE_CHECKING,
Any,
Literal,
final,
)

import numpy as np

from pandas._config import using_copy_on_write

from pandas._libs import (
algos as libalgos,
lib,
)
from pandas.errors import AbstractMethodError
from pandas.util._validators import validate_bool_kwarg

from pandas.core.dtypes.cast import (
find_common_type,
np_can_hold_element,
)

from pandas.core.base import PandasObject
from pandas.core.construction import extract_array
from pandas.core.indexes.api import (
Index,
default_index,
Expand Down Expand Up @@ -138,10 +147,128 @@ def apply(
) -> Self:
raise AbstractMethodError(self)

def apply_with_block(
self,
f,
align_keys: list[str] | None = None,
**kwargs,
) -> Self:
raise AbstractMethodError(self)

@final
def isna(self, func) -> Self:
return self.apply("apply", func=func)

@final
def fillna(self, value, limit: int | None, inplace: bool, downcast) -> Self:
if limit is not None:
# Do this validation even if we go through one of the no-op paths
limit = libalgos.validate_limit(None, limit=limit)

return self.apply_with_block(
"fillna",
value=value,
limit=limit,
inplace=inplace,
downcast=downcast,
using_cow=using_copy_on_write(),
)

@final
def where(self, other, cond, align: bool) -> Self:
if align:
align_keys = ["other", "cond"]
else:
align_keys = ["cond"]
other = extract_array(other, extract_numpy=True)

return self.apply_with_block(
"where",
align_keys=align_keys,
other=other,
cond=cond,
using_cow=using_copy_on_write(),
)

@final
def putmask(self, mask, new, align: bool = True) -> Self:
if align:
align_keys = ["new", "mask"]
else:
align_keys = ["mask"]
new = extract_array(new, extract_numpy=True)

return self.apply_with_block(
"putmask",
align_keys=align_keys,
mask=mask,
new=new,
using_cow=using_copy_on_write(),
)

@final
def round(self, decimals: int, using_cow: bool = False) -> Self:
return self.apply_with_block(
"round",
decimals=decimals,
using_cow=using_cow,
)

@final
def replace(self, to_replace, value, inplace: bool) -> Self:
inplace = validate_bool_kwarg(inplace, "inplace")
# NDFrame.replace ensures the not-is_list_likes here
assert not lib.is_list_like(to_replace)
assert not lib.is_list_like(value)
return self.apply_with_block(
"replace",
to_replace=to_replace,
value=value,
inplace=inplace,
using_cow=using_copy_on_write(),
)

@final
def replace_regex(self, **kwargs) -> Self:
return self.apply_with_block(
"_replace_regex", **kwargs, using_cow=using_copy_on_write()
)

@final
def replace_list(
self,
src_list: list[Any],
dest_list: list[Any],
inplace: bool = False,
regex: bool = False,
) -> Self:
"""do a list replace"""
inplace = validate_bool_kwarg(inplace, "inplace")

bm = self.apply_with_block(
"replace_list",
src_list=src_list,
dest_list=dest_list,
inplace=inplace,
regex=regex,
using_cow=using_copy_on_write(),
)
bm._consolidate_inplace()
return bm

def interpolate(self, inplace: bool, **kwargs) -> Self:
return self.apply_with_block(
"interpolate", inplace=inplace, **kwargs, using_cow=using_copy_on_write()
)

def pad_or_backfill(self, inplace: bool, **kwargs) -> Self:
return self.apply_with_block(
"pad_or_backfill",
inplace=inplace,
**kwargs,
using_cow=using_copy_on_write(),
)

# --------------------------------------------------------------------
# Consolidation: No-ops for all but BlockManager

Expand Down
Loading