@@ -218,8 +218,27 @@ def read_gbq(
218
218
read_gbq .__doc__ = inspect .getdoc (bigframes .session .Session .read_gbq )
219
219
220
220
221
+ def _run_read_gbq_colab_sessionless_dry_run (
222
+ query : str ,
223
+ * ,
224
+ pyformat_args : Dict [str , Any ],
225
+ ) -> pandas .Series :
226
+ """Run a dry_run without a session."""
227
+
228
+ query_formatted = bigframes .core .pyformat .pyformat (
229
+ query ,
230
+ pyformat_args = pyformat_args ,
231
+ dry_run = True ,
232
+ )
233
+ bqclient = _get_bqclient ()
234
+ job = _dry_run (query_formatted , bqclient )
235
+ return dry_runs .get_query_stats_with_inferred_dtypes (job , (), ())
236
+
237
+
221
238
def _try_read_gbq_colab_sessionless_dry_run (
222
- create_query : Callable [[], str ],
239
+ query : str ,
240
+ * ,
241
+ pyformat_args : Dict [str , Any ],
223
242
) -> Optional [pandas .Series ]:
224
243
"""Run a dry_run without a session, only if the session hasn't yet started."""
225
244
@@ -230,10 +249,9 @@ def _try_read_gbq_colab_sessionless_dry_run(
230
249
# to local data and not any BigQuery tables.
231
250
with _default_location_lock :
232
251
if not config .options .bigquery ._session_started :
233
- bqclient = _get_bqclient ()
234
- query = create_query ()
235
- job = _dry_run (query , bqclient )
236
- return dry_runs .get_query_stats_with_inferred_dtypes (job , (), ())
252
+ return _run_read_gbq_colab_sessionless_dry_run (
253
+ query , pyformat_args = pyformat_args
254
+ )
237
255
238
256
# Explicitly return None to indicate that we didn't run the dry run query.
239
257
return None
@@ -286,21 +304,13 @@ def _read_gbq_colab(
286
304
if pyformat_args is None :
287
305
pyformat_args = {}
288
306
289
- # Delay formatting the query with the special "session-less" logic. This
290
- # avoids doing unnecessary work if the session already has a location or has
291
- # already started.
292
- create_query = functools .partial (
293
- bigframes .core .pyformat .pyformat ,
294
- query_or_table ,
295
- pyformat_args = pyformat_args ,
296
- dry_run = True ,
297
- )
298
-
299
307
# Only try to set the global location if it's not a dry run. We don't want
300
308
# to bind to a location too early. This is especially important if the query
301
309
# only refers to local data and not any BigQuery tables.
302
310
if dry_run :
303
- result = _try_read_gbq_colab_sessionless_dry_run (create_query )
311
+ result = _try_read_gbq_colab_sessionless_dry_run (
312
+ query_or_table , pyformat_args = pyformat_args
313
+ )
304
314
305
315
if result is not None :
306
316
return result
@@ -309,6 +319,15 @@ def _read_gbq_colab(
309
319
# started. That means we can safely call the "real" _read_gbq_colab,
310
320
# which generates slightly nicer SQL.
311
321
else :
322
+ # Delay formatting the query with the special "session-less" logic. This
323
+ # avoids doing unnecessary work if the session already has a location or has
324
+ # already started.
325
+ create_query = functools .partial (
326
+ bigframes .core .pyformat .pyformat ,
327
+ query_or_table ,
328
+ pyformat_args = pyformat_args ,
329
+ dry_run = True ,
330
+ )
312
331
_set_default_session_location_if_possible_deferred_query (create_query )
313
332
314
333
return global_session .with_default_session (
0 commit comments