diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index c65bbdd2c8..5a93646677 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..15f800097b 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -1197,6 +1197,11 @@ 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):