diff --git a/bigframes/core/indexers.py b/bigframes/core/indexers.py index 69048b6845..12a1303d29 100644 --- a/bigframes/core/indexers.py +++ b/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, diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 45490e00ca..57115335dc 100644 --- a/tests/system/small/test_dataframe.py +++ b/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(