Skip to content

Commit a272fc8

Browse files
committed
BUG: Accept dict or Series in fillna for categorical Series
1 parent 63e8527 commit a272fc8

File tree

2 files changed

+58
-9
lines changed

2 files changed

+58
-9
lines changed

pandas/core/categorical.py

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1660,16 +1660,26 @@ def fillna(self, value=None, method=None, limit=None):
16601660

16611661
else:
16621662

1663-
if not isna(value) and value not in self.categories:
1664-
raise ValueError("fill value must be in categories")
1663+
if isinstance(value, ABCSeries):
1664+
if not value[~value.isin(self.categories)].isna().all():
1665+
raise ValueError("fill value must be in categories")
16651666

1666-
mask = values == -1
1667-
if mask.any():
1668-
values = values.copy()
1669-
if isna(value):
1670-
values[mask] = -1
1671-
else:
1672-
values[mask] = self.categories.get_loc(value)
1667+
values_codes = _get_codes_for_values(value, self.categories)
1668+
indexer = np.where(values_codes != -1)
1669+
values[indexer] = values_codes[values_codes != -1]
1670+
1671+
# Scalar value
1672+
else:
1673+
if not isna(value) and value not in self.categories:
1674+
raise ValueError("fill value must be in categories")
1675+
1676+
mask = values == -1
1677+
if mask.any():
1678+
values = values.copy()
1679+
if isna(value):
1680+
values[mask] = -1
1681+
else:
1682+
values[mask] = self.categories.get_loc(value)
16731683

16741684
return self._constructor(values, categories=self.categories,
16751685
ordered=self.ordered, fastpath=True)

pandas/tests/test_categorical.py

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4601,6 +4601,45 @@ def f():
46014601
df = pd.DataFrame({'a': pd.Categorical(idx)})
46024602
tm.assert_frame_equal(df.fillna(value=pd.NaT), df)
46034603

4604+
@pytest.mark.parametrize('fill_value expected_output', [
4605+
('a', ['a', 'a', 'b', 'a', 'a']),
4606+
({1: 'a', 3: 'b', 4: 'b'}, ['a', 'a', 'b', 'b', 'b']),
4607+
({1: 'a'}, ['a', 'a', 'b', np.nan, np.nan]),
4608+
({1: 'a', 3: 'b'}, ['a', 'a', 'b', 'b', np.nan]),
4609+
(pd.Series('a'), ['a', np.nan, 'b', np.nan, np.nan]),
4610+
(pd.Series('a', index=[1]), ['a', 'a', 'b', np.nan, np.nan]),
4611+
(pd.Series({1: 'a', 3: 'b'}), ['a', 'a', 'b', 'b', np.nan]),
4612+
(pd.Series(['a', 'b'], index=[3, 4]))
4613+
])
4614+
def fillna_series_categorical(self, fill_value, expected_output):
4615+
# GH 17033
4616+
# Test fillna for a Categorical series
4617+
data = ['a', np.nan, 'b', np.nan, np.nan]
4618+
s = pd.Series(pd.Categorical(data, categories=['a', 'b']))
4619+
exp = pd.Series(pd.Categorical(expected_output, categories=['a', 'b']))
4620+
tm.assert_series_equal(s.fillna(fill_value), exp)
4621+
4622+
def fillna_series_categorical_errormsg(self):
4623+
data = ['a', np.nan, 'b', np.nan, np.nan]
4624+
s = pd.Series(pd.Categorical(data, categories=['a', 'b']))
4625+
4626+
with tm.assert_raises_regex(ValueError,
4627+
"fill value must be in categories"):
4628+
s.fillna('d')
4629+
4630+
with tm.assert_raises_regex(ValueError,
4631+
"fill value must be in categories"):
4632+
s.fillna(pd.Series('d'))
4633+
4634+
with tm.assert_raises_regex(ValueError,
4635+
"fill value must be in categories"):
4636+
s.fillna({1: 'd', 3: 'a'})
4637+
4638+
with tm.assert_raises_regex(TypeError,
4639+
'"value" parameter must be a scalar or '
4640+
'dict but you passed a "list"'):
4641+
s.fillna(['a', 'b'])
4642+
46044643
def test_astype_to_other(self):
46054644

46064645
s = self.cat['value_group']

0 commit comments

Comments
 (0)