From a3852190a4b8542ef9d81a23e253a5cdac94ad89 Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Tue, 9 Sep 2025 20:29:02 +0000 Subject: [PATCH 1/2] feat: support pandas.Index as key for DataFrame.__setitem__() --- bigframes/dataframe.py | 8 +++++--- tests/system/small/test_dataframe.py | 3 +++ 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index c65bbdd2c8..6ebb8dc2dc 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -890,9 +890,11 @@ def __delitem__(self, key: str): self._set_block(df._get_block()) def __setitem__( - self, key: str | list[str], value: SingleItemValue | MultiItemValue + self, + key: str | list[str] | pandas.Index, + value: SingleItemValue | MultiItemValue, ): - if isinstance(key, list): + if isinstance(key, (list, pandas.Index)): df = self._assign_multi_items(key, value) else: df = self._assign_single_item(key, value) @@ -2246,7 +2248,7 @@ def _assign_single_item( def _assign_multi_items( self, - k: list[str], + k: list[str] | pandas.Index, v: SingleItemValue | MultiItemValue, ) -> DataFrame: value_sources: Sequence[Any] = [] diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 323956b038..41420cfa57 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -1197,6 +1197,9 @@ def test_assign_new_column_w_setitem_list_error(scalars_dfs): pytest.param( ["new_col", "new_col_too"], [1, 2], id="sequence_to_full_new_column" ), + pytest.param( + pd.Index(("new_col", "new_col_too")), [1, 2], id="sequence_to_full_new_column_as_index" + ) ], ) def test_setitem_multicolumn_with_literals(scalars_dfs, key, value): From 6f519475dfa8305cf43496be62e751a1447d738e Mon Sep 17 00:00:00 2001 From: Shenyang Cai Date: Tue, 9 Sep 2025 20:31:56 +0000 Subject: [PATCH 2/2] fix format --- bigframes/dataframe.py | 2 +- tests/system/small/test_dataframe.py | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 6ebb8dc2dc..5a93646677 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -891,7 +891,7 @@ def __delitem__(self, key: str): def __setitem__( self, - key: str | list[str] | pandas.Index, + key: str | list[str] | pandas.Index, value: SingleItemValue | MultiItemValue, ): if isinstance(key, (list, pandas.Index)): diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 41420cfa57..15f800097b 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -1198,8 +1198,10 @@ def test_assign_new_column_w_setitem_list_error(scalars_dfs): ["new_col", "new_col_too"], [1, 2], id="sequence_to_full_new_column" ), pytest.param( - pd.Index(("new_col", "new_col_too")), [1, 2], id="sequence_to_full_new_column_as_index" - ) + pd.Index(("new_col", "new_col_too")), + [1, 2], + id="sequence_to_full_new_column_as_index", + ), ], ) def test_setitem_multicolumn_with_literals(scalars_dfs, key, value):