diff --git a/bigframes/core/indexers.py b/bigframes/core/indexers.py index 6258eb00d5..c60e40880b 100644 --- a/bigframes/core/indexers.py +++ b/bigframes/core/indexers.py @@ -155,8 +155,8 @@ def __getitem__(self, key): # row key. We must choose one, so bias towards treating as multi-part row label if isinstance(key, tuple) and len(key) == 2: is_row_multi_index = self._dataframe.index.nlevels > 1 - is_first_item_tuple = isinstance(key[0], tuple) - if not is_row_multi_index or is_first_item_tuple: + is_first_item_list_or_tuple = isinstance(key[0], (tuple, list)) + if not is_row_multi_index or is_first_item_list_or_tuple: df = typing.cast( bigframes.dataframe.DataFrame, _loc_getitem_series_or_dataframe(self._dataframe, key[0]), diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 582d164540..596b9b17f1 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -3638,9 +3638,7 @@ def test_iat_errors(scalars_df_index, scalars_pandas_df_index, index, error): scalars_df_index.iat[index] -def test_iloc_single_integer_out_of_bound_error( - scalars_df_index, scalars_pandas_df_index -): +def test_iloc_single_integer_out_of_bound_error(scalars_df_index): with pytest.raises(IndexError, match="single positional indexer is out-of-bounds"): scalars_df_index.iloc[99] @@ -3655,6 +3653,17 @@ def test_loc_bool_series(scalars_df_index, scalars_pandas_df_index): ) +def test_loc_list_select_rows_and_columns(scalars_df_index, scalars_pandas_df_index): + idx_list = [0, 3, 5] + bf_result = scalars_df_index.loc[idx_list, ["bool_col", "int64_col"]].to_pandas() + pd_result = scalars_pandas_df_index.loc[idx_list, ["bool_col", "int64_col"]] + + pd.testing.assert_frame_equal( + bf_result, + pd_result, + ) + + def test_loc_select_column(scalars_df_index, scalars_pandas_df_index): bf_result = scalars_df_index.loc[:, "int64_col"].to_pandas() pd_result = scalars_pandas_df_index.loc[:, "int64_col"]