This repository was archived by the owner on Apr 1, 2026. It is now read-only.
File tree Expand file tree Collapse file tree
tests/system/small/session Expand file tree Collapse file tree Original file line number Diff line number Diff 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 ,
Original file line number Diff line number Diff line change 1919import pandas .testing
2020import pytest
2121
22+ import bigframes
23+ import bigframes .pandas
24+
2225
2326def 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+
6290def 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 (
You can’t perform that action at this time.
0 commit comments