From 0fd753345c8ab01f8fd1ea5ed9d95d03df1fd306 Mon Sep 17 00:00:00 2001 From: Huan Chen Date: Tue, 28 Nov 2023 01:19:47 +0000 Subject: [PATCH] feat: support dataframe.loc with conditional columns selection --- bigframes/core/indexers.py | 7 ++++++- tests/system/small/test_dataframe.py | 11 +++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) 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 e25e9ce501..106e2456d9 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -2412,6 +2412,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(