feat(bigquery): read Arrow results from explicit destinations

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:

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.

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

@@ -1983,6 +1983,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

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

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()