feat: _read_gbq_colab creates hybrid session

bigframes/pandas/io/api.py

@@ -279,6 +279,7 @@ def _read_gbq_colab(  # type: ignore[overload-overlap]
     *,
     pyformat_args: Optional[Dict[str, Any]] = ...,
     dry_run: Literal[False] = ...,
+    use_hybrid_engine: bool = ...,
 ) -> bigframes.dataframe.DataFrame:
     ...
 
@@ -289,6 +290,7 @@ def _read_gbq_colab(
     *,
     pyformat_args: Optional[Dict[str, Any]] = ...,
     dry_run: Literal[True] = ...,
+    use_hybrid_engine: bool = ...,
 ) -> pandas.Series:
     ...
 
@@ -298,6 +300,7 @@ def _read_gbq_colab(
     *,
     pyformat_args: Optional[Dict[str, Any]] = None,
     dry_run: bool = False,
+    use_hybrid_engine: bool = False,
 ) -> bigframes.dataframe.DataFrame | pandas.Series:
     """A Colab-specific version of read_gbq.
 
@@ -312,7 +315,9 @@ def _read_gbq_colab(
         dry_run (bool):
             If True, estimates the query results size without returning data.
             The return will be a pandas Series with query metadata.
-
+        use_hybrid_engine (bool):
+            If True, and session not started, new session started will use
+            hybrid execution which pushes some execution to local cpu.
     Returns:
         Union[bigframes.dataframe.DataFrame, pandas.Series]:
             A BigQuery DataFrame if `dry_run` is False, otherwise a pandas Series.
@@ -345,6 +350,8 @@ def _read_gbq_colab(
             dry_run=True,
         )
         _set_default_session_location_if_possible_deferred_query(create_query)
+        if use_hybrid_engine and not config.options.bigquery._session_started:
+            config.options.bigquery.enable_polars_execution = True
 
     return global_session.with_default_session(
         bigframes.session.Session._read_gbq_colab,

tests/system/small/session/test_read_gbq_colab.py

@@ -19,6 +19,11 @@
 import pandas.testing
 import pytest
 
+import bigframes
+import bigframes.pandas
+
+pytest.importorskip("polars")
+
 
 def test_read_gbq_colab_to_pandas_batches_preserves_order_by(maybe_ordered_session):
     # This query should return enough results to be too big to fit in a single
@@ -59,6 +64,32 @@ def test_read_gbq_colab_to_pandas_batches_preserves_order_by(maybe_ordered_sessi
     assert executions_after == executions_before_python == executions_before_sql + 1
 
 
+def test_read_gbq_colab_fresh_session_is_hybrid():
+    bigframes.close_session()
+    df = bigframes.pandas._read_gbq_colab(
+        """
+        SELECT
+            name,
+            SUM(number) AS total
+        FROM
+            `bigquery-public-data.usa_names.usa_1910_2013`
+        WHERE state LIKE 'W%'
+        GROUP BY name
+        ORDER BY total DESC
+        LIMIT 300
+        """,
+        use_hybrid_engine=True,
+    )
+    session = df._session
+    executions_before_python = session._metrics.execution_count
+    result = df.sort_values("name").peek(100)
+    executions_after = session._metrics.execution_count
+
+    assert len(result) == 100
+    assert session._executor._enable_polars_execution is True  # type: ignore
+    assert executions_after == executions_before_python == 1
+
+
 def test_read_gbq_colab_peek_avoids_requery(maybe_ordered_session):
     executions_before_sql = maybe_ordered_session._metrics.execution_count
     df = maybe_ordered_session._read_gbq_colab(