Merge remote-tracking branch 'github/main' into prerelease_and_pandas_3

Author: TrevorBergeron

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/ml/llm.py

@@ -877,7 +877,7 @@ class Claude3TextGenerator(base.RetriableRemotePredictor):
             The model for natural language tasks. Possible values are "claude-3-sonnet", "claude-3-haiku", "claude-3-5-sonnet" and "claude-3-opus".
             "claude-3-sonnet" (deprecated) is Anthropic's dependable combination of skills and speed. It is engineered to be dependable for scaled AI deployments across a variety of use cases.
             "claude-3-haiku" is Anthropic's fastest, most compact vision and text model for near-instant responses to simple queries, meant for seamless AI experiences mimicking human interactions.
-            "claude-3-5-sonnet" is Anthropic's most powerful AI model and maintains the speed and cost of Claude 3 Sonnet, which is a mid-tier model.
+            "claude-3-5-sonnet" (deprecated) is Anthropic's most powerful AI model and maintains the speed and cost of Claude 3 Sonnet, which is a mid-tier model.
             "claude-3-opus" (deprecated) is Anthropic's second-most powerful AI model, with strong performance on highly complex tasks.
             https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude#available-claude-models
             If no setting is provided, "claude-3-sonnet" will be used by default

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,