From dbc30a07ee106c50f273549e10839a114dd1535a Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 13 Oct 2026 09:32:49 +0000 Subject: [PATCH] fix(bigquery): keep the Arrow schema for empty results A query that returned no rows with ARROW yielded no batches, so to_arrow() fell back to the iterator's schema, which jobs.query omits for ARROW, and returned a table without columns. jobs.query sends arrowSchema even for zero rows, but the default job stream sends nothing for them. to_arrow_iterable() now yields one empty RecordBatch if no rows arrive, like _EmptyRowIterator already does. It has the Arrow schema from the first page or the read session; otherwise the schema from jobs.getQueryResults, which the jobs.insert path has; otherwise no columns. to_arrow() and pyarrow.Table.from_batches() keep that schema. --- .../google/cloud/bigquery/_arrow_query.py | 35 +++++- .../tests/system/test_arrow.py | 48 ++++++++ .../unit/test_query_results_format_arrow.py | 115 +++++++++++++++++- 3 files changed, 193 insertions(+), 5 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 6ce36bd71c21..fe705e9e51ba 100644 --- a/packages/google-cloud-bigquery/google/cloud/bigquery/_arrow_query.py +++ b/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) diff --git a/packages/google-cloud-bigquery/tests/system/test_arrow.py b/packages/google-cloud-bigquery/tests/system/test_arrow.py index 1ee4bbd0f2c0..03ce2ad340dd 100644 --- a/packages/google-cloud-bigquery/tests/system/test_arrow.py +++ b/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)) 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] 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 d730e464f70e..6c6bcfb4bf27 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 @@ -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()