Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -160,23 +160,13 @@ def download_arrow_from_job(
if pyarrow is None:
raise ValueError(_NO_PYARROW_ERROR)

# Step 1: Ensure BigQuery Read API client is available upfront.
if bqstorage_client is None:
if row_iterator.client is None:
raise ValueError("RowIterator client is None.")
bqstorage_client = row_iterator.client._ensure_bqstorage_client()
if bqstorage_client is None:
raise ValueError(
"The google-cloud-bigquery-storage library is required to read Arrow results."
)

offset = 0
pa_schema = None
total_rows = row_iterator.total_rows
job_complete = False
remaining = row_iterator.max_results

# Step 2: Process initial inline Arrow response from jobs.query if available.
# Step 1: Process initial inline Arrow response from jobs.query if available.
if row_iterator._first_page_response:
first_page = row_iterator._first_page_response

Expand Down Expand Up @@ -206,12 +196,22 @@ def download_arrow_from_job(
batch, remaining = _slice_batch(batch, remaining)
yield batch

# Step 3: Return early if all results were delivered in the first page
# Step 2: Return early if all results were delivered in the first page
# response, or if the first page already reached max_results.
if (job_complete and offset >= total_rows) or remaining == 0:
return

# Step 4: Stream remaining Arrow record batches from the job default stream.
# Step 3: Stream remaining Arrow record batches from the job default stream,
# with the BigQuery client's shared BigQuery Storage client unless the
# caller passed one. Results that fit in the first page never create it.
if bqstorage_client is None:
if row_iterator.client is None:
raise ValueError("RowIterator client is None.")
bqstorage_client = row_iterator.client._get_shared_bqstorage_client()
if bqstorage_client is None:
raise ValueError(
"The google-cloud-bigquery-storage library is required to read Arrow results."
)
project = row_iterator._project or (
row_iterator.client.project if row_iterator.client else None
)
Expand Down
77 changes: 77 additions & 0 deletions packages/google-cloud-bigquery/google/cloud/bigquery/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import math
import os
import tempfile
import threading
import typing
import uuid
import warnings
Expand Down Expand Up @@ -166,6 +167,35 @@
# https://github.com/googleapis/python-bigquery/issues/438
_MIN_GET_QUERY_RESULTS_TIMEOUT = 120

# While a stream is open and the connection is idle, the shared BigQuery
# Storage client pings the server after a minute, and closes the connection
# if a ping isn't answered within a minute. A dropped ReadRows stream is
# resumed at the next unread row.
_BQSTORAGE_KEEPALIVE_TIME_MS = 60 * 1000
_BQSTORAGE_KEEPALIVE_TIMEOUT_MS = 60 * 1000


def _create_bqstorage_transport(**kwargs):
"""Create a BigQuery Storage gRPC transport whose channel sends keepalive
pings.

``BigQueryReadClient`` calls this with the arguments it passes to its
default transport.
"""
from google.cloud import bigquery_storage

transport_class = bigquery_storage.BigQueryReadClient.get_transport_class("grpc")

def create_channel(*args, options=(), **channel_kwargs):
options = list(options) + [
("grpc.keepalive_time_ms", _BQSTORAGE_KEEPALIVE_TIME_MS),
("grpc.keepalive_timeout_ms", _BQSTORAGE_KEEPALIVE_TIMEOUT_MS),
]
return transport_class.create_channel(*args, options=options, **channel_kwargs)

return transport_class(channel=create_channel, **kwargs)


_LOAD_TABLE_FROM_DATAFRAME_DEPRECATED = (
"Loading DataFrames via google-cloud-bigquery is deprecated. "
"For direct, optimized loading, please call 'pandas_gbq.to_gbq()' directly."
Expand Down Expand Up @@ -300,6 +330,14 @@ def __init__(
self._default_load_job_config = copy.deepcopy(default_load_job_config)
self.default_job_creation_mode = default_job_creation_mode

# ARROW query results share one BigQuery Storage client. It's created
# the first time a query needs a stream, and close() closes it.
self._universe_domain = _get_client_universe(client_options)
self._shared_bqstorage_client: Optional[
"google.cloud.bigquery_storage.BigQueryReadClient"
] = None
self._shared_bqstorage_client_lock = threading.Lock()

# Use property setter so validation can run.
self.default_query_job_config = default_query_job_config

Expand Down Expand Up @@ -348,6 +386,9 @@ def default_load_job_config(self, value: LoadJobConfig):
def close(self):
"""Close the underlying transport objects, releasing system resources.

This also closes the BigQuery Storage client that ``ARROW`` query
results share.

.. note::

The client instance can be used for making additional requests even
Expand All @@ -356,6 +397,11 @@ def close(self):
"""
self._http._auth_request.session.close()
self._http.close()
with self._shared_bqstorage_client_lock:
bqstorage_client = self._shared_bqstorage_client
self._shared_bqstorage_client = None
if bqstorage_client is not None:
bqstorage_client._transport.close()

def get_service_account_email(
self,
Expand Down Expand Up @@ -600,6 +646,7 @@ def _ensure_bqstorage_client(
client_info: Optional[
"google.api_core.gapic_v1.client_info.ClientInfo"
] = DEFAULT_BQSTORAGE_CLIENT_INFO,
transport: Optional[Callable[..., Any]] = None,
) -> Optional["google.cloud.bigquery_storage.BigQueryReadClient"]:
"""Create a BigQuery Storage API client using this client's credentials.

Expand All @@ -613,6 +660,9 @@ def _ensure_bqstorage_client(
client_info:
The client info used with a new BigQuery Storage client
instance if one is created.
transport:
A callable that creates the transport of a new BigQuery Storage
client instance. If ``None``, the default transport is used.

Returns:
A BigQuery Storage API client.
Expand Down Expand Up @@ -680,10 +730,37 @@ def _ensure_bqstorage_client(
credentials=self._credentials,
client_options=client_options,
client_info=amended_client_info,
transport=transport,
)

return bqstorage_client

def _get_shared_bqstorage_client(
self,
) -> Optional["google.cloud.bigquery_storage.BigQueryReadClient"]:
"""Return the BigQuery Storage client that ``ARROW`` query results share.

The client is created on first use, with this client's credentials and
universe domain and a gRPC channel that sends keepalive pings.
:meth:`close` closes it, and the next call creates a new one.

Returns:
A BigQuery Storage API client, or ``None`` if
google-cloud-bigquery-storage is missing or outdated.
"""
with self._shared_bqstorage_client_lock:
if self._shared_bqstorage_client is None:
client_options = None
if self._universe_domain != _DEFAULT_UNIVERSE:
client_options = google.api_core.client_options.ClientOptions(
universe_domain=self._universe_domain
)
self._shared_bqstorage_client = self._ensure_bqstorage_client(
client_options=client_options,
transport=_create_bqstorage_transport,
)
return self._shared_bqstorage_client

def _dataset_from_arg(self, dataset) -> Union[Dataset, DatasetReference]:
if isinstance(dataset, str):
dataset = DatasetReference.from_string(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2323,7 +2323,8 @@ def _download_arrow_from_job_id(

Args:
bqstorage_client (Optional[bigquery_storage.BigQueryReadClient]):
Client for BigQuery Storage Read API. If None, one will be created.
Client for BigQuery Storage Read API. If None, the BigQuery
Storage client that the BigQuery client shares is used.
timeout (Optional[float]):
Timeout in seconds for read operations.

Expand Down
66 changes: 66 additions & 0 deletions packages/google-cloud-bigquery/tests/system/test_arrow.py
Original file line number Diff line number Diff line change
Expand Up @@ -755,3 +755,69 @@ def test_query_and_wait_arrow_max_results(
table = pyarrow.Table.from_batches(batches)
assert table.column("x").to_pylist() == list(range(1, expected_rows + 1))
assert read_rows.called == streamed


def _query_arrow(client, row_count):
return client.query_and_wait(
f"SELECT x FROM UNNEST(GENERATE_ARRAY(1, {row_count})) AS x ORDER BY x",
job_config=bigquery.QueryJobConfig(use_query_cache=False),
page_size=100,
query_results_format=enums.QueryResultsFormat.ARROW,
)


def test_query_and_wait_arrow_queries_share_one_bqstorage_client():
from google.cloud.bigquery import client as client_module

with bigquery.Client() as client:
# Wrap the client factory to count how many clients get created.
ensure = mock.Mock(wraps=client._ensure_bqstorage_client)
with mock.patch.object(client, "_ensure_bqstorage_client", ensure):
tables = [
pyarrow.Table.from_batches(
list(_query_arrow(client, 1000).to_arrow_iterable())
),
_query_arrow(client, 1000).to_arrow(),
pyarrow.Table.from_batches(
list(_query_arrow(client, 1000).to_arrow_iterable())
),
]
shared = client._shared_bqstorage_client

for table in tables:
assert table.column("x").to_pylist() == list(range(1, 1001))
ensure.assert_called_once_with(
client_options=None, transport=client_module._create_bqstorage_transport
)
assert shared is not None
# Leaving the with block calls close(), which closes the shared channel.
assert client._shared_bqstorage_client is None
with pytest.raises(ValueError, match="(?i)closed"):
shared.read_rows("projects/p/locations/US/jobs/j/streams/_default")


def test_query_and_wait_arrow_ending_a_read_early_cancels_only_its_stream():
from google.cloud.bigquery_storage_v1 import reader

# Wrap ReadRowsStream.cancel() to count calls without changing it.
cancel_stream = reader.ReadRowsStream.cancel
with bigquery.Client() as client, mock.patch.object(
reader.ReadRowsStream, "cancel", autospec=True, side_effect=cancel_stream
) as cancel:
ended_early_rows = _query_arrow(client, 100000)
ended_early = ended_early_rows.to_arrow_iterable()
read_fully = _query_arrow(client, 100000).to_arrow_iterable()
# Two reads reach the stream, whether or not jobs.query returned a
# first page.
next(ended_early)
next(ended_early)
rows = [next(read_fully), next(read_fully)]

ended_early.close()
rows.extend(read_fully)

(stream,) = cancel.call_args.args
cancel.assert_called_once_with(stream)
assert stream._name.endswith(f"/jobs/{ended_early_rows.job_id}/streams/_default")
table = pyarrow.Table.from_batches(rows)
assert table.column("x").to_pylist() == list(range(1, 100001))
Loading
Loading