Skip to content

Commit 25684ff

Browse files
authored
feat: support item assignment in series (#1859)
* Add item assignment for bigframes dataframe * testcase update
1 parent daf0c3b commit 25684ff

File tree

2 files changed

+67
-0
lines changed

2 files changed

+67
-0
lines changed

bigframes/series.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1598,6 +1598,10 @@ def __getattr__(self, key: str):
15981598
else:
15991599
raise AttributeError(key)
16001600

1601+
def __setitem__(self, key, value) -> None:
1602+
"""Set item using direct assignment, delegating to .loc indexer."""
1603+
self.loc[key] = value
1604+
16011605
def _apply_aggregation(
16021606
self, op: agg_ops.UnaryAggregateOp | agg_ops.NullaryAggregateOp
16031607
) -> Any:

tests/system/small/test_series.py

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -510,6 +510,69 @@ def test_series___getitem___with_default_index(scalars_dfs):
510510
assert bf_result == pd_result
511511

512512

513+
@pytest.mark.parametrize(
514+
("index_col", "key", "value"),
515+
(
516+
("int64_too", 2, "new_string_value"),
517+
("string_col", "Hello, World!", "updated_value"),
518+
("int64_too", 0, None),
519+
),
520+
)
521+
def test_series___setitem__(scalars_dfs, index_col, key, value):
522+
col_name = "string_col"
523+
scalars_df, scalars_pandas_df = scalars_dfs
524+
scalars_df = scalars_df.set_index(index_col, drop=False)
525+
scalars_pandas_df = scalars_pandas_df.set_index(index_col, drop=False)
526+
527+
bf_series = scalars_df[col_name]
528+
pd_series = scalars_pandas_df[col_name].copy()
529+
530+
bf_series[key] = value
531+
pd_series[key] = value
532+
533+
pd.testing.assert_series_equal(bf_series.to_pandas(), pd_series)
534+
535+
536+
@pytest.mark.parametrize(
537+
("key", "value"),
538+
(
539+
(0, 999),
540+
(1, 888),
541+
(0, None),
542+
(-2345, 777),
543+
),
544+
)
545+
def test_series___setitem___with_int_key_numeric(scalars_dfs, key, value):
546+
col_name = "int64_col"
547+
index_col = "int64_too"
548+
scalars_df, scalars_pandas_df = scalars_dfs
549+
scalars_df = scalars_df.set_index(index_col, drop=False)
550+
scalars_pandas_df = scalars_pandas_df.set_index(index_col, drop=False)
551+
552+
bf_series = scalars_df[col_name]
553+
pd_series = scalars_pandas_df[col_name].copy()
554+
555+
bf_series[key] = value
556+
pd_series[key] = value
557+
558+
pd.testing.assert_series_equal(bf_series.to_pandas(), pd_series)
559+
560+
561+
def test_series___setitem___with_default_index(scalars_dfs):
562+
col_name = "float64_col"
563+
key = 2
564+
value = 123.456
565+
scalars_df, scalars_pandas_df = scalars_dfs
566+
567+
bf_series = scalars_df[col_name]
568+
pd_series = scalars_pandas_df[col_name].copy()
569+
570+
bf_series[key] = value
571+
pd_series[key] = value
572+
573+
assert bf_series.to_pandas().iloc[key] == pd_series.iloc[key]
574+
575+
513576
@pytest.mark.parametrize(
514577
("col_name",),
515578
(

0 commit comments

Comments
 (0)