feat: support dataframe.loc with conditional columns selection (#233)

Co-authored-by: Tim Swast <[email protected]>

Author: Genesis929

bigframes/core/indexers.py

@@ -156,7 +156,12 @@ def __getitem__(self, key):
                 bigframes.dataframe.DataFrame,
                 _loc_getitem_series_or_dataframe(self._dataframe, key[0]),
             )
-            return df[key[1]]
+
+            columns = key[1]
+            if isinstance(columns, pd.Series) and columns.dtype == "bool":
+                columns = df.columns[columns]
+
+            return df[columns]
 
         return typing.cast(
             bigframes.dataframe.DataFrame,

tests/system/small/test_dataframe.py

@@ -2474,6 +2474,17 @@ def test_loc_select_column(scalars_df_index, scalars_pandas_df_index):
     )
 
 
+def test_loc_select_with_column_condition(scalars_df_index, scalars_pandas_df_index):
+    bf_result = scalars_df_index.loc[:, scalars_df_index.dtypes == "Int64"].to_pandas()
+    pd_result = scalars_pandas_df_index.loc[
+        :, scalars_pandas_df_index.dtypes == "Int64"
+    ]
+    pd.testing.assert_frame_equal(
+        bf_result,
+        pd_result,
+    )
+
+
 def test_loc_single_index_with_duplicate(scalars_df_index, scalars_pandas_df_index):
     scalars_df_index = scalars_df_index.set_index("string_col", drop=False)
     scalars_pandas_df_index = scalars_pandas_df_index.set_index(