feat: add support for df.loc[list, column(s)] (#1761)

Author: Genesis929

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]),

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"]