Skip to content

CLN: use ensure_block_shape #40792

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
Apr 7, 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
4 changes: 2 additions & 2 deletions pandas/core/internals/array_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -479,8 +479,7 @@ def apply_with_block(self: T, f, align_keys=None, swap_axis=True, **kwargs) -> T
arr = arr._data # type: ignore[attr-defined]

if self.ndim == 2:
if isinstance(arr, np.ndarray):
arr = np.atleast_2d(arr)
arr = ensure_block_shape(arr, 2)
block = new_block(arr, placement=slice(0, 1, 1), ndim=2)
else:
block = new_block(arr, placement=slice(0, len(self), 1), ndim=1)
Expand All @@ -490,6 +489,7 @@ def apply_with_block(self: T, f, align_keys=None, swap_axis=True, **kwargs) -> T
applied = applied[0]
arr = applied.values
if self.ndim == 2 and arr.ndim == 2:
# 2D for np.ndarray or DatetimeArray/TimedeltaArray
assert len(arr) == 1
# error: Invalid index type "Tuple[int, slice]" for
# "Union[ndarray, ExtensionArray]"; expected type
Expand Down
22 changes: 13 additions & 9 deletions pandas/core/internals/blocks.py
Original file line number Diff line number Diff line change
Expand Up @@ -925,15 +925,6 @@ def setitem(self, indexer, value):
# current dtype cannot store value, coerce to common dtype
return self.coerce_to_target_dtype(value).setitem(indexer, value)

if self.dtype.kind in ["m", "M"]:
arr = self.values
if self.ndim > 1:
# Dont transpose with ndim=1 bc we would fail to invalidate
# arr.freq
arr = arr.T
arr[indexer] = value
return self

# value must be storable at this moment
if is_extension_array_dtype(getattr(value, "dtype", None)):
# We need to be careful not to allow through strings that
Expand Down Expand Up @@ -1697,6 +1688,19 @@ def iget(self, key):
# TODO(EA2D): this can be removed if we ever have 2D EA
return self.values.reshape(self.shape)[key]

def setitem(self, indexer, value):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

typing when you can

if not self._can_hold_element(value):
# TODO: general case needs casting logic.
return self.astype(object).setitem(indexer, value)

values = self.values
if self.ndim > 1:
# Dont transpose with ndim=1 bc we would fail to invalidate
# arr.freq
values = values.T
values[indexer] = value
return self

def putmask(self, mask, new) -> list[Block]:
mask = extract_bool_array(mask)

Expand Down
12 changes: 9 additions & 3 deletions pandas/core/internals/managers.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@

from pandas.core.dtypes.cast import infer_dtype_from_scalar
from pandas.core.dtypes.common import (
ensure_int64,
ensure_platform_int,
is_dtype_equal,
is_extension_array_dtype,
is_list_like,
Expand Down Expand Up @@ -346,6 +346,13 @@ def unpickle_block(values, mgr_locs, ndim: int):
state = state[3]["0.14.1"]
self.axes = [ensure_index(ax) for ax in state["axes"]]
ndim = len(self.axes)

for blk in state["blocks"]:
vals = blk["values"]
# older versions may hold e.g. DatetimeIndex instead of DTA
vals = extract_array(vals, extract_numpy=True)
blk["values"] = ensure_block_shape(vals, ndim=ndim)

self.blocks = tuple(
unpickle_block(b["values"], b["mgr_locs"], ndim=ndim)
for b in state["blocks"]
Expand Down Expand Up @@ -1975,8 +1982,7 @@ def _preprocess_slice_or_indexer(
dtype = getattr(slice_or_indexer, "dtype", None)
raise TypeError(type(slice_or_indexer), dtype)

# TODO: np.intp?
indexer = ensure_int64(slice_or_indexer)
indexer = ensure_platform_int(slice_or_indexer)
if not allow_fill:
indexer = maybe_convert_indices(indexer, length)
return "fancy", indexer, len(indexer)
6 changes: 5 additions & 1 deletion pandas/tests/internals/test_internals.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,10 @@
SingleBlockManager,
make_block,
)
from pandas.core.internals.blocks import new_block
from pandas.core.internals.blocks import (
ensure_block_shape,
new_block,
)

# this file contains BlockManager specific tests
# TODO(ArrayManager) factor out interleave_dtype tests
Expand Down Expand Up @@ -138,6 +141,7 @@ def create_block(typestr, placement, item_shape=None, num_offset=0, maker=new_bl
tz = m.groups()[0]
assert num_items == 1, "must have only 1 num items for a tz-aware"
values = DatetimeIndex(np.arange(N) * 1e9, tz=tz)._data
values = ensure_block_shape(values, ndim=len(shape))
elif typestr in ("timedelta", "td", "m8[ns]"):
values = (mat * 1).astype("m8[ns]")
elif typestr in ("category",):
Expand Down
2 changes: 1 addition & 1 deletion pandas/tests/series/methods/test_isin.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ def test_isin_with_i8(self):
tm.assert_series_equal(result, expected)

# fails on dtype conversion in the first place
result = s.isin(s[0:2].values.astype("datetime64[D]"))
result = s.isin(np.asarray(s[0:2].values).astype("datetime64[D]"))
tm.assert_series_equal(result, expected)

result = s.isin([s[1]])
Expand Down