Skip to content

Commit b3cf824

Browse files
authored
feat: support pandas.Index as key for DataFrame.__setitem__() (#2062)
* feat: support pandas.Index as key for DataFrame.__setitem__() * fix format
1 parent dc54f58 commit b3cf824

File tree

2 files changed

+10
-3
lines changed

2 files changed

+10
-3
lines changed

bigframes/dataframe.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -890,9 +890,11 @@ def __delitem__(self, key: str):
890890
self._set_block(df._get_block())
891891

892892
def __setitem__(
893-
self, key: str | list[str], value: SingleItemValue | MultiItemValue
893+
self,
894+
key: str | list[str] | pandas.Index,
895+
value: SingleItemValue | MultiItemValue,
894896
):
895-
if isinstance(key, list):
897+
if isinstance(key, (list, pandas.Index)):
896898
df = self._assign_multi_items(key, value)
897899
else:
898900
df = self._assign_single_item(key, value)
@@ -2246,7 +2248,7 @@ def _assign_single_item(
22462248

22472249
def _assign_multi_items(
22482250
self,
2249-
k: list[str],
2251+
k: list[str] | pandas.Index,
22502252
v: SingleItemValue | MultiItemValue,
22512253
) -> DataFrame:
22522254
value_sources: Sequence[Any] = []

tests/system/small/test_dataframe.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1197,6 +1197,11 @@ def test_assign_new_column_w_setitem_list_error(scalars_dfs):
11971197
pytest.param(
11981198
["new_col", "new_col_too"], [1, 2], id="sequence_to_full_new_column"
11991199
),
1200+
pytest.param(
1201+
pd.Index(("new_col", "new_col_too")),
1202+
[1, 2],
1203+
id="sequence_to_full_new_column_as_index",
1204+
),
12001205
],
12011206
)
12021207
def test_setitem_multicolumn_with_literals(scalars_dfs, key, value):

0 commit comments

Comments
 (0)