Skip to content

Commit bb3452b

Browse files
committed
FIX-#2362: fix key handling in 'Series.__setitem__'
Signed-off-by: Dmitry Chigarev <[email protected]>
1 parent f14a3c1 commit bb3452b

File tree

2 files changed

+31
-10
lines changed

2 files changed

+31
-10
lines changed

modin/pandas/series.py

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -285,11 +285,16 @@ def __round__(self, decimals=0):
285285
)
286286

287287
def __setitem__(self, key, value):
288-
if key not in self.keys():
289-
raise KeyError(key)
290-
self._create_or_update_from_compiler(
291-
self._query_compiler.setitem(1, key, value), inplace=True
292-
)
288+
if isinstance(key, slice) and (
289+
isinstance(key.start, int) or isinstance(key.stop, int)
290+
):
291+
# There could be two type of slices:
292+
# - Location based slice (1:5)
293+
# - Labels based slice ("a":"e")
294+
# For location based slice we're going to `iloc`, since `loc` can't manage it.
295+
self.iloc[key] = value
296+
else:
297+
self.loc[key] = value
293298

294299
def __sub__(self, right):
295300
return self.sub(right)

modin/pandas/test/test_series.py

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -180,13 +180,13 @@ def inter_df_math_helper_one_side(modin_series, pandas_series, op):
180180
pass
181181

182182

183-
def create_test_series(vals, sort=False):
183+
def create_test_series(vals, sort=False, **kwargs):
184184
if isinstance(vals, dict):
185-
modin_series = pd.Series(vals[next(iter(vals.keys()))])
186-
pandas_series = pandas.Series(vals[next(iter(vals.keys()))])
185+
modin_series = pd.Series(vals[next(iter(vals.keys()))], **kwargs)
186+
pandas_series = pandas.Series(vals[next(iter(vals.keys()))], **kwargs)
187187
else:
188-
modin_series = pd.Series(vals)
189-
pandas_series = pandas.Series(vals)
188+
modin_series = pd.Series(vals, **kwargs)
189+
pandas_series = pandas.Series(vals, **kwargs)
190190
if sort:
191191
modin_series = modin_series.sort_values().reset_index(drop=True)
192192
pandas_series = pandas_series.sort_values().reset_index(drop=True)
@@ -526,6 +526,22 @@ def test___setitem__(data):
526526
df_equals(modin_series, pandas_series)
527527

528528

529+
@pytest.mark.parametrize(
530+
"key",
531+
[
532+
pytest.param(slice(1, 3), id="numeric_slice"),
533+
pytest.param(slice("a", "c"), id="index_based_slice"),
534+
pytest.param(["a", "c", "e"], id="list_of_labels"),
535+
pytest.param([True, False, True, False, True], id="boolean_mask"),
536+
],
537+
)
538+
def test___setitem___non_hashable(key):
539+
md_sr, pd_sr = create_test_series([1, 2, 3, 4, 5], index=["a", "b", "c", "d", "e"])
540+
md_sr[key] = 10
541+
pd_sr[key] = 10
542+
df_equals(md_sr, pd_sr)
543+
544+
529545
@pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)
530546
def test___sizeof__(data):
531547
modin_series, pandas_series = create_test_series(data)

0 commit comments

Comments
 (0)