BigQuery: make jobs awaitable

[dkapitan]

I know BigQuery jobs are asynchronous by default. However, I am struggling to make my datapipeline async end-to-end.

Looking at this JS example, I thought it would be the most Pythonic to make a BigQuery job awaitable. However, I can't get that to work in Python i.e. errors when await client.query(query). Looking at the source code, I don't see which method returns an awaitable object.

I have little experience in writing async Python code and found this example that wraps jobs in a async def coroutine.

class BQApi(object):                                                                                                 
    def __init__(self):                                                                                              
        self.api = bigquery.Client.from_service_account_json(BQ_CONFIG["credentials"])                               

    async def exec_query(self, query, **kwargs) -> bigquery.table.RowIterator:                                       
        job = self.api.query(query, **kwargs)                                                                        
        task = asyncio.create_task(self.coroutine_job(job))                                                          
        return await task                                                                                            

    @staticmethod                                                                                                    
    async def coroutine_job(job):                                                                                    
        return job.result()   

The google.api_core.operation.Operation shows how to use add_done_callback to asynchronously wait for long-running operations. I have tried that, but the following yields AttributeError: 'QueryJob' object has no attribute '_condition' :

from concurrent.futures import ThreadPoolExecutor, as_completed
query1 = 'SELECT 1'
query2 = 'SELECT 2'

def my_callback(future):
    result = future.result()

operations = [bq.query(query1), bq.query(query2)]
[operation.add_done_callback(my_callback) for operation in operations]
results2 = []
for future in as_completed(operations):
  results2.append(list(future.result()))

Given that jobs are already asynchronous, would it make sense to add a method that returns an awaitable?

Or am I missing something and is there an Pythonic way to use the BigQuery client with the async/await pattern?

[tseaver]

@dkapitan Re: the async await / async def keywords: our codebase straddles both Python2 and Python3 until at least the end of this year, and thus cannot adopt the syntax directly in the codebase until we drop Python2 support.

We might be able to tweak the google.api_core.operation.Operation class to support that usage at a low level. Can you please post the traceback for the AttributeError?

Also, for your example above, your callback isn't doing anything with result: it is being called from a helper thread, and would need to do something to signal the main thread

[dkapitan]

@tseaver thanks for clarifying. I have read up more on concurrent.futures and asyncio. I understand the latter is indeed quite new and Python 3 only. Will investigate the sample code a bit more, and get to you with the results and/or the traceback.

[dkapitan]

@tseaver here's traceback from the AttributeError

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-19-e9025f68a561> in <module>
      9 [operation.add_done_callback(my_callback) for operation in operations]
     10 results2 = []
---> 11 for future in as_completed(operations):
     12   results2.append(list(future.result()))

/usr/local/Cellar/python/3.7.4/Frameworks/Python.framework/Versions/3.7/lib/python3.7/concurrent/futures/_base.py in as_completed(fs, timeout)
    217     fs = set(fs)
    218     total_futures = len(fs)
--> 219     with _AcquireFutures(fs):
    220         finished = set(
    221                 f for f in fs

/usr/local/Cellar/python/3.7.4/Frameworks/Python.framework/Versions/3.7/lib/python3.7/concurrent/futures/_base.py in __enter__(self)
    144     def __enter__(self):
    145         for future in self.futures:
--> 146             future._condition.acquire()
    147 
    148     def __exit__(self, *args):

AttributeError: 'QueryJob' object has no attribute '_condition'

[dkapitan]

@tseaver
For now, I have decided to do the following:

import numpy as np
from time import sleep


query1 = """
SELECT
  language.name,
  average(language.bytes)
FROM `bigquery-public-data.github_repos.languages` 
, UNNEST(language) AS language
GROUP BY language.name"""
query2 = 'SELECT 2'


def dummy_callback(future):
    global jobs_done
    jobs_done[future.job_id] = True


jobs = [bq.query(query1), bq.query(query2)]
jobs_done = {job.job_id: False for job in jobs}
[job.add_done_callback(dummy_callback) for job in jobs]

# blocking loop to wait for jobs to finish
while not (np.all(list(jobs_done.values()))):
    print('waiting for jobs to finish ... sleeping for 1s')
    sleep(1)

print('all jobs done, do your stuff')

Rather than using as_completed I prefer to use the built-in async functionality from the bigquery jobs themselves. This also makes it possible for me to decompose the datapipeline into separate Cloud Functions, without having to keep the main ThreadPoolExecutor live for the duration of the whole pipeline. Incidentally, this was the reason why I was looking into this: my pipelines are longer than the max timeout of 9 minutes for Cloud Functions (or even 15 minutes for Cloud Run).

Downside is I need to keep track of all the job_ids across the various functions, but that is relatively easy to solve when configuring the pipeline by specifying inputs and outputs such that they form a directed acyclic graph.

Any suggestions are welcome. Would be nice if at some point in the future, Google's api future plays nicely with Python's future (no pun intended).

[tseaver]

@dkapitan Thanks for the follow-up. I agree that it would be ideal if jobs could be awaited.

You could simplify your example code using a set, e.g.:

from time import sleep

query1 = """
SELECT
  language.name,
  average(language.bytes)
FROM `bigquery-public-data.github_repos.languages` 
, UNNEST(language) AS language
GROUP BY language.name"""

query2 = 'SELECT 2'

queries = [query_1, query_2]

awaiting_jobs = set()

def callback(future):
    awaiting_jobs.discard(future.job_id)

for query in queries:
    job = bq.query(query)
    awaiting_jobs.add(job.job_id)
    job.add_done_callback(callback)

while awaiting_jobs:
    print('waiting for jobs to finish ... sleeping for 1s')
    sleep(1)

print('all jobs done, do your stuff')

[dkapitan]

@tseaver
Nice, thanks. Never thought of using a set to do a kind of countdown ... 😄

[northtree]

@tseaver Thanks for the simplify example. Could you fix two typos?

awaiting_jobs.discard(future.job_id)

job = bq.query(query)

[tseaver]

@northtree Thanks for catching those: updated.

[sambcom]

How do you get the rows from your different queries? Won't this block on job.result() What am I missing?

[dkapitan]

@sambcom
apologies for the late response; forgot to subscribe to this issue.
In my use case, the results of the query are explicitly written to new tables, as part of an ELT pipeline. So that's no issue.

I understand that generically speaking, you could write query results to (temporary) tables in BigQuery, so this should not be a blocking issue.

[zbarnett]

I just ran into this problem myself and found out that even though the job itself is "asynchronous" there are actually two places where synchronous I/O is happening.

From the example usage on https://googleapis.dev/python/bigquery/latest/index.html

from google.cloud import bigquery

client = bigquery.Client()

# Perform a query.
QUERY = (
    'SELECT name FROM `bigquery-public-data.usa_names.usa_1910_2013` '
    'WHERE state = "TX" '
    'LIMIT 100')
query_job = client.query(QUERY)  # API request
rows = query_job.result()  # Waits for query to finish

for row in rows:
    print(row.name)

client.query(QUERY) makes a synchronous network request as does query_job.result()

Usually, blocking on result() takes the majority of the time, but I've seen cases where query() can take over a second to complete by itself.

As a workaround for now, I'm running all the BigQuery data fetching code in a separate thread to free up the event loop but it would be fantastic if the BigQuery API supported async I/O.

Example workaround (code fragment):

def get_data(client, query):
    return client.query(query).result()

loop = asyncio.get_running_loop()
data = await loop.run_in_executor(None, get_data, client, query)

[yan-hic]

We too have hundreds of load and dedupe queries running async. We do use futures but only to submit the jobs at once and get the id's.

From there we poke regularly the Stackdriver logging (using airflow) for the status of all in one API call.
Waiting for blocking result() is inefficient if you only want to know when jobs are done.

I wish get_job() could support multiple jobids but the API doesn't and an enhancement request has stalled.

[tswast]

We've identified EOL for new versions of the BigQuery client library as January 1, 2020 in #9036 We can revisit this request after that date.