From 8cf0e3d7d13ecc1e3f05f3e26a1dc25ee32b37c2 Mon Sep 17 00:00:00 2001 From: Michael Tiemann <72577720+MichaelTiemannOSC@users.noreply.github.com> Date: Sun, 13 Aug 2023 22:37:41 -0400 Subject: [PATCH 1/2] Update managers.py When a loc indexer goes to create a fresh array to hold values from an Extension Array, it makes the array allocation based on the na_value of the EA, but that na_value may be smaller than the size of the things that the EA can hold (such as complex128). Note that np.nan is itself a legitimate complex128 value. If the allocated array cannot hold the values from the block manager, and if the EA is not immutable, it will try a second strategy of allocating an array based on the dtype of the values of the blocks. If the blocks hold complex numbers, the array will be properly sized. This should close #54445. Signed-off-by: Michael Tiemann <72577720+MichaelTiemannOSC@users.noreply.github.com> --- pandas/core/internals/managers.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/pandas/core/internals/managers.py b/pandas/core/internals/managers.py index 7bb579b22aeed..a8e1a174d3295 100644 --- a/pandas/core/internals/managers.py +++ b/pandas/core/internals/managers.py @@ -983,11 +983,18 @@ def fast_xs(self, loc: int) -> SingleBlockManager: ) result = ensure_wrapped_if_datetimelike(result) - for blk in self.blocks: - # Such assignment may incorrectly coerce NaT to None - # result[blk.mgr_locs] = blk._slice((slice(None), loc)) - for i, rl in enumerate(blk.mgr_locs): - result[rl] = blk.iget((i, loc)) + try: + for blk in self.blocks: + # Such assignment may incorrectly coerce NaT to None + # result[blk.mgr_locs] = blk._slice((slice(None), loc)) + for i, rl in enumerate(blk.mgr_locs): + result[rl] = blk.iget((i, loc)) + except TypeError: + if isinstance(dtype, ExtensionDtype) and not immutable_ea: + values = [v[loc] for v in self.arrays] + result = cls._from_sequence(values, dtype) + else: + raise TypeError if immutable_ea: dtype = cast(ExtensionDtype, dtype) From 7b9c68a3aba1a9f7e2c6870ea2127fb70e7fdcf8 Mon Sep 17 00:00:00 2001 From: Michael Tiemann <72577720+MichaelTiemannOSC@users.noreply.github.com> Date: Mon, 14 Aug 2023 04:04:48 -0400 Subject: [PATCH 2/2] Update managers.py Make mypy happy. Signed-off-by: Michael Tiemann <72577720+MichaelTiemannOSC@users.noreply.github.com> --- pandas/core/internals/managers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pandas/core/internals/managers.py b/pandas/core/internals/managers.py index a8e1a174d3295..7904972a84dfa 100644 --- a/pandas/core/internals/managers.py +++ b/pandas/core/internals/managers.py @@ -992,7 +992,7 @@ def fast_xs(self, loc: int) -> SingleBlockManager: except TypeError: if isinstance(dtype, ExtensionDtype) and not immutable_ea: values = [v[loc] for v in self.arrays] - result = cls._from_sequence(values, dtype) + result = cls._from_sequence(values, dtype=dtype) else: raise TypeError