From 02dd36a5c0884b36632fbcd289c53b92285f29e8 Mon Sep 17 00:00:00 2001 From: Brock Date: Tue, 23 Nov 2021 12:48:46 -0800 Subject: [PATCH 1/2] BUG: DataFrame(EA2D) --- pandas/core/internals/construction.py | 2 +- pandas/tests/extension/base/dim2.py | 7 +++++++ pandas/tests/frame/test_constructors.py | 8 ++++++++ 3 files changed, 16 insertions(+), 1 deletion(-) diff --git a/pandas/core/internals/construction.py b/pandas/core/internals/construction.py index a766f8321a641..c6fb293b3452a 100644 --- a/pandas/core/internals/construction.py +++ b/pandas/core/internals/construction.py @@ -294,7 +294,7 @@ def ndarray_to_mgr( if is_1d_only_ea_dtype(vdtype) or isinstance(dtype, ExtensionDtype): # GH#19157 - if isinstance(values, np.ndarray) and values.ndim > 1: + if isinstance(values, (np.ndarray, ExtensionArray)) and values.ndim > 1: # GH#12513 a EA dtype passed with a 2D array, split into # multiple EAs that view the values values = [values[:, n] for n in range(values.shape[1])] diff --git a/pandas/tests/extension/base/dim2.py b/pandas/tests/extension/base/dim2.py index b4a817cbc37ec..1ac06d4c6932f 100644 --- a/pandas/tests/extension/base/dim2.py +++ b/pandas/tests/extension/base/dim2.py @@ -14,6 +14,13 @@ class Dim2CompatTests(BaseExtensionTests): + def test_frame_from_2d_array(self, data): + arr2d = data.repeat(2).reshape(-1, 2) + + df = pd.DataFrame(arr2d) + expected = pd.DataFrame({0: arr2d[:, 0], 1: arr2d[:, 1]}) + self.assert_frame_equal(df, expected) + def test_swapaxes(self, data): arr2d = data.repeat(2).reshape(-1, 2) diff --git a/pandas/tests/frame/test_constructors.py b/pandas/tests/frame/test_constructors.py index fc5bffab118af..7347640fc05a7 100644 --- a/pandas/tests/frame/test_constructors.py +++ b/pandas/tests/frame/test_constructors.py @@ -70,6 +70,14 @@ class TestDataFrameConstructors: + def test_constructor_from_2d_datetimearray(self): + dti = date_range("2016-01-01", periods=6, tz="US/Pacific") + dta = dti._data.reshape(3, 2) + + df = DataFrame(dta) + expected = DataFrame({0: dta[:, 0], 1: dta[:, 1]}) + tm.assert_frame_equal(df, expected) + def test_constructor_dict_with_tzaware_scalar(self): # GH#42505 dt = Timestamp("2019-11-03 01:00:00-0700").tz_convert("America/Los_Angeles") From 6ca443036b25a8b010312a43b2fcb4a9c572d275 Mon Sep 17 00:00:00 2001 From: Brock Date: Tue, 23 Nov 2021 14:50:26 -0800 Subject: [PATCH 2/2] mypy fixup --- pandas/core/internals/construction.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/pandas/core/internals/construction.py b/pandas/core/internals/construction.py index c6fb293b3452a..77f3db0d09df5 100644 --- a/pandas/core/internals/construction.py +++ b/pandas/core/internals/construction.py @@ -297,7 +297,12 @@ def ndarray_to_mgr( if isinstance(values, (np.ndarray, ExtensionArray)) and values.ndim > 1: # GH#12513 a EA dtype passed with a 2D array, split into # multiple EAs that view the values - values = [values[:, n] for n in range(values.shape[1])] + # error: No overload variant of "__getitem__" of "ExtensionArray" + # matches argument type "Tuple[slice, int]" + values = [ + values[:, n] # type: ignore[call-overload] + for n in range(values.shape[1]) + ] else: values = [values]