fix(bigquery): keep the Arrow schema for empty results

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

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 import base64
-from typing import Any, Dict, Iterator, Optional, Tuple, TYPE_CHECKING
+from typing import Any, Dict, Iterator, Optional, Sequence, Tuple, TYPE_CHECKING
 
 from google.cloud.bigquery import _versions_helpers
 from google.cloud.bigquery import enums
@@ -138,6 +138,29 @@ def _slice_batch(
     return batch, remaining - batch.num_rows
 
 
+def _empty_batch(
+    pa_schema: Optional["pyarrow.Schema"], bq_schema: Sequence[Any]
+) -> "pyarrow.RecordBatch":
+    """Build a batch without rows that keeps the result's schema.
+
+    Use the Arrow schema BigQuery sent if there is one. The default job
+    stream sends nothing for zero rows, so fall back to the schema from
+    ``jobs.getQueryResults``, and else to no columns.
+    """
+    if pa_schema is None and bq_schema:
+        # Imported here: _pandas_helpers imports pandas_gbq, which imports
+        # this package's table module.
+        from google.cloud.bigquery import _pandas_helpers
+
+        pa_schema = _pandas_helpers.bq_to_arrow_schema(bq_schema)
+    if pa_schema is None:
+        pa_schema = pyarrow.schema([])
+    return pyarrow.RecordBatch.from_arrays(
+        [pyarrow.array([], type=field.type) for field in pa_schema],
+        schema=pa_schema,
+    )
+
+
 def _read_rows_options(options: Optional[Dict[str, str]]) -> Dict[str, str]:
     """Convert ``jobs.query`` ``arrowSerializationOptions`` to the
     ``ArrowSerializationOptions`` fields that ``ReadRows`` takes."""
@@ -235,6 +258,9 @@ def download_arrow_from_job(
     total_rows = row_iterator.total_rows
     job_complete = False
     remaining = row_iterator.max_results
+    # Results without rows still yield one empty batch, so callers get the
+    # schema.
+    yielded = False
 
     # Step 1: Process initial inline Arrow response from jobs.query if available.
     if row_iterator._first_page_response:
@@ -264,11 +290,14 @@ def download_arrow_from_job(
                 )
                 offset += batch.num_rows
                 batch, remaining = _slice_batch(batch, remaining)
+                yielded = True
                 yield batch
 
     # 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:
+        if not yielded:
+            yield _empty_batch(pa_schema, row_iterator._schema)
         return
 
     # Step 3: Stream remaining Arrow record batches from the job default stream,
@@ -291,6 +320,7 @@ def download_arrow_from_job(
         reader = _open_job_stream(row_iterator, bqstorage_client, offset, timeout)
     if reader is None:
         # An empty destination table has no streams.
+        yield _empty_batch(pa_schema, row_iterator._schema)
         return
     finished = False
     try:
@@ -315,6 +345,7 @@ def download_arrow_from_job(
                     pa_schema,
                 )
                 batch, remaining = _slice_batch(batch, remaining)
+                yielded = True
                 yield batch
                 if remaining == 0:
                     break
@@ -325,3 +356,5 @@ def download_arrow_from_job(
         # an error, or when the caller closes the iterator.
         if not finished:
             reader.cancel()
+    if not yielded:
+        yield _empty_batch(pa_schema, row_iterator._schema)

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

@@ -860,3 +860,51 @@ def test_query_and_wait_arrow_explicit_destination(
         f"projects/{bigquery_client.project}/datasets/{dataset_id}/tables/{table_name}"
     )
     assert kwargs["max_stream_count"] == 1
+
+
+@pytest.mark.parametrize(
+    "priority",
+    [None, enums.QueryPriority.INTERACTIVE],
+    ids=["jobs_query", "jobs_insert"],
+)
+def test_query_and_wait_arrow_zero_rows_keep_the_schema(bigquery_client, priority):
+    job_config = bigquery.QueryJobConfig(use_query_cache=False)
+    if priority is not None:
+        # jobs.query doesn't accept a priority, so this runs through jobs.insert.
+        job_config.priority = priority
+
+    def query():
+        return bigquery_client.query_and_wait(
+            "SELECT x, CAST(x AS STRING) AS s "
+            "FROM UNNEST(CAST([] AS ARRAY<INT64>)) AS x",
+            job_config=job_config,
+            query_results_format=enums.QueryResultsFormat.ARROW,
+        )
+
+    batches = list(query().to_arrow_iterable())
+
+    assert len(batches) == 1
+    for table in (pyarrow.Table.from_batches(batches), query().to_arrow()):
+        assert table.num_rows == 0
+        assert table.schema.names == ["x", "s"]
+        assert table.schema.field("x").type == pyarrow.int64()
+        assert table.schema.field("s").type == pyarrow.string()
+
+
+def test_query_and_wait_arrow_statements_without_results(bigquery_client, dataset_id):
+    table_id = f"{bigquery_client.project}.{dataset_id}.arrow_statements"
+
+    for statement in (
+        f"CREATE OR REPLACE TABLE `{table_id}` AS SELECT 1 AS x",
+        f"UPDATE `{table_id}` SET x = 2 WHERE TRUE",
+    ):
+        rows = bigquery_client.query_and_wait(
+            statement, query_results_format=enums.QueryResultsFormat.ARROW
+        )
+        (batch,) = rows.to_arrow_iterable()
+        assert batch.num_rows == 0
+
+    assert rows.num_dml_affected_rows == 1
+    assert [
+        row.x for row in bigquery_client.query_and_wait(f"SELECT x FROM `{table_id}`")
+    ] == [2]

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

@@ -232,7 +232,11 @@ def test_download_arrow_from_job_id_constructs_stream_and_reads(self):
         with mock.patch("google.cloud.bigquery._arrow_query.pyarrow") as mock_pyarrow:
             mock_pyarrow.py_buffer = lambda x: x
             batches = list(iterator._download_arrow_from_job_id(timeout=5.0))
-            self.assertEqual(batches, [])
+            # No rows and no schema: one empty batch without columns.
+            self.assertEqual(
+                batches, [mock_pyarrow.RecordBatch.from_arrays.return_value]
+            )
+            mock_pyarrow.schema.assert_called_once_with([])
 
             expected_stream_name = (
                 "projects/test-proj/locations/US/jobs/test-job-456/streams/_default"
@@ -1357,12 +1361,115 @@ def test_download_arrow_explicit_destination_reads_a_read_session():
     )
 
 
-def test_download_arrow_explicit_destination_without_streams_reads_nothing():
-    iterator, bqstorage_client = _destination_iterator([], serialized_schema=b"")
+@pytest.mark.parametrize(
+    ("serialized_schema", "expected_names"),
+    [(None, ["x"]), (b"", [])],
+    ids=["with_schema", "without_schema"],
+)
+def test_download_arrow_explicit_destination_without_rows(
+    serialized_schema, expected_names
+):
+    iterator, bqstorage_client = _destination_iterator(
+        [], serialized_schema=serialized_schema
+    )
 
-    assert list(iterator._download_arrow_from_job_id(bqstorage_client)) == []
+    (batch,) = iterator._download_arrow_from_job_id(bqstorage_client)
+
+    assert batch.num_rows == 0
+    assert batch.schema.names == expected_names
     bqstorage_client.read_rows.assert_not_called()
 
 
+def _zero_row_iterator():
+    """A RowIterator whose jobs.query response has a schema and no rows."""
+    schema, _ = _arrow_pages()
+    return RowIterator(
+        client=mock.MagicMock(),
+        api_request=mock.MagicMock(),
+        path=None,
+        schema=(),
+        project="p",
+        location="US",
+        job_id="j",
+        query_results_format="ARROW",
+        first_page_response={
+            "jobComplete": True,
+            "totalRows": "0",
+            "arrowSchema": {"serializedSchema": base64.b64encode(schema).decode()},
+            # A response without rows carries no serialized batch.
+            "arrowRecordBatch": {},
+        },
+    )
+
+
+def test_download_arrow_zero_rows_inline_keep_the_schema():
+    iterator = _zero_row_iterator()
+
+    (batch,) = iterator._download_arrow_from_job_id()
+
+    assert batch.num_rows == 0
+    assert batch.schema.names == ["x"]
+    iterator.client._get_shared_bqstorage_client.assert_not_called()
+
+
+def test_to_arrow_zero_rows_keep_the_schema():
+    table = _zero_row_iterator().to_arrow()
+
+    assert table.num_rows == 0
+    assert table.schema.names == ["x"]
+
+
+def test_download_arrow_zero_rows_streamed_use_the_query_results_schema():
+    pyarrow = pytest.importorskip("pyarrow")
+    from google.cloud.bigquery.schema import SchemaField
+
+    bqstorage_client = mock.MagicMock()
+    # The default job stream sends no responses for zero rows.
+    bqstorage_client.read_rows.return_value = []
+    iterator = RowIterator(
+        client=mock.MagicMock(),
+        api_request=mock.MagicMock(),
+        path=None,
+        schema=[SchemaField("x", "INTEGER")],
+        project="p",
+        location="US",
+        job_id="j",
+        query_results_format="ARROW",
+        first_page_response={"jobComplete": False},
+    )
+
+    (batch,) = iterator._download_arrow_from_job_id(bqstorage_client)
+
+    assert batch.num_rows == 0
+    assert batch.schema.names == ["x"]
+    assert batch.schema.field("x").type == pyarrow.int64()
+
+
+def test_download_arrow_rows_yield_no_extra_empty_batch():
+    schema, batches = _arrow_pages(3)
+    iterator = RowIterator(
+        client=mock.MagicMock(),
+        api_request=mock.MagicMock(),
+        path=None,
+        schema=(),
+        project="p",
+        location="US",
+        job_id="j",
+        query_results_format="ARROW",
+        first_page_response={
+            "jobComplete": True,
+            "totalRows": "3",
+            "arrowSchema": {"serializedSchema": base64.b64encode(schema).decode()},
+            "arrowRecordBatch": {
+                "serializedRecordBatch": base64.b64encode(batches[0]).decode()
+            },
+        },
+    )
+
+    (batch,) = iterator._download_arrow_from_job_id()
+
+    assert batch.num_rows == 3
+
+
 if __name__ == "__main__":
     unittest.main()