From a158512a7bb37bf0b6b006ecbc0bc3eae550e240 Mon Sep 17 00:00:00 2001 From: Brock Date: Thu, 18 Feb 2021 10:24:41 -0800 Subject: [PATCH 1/2] REF: de-duplicate block_shape vs safe_reshape --- pandas/core/internals/__init__.py | 2 -- pandas/core/internals/blocks.py | 38 ++++++++----------------------- pandas/core/internals/managers.py | 6 ++--- 3 files changed, 13 insertions(+), 33 deletions(-) diff --git a/pandas/core/internals/__init__.py b/pandas/core/internals/__init__.py index ff4e186e147d7..132598e03d6c0 100644 --- a/pandas/core/internals/__init__.py +++ b/pandas/core/internals/__init__.py @@ -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 ( @@ -31,7 +30,6 @@ "FloatBlock", "ObjectBlock", "TimeDeltaBlock", - "safe_reshape", "make_block", "DataManager", "ArrayManager", diff --git a/pandas/core/internals/blocks.py b/pandas/core/internals/blocks.py index ae019c2f853a1..1d8471e76e10f 100644 --- a/pandas/core/internals/blocks.py +++ b/pandas/core/internals/blocks.py @@ -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) @@ -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 @@ -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)] @@ -2431,36 +2433,16 @@ 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) # type: ignore[attr-defined] 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 diff --git a/pandas/core/internals/managers.py b/pandas/core/internals/managers.py index b3f0466f236b6..4ad094f315c78 100644 --- a/pandas/core/internals/managers.py +++ b/pandas/core/internals/managers.py @@ -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, @@ -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 @@ -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)) From 19e8d5576e63c7596585fdc9990829ba367cc6c6 Mon Sep 17 00:00:00 2001 From: Brock Date: Thu, 18 Feb 2021 12:39:23 -0800 Subject: [PATCH 2/2] mypy fixup --- pandas/core/internals/blocks.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pandas/core/internals/blocks.py b/pandas/core/internals/blocks.py index 1d8471e76e10f..dee79efc538e3 100644 --- a/pandas/core/internals/blocks.py +++ b/pandas/core/internals/blocks.py @@ -2443,6 +2443,5 @@ def ensure_block_shape(values: ArrayLike, ndim: int = 1) -> ArrayLike: # block.shape is incorrect for "2D" ExtensionArrays # We can't, and don't need to, reshape. - # error: "ExtensionArray" has no attribute "reshape" - values = np.asarray(values).reshape(1, -1) # type: ignore[attr-defined] + values = np.asarray(values).reshape(1, -1) return values