test: parametrize to_arrow empty stream tests and update docstrings

Author: chalmerlowe

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

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)