feat: Support pd.col expressions with .loc and getitem

Author: TrevorBergeron

bigframes/core/array_value.py

@@ -204,7 +204,12 @@ def filter_by_id(self, predicate_id: str, keep_null: bool = False) -> ArrayValue
         return self.filter(predicate)
 
     def filter(self, predicate: ex.Expression):
-        return ArrayValue(nodes.FilterNode(child=self.node, predicate=predicate))
+        if predicate.is_scalar_expr:
+            return ArrayValue(nodes.FilterNode(child=self.node, predicate=predicate))
+        else:
+            arr, filter_ids = self.compute_general_expression([predicate])
+            arr = arr.filter_by_id(filter_ids[0])
+            return arr.drop_columns(filter_ids)
 
     def order_by(
         self, by: Sequence[OrderingExpression], is_total_order: bool = False

bigframes/core/indexers.py

@@ -23,6 +23,7 @@
 import pandas as pd
 
 import bigframes.core.blocks
+import bigframes.core.col
 import bigframes.core.expression as ex
 import bigframes.core.guid as guid
 import bigframes.core.indexes as indexes
@@ -36,7 +37,11 @@
 
 if typing.TYPE_CHECKING:
     LocSingleKey = Union[
-        bigframes.series.Series, indexes.Index, slice, bigframes.core.scalar.Scalar
+        bigframes.series.Series,
+        indexes.Index,
+        slice,
+        bigframes.core.scalar.Scalar,
+        bigframes.core.col.Expression,
     ]
 
 
@@ -309,6 +314,15 @@ def _loc_getitem_series_or_dataframe(
         raise NotImplementedError(
             f"loc does not yet support indexing with a slice. {constants.FEEDBACK_LINK}"
         )
+    if isinstance(key, bigframes.core.col.Expression):
+        label_to_col_ref = {
+            label: ex.deref(id)
+            for id, label in series_or_dataframe._block.col_id_to_label.items()
+        }
+        resolved_expr = key._value.bind_variables(label_to_col_ref)
+        result = series_or_dataframe.copy()
+        result._set_block(series_or_dataframe._block.filter(resolved_expr))
+        return result
     if callable(key):
         raise NotImplementedError(
             f"loc does not yet support indexing with a callable. {constants.FEEDBACK_LINK}"

bigframes/dataframe.py

@@ -623,13 +623,18 @@ def __getitem__(
     ):  # No return type annotations (like pandas) as type cannot always be determined statically
         # NOTE: This implements the operations described in
         # https://pandas.pydata.org/docs/getting_started/intro_tutorials/03_subset_data.html
+        import bigframes.core.col
+        import bigframes.pandas
 
-        if isinstance(key, bigframes.series.Series):
+        if isinstance(key, bigframes.pandas.Series):
             return self._getitem_bool_series(key)
 
         if isinstance(key, slice):
             return self.iloc[key]
 
+        if isinstance(key, bigframes.core.col.Expression):
+            return self.loc[key]
+
         # TODO(tswast): Fix this pylance warning: Class overlaps "Hashable"
         # unsafely and could produce a match at runtime
         if isinstance(key, blocks.Label):

tests/unit/test_col.py

@@ -158,3 +158,21 @@ def test_pd_col_binary_bool_operators(scalars_dfs, op):
     pd_result = scalars_pandas_df.assign(**pd_kwargs)
 
     assert_frame_equal(bf_result, pd_result)
+
+
+def test_loc_with_pd_col(scalars_dfs):
+    scalars_df, scalars_pandas_df = scalars_dfs
+
+    bf_result = scalars_df.loc[bpd.col("float64_col") > 4].to_pandas()
+    pd_result = scalars_pandas_df.loc[pd.col("float64_col") > 4]  # type: ignore
+
+    assert_frame_equal(bf_result, pd_result)
+
+
+def test_getitem_with_pd_col(scalars_dfs):
+    scalars_df, scalars_pandas_df = scalars_dfs
+
+    bf_result = scalars_df[bpd.col("float64_col") > 4].to_pandas()
+    pd_result = scalars_pandas_df[pd.col("float64_col") > 4]  # type: ignore
+
+    assert_frame_equal(bf_result, pd_result)