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..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 @@ -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,32 @@ 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. + + 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, {}) + + 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)