From 73fc241982926c2f75f0c6a3bb12d9f11d4431a7 Mon Sep 17 00:00:00 2001 From: chalmer lowe Date: Mon, 27 Jul 2026 18:22:48 -0400 Subject: [PATCH 1/3] test: parametrize to_arrow empty stream tests and update docstrings --- .../cloud/bigquery_storage_v1/reader.py | 10 +++++++++ .../tests/unit/test_reader_v1_arrow.py | 22 ++++++++++++++++++- 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py index a2222429240a..c4caa536da1d 100644 --- a/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py +++ b/packages/google-cloud-bigquery-storage/google/cloud/bigquery_storage_v1/reader.py @@ -349,6 +349,13 @@ def __iter__(self): def to_arrow(self): """Create a :class:`pyarrow.Table` of all rows in the stream. + Note: This is the :class:`ReadRowsIterable` version of ``to_arrow``. It is + typically invoked by calling :meth:`ReadRowsStream.to_arrow`, which + delegates here after handling optional session context. The key difference + is that :meth:`ReadRowsStream.to_arrow` accepts a ``read_session`` argument + to provide schema hints for empty streams, whereas this method relies on + the parser initialized during :class:`ReadRowsIterable` construction. + This method requires the pyarrow library and a stream using the Arrow format. @@ -365,6 +372,9 @@ def to_arrow(self): # No data, return an empty Table. if self._stream_parser is None: + # Note: This returns a table with an empty schema (no columns). + # Downstream consumers (like Vertex Ray) might fail if they expect specific columns. + # To guarantee the correct schema, provide 'read_session' to `ReadRowsStream.to_arrow()`. return pyarrow.Table.from_batches([], schema=pyarrow.schema([])) self._stream_parser._parse_arrow_schema() diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1_arrow.py b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1_arrow.py index 29af246491bf..db494ec8b285 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1_arrow.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1_arrow.py @@ -27,7 +27,6 @@ import importlib_metadata as metadata import google.api_core.exceptions - from google.cloud.bigquery_storage import types from .helpers import SCALAR_BLOCKS, SCALAR_COLUMN_NAMES, SCALAR_COLUMNS @@ -208,6 +207,27 @@ def test_rows_w_empty_stream_arrow(class_under_test, mock_gapic_client): assert tuple(got) == () +@pytest.mark.parametrize( + "use_session", + [False, True], + ids=["no_session", "with_session"], +) +def test_to_arrow_empty_stream(class_under_test, mock_gapic_client, use_session): + """Verify that to_arrow() handles empty streams safely.""" + arrow_schema = _bq_to_arrow_schema(SCALAR_COLUMNS) + mock_gapic_client.read_rows.return_value = iter([]) + + reader = class_under_test(mock_gapic_client, "name", 0, {}) + + read_session = _generate_arrow_read_session(arrow_schema) if use_session else None + expected_schema = arrow_schema if use_session else pyarrow.schema([]) + + table = reader.to_arrow(read_session) + + assert len(table) == 0 + assert table.schema == expected_schema + + def test_rows_w_scalars_arrow(class_under_test, mock_gapic_client): arrow_schema = _bq_to_arrow_schema(SCALAR_COLUMNS) arrow_batches = _bq_to_arrow_batches(SCALAR_BLOCKS, arrow_schema) From 80f27f0942b8e1355c5ad365caa2e45dc6d769aa Mon Sep 17 00:00:00 2001 From: chalmer lowe Date: Tue, 28 Jul 2026 06:34:06 -0400 Subject: [PATCH 2/3] test: apply reviewer suggestion for robustness in to_arrow tests --- .../tests/unit/test_reader_v1_arrow.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1_arrow.py b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1_arrow.py index db494ec8b285..e455388401b2 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1_arrow.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1_arrow.py @@ -219,10 +219,14 @@ def test_to_arrow_empty_stream(class_under_test, mock_gapic_client, use_session) reader = class_under_test(mock_gapic_client, "name", 0, {}) + if use_session and class_under_test.__name__ == "ReadRowsIterable": + return + read_session = _generate_arrow_read_session(arrow_schema) if use_session else None expected_schema = arrow_schema if use_session else pyarrow.schema([]) - table = reader.to_arrow(read_session) + kwargs = {"read_session": read_session} if use_session else {} + table = reader.to_arrow(**kwargs) assert len(table) == 0 assert table.schema == expected_schema From 3fdb263a9ef84cf7e72c274d0f488b2de7cbdac7 Mon Sep 17 00:00:00 2001 From: chalmer lowe Date: Tue, 28 Jul 2026 06:39:53 -0400 Subject: [PATCH 3/3] refactor(test): clean up to_arrow empty stream test and clarify focus --- .../tests/unit/test_reader_v1_arrow.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1_arrow.py b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1_arrow.py index e455388401b2..8b6ec1f95d3e 100644 --- a/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1_arrow.py +++ b/packages/google-cloud-bigquery-storage/tests/unit/test_reader_v1_arrow.py @@ -213,20 +213,21 @@ def test_rows_w_empty_stream_arrow(class_under_test, mock_gapic_client): ids=["no_session", "with_session"], ) def test_to_arrow_empty_stream(class_under_test, mock_gapic_client, use_session): - """Verify that to_arrow() handles empty streams safely.""" + """Verify that to_arrow() handles empty streams safely. + + Note: This test focuses specifically on ReadRowsStream.to_arrow(), which + accepts a read_session argument to provide schema hints for empty streams, + unlike ReadRowsIterable.to_arrow(). + """ arrow_schema = _bq_to_arrow_schema(SCALAR_COLUMNS) mock_gapic_client.read_rows.return_value = iter([]) reader = class_under_test(mock_gapic_client, "name", 0, {}) - if use_session and class_under_test.__name__ == "ReadRowsIterable": - return - read_session = _generate_arrow_read_session(arrow_schema) if use_session else None expected_schema = arrow_schema if use_session else pyarrow.schema([]) - kwargs = {"read_session": read_session} if use_session else {} - table = reader.to_arrow(**kwargs) + table = reader.to_arrow(read_session) assert len(table) == 0 assert table.schema == expected_schema