Skip to content

Commit f2e76e6

Browse files
ohmayrtswast
authored andcommitted
feat(bigquery): read Arrow results from explicit destinations
The job's default read stream rejects queries whose job config sets a destination table ("This API does not support explicitly-specified destination tables"), so query_and_wait() with ARROW results failed for those queries. When the query wrote to an explicit destination, read the table through create_read_session() instead, with data_format=ARROW and max_stream_count=1 so rows stay in order. The session's read options carry the same serialization options as jobs.query. An empty table has no streams and yields nothing. WRITE_APPEND destinations are already rejected by validate_query().
1 parent 68d6bfb commit f2e76e6

5 files changed

Lines changed: 252 additions & 18 deletions

File tree

‎packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py‎

Lines changed: 82 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,76 @@ def _read_rows_options(options: Optional[Dict[str, str]]) -> Dict[str, str]:
144144
}
145145

146146

147+
def _open_job_stream(
148+
row_iterator: "RowIterator",
149+
bqstorage_client: "bigquery_storage.BigQueryReadClient",
150+
offset: int,
151+
timeout: Optional[float],
152+
):
153+
"""Open the job's default stream, starting at row ``offset``."""
154+
project = row_iterator._project or (
155+
row_iterator.client.project if row_iterator.client else None
156+
)
157+
location = row_iterator._location or (
158+
row_iterator.client.location if row_iterator.client else None
159+
)
160+
stream_name = f"projects/{project}/locations/{location}/jobs/{row_iterator._job_id}/streams/_default"
161+
return bqstorage_client.read_rows(
162+
stream_name,
163+
offset=offset,
164+
timeout=timeout,
165+
arrow_serialization_options=_read_rows_options(
166+
row_iterator._arrow_serialization_options
167+
),
168+
)
169+
170+
171+
def _open_destination_stream(
172+
row_iterator: "RowIterator",
173+
bqstorage_client: "bigquery_storage.BigQueryReadClient",
174+
timeout: Optional[float],
175+
pa_schema: Optional["pyarrow.Schema"],
176+
):
177+
"""Open a read session on the query's explicit destination table.
178+
179+
The default job stream rejects jobs with an explicit destination table
180+
("This API does not support explicitly-specified destination tables"), so
181+
read the table instead. One stream keeps the rows in order, and the
182+
session's read options carry the same serialization options as
183+
``jobs.query``. Explicit destinations always go through ``jobs.insert``,
184+
so there's no first page and the stream starts at row 0.
185+
186+
Returns ``(reader, pa_schema)``, with ``reader`` set to ``None`` if the
187+
table is empty and the session has no streams.
188+
"""
189+
from google.cloud import bigquery_storage
190+
191+
session = bqstorage_client.create_read_session(
192+
parent=f"projects/{row_iterator._billing_project}",
193+
read_session=bigquery_storage.types.stream.ReadSession(
194+
table=row_iterator._table.to_bqstorage(),
195+
data_format=bigquery_storage.types.stream.DataFormat.ARROW,
196+
read_options=bigquery_storage.types.stream.ReadSession.TableReadOptions(
197+
arrow_serialization_options=_read_rows_options(
198+
row_iterator._arrow_serialization_options
199+
)
200+
),
201+
),
202+
max_stream_count=1,
203+
timeout=timeout,
204+
)
205+
if pa_schema is None and session.arrow_schema.serialized_schema:
206+
pa_schema = pyarrow.ipc.read_schema(
207+
pyarrow.py_buffer(session.arrow_schema.serialized_schema)
208+
)
209+
if not session.streams:
210+
return None, pa_schema
211+
return (
212+
bqstorage_client.read_rows(session.streams[0].name, offset=0, timeout=timeout),
213+
pa_schema,
214+
)
215+
216+
147217
def download_arrow_from_job(
148218
row_iterator: "RowIterator",
149219
bqstorage_client: Optional["bigquery_storage.BigQueryReadClient"] = None,
@@ -204,8 +274,9 @@ def download_arrow_from_job(
204274
return
205275

206276
# Step 3: Stream remaining Arrow record batches from the job default stream,
207-
# with the BigQuery client's shared BigQuery Storage client unless the
208-
# caller passed one. Results that fit in the first page never create it.
277+
# or from the destination table if the job config set one, with the
278+
# BigQuery client's shared BigQuery Storage client unless the caller passed
279+
# one. Results that fit in the first page never create it.
209280
if bqstorage_client is None:
210281
if row_iterator.client is None:
211282
raise ValueError("RowIterator client is None.")
@@ -214,21 +285,15 @@ def download_arrow_from_job(
214285
raise ValueError(
215286
"The google-cloud-bigquery-storage library is required to read Arrow results."
216287
)
217-
project = row_iterator._project or (
218-
row_iterator.client.project if row_iterator.client else None
219-
)
220-
location = row_iterator._location or (
221-
row_iterator.client.location if row_iterator.client else None
222-
)
223-
stream_name = f"projects/{project}/locations/{location}/jobs/{row_iterator._job_id}/streams/_default"
224-
reader = bqstorage_client.read_rows(
225-
stream_name,
226-
offset=offset,
227-
timeout=timeout,
228-
arrow_serialization_options=_read_rows_options(
229-
row_iterator._arrow_serialization_options
230-
),
231-
)
288+
if row_iterator._explicit_destination and row_iterator._table is not None:
289+
reader, pa_schema = _open_destination_stream(
290+
row_iterator, bqstorage_client, timeout, pa_schema
291+
)
292+
else:
293+
reader = _open_job_stream(row_iterator, bqstorage_client, offset, timeout)
294+
if reader is None:
295+
# An empty destination table has no streams.
296+
return
232297
finished = False
233298
try:
234299
for response in reader:

‎packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -561,6 +561,9 @@ def query_and_wait(
561561
max_results=max_results,
562562
query_results_format=query_results_format,
563563
arrow_serialization_options=arrow_serialization_options,
564+
explicit_destination=(
565+
job_config is not None and job_config.destination is not None
566+
),
564567
callback=callback,
565568
)
566569

@@ -739,6 +742,7 @@ def _wait_or_cancel(
739742
*,
740743
query_results_format: Optional[str] = None,
741744
arrow_serialization_options: Optional[Dict[str, str]] = None,
745+
explicit_destination: bool = False,
742746
callback: Callable = lambda _: None,
743747
) -> table.RowIterator:
744748
"""Wait for a job to complete and return the results.
@@ -785,6 +789,7 @@ def _wait_or_cancel(
785789
)
786790
query_results._query_results_format = query_results_format
787791
query_results._arrow_serialization_options = arrow_serialization_options
792+
query_results._explicit_destination = explicit_destination
788793
return query_results
789794
except Exception:
790795
# Attempt to cancel the job since we can't return the results.

‎packages/google-cloud-bigquery/google/cloud/bigquery/table.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1978,6 +1978,9 @@ def __init__(
19781978
self._job_ended = ended
19791979
self._query_results_format = query_results_format
19801980
self._arrow_serialization_options = arrow_serialization_options
1981+
# Set by query_and_wait() if the query wrote to a destination table
1982+
# set in its job config.
1983+
self._explicit_destination = False
19811984

19821985
@property
19831986
@_disallow_when_arrow

‎packages/google-cloud-bigquery/tests/system/test_arrow.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -841,3 +841,42 @@ def record_open_cancel(stream):
841841
assert stream._name.endswith(f"/jobs/{ended_early_rows.job_id}/streams/_default")
842842
table = pyarrow.Table.from_batches(rows)
843843
assert table.column("x").to_pylist() == list(range(1, 100001))
844+
845+
846+
@pytest.mark.parametrize(
847+
"write_disposition",
848+
[None, enums.WriteDisposition.WRITE_TRUNCATE],
849+
ids=["default", "overwrite"],
850+
)
851+
def test_query_and_wait_arrow_explicit_destination(
852+
bigquery_client, bqstorage_client, dataset_id, write_disposition
853+
):
854+
table_name = f"arrow_destination_{write_disposition or 'default'}".lower()
855+
table_id = f"{bigquery_client.project}.{dataset_id}.{table_name}"
856+
job_config = bigquery.QueryJobConfig(destination=table_id)
857+
if write_disposition is not None:
858+
# The rows already in the table are replaced, not read.
859+
bigquery_client.query_and_wait(
860+
f"CREATE OR REPLACE TABLE `{table_id}` AS SELECT -1 AS x"
861+
)
862+
job_config.write_disposition = write_disposition
863+
# Wrap create_read_session() to check how the table is read.
864+
create_read_session = mock.Mock(wraps=bqstorage_client.create_read_session)
865+
866+
rows = bigquery_client.query_and_wait(
867+
"SELECT x FROM UNNEST(GENERATE_ARRAY(1, 1000)) AS x ORDER BY x",
868+
job_config=job_config,
869+
query_results_format=enums.QueryResultsFormat.ARROW,
870+
)
871+
with mock.patch.object(
872+
bqstorage_client, "create_read_session", create_read_session
873+
):
874+
batches = list(rows.to_arrow_iterable(bqstorage_client=bqstorage_client))
875+
876+
table = pyarrow.Table.from_batches(batches)
877+
assert table.column("x").to_pylist() == list(range(1, 1001))
878+
kwargs = create_read_session.call_args.kwargs
879+
assert kwargs["read_session"].table == (
880+
f"projects/{bigquery_client.project}/datasets/{dataset_id}/tables/{table_name}"
881+
)
882+
assert kwargs["max_stream_count"] == 1

‎packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py‎

Lines changed: 123 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@
3030
from google.cloud.bigquery.enums import TimestampPrecision
3131
from google.cloud.bigquery.enums import WriteDisposition
3232
from google.cloud.bigquery.job import QueryJobConfig
33-
from google.cloud.bigquery.table import RowIterator, _EmptyRowIterator
33+
from google.cloud.bigquery.table import RowIterator, TableReference, _EmptyRowIterator
3434

3535
from .helpers import make_client
3636

@@ -954,6 +954,7 @@ def test_query_and_wait_arrow_jobs_insert_omits_arrow_fields():
954954
max_results=mock.ANY,
955955
query_results_format="ARROW",
956956
arrow_serialization_options=_DEFAULT_QUERY_OPTIONS,
957+
explicit_destination=True,
957958
callback=mock.ANY,
958959
)
959960
assert wait_or_cancel.call_args.kwargs["page_size"] == 0
@@ -1356,6 +1357,127 @@ def test_to_arrow_with_arrow_format_reads_with_the_shared_client():
13561357
assert to_arrow_iterable.call_args.kwargs["bqstorage_client"] is None
13571358

13581359

1360+
@pytest.mark.parametrize(
1361+
("job_config", "expected"),
1362+
[
1363+
(QueryJobConfig(destination="p.d.t"), True),
1364+
# jobs.query doesn't accept a priority either, so this also runs
1365+
# through jobs.insert, without a destination.
1366+
(QueryJobConfig(priority="INTERACTIVE"), False),
1367+
],
1368+
ids=["destination", "no_destination"],
1369+
)
1370+
def test_query_and_wait_arrow_marks_explicit_destination(job_config, expected):
1371+
client = mock.MagicMock(spec=Client)
1372+
1373+
with (
1374+
mock.patch.object(_job_helpers, "query_jobs_insert") as query_jobs_insert,
1375+
mock.patch.object(_job_helpers, "_wait_or_cancel") as wait_or_cancel,
1376+
):
1377+
_query_and_wait(client, job_config=job_config, query_results_format="ARROW")
1378+
1379+
wait_or_cancel.assert_called_once_with(
1380+
query_jobs_insert.return_value,
1381+
api_timeout=mock.ANY,
1382+
wait_timeout=mock.ANY,
1383+
retry=mock.ANY,
1384+
page_size=mock.ANY,
1385+
max_results=mock.ANY,
1386+
query_results_format="ARROW",
1387+
arrow_serialization_options=mock.ANY,
1388+
explicit_destination=expected,
1389+
callback=mock.ANY,
1390+
)
1391+
1392+
1393+
def test_wait_or_cancel_sets_explicit_destination():
1394+
job = mock.MagicMock(dry_run=False)
1395+
job.result.return_value = RowIterator(
1396+
client=mock.MagicMock(), api_request=mock.MagicMock(), path=None, schema=()
1397+
)
1398+
assert not job.result.return_value._explicit_destination
1399+
1400+
rows = _job_helpers._wait_or_cancel(
1401+
job,
1402+
api_timeout=None,
1403+
wait_timeout=None,
1404+
retry=None,
1405+
page_size=0,
1406+
max_results=None,
1407+
explicit_destination=True,
1408+
)
1409+
1410+
assert rows._explicit_destination
1411+
1412+
1413+
def _destination_iterator(row_counts, serialized_schema=None):
1414+
"""A RowIterator for a query with an explicit destination table, plus its
1415+
mock BigQuery Storage client.
1416+
"""
1417+
bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage")
1418+
options = bigquery_storage.types.ArrowSerializationOptions
1419+
if "picos_timestamp_precision" not in options.meta.fields:
1420+
# query_and_wait() rejects ARROW with this release before it runs.
1421+
pytest.skip("google-cloud-bigquery-storage has no timestamp precision")
1422+
schema, batches = _arrow_pages(*row_counts)
1423+
reader = mock.MagicMock()
1424+
reader.__iter__.return_value = iter([_stream_response(batch) for batch in batches])
1425+
stream = mock.Mock()
1426+
stream.name = "projects/p/locations/us/sessions/s/streams/0"
1427+
session = mock.Mock(streams=[stream] if row_counts else [])
1428+
session.arrow_schema.serialized_schema = (
1429+
schema if serialized_schema is None else serialized_schema
1430+
)
1431+
bqstorage_client = mock.MagicMock()
1432+
bqstorage_client.create_read_session.return_value = session
1433+
bqstorage_client.read_rows.return_value = reader
1434+
iterator = RowIterator(
1435+
client=mock.MagicMock(project="billing-project"),
1436+
api_request=mock.MagicMock(),
1437+
path=None,
1438+
schema=(),
1439+
table=TableReference.from_string("p.d.t"),
1440+
project="p",
1441+
location="US",
1442+
job_id="j",
1443+
query_results_format="ARROW",
1444+
)
1445+
iterator._explicit_destination = True
1446+
return iterator, bqstorage_client
1447+
1448+
1449+
def test_download_arrow_explicit_destination_reads_a_read_session():
1450+
iterator, bqstorage_client = _destination_iterator([10, 15])
1451+
from google.cloud import bigquery_storage
1452+
1453+
batches = list(iterator._download_arrow_from_job_id(bqstorage_client))
1454+
1455+
assert _values(batches) == list(range(1, 26))
1456+
bqstorage_client.create_read_session.assert_called_once_with(
1457+
parent="projects/billing-project",
1458+
read_session=bigquery_storage.types.stream.ReadSession(
1459+
table="projects/p/datasets/d/tables/t",
1460+
data_format=bigquery_storage.types.stream.DataFormat.ARROW,
1461+
read_options=bigquery_storage.types.stream.ReadSession.TableReadOptions(
1462+
arrow_serialization_options=_DEFAULT_READ_OPTIONS
1463+
),
1464+
),
1465+
max_stream_count=1,
1466+
timeout=None,
1467+
)
1468+
# The read session carries the serialization options, so ReadRows doesn't.
1469+
bqstorage_client.read_rows.assert_called_once_with(
1470+
"projects/p/locations/us/sessions/s/streams/0", offset=0, timeout=None
1471+
)
1472+
1473+
1474+
def test_download_arrow_explicit_destination_without_streams_reads_nothing():
1475+
iterator, bqstorage_client = _destination_iterator([], serialized_schema=b"")
1476+
1477+
assert list(iterator._download_arrow_from_job_id(bqstorage_client)) == []
1478+
bqstorage_client.read_rows.assert_not_called()
1479+
1480+
13591481
def test_to_arrow_iterable_reads_inline_result_twice():
13601482
# With JOB_CREATION_OPTIONAL and every row inline there's no job to
13611483
# stream from, so the first page has to stay on the iterator for the

0 commit comments

Comments
 (0)