Merge branch 'main' into shuowei-anywidget-cell-visual

Author: shuoweil

bigframes/bigquery/_operations/ai.py

@@ -1018,7 +1018,7 @@ def _convert_series(
 
 def _resolve_connection_id(series: series.Series, connection_id: str | None):
     return clients.get_canonical_bq_connection_id(
-        connection_id or series._session._bq_connection,
+        connection_id or series._session.bq_connection,
         series._session._project,
         series._session._location,
     )

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):

bigframes/functions/_function_session.py

@@ -162,7 +162,7 @@ def _resolve_bigquery_connection_id(
     ) -> str:
         """Resolves BigQuery connection id."""
         if not bigquery_connection:
-            bigquery_connection = session._bq_connection  # type: ignore
+            bigquery_connection = session.bq_connection  # type: ignore
 
         bigquery_connection = clients.get_canonical_bq_connection_id(
             bigquery_connection,

bigframes/operations/blob.py

@@ -311,7 +311,7 @@ def _resolve_connection(self, connection: Optional[str] = None) -> str:
         Raises:
             ValueError: If the connection cannot be resolved to a valid string.
         """
-        connection = connection or self._data._block.session._bq_connection
+        connection = connection or self._data._block.session.bq_connection
         return clients.get_canonical_bq_connection_id(
             connection,
             default_project=self._data._block.session._project,

bigframes/session/__init__.py

@@ -374,6 +374,16 @@ def _allows_ambiguity(self) -> bool:
     def _anonymous_dataset(self):
         return self._anon_dataset_manager.dataset
 
+    @property
+    def bq_connection(self) -> str:
+        msg = bfe.format_message(
+            f"""You are using the BigFrames session default connection: {self._bq_connection},
+            which can be different from the BigQuery project default connection.
+            This default connection may change in the future."""
+        )
+        warnings.warn(msg, category=FutureWarning)
+        return self._bq_connection
+
     def __hash__(self):
         # Stable hash needed to use in expression tree
         return hash(str(self._session_id))
@@ -2253,7 +2263,7 @@ def _create_bq_connection(
     ) -> str:
         """Create the connection with the session settings and try to attach iam role to the connection SA.
         If any of project, location or connection isn't specified, use the session defaults. Returns fully-qualified connection name."""
-        connection = self._bq_connection if not connection else connection
+        connection = self.bq_connection if not connection else connection
         connection = bigframes.clients.get_canonical_bq_connection_id(
             connection_id=connection,
             default_project=self._project,

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)