feat: support pandas.Index as key for DataFrame.__setitem__()

Author: sycai

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] = []

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):