Skip to content

Commit e8b9749

Browse files
authored
PERF: DataFrame.iloc[int] for EA dtypes (#54508)
* improve perf of fast_xs for EA dtypes * whatsnew
1 parent 65f22a0 commit e8b9749

File tree

2 files changed

+7
-15
lines changed

2 files changed

+7
-15
lines changed

doc/source/whatsnew/v2.1.0.rst

+1
Original file line numberDiff line numberDiff line change
@@ -622,6 +622,7 @@ Performance improvements
622622
- Performance improvement in :meth:`.ArrowExtensionArray.to_numpy` (:issue:`52525`)
623623
- Performance improvement in :meth:`.DataFrameGroupBy.groups` (:issue:`53088`)
624624
- Performance improvement in :meth:`DataFrame.astype` when ``dtype`` is an extension dtype (:issue:`54299`)
625+
- Performance improvement in :meth:`DataFrame.iloc` when input is an single integer and dataframe is backed by extension dtypes (:issue:`54508`)
625626
- Performance improvement in :meth:`DataFrame.isin` for extension dtypes (:issue:`53514`)
626627
- Performance improvement in :meth:`DataFrame.loc` when selecting rows and columns (:issue:`53014`)
627628
- Performance improvement in :meth:`DataFrame.transpose` when transposing a DataFrame with a single PyArrow dtype (:issue:`54224`)

pandas/core/internals/managers.py

+6-15
Original file line numberDiff line numberDiff line change
@@ -966,19 +966,10 @@ def fast_xs(self, loc: int) -> SingleBlockManager:
966966

967967
n = len(self)
968968

969-
# GH#46406
970-
immutable_ea = isinstance(dtype, ExtensionDtype) and dtype._is_immutable
971-
972-
if isinstance(dtype, ExtensionDtype) and not immutable_ea:
973-
cls = dtype.construct_array_type()
974-
result = cls._empty((n,), dtype=dtype)
969+
if isinstance(dtype, ExtensionDtype):
970+
result = np.empty(n, dtype=object)
975971
else:
976-
# error: Argument "dtype" to "empty" has incompatible type
977-
# "Union[Type[object], dtype[Any], ExtensionDtype, None]"; expected
978-
# "None"
979-
result = np.empty(
980-
n, dtype=object if immutable_ea else dtype # type: ignore[arg-type]
981-
)
972+
result = np.empty(n, dtype=dtype)
982973
result = ensure_wrapped_if_datetimelike(result)
983974

984975
for blk in self.blocks:
@@ -987,9 +978,9 @@ def fast_xs(self, loc: int) -> SingleBlockManager:
987978
for i, rl in enumerate(blk.mgr_locs):
988979
result[rl] = blk.iget((i, loc))
989980

990-
if immutable_ea:
991-
dtype = cast(ExtensionDtype, dtype)
992-
result = dtype.construct_array_type()._from_sequence(result, dtype=dtype)
981+
if isinstance(dtype, ExtensionDtype):
982+
cls = dtype.construct_array_type()
983+
result = cls._from_sequence(result, dtype=dtype)
993984

994985
bp = BlockPlacement(slice(0, len(result)))
995986
block = new_block(result, placement=bp, ndim=1)

0 commit comments

Comments
 (0)