Skip to content

Commit dbef14d

Browse files
plamutplamut
committed
Support both v1 stable and beta1 BQ Storage client
1 parent 7d7dc73 commit dbef14d

File tree

6 files changed

+498
-48
lines changed

6 files changed

+498
-48
lines changed

‎google/cloud/bigquery/_pandas_helpers.py‎

Lines changed: 55 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -577,7 +577,19 @@ def _bqstorage_page_to_dataframe(column_names, dtypes, page):
577577
def _download_table_bqstorage_stream(
578578
download_state, bqstorage_client, session, stream, worker_queue, page_to_item
579579
):
580-
rowstream = bqstorage_client.read_rows(stream.name).rows(session)
580+
# Passing a BQ Storage client in implies that the BigQuery Storage library
581+
# is available and can be imported.
582+
from google.cloud import bigquery_storage_v1beta1
583+
584+
# We want to preserve comaptibility with the v1beta1 BQ Storage clients,
585+
# thus adjust constructing the rowstream if needed.
586+
# The assumption is that the caller provides a BQ Storage `session` that is
587+
# compatible with the version of the BQ Storage client passed in.
588+
if isinstance(bqstorage_client, bigquery_storage_v1beta1.BigQueryStorageClient):
589+
position = bigquery_storage_v1beta1.types.StreamPosition(stream=stream)
590+
rowstream = bqstorage_client.read_rows(position).rows(session)
591+
else:
592+
rowstream = bqstorage_client.read_rows(stream.name).rows(session)
581593

582594
for page in rowstream.pages:
583595
if download_state.done:
@@ -609,29 +621,57 @@ def _download_table_bqstorage(
609621
page_to_item=None,
610622
):
611623
"""Use (faster, but billable) BQ Storage API to construct DataFrame."""
624+
625+
# Passing a BQ Storage client in implies that the BigQuery Storage library
626+
# is available and can be imported.
627+
from google.cloud import bigquery_storage_v1
628+
from google.cloud import bigquery_storage_v1beta1
629+
612630
if "$" in table.table_id:
613631
raise ValueError(
614632
"Reading from a specific partition is not currently supported."
615633
)
616634
if "@" in table.table_id:
617635
raise ValueError("Reading from a specific snapshot is not currently supported.")
618636

619-
requested_session = bigquery_storage_v1.types.ReadSession(
620-
table=table.to_bqstorage(),
621-
data_format=bigquery_storage_v1.enums.DataFormat.ARROW,
622-
)
623-
624-
if selected_fields is not None:
625-
for field in selected_fields:
626-
requested_session.read_options.selected_fields.append(field.name)
627-
628637
requested_streams = 1 if preserve_order else 0
629638

630-
session = bqstorage_client.create_read_session(
631-
parent="projects/{}".format(project_id),
632-
read_session=requested_session,
633-
max_stream_count=requested_streams,
634-
)
639+
# We want to preserve comaptibility with the v1beta1 BQ Storage clients,
640+
# thus adjust the session creation if needed.
641+
if isinstance(bqstorage_client, bigquery_storage_v1beta1.BigQueryStorageClient):
642+
warnings.warn(
643+
"Support for BigQuery Storage v1beta1 clients is deprecated, please "
644+
"consider upgrading the client to BigQuery Storage v1 stable version.",
645+
category=DeprecationWarning,
646+
)
647+
read_options = bigquery_storage_v1beta1.types.TableReadOptions()
648+
649+
if selected_fields is not None:
650+
for field in selected_fields:
651+
read_options.selected_fields.append(field.name)
652+
653+
session = bqstorage_client.create_read_session(
654+
table.to_bqstorage(v1beta1=True),
655+
"projects/{}".format(project_id),
656+
format_=bigquery_storage_v1beta1.enums.DataFormat.ARROW,
657+
read_options=read_options,
658+
requested_streams=requested_streams,
659+
)
660+
else:
661+
requested_session = bigquery_storage_v1.types.ReadSession(
662+
table=table.to_bqstorage(),
663+
data_format=bigquery_storage_v1.enums.DataFormat.ARROW,
664+
)
665+
if selected_fields is not None:
666+
for field in selected_fields:
667+
requested_session.read_options.selected_fields.append(field.name)
668+
669+
session = bqstorage_client.create_read_session(
670+
parent="projects/{}".format(project_id),
671+
read_session=requested_session,
672+
max_stream_count=requested_streams,
673+
)
674+
635675
_LOGGER.debug(
636676
"Started reading table '{}.{}.{}' with BQ Storage API session '{}'.".format(
637677
table.project, table.dataset_id, table.table_id, session.name

‎google/cloud/bigquery/dbapi/cursor.py‎

Lines changed: 39 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
"""Cursor for the Google BigQuery DB-API."""
1616

1717
import collections
18+
import warnings
1819

1920
try:
2021
from collections import abc as collections_abc
@@ -268,28 +269,54 @@ def _bqstorage_fetch(self, bqstorage_client):
268269
A sequence of rows, represented as dictionaries.
269270
"""
270271
# Hitting this code path with a BQ Storage client instance implies that
271-
# bigquery_storage_v1 can indeed be imported here without errors.
272+
# bigquery_storage_v1* can indeed be imported here without errors.
272273
from google.cloud import bigquery_storage_v1
274+
from google.cloud import bigquery_storage_v1beta1
273275

274276
table_reference = self._query_job.destination
275277

276-
requested_session = bigquery_storage_v1.types.ReadSession(
277-
table=table_reference.to_bqstorage(),
278-
data_format=bigquery_storage_v1.enums.DataFormat.AVRO,
278+
is_v1beta1_client = isinstance(
279+
bqstorage_client, bigquery_storage_v1beta1.BigQueryStorageClient
279280
)
280281

281-
read_session = bqstorage_client.create_read_session(
282-
parent="projects/{}".format(table_reference.project),
283-
read_session=requested_session,
284-
# a single stream only, as DB API is not well-suited for multithreading
285-
max_stream_count=1,
286-
)
282+
# We want to preserve comaptibility with the v1beta1 BQ Storage clients,
283+
# thus adjust the session creation if needed.
284+
if is_v1beta1_client:
285+
warnings.warn(
286+
"Support for BigQuery Storage v1beta1 clients is deprecated, please "
287+
"consider upgrading the client to BigQuery Storage v1 stable version.",
288+
category=DeprecationWarning,
289+
)
290+
read_session = bqstorage_client.create_read_session(
291+
table_reference.to_bqstorage(v1beta1=True),
292+
"projects/{}".format(table_reference.project),
293+
# a single stream only, as DB API is not well-suited for multithreading
294+
requested_streams=1,
295+
)
296+
else:
297+
requested_session = bigquery_storage_v1.types.ReadSession(
298+
table=table_reference.to_bqstorage(),
299+
data_format=bigquery_storage_v1.enums.DataFormat.AVRO,
300+
)
301+
read_session = bqstorage_client.create_read_session(
302+
parent="projects/{}".format(table_reference.project),
303+
read_session=requested_session,
304+
# a single stream only, as DB API is not well-suited for multithreading
305+
max_stream_count=1,
306+
)
287307

288308
if not read_session.streams:
289309
return iter([]) # empty table, nothing to read
290310

291-
stream_name = read_session.streams[0].name
292-
read_rows_stream = bqstorage_client.read_rows(stream_name)
311+
if is_v1beta1_client:
312+
read_position = bigquery_storage_v1beta1.types.StreamPosition(
313+
stream=read_session.streams[0],
314+
)
315+
read_rows_stream = bqstorage_client.read_rows(read_position)
316+
else:
317+
stream_name = read_session.streams[0].name
318+
read_rows_stream = bqstorage_client.read_rows(stream_name)
319+
293320
rows_iterable = read_rows_stream.rows(read_session)
294321
return rows_iterable
295322

‎google/cloud/bigquery/table.py‎

Lines changed: 51 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,12 @@
2525

2626
import six
2727

28+
try:
29+
# Needed for the to_bqstorage() method.
30+
from google.cloud import bigquery_storage_v1beta1
31+
except ImportError: # pragma: NO COVER
32+
bigquery_storage_v1beta1 = None
33+
2834
try:
2935
import pandas
3036
except ImportError: # pragma: NO COVER
@@ -221,7 +227,7 @@ def to_api_repr(self):
221227
"tableId": self._table_id,
222228
}
223229

224-
def to_bqstorage(self):
230+
def to_bqstorage(self, v1beta1=False):
225231
"""Construct a BigQuery Storage API representation of this table.
226232
227233
Install the ``google-cloud-bigquery-storage`` package to use this
@@ -235,15 +241,37 @@ def to_bqstorage(self):
235241
:class:`google.cloud.bigquery_storage_v1.types.ReadSession.TableModifiers`
236242
to select a specific snapshot to read from.
237243
244+
Args:
245+
v1beta1 (Optiona[bool]):
246+
If :data:`True`, return representation compatible with BigQuery
247+
Storage ``v1beta1`` version. Defaults to :data:`False`.
248+
238249
Returns:
239-
str: A reference to this table in the BigQuery Storage API.
250+
Union[str, google.cloud.bigquery_storage_v1beta1.types.TableReference:]:
251+
A reference to this table in the BigQuery Storage API.
252+
253+
Raises:
254+
ValueError:
255+
If ``v1beta1`` compatibility is requested, but the
256+
:mod:`google.cloud.bigquery_storage_v1beta1` module cannot be imported.
240257
"""
258+
if v1beta1 and bigquery_storage_v1beta1 is None:
259+
raise ValueError(_NO_BQSTORAGE_ERROR)
260+
241261
table_id, _, _ = self._table_id.partition("@")
242262
table_id, _, _ = table_id.partition("$")
243263

244-
table_ref = "projects/{}/datasets/{}/tables/{}".format(
245-
self._project, self._dataset_id, table_id,
246-
)
264+
if v1beta1:
265+
table_ref = bigquery_storage_v1beta1.types.TableReference(
266+
project_id=self._project,
267+
dataset_id=self._dataset_id,
268+
table_id=table_id,
269+
)
270+
else:
271+
table_ref = "projects/{}/datasets/{}/tables/{}".format(
272+
self._project, self._dataset_id, table_id,
273+
)
274+
247275
return table_ref
248276

249277
def _key(self):
@@ -849,13 +877,19 @@ def to_api_repr(self):
849877
"""
850878
return copy.deepcopy(self._properties)
851879

852-
def to_bqstorage(self):
880+
def to_bqstorage(self, v1beta1=False):
853881
"""Construct a BigQuery Storage API representation of this table.
854882
883+
Args:
884+
v1beta1 (Optiona[bool]):
885+
If :data:`True`, return representation compatible with BigQuery
886+
Storage ``v1beta1`` version. Defaults to :data:`False`.
887+
855888
Returns:
856-
str: A reference to this table in the BigQuery Storage API.
889+
Union[str, google.cloud.bigquery_storage_v1beta1.types.TableReference:]:
890+
A reference to this table in the BigQuery Storage API.
857891
"""
858-
return self.reference.to_bqstorage()
892+
return self.reference.to_bqstorage(v1beta1=v1beta1)
859893

860894
def _build_resource(self, filter_fields):
861895
"""Generate a resource for ``update``."""
@@ -1063,13 +1097,19 @@ def from_string(cls, full_table_id):
10631097
{"tableReference": TableReference.from_string(full_table_id).to_api_repr()}
10641098
)
10651099

1066-
def to_bqstorage(self):
1100+
def to_bqstorage(self, v1beta1=False):
10671101
"""Construct a BigQuery Storage API representation of this table.
10681102
1103+
Args:
1104+
v1beta1 (Optiona[bool]):
1105+
If :data:`True`, return representation compatible with BigQuery
1106+
Storage ``v1beta1`` version. Defaults to :data:`False`.
1107+
10691108
Returns:
1070-
str: A reference to this table in the BigQuery Storage API.
1109+
Union[str, google.cloud.bigquery_storage_v1beta1.types.TableReference:]:
1110+
A reference to this table in the BigQuery Storage API.
10711111
"""
1072-
return self.reference.to_bqstorage()
1112+
return self.reference.to_bqstorage(v1beta1=v1beta1)
10731113

10741114

10751115
def _row_from_mapping(mapping, schema):

0 commit comments

Comments
 (0)