Skip to content

Commit 8306605

Browse files
jbrockmendelJulianWgs
authored andcommitted
REGR: preserve Int32 dtype on setitem (pandas-dev#42166)
1 parent b6b721a commit 8306605

File tree

2 files changed

+37
-1
lines changed

2 files changed

+37
-1
lines changed

pandas/core/indexing.py

+20-1
Original file line numberDiff line numberDiff line change
@@ -42,8 +42,12 @@
4242
isna,
4343
)
4444

45+
from pandas.core import algorithms as algos
4546
import pandas.core.common as com
46-
from pandas.core.construction import array as pd_array
47+
from pandas.core.construction import (
48+
array as pd_array,
49+
extract_array,
50+
)
4751
from pandas.core.indexers import (
4852
check_array_indexer,
4953
is_empty_indexer,
@@ -1660,6 +1664,21 @@ def _setitem_with_indexer(self, indexer, value, name="iloc"):
16601664
if com.is_null_slice(indexer[0]):
16611665
# We are setting an entire column
16621666
self.obj[key] = value
1667+
return
1668+
elif is_array_like(value):
1669+
# GH#42099
1670+
arr = extract_array(value, extract_numpy=True)
1671+
taker = -1 * np.ones(len(self.obj), dtype=np.intp)
1672+
empty_value = algos.take_nd(arr, taker)
1673+
if not isinstance(value, ABCSeries):
1674+
# if not Series (in which case we need to align),
1675+
# we can short-circuit
1676+
empty_value[indexer[0]] = arr
1677+
self.obj[key] = empty_value
1678+
return
1679+
1680+
self.obj[key] = empty_value
1681+
16631682
else:
16641683
self.obj[key] = infer_fill_value(value)
16651684

pandas/tests/indexing/test_loc.py

+17
Original file line numberDiff line numberDiff line change
@@ -1830,6 +1830,23 @@ def test_loc_setitem_with_expansion_nonunique_index(self, index, request):
18301830
)
18311831
tm.assert_frame_equal(df, expected)
18321832

1833+
@pytest.mark.parametrize(
1834+
"dtype", ["Int32", "Int64", "UInt32", "UInt64", "Float32", "Float64"]
1835+
)
1836+
def test_loc_setitem_with_expansion_preserves_nullable_int(self, dtype):
1837+
# GH#42099
1838+
ser = Series([0, 1, 2, 3], dtype=dtype)
1839+
df = DataFrame({"data": ser})
1840+
1841+
result = DataFrame(index=df.index)
1842+
result.loc[df.index, "data"] = ser
1843+
1844+
tm.assert_frame_equal(result, df)
1845+
1846+
result = DataFrame(index=df.index)
1847+
result.loc[df.index, "data"] = ser._values
1848+
tm.assert_frame_equal(result, df)
1849+
18331850

18341851
class TestLocCallable:
18351852
def test_frame_loc_getitem_callable(self):

0 commit comments

Comments
 (0)