diff --git a/pandas/core/internals/array_manager.py b/pandas/core/internals/array_manager.py index c3d17d3b400a5..76d247d9804ee 100644 --- a/pandas/core/internals/array_manager.py +++ b/pandas/core/internals/array_manager.py @@ -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) @@ -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 diff --git a/pandas/core/internals/blocks.py b/pandas/core/internals/blocks.py index c4c70851bd54a..5529f9f430375 100644 --- a/pandas/core/internals/blocks.py +++ b/pandas/core/internals/blocks.py @@ -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 @@ -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): + 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) diff --git a/pandas/core/internals/managers.py b/pandas/core/internals/managers.py index 05d65537c69ba..3652c3a18dbcc 100644 --- a/pandas/core/internals/managers.py +++ b/pandas/core/internals/managers.py @@ -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, @@ -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"] @@ -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) diff --git a/pandas/tests/internals/test_internals.py b/pandas/tests/internals/test_internals.py index c67ef9177ca96..7d4f3b62d183a 100644 --- a/pandas/tests/internals/test_internals.py +++ b/pandas/tests/internals/test_internals.py @@ -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 @@ -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",): diff --git a/pandas/tests/series/methods/test_isin.py b/pandas/tests/series/methods/test_isin.py index 320179c0a8b0a..898a769dfac48 100644 --- a/pandas/tests/series/methods/test_isin.py +++ b/pandas/tests/series/methods/test_isin.py @@ -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]])