From 35c581b0f19e1678e6617e7bbd154201e72edf90 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:32:49 +0000 Subject: [PATCH] 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(). --- .../google/cloud/bigquery/_arrow_query.py | 99 +++++++++++--- .../google/cloud/bigquery/_job_helpers.py | 5 + .../google/cloud/bigquery/table.py | 3 + .../tests/system/test_arrow.py | 39 ++++++ .../unit/test_query_results_format_arrow.py | 124 +++++++++++++++++- 5 files changed, 252 insertions(+), 18 deletions(-) diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py index 30534cbe4a3e..6ce36bd71c21 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py @@ -148,6 +148,76 @@ def _read_rows_options(options: Optional[Dict[str, str]]) -> Dict[str, str]: } +def _open_job_stream( + row_iterator: "RowIterator", + bqstorage_client: "bigquery_storage.BigQueryReadClient", + offset: int, + timeout: Optional[float], +): + """Open the job's default stream, starting at row ``offset``.""" + project = row_iterator._project or ( + row_iterator.client.project if row_iterator.client else None + ) + location = row_iterator._location or ( + row_iterator.client.location if row_iterator.client else None + ) + stream_name = f"projects/{project}/locations/{location}/jobs/{row_iterator._job_id}/streams/_default" + return bqstorage_client.read_rows( + stream_name, + offset=offset, + timeout=timeout, + arrow_serialization_options=_read_rows_options( + row_iterator._arrow_serialization_options + ), + ) + + +def _open_destination_stream( + row_iterator: "RowIterator", + bqstorage_client: "bigquery_storage.BigQueryReadClient", + timeout: Optional[float], + pa_schema: Optional["pyarrow.Schema"], +): + """Open a read session on the query's explicit destination table. + + The default job stream rejects jobs with an explicit destination table + ("This API does not support explicitly-specified destination tables"), so + read the table instead. One stream keeps the rows in order, and the + session's read options carry the same serialization options as + ``jobs.query``. Explicit destinations always go through ``jobs.insert``, + so there's no first page and the stream starts at row 0. + + Returns ``(reader, pa_schema)``, with ``reader`` set to ``None`` if the + table is empty and the session has no streams. + """ + from google.cloud import bigquery_storage + + session = bqstorage_client.create_read_session( + parent=f"projects/{row_iterator._billing_project}", + read_session=bigquery_storage.types.stream.ReadSession( + table=row_iterator._table.to_bqstorage(), + data_format=bigquery_storage.types.stream.DataFormat.ARROW, + read_options=bigquery_storage.types.stream.ReadSession.TableReadOptions( + arrow_serialization_options=_read_rows_options( + row_iterator._arrow_serialization_options + ) + ), + ), + max_stream_count=1, + timeout=timeout, + ) + if pa_schema is None and session.arrow_schema.serialized_schema: + pa_schema = pyarrow.ipc.read_schema( + pyarrow.py_buffer(session.arrow_schema.serialized_schema) + ) + if not session.streams: + return None, pa_schema + return ( + bqstorage_client.read_rows(session.streams[0].name, offset=0, timeout=timeout), + pa_schema, + ) + + def download_arrow_from_job( row_iterator: "RowIterator", bqstorage_client: Optional["bigquery_storage.BigQueryReadClient"] = None, @@ -202,8 +272,9 @@ def download_arrow_from_job( return # 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. + # or from the destination table if the job config set one, 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.") @@ -212,21 +283,15 @@ def download_arrow_from_job( 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 - ) - location = row_iterator._location or ( - row_iterator.client.location if row_iterator.client else None - ) - stream_name = f"projects/{project}/locations/{location}/jobs/{row_iterator._job_id}/streams/_default" - reader = bqstorage_client.read_rows( - stream_name, - offset=offset, - timeout=timeout, - arrow_serialization_options=_read_rows_options( - row_iterator._arrow_serialization_options - ), - ) + if row_iterator._explicit_destination and row_iterator._table is not None: + reader, pa_schema = _open_destination_stream( + row_iterator, bqstorage_client, timeout, pa_schema + ) + else: + reader = _open_job_stream(row_iterator, bqstorage_client, offset, timeout) + if reader is None: + # An empty destination table has no streams. + return finished = False try: for response in reader: diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py index a916245891f4..da6a806ba20d 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/_job_helpers.py @@ -561,6 +561,9 @@ def query_and_wait( max_results=max_results, query_results_format=query_results_format, arrow_serialization_options=arrow_serialization_options, + explicit_destination=( + job_config is not None and job_config.destination is not None + ), callback=callback, ) @@ -736,6 +739,7 @@ def _wait_or_cancel( *, query_results_format: Optional[str] = None, arrow_serialization_options: Optional[Dict[str, str]] = None, + explicit_destination: bool = False, callback: Callable = lambda _: None, ) -> table.RowIterator: """Wait for a job to complete and return the results. @@ -782,6 +786,7 @@ def _wait_or_cancel( ) query_results._query_results_format = query_results_format query_results._arrow_serialization_options = arrow_serialization_options + query_results._explicit_destination = explicit_destination return query_results except Exception: # Attempt to cancel the job since we can't return the results. diff --git a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py index b7e2bd55fcea..71d2148d828d 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/table.py +++ b/packages/google-cloud-bigquery/google/cloud/bigquery/table.py @@ -1977,6 +1977,9 @@ def __init__( self._job_ended = ended self._query_results_format = query_results_format self._arrow_serialization_options = arrow_serialization_options + # Set by query_and_wait() if the query wrote to a destination table + # set in its job config. + self._explicit_destination = False @property @_disallow_when_arrow diff --git a/packages/google-cloud-bigquery/tests/system/test_arrow.py b/packages/google-cloud-bigquery/tests/system/test_arrow.py index 339e6475a1ac..1ee4bbd0f2c0 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/packages/google-cloud-bigquery/tests/system/test_arrow.py @@ -821,3 +821,42 @@ def test_query_and_wait_arrow_ending_a_read_early_cancels_only_its_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)) + + +@pytest.mark.parametrize( + "write_disposition", + [None, enums.WriteDisposition.WRITE_TRUNCATE], + ids=["default", "overwrite"], +) +def test_query_and_wait_arrow_explicit_destination( + bigquery_client, bqstorage_client, dataset_id, write_disposition +): + table_name = f"arrow_destination_{write_disposition or 'default'}".lower() + table_id = f"{bigquery_client.project}.{dataset_id}.{table_name}" + job_config = bigquery.QueryJobConfig(destination=table_id) + if write_disposition is not None: + # The rows already in the table are replaced, not read. + bigquery_client.query_and_wait( + f"CREATE OR REPLACE TABLE `{table_id}` AS SELECT -1 AS x" + ) + job_config.write_disposition = write_disposition + # Wrap create_read_session() to check how the table is read. + create_read_session = mock.Mock(wraps=bqstorage_client.create_read_session) + + rows = bigquery_client.query_and_wait( + "SELECT x FROM UNNEST(GENERATE_ARRAY(1, 1000)) AS x ORDER BY x", + job_config=job_config, + query_results_format=enums.QueryResultsFormat.ARROW, + ) + with mock.patch.object( + bqstorage_client, "create_read_session", create_read_session + ): + batches = list(rows.to_arrow_iterable(bqstorage_client=bqstorage_client)) + + table = pyarrow.Table.from_batches(batches) + assert table.column("x").to_pylist() == list(range(1, 1001)) + kwargs = create_read_session.call_args.kwargs + assert kwargs["read_session"].table == ( + f"projects/{bigquery_client.project}/datasets/{dataset_id}/tables/{table_name}" + ) + assert kwargs["max_stream_count"] == 1 diff --git a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py index 46489b355585..d730e464f70e 100644 --- a/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py +++ b/packages/google-cloud-bigquery/tests/unit/test_query_results_format_arrow.py @@ -30,7 +30,7 @@ from google.cloud.bigquery.enums import TimestampPrecision from google.cloud.bigquery.enums import WriteDisposition from google.cloud.bigquery.job import QueryJobConfig -from google.cloud.bigquery.table import RowIterator, _EmptyRowIterator +from google.cloud.bigquery.table import RowIterator, TableReference, _EmptyRowIterator from .helpers import make_client @@ -847,6 +847,7 @@ def test_query_and_wait_arrow_jobs_insert_omits_arrow_fields(): max_results=mock.ANY, query_results_format="ARROW", arrow_serialization_options=_DEFAULT_QUERY_OPTIONS, + explicit_destination=True, callback=mock.ANY, ) assert wait_or_cancel.call_args.kwargs["page_size"] == 0 @@ -1242,5 +1243,126 @@ def test_to_arrow_with_arrow_format_reads_with_the_shared_client(): assert to_arrow_iterable.call_args.kwargs["bqstorage_client"] is None +@pytest.mark.parametrize( + ("job_config", "expected"), + [ + (QueryJobConfig(destination="p.d.t"), True), + # jobs.query doesn't accept a priority either, so this also runs + # through jobs.insert, without a destination. + (QueryJobConfig(priority="INTERACTIVE"), False), + ], + ids=["destination", "no_destination"], +) +def test_query_and_wait_arrow_marks_explicit_destination(job_config, expected): + client = mock.MagicMock(spec=Client) + + with ( + mock.patch.object(_job_helpers, "query_jobs_insert") as query_jobs_insert, + mock.patch.object(_job_helpers, "_wait_or_cancel") as wait_or_cancel, + ): + _query_and_wait(client, job_config=job_config, query_results_format="ARROW") + + wait_or_cancel.assert_called_once_with( + query_jobs_insert.return_value, + api_timeout=mock.ANY, + wait_timeout=mock.ANY, + retry=mock.ANY, + page_size=mock.ANY, + max_results=mock.ANY, + query_results_format="ARROW", + arrow_serialization_options=mock.ANY, + explicit_destination=expected, + callback=mock.ANY, + ) + + +def test_wait_or_cancel_sets_explicit_destination(): + job = mock.MagicMock(dry_run=False) + job.result.return_value = RowIterator( + client=mock.MagicMock(), api_request=mock.MagicMock(), path=None, schema=() + ) + assert not job.result.return_value._explicit_destination + + rows = _job_helpers._wait_or_cancel( + job, + api_timeout=None, + wait_timeout=None, + retry=None, + page_size=0, + max_results=None, + explicit_destination=True, + ) + + assert rows._explicit_destination + + +def _destination_iterator(row_counts, serialized_schema=None): + """A RowIterator for a query with an explicit destination table, plus its + mock BigQuery Storage client. + """ + bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage") + options = bigquery_storage.types.ArrowSerializationOptions + if "picos_timestamp_precision" not in options.meta.fields: + # query_and_wait() rejects ARROW with this release before it runs. + pytest.skip("google-cloud-bigquery-storage has no timestamp precision") + schema, batches = _arrow_pages(*row_counts) + reader = mock.MagicMock() + reader.__iter__.return_value = iter([_stream_response(batch) for batch in batches]) + stream = mock.Mock() + stream.name = "projects/p/locations/us/sessions/s/streams/0" + session = mock.Mock(streams=[stream] if row_counts else []) + session.arrow_schema.serialized_schema = ( + schema if serialized_schema is None else serialized_schema + ) + bqstorage_client = mock.MagicMock() + bqstorage_client.create_read_session.return_value = session + bqstorage_client.read_rows.return_value = reader + iterator = RowIterator( + client=mock.MagicMock(project="billing-project"), + api_request=mock.MagicMock(), + path=None, + schema=(), + table=TableReference.from_string("p.d.t"), + project="p", + location="US", + job_id="j", + query_results_format="ARROW", + ) + iterator._explicit_destination = True + return iterator, bqstorage_client + + +def test_download_arrow_explicit_destination_reads_a_read_session(): + iterator, bqstorage_client = _destination_iterator([10, 15]) + from google.cloud import bigquery_storage + + batches = list(iterator._download_arrow_from_job_id(bqstorage_client)) + + assert _values(batches) == list(range(1, 26)) + bqstorage_client.create_read_session.assert_called_once_with( + parent="projects/billing-project", + read_session=bigquery_storage.types.stream.ReadSession( + table="projects/p/datasets/d/tables/t", + data_format=bigquery_storage.types.stream.DataFormat.ARROW, + read_options=bigquery_storage.types.stream.ReadSession.TableReadOptions( + arrow_serialization_options=_DEFAULT_READ_OPTIONS + ), + ), + max_stream_count=1, + timeout=None, + ) + # The read session carries the serialization options, so ReadRows doesn't. + bqstorage_client.read_rows.assert_called_once_with( + "projects/p/locations/us/sessions/s/streams/0", offset=0, timeout=None + ) + + +def test_download_arrow_explicit_destination_without_streams_reads_nothing(): + iterator, bqstorage_client = _destination_iterator([], serialized_schema=b"") + + assert list(iterator._download_arrow_from_job_id(bqstorage_client)) == [] + bqstorage_client.read_rows.assert_not_called() + + if __name__ == "__main__": unittest.main()