googleapis · Genesis929 · Dec 11, 2023 · Nov 28, 2023 · Nov 28, 2023 · Dec 8, 2023
@@ -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,

@@ -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(