Skip to content

Commit 0197e0c

Browse files
committed
BUG: Preserve dtype on homogeneous EA xs
1 parent e8b37da commit 0197e0c

File tree

3 files changed

+54
-8
lines changed

3 files changed

+54
-8
lines changed

doc/source/whatsnew/v0.24.0.txt

+1
Original file line numberDiff line numberDiff line change
@@ -545,6 +545,7 @@ Other API Changes
545545
- :class:`pandas.io.formats.style.Styler` supports a ``number-format`` property when using :meth:`~pandas.io.formats.style.Styler.to_excel` (:issue:`22015`)
546546
- :meth:`DataFrame.corr` and :meth:`Series.corr` now raise a ``ValueError`` along with a helpful error message instead of a ``KeyError`` when supplied with an invalid method (:issue:`22298`)
547547
- :meth:`shift` will now always return a copy, instead of the previous behaviour of returning self when shifting by 0 (:issue:`22397`)
548+
- Slicing a single row of a DataFrame with multiple ExtensionArrays of the same type now preserves the dtype, rather than coercing to object (:issue:`22784`)
548549

549550
.. _whatsnew_0240.deprecations:
550551

pandas/core/internals/managers.py

+25-8
Original file line numberDiff line numberDiff line change
@@ -906,14 +906,25 @@ def fast_xs(self, loc):
906906

907907
# unique
908908
dtype = _interleaved_dtype(self.blocks)
909+
909910
n = len(items)
910-
result = np.empty(n, dtype=dtype)
911+
if is_extension_array_dtype(dtype):
912+
# we'll eventually construct an ExtensionArray.
913+
result = np.empty(n, dtype=object)
914+
else:
915+
result = np.empty(n, dtype=dtype)
916+
911917
for blk in self.blocks:
912918
# Such assignment may incorrectly coerce NaT to None
913919
# result[blk.mgr_locs] = blk._slice((slice(None), loc))
914920
for i, rl in enumerate(blk.mgr_locs):
915921
result[rl] = blk._try_coerce_result(blk.iget((i, loc)))
916922

923+
if is_extension_array_dtype(dtype):
924+
result = dtype.construct_array_type()._from_sequence(
925+
result, dtype=dtype
926+
)
927+
917928
return result
918929

919930
def consolidate(self):
@@ -1855,16 +1866,22 @@ def _shape_compat(x):
18551866

18561867

18571868
def _interleaved_dtype(blocks):
1858-
if not len(blocks):
1859-
return None
1869+
# type: (List[Block]) -> Optional[Union[np.dtype, ExtensionDtype]]
1870+
"""Find the common dtype for `blocks`.
18601871
1861-
dtype = find_common_type([b.dtype for b in blocks])
1872+
Parameters
1873+
----------
1874+
blocks : List[Block]
18621875
1863-
# only numpy compat
1864-
if isinstance(dtype, (PandasExtensionDtype, ExtensionDtype)):
1865-
dtype = np.object
1876+
Returns
1877+
-------
1878+
dtype : Optional[Union[np.dtype, ExtensionDtype]]
1879+
None is returned when `blocks` is empty.
1880+
"""
1881+
if not len(blocks):
1882+
return None
18661883

1867-
return dtype
1884+
return find_common_type([b.dtype for b in blocks])
18681885

18691886

18701887
def _consolidate(blocks):

pandas/tests/indexing/test_indexing.py

+28
Original file line numberDiff line numberDiff line change
@@ -1079,3 +1079,31 @@ def test_validate_indices_high():
10791079
def test_validate_indices_empty():
10801080
with tm.assert_raises_regex(IndexError, "indices are out"):
10811081
validate_indices(np.array([0, 1]), 0)
1082+
1083+
1084+
def test_extension_array_cross_section():
1085+
# A cross-section of a homogeneous EA should be an EA
1086+
df = pd.DataFrame({
1087+
"A": pd.core.arrays.integer_array([1, 2]),
1088+
"B": pd.core.arrays.integer_array([3, 4])
1089+
}, index=['a', 'b'])
1090+
expected = pd.Series(pd.core.arrays.integer_array([1, 3]),
1091+
index=['A', 'B'], name='a')
1092+
result = df.loc['a']
1093+
tm.assert_series_equal(result, expected)
1094+
1095+
result = df.iloc[0]
1096+
tm.assert_series_equal(result, expected)
1097+
1098+
1099+
def test_extension_array_cross_section_converts():
1100+
df = pd.DataFrame({
1101+
"A": pd.core.arrays.integer_array([1, 2]),
1102+
"B": np.array([1, 2]),
1103+
}, index=['a', 'b'])
1104+
result = df.loc['a']
1105+
expected = pd.Series([1, 1], dtype=object, index=['A', 'B'], name='a')
1106+
tm.assert_series_equal(result, expected)
1107+
1108+
result = df.iloc[0]
1109+
tm.assert_series_equal(result, expected)

0 commit comments

Comments
 (0)