diff --git a/doc/source/whatsnew/v2.1.0.rst b/doc/source/whatsnew/v2.1.0.rst index 11533647ca124..36a24c93600b8 100644 --- a/doc/source/whatsnew/v2.1.0.rst +++ b/doc/source/whatsnew/v2.1.0.rst @@ -432,6 +432,8 @@ Groupby/resample/rolling Reshaping ^^^^^^^^^ - Bug in :func:`crosstab` when ``dropna=False`` would not keep ``np.nan`` in the result (:issue:`10772`) +- Bug in :meth:`BlockManager.insert` allows assignments for multiple-dimension ndarray into one column (:issue:`53366`) +- Bug in :meth:`DataFrame.__setitem__` allows assignments for multiple-dimension ndarray into one column (:issue:`51925`) - Bug in :meth:`DataFrame.agg` and :meth:`Series.agg` on non-unique columns would return incorrect type when dist-like argument passed in (:issue:`51099`) - Bug in :meth:`DataFrame.idxmin` and :meth:`DataFrame.idxmax`, where the axis dtype would be lost for empty frames (:issue:`53265`) - Bug in :meth:`DataFrame.merge` not merging correctly when having ``MultiIndex`` with single level (:issue:`52331`) diff --git a/pandas/core/frame.py b/pandas/core/frame.py index 8afb3ee96ba94..59282f2e8fe12 100644 --- a/pandas/core/frame.py +++ b/pandas/core/frame.py @@ -3966,6 +3966,24 @@ def __setitem__(self, key, value): ): # Column to set is duplicated self._setitem_array([key], value) + elif ( + isinstance(value, np.ndarray) and value.ndim > 1 and self.columns.is_unique + ): + # TODO: a check for MultiIndex should be added + if isinstance(self.columns, MultiIndex): + self._set_item(key, value) + return + # squeeze 2d ndarray to 1d if possible + # this keeps the backward compatibility + if np.ndim(value) == 2 and (1 in np.shape(value)): + # if value is np.matrix, convert to np.ndarray + value = np.asarray(value).flatten() + self._set_item(key, value) + else: + # avoid assign non 1d array to column + raise ValueError( + f"Expected a 1D array, got an array with shape {value.shape}" + ) else: # set column self._set_item(key, value) diff --git a/pandas/core/internals/managers.py b/pandas/core/internals/managers.py index 2a7c0536c66a4..b5cbde6dbd65c 100644 --- a/pandas/core/internals/managers.py +++ b/pandas/core/internals/managers.py @@ -1401,7 +1401,7 @@ def insert(self, loc: int, item: Hashable, value: ArrayLike, refs=None) -> None: # insert to the axis; this could possibly raise a TypeError new_axis = self.items.insert(loc, item) - if value.ndim == 2: + if value.ndim >= 2: value = value.T if len(value) > 1: raise ValueError( diff --git a/pandas/tests/frame/test_constructors.py b/pandas/tests/frame/test_constructors.py index 47e307f561cf4..a011da84804b2 100644 --- a/pandas/tests/frame/test_constructors.py +++ b/pandas/tests/frame/test_constructors.py @@ -283,6 +283,11 @@ def test_constructor_cast_failure(self): with pytest.raises(ValueError, match=msg): df["test"] = np.ones((4, 2)) + # this is not ok + msg = "Expected a 1D array, got an array with shape \\(4, 2, 3\\)" + with pytest.raises(ValueError, match=msg): + df["test"] = np.ones((4, 2, 3)) + # this is ok df["foo2"] = np.ones((4, 2)).tolist()