Skip to content
This repository was archived by the owner on Apr 1, 2026. It is now read-only.

Commit 921fb7a

Browse files
feat: _read_gbq_colab creates hybrid session
1 parent 05cb7d0 commit 921fb7a

2 files changed

Lines changed: 30 additions & 0 deletions

File tree

bigframes/pandas/io/api.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -345,6 +345,8 @@ def _read_gbq_colab(
345345
dry_run=True,
346346
)
347347
_set_default_session_location_if_possible_deferred_query(create_query)
348+
if not config.options.bigquery._session_started:
349+
config.options.bigquery.enable_polars_execution = True
348350

349351
return global_session.with_default_session(
350352
bigframes.session.Session._read_gbq_colab,

tests/system/small/session/test_read_gbq_colab.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,9 @@
1919
import pandas.testing
2020
import pytest
2121

22+
import bigframes
23+
import bigframes.pandas
24+
2225

2326
def test_read_gbq_colab_to_pandas_batches_preserves_order_by(maybe_ordered_session):
2427
# This query should return enough results to be too big to fit in a single
@@ -59,6 +62,31 @@ def test_read_gbq_colab_to_pandas_batches_preserves_order_by(maybe_ordered_sessi
5962
assert executions_after == executions_before_python == executions_before_sql + 1
6063

6164

65+
def test_read_gbq_colab_fresh_session_is_hybrid():
66+
bigframes.close_session()
67+
df = bigframes.pandas._read_gbq_colab(
68+
"""
69+
SELECT
70+
name,
71+
SUM(number) AS total
72+
FROM
73+
`bigquery-public-data.usa_names.usa_1910_2013`
74+
WHERE state LIKE 'W%'
75+
GROUP BY name
76+
ORDER BY total DESC
77+
LIMIT 300
78+
"""
79+
)
80+
session = df._session
81+
executions_before_python = session._metrics.execution_count
82+
result = df.sort_values("name").peek(100)
83+
executions_after = session._metrics.execution_count
84+
85+
assert len(result) == 100
86+
assert session._executor._enable_polars_execution is True # type: ignore
87+
assert executions_after == executions_before_python == 1
88+
89+
6290
def test_read_gbq_colab_peek_avoids_requery(maybe_ordered_session):
6391
executions_before_sql = maybe_ordered_session._metrics.execution_count
6492
df = maybe_ordered_session._read_gbq_colab(

0 commit comments

Comments
 (0)