diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2d8f6468aca83..aec19e27a33e2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -153,3 +153,9 @@ jobs: run: | source activate pandas-dev pytest pandas/tests/frame/methods --array-manager + + # indexing iset related (temporary since other tests don't pass yet) + pytest pandas/tests/frame/indexing/test_indexing.py::TestDataFrameIndexing::test_setitem_multi_index --array-manager + pytest pandas/tests/frame/indexing/test_setitem.py::TestDataFrameSetItem::test_setitem_listlike_indexer_duplicate_columns --array-manager + pytest pandas/tests/indexing/multiindex/test_setitem.py::TestMultiIndexSetItem::test_astype_assignment_with_dups --array-manager + pytest pandas/tests/indexing/multiindex/test_setitem.py::TestMultiIndexSetItem::test_frame_setitem_multi_column --array-manager diff --git a/pandas/core/frame.py b/pandas/core/frame.py index 63d238da12101..782539c2ba4b2 100644 --- a/pandas/core/frame.py +++ b/pandas/core/frame.py @@ -9723,12 +9723,3 @@ def _reindex_for_setitem(value: FrameOrSeriesUnion, index: Index) -> ArrayLike: "incompatible index of inserted column with frame index" ) from err return reindexed_value - - -def _maybe_atleast_2d(value): - # TODO(EA2D): not needed with 2D EAs - - if is_extension_array_dtype(value): - return value - - return np.atleast_2d(np.asarray(value)) diff --git a/pandas/core/internals/array_manager.py b/pandas/core/internals/array_manager.py index a8493e647f39a..083f32488acd4 100644 --- a/pandas/core/internals/array_manager.py +++ b/pandas/core/internals/array_manager.py @@ -659,24 +659,53 @@ def idelete(self, indexer): def iset(self, loc: Union[int, slice, np.ndarray], value): """ - Set new item in-place. Does not consolidate. Adds new Block if not - contained in the current set of items + Set new column(s). + + This changes the ArrayManager in-place, but replaces (an) existing + column(s), not changing column values in-place). + + Parameters + ---------- + loc : integer, slice or boolean mask + Positional location (already bounds checked) + value : array-like """ + # single column -> single integer index if lib.is_integer(loc): - # TODO normalize array -> this should in theory not be needed? + # TODO the extract array should in theory not be needed? value = extract_array(value, extract_numpy=True) + + # TODO can we avoid needing to unpack this here? That means converting + # DataFrame into 1D array when loc is an integer if isinstance(value, np.ndarray) and value.ndim == 2: + assert value.shape[1] == 1 value = value[0, :] assert isinstance(value, (np.ndarray, ExtensionArray)) - # value = np.asarray(value) - # assert isinstance(value, np.ndarray) + assert value.ndim == 1 assert len(value) == len(self._axes[0]) self.arrays[loc] = value return - # TODO - raise Exception + # multiple columns -> convert slice or array to integer indices + elif isinstance(loc, slice): + indices = range( + loc.start if loc.start is not None else 0, + loc.stop if loc.stop is not None else self.shape_proper[1], + loc.step if loc.step is not None else 1, + ) + else: + assert isinstance(loc, np.ndarray) + assert loc.dtype == "bool" + indices = np.nonzero(loc)[0] + + assert value.ndim == 2 + assert value.shape[0] == len(self._axes[0]) + + for value_idx, mgr_idx in enumerate(indices): + value_arr = value[:, value_idx] + self.arrays[mgr_idx] = value_arr + return def insert(self, loc: int, item: Hashable, value, allow_duplicates: bool = False): """ diff --git a/pandas/core/internals/managers.py b/pandas/core/internals/managers.py index 1c45b39ba990a..ccac2696b34c5 100644 --- a/pandas/core/internals/managers.py +++ b/pandas/core/internals/managers.py @@ -1140,8 +1140,7 @@ def insert(self, loc: int, item: Hashable, value, allow_duplicates: bool = False if value.ndim == 2: value = value.T - - if value.ndim == self.ndim - 1 and not is_extension_array_dtype(value.dtype): + 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)