Skip to content

REF: de-duplicate block_shape vs safe_reshape #39892

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 2 commits into from
Feb 19, 2021
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
2 changes: 0 additions & 2 deletions pandas/core/internals/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@
ObjectBlock,
TimeDeltaBlock,
make_block,
safe_reshape,
)
from pandas.core.internals.concat import concatenate_block_managers
from pandas.core.internals.managers import (
Expand All @@ -31,7 +30,6 @@
"FloatBlock",
"ObjectBlock",
"TimeDeltaBlock",
"safe_reshape",
"make_block",
"DataManager",
"ArrayManager",
Expand Down
39 changes: 10 additions & 29 deletions pandas/core/internals/blocks.py
Original file line number Diff line number Diff line change
Expand Up @@ -308,7 +308,7 @@ def make_block(self, values, placement=None) -> Block:
if placement is None:
placement = self.mgr_locs
if self.is_extension:
values = _block_shape(values, ndim=self.ndim)
values = ensure_block_shape(values, ndim=self.ndim)

return make_block(values, placement=placement, ndim=self.ndim)

Expand Down Expand Up @@ -533,7 +533,7 @@ def make_a_block(nv, ref_loc):
else:
# Put back the dimension that was taken from it and make
# a block out of the result.
nv = _block_shape(nv, ndim=self.ndim)
nv = ensure_block_shape(nv, ndim=self.ndim)
block = self.make_block(values=nv, placement=ref_loc)
return block

Expand Down Expand Up @@ -1569,7 +1569,9 @@ def putmask(self, mask, new) -> List[Block]:
if isinstance(new, (np.ndarray, ExtensionArray)) and len(new) == len(mask):
new = new[mask]

mask = safe_reshape(mask, new_values.shape)
if mask.ndim == new_values.ndim + 1:
# TODO(EA2D): unnecessary with 2D EAs
mask = mask.reshape(new_values.shape)

new_values[mask] = new
return [self.make_block(values=new_values)]
Expand Down Expand Up @@ -2431,36 +2433,15 @@ def extend_blocks(result, blocks=None) -> List[Block]:
return blocks


def _block_shape(values: ArrayLike, ndim: int = 1) -> ArrayLike:
""" guarantee the shape of the values to be at least 1 d """
def ensure_block_shape(values: ArrayLike, ndim: int = 1) -> ArrayLike:
"""
Reshape if possible to have values.ndim == ndim.
"""
if values.ndim < ndim:
shape = values.shape
if not is_extension_array_dtype(values.dtype):
# TODO(EA2D): https://github.com/pandas-dev/pandas/issues/23023
# block.shape is incorrect for "2D" ExtensionArrays
# We can't, and don't need to, reshape.

# error: "ExtensionArray" has no attribute "reshape"
values = values.reshape(tuple((1,) + shape)) # type: ignore[attr-defined]
values = np.asarray(values).reshape(1, -1)
return values


def safe_reshape(arr: ArrayLike, new_shape: Shape) -> ArrayLike:
"""
Reshape `arr` to have shape `new_shape`, unless it is an ExtensionArray,
in which case it will be returned unchanged (see gh-13012).

Parameters
----------
arr : np.ndarray or ExtensionArray
new_shape : Tuple[int]

Returns
-------
np.ndarray or ExtensionArray
"""
if not is_extension_array_dtype(arr.dtype):
# Note: this will include TimedeltaArray and tz-naive DatetimeArray
# TODO(EA2D): special case will be unnecessary with 2D EAs
arr = np.asarray(arr).reshape(new_shape)
return arr
6 changes: 3 additions & 3 deletions pandas/core/internals/managers.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,10 +69,10 @@
DatetimeTZBlock,
ExtensionBlock,
ObjectValuesExtensionBlock,
ensure_block_shape,
extend_blocks,
get_block_type,
make_block,
safe_reshape,
)
from pandas.core.internals.ops import (
blockwise_all,
Expand Down Expand Up @@ -1042,7 +1042,7 @@ def value_getitem(placement):
value = value.T

if value.ndim == self.ndim - 1:
value = safe_reshape(value, (1,) + value.shape)
value = ensure_block_shape(value, ndim=2)

def value_getitem(placement):
return value
Expand Down Expand Up @@ -1167,7 +1167,7 @@ def insert(self, loc: int, item: Hashable, value, allow_duplicates: bool = False
value = value.T
elif value.ndim == self.ndim - 1 and not is_extension_array_dtype(value.dtype):
# TODO(EA2D): special case not needed with 2D EAs
value = safe_reshape(value, (1,) + value.shape)
value = ensure_block_shape(value, ndim=2)

block = make_block(values=value, ndim=self.ndim, placement=slice(loc, loc + 1))

Expand Down