From 15aa0a6b4d1a8f9311d88991f525768d8c3949b3 Mon Sep 17 00:00:00 2001 From: Brock Date: Tue, 24 Aug 2021 19:46:57 -0700 Subject: [PATCH 1/2] REF: fast_xs avoid object dtype --- pandas/core/internals/managers.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/pandas/core/internals/managers.py b/pandas/core/internals/managers.py index d0dad26cfb275..30a9d43c5d150 100644 --- a/pandas/core/internals/managers.py +++ b/pandas/core/internals/managers.py @@ -941,13 +941,11 @@ def fast_xs(self, loc: int) -> ArrayLike: n = len(self) if isinstance(dtype, ExtensionDtype): - # we'll eventually construct an ExtensionArray. - result = np.empty(n, dtype=object) - # TODO: let's just use dtype.empty? + cls = dtype.construct_array_type() + result = cls._empty(n, dtype=dtype) else: result = np.empty(n, dtype=dtype) - - result = ensure_wrapped_if_datetimelike(result) + result = ensure_wrapped_if_datetimelike(result) for blk in self.blocks: # Such assignment may incorrectly coerce NaT to None @@ -955,9 +953,6 @@ def fast_xs(self, loc: int) -> ArrayLike: for i, rl in enumerate(blk.mgr_locs): result[rl] = blk.iget((i, loc)) - if isinstance(dtype, ExtensionDtype): - result = dtype.construct_array_type()._from_sequence(result, dtype=dtype) - return result def iget(self, i: int) -> SingleBlockManager: From 6470a353c8c75b8b469fa0d59f766ce44b57a691 Mon Sep 17 00:00:00 2001 From: Brock Date: Wed, 25 Aug 2021 11:27:15 -0700 Subject: [PATCH 2/2] mypy fixup --- 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 30a9d43c5d150..e558e18d183e1 100644 --- a/pandas/core/internals/managers.py +++ b/pandas/core/internals/managers.py @@ -942,7 +942,7 @@ def fast_xs(self, loc: int) -> ArrayLike: n = len(self) if isinstance(dtype, ExtensionDtype): cls = dtype.construct_array_type() - result = cls._empty(n, dtype=dtype) + result = cls._empty((n,), dtype=dtype) else: result = np.empty(n, dtype=dtype) result = ensure_wrapped_if_datetimelike(result)