test: parametrize to_arrow empty stream tests and update docstrings - #17919
test: parametrize to_arrow empty stream tests and update docstrings#17919chalmerlowe wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
Code Review
This pull request adds documentation and inline comments to clarify the behavior of to_arrow() on empty streams in ReadRowsIterable and ReadRowsStream. It also introduces a new unit test, test_to_arrow_empty_stream, to verify empty stream handling. Feedback on the unit test highlights a potential TypeError when testing ReadRowsIterable because it does not accept a read_session argument, and suggests a code modification to conditionally handle this case.
| 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) |
There was a problem hiding this comment.
The class_under_test fixture parametrizes both ReadRowsStream and ReadRowsIterable. Since ReadRowsIterable.to_arrow() does not accept any arguments (unlike ReadRowsStream.to_arrow()), calling reader.to_arrow(read_session) will raise a TypeError when testing ReadRowsIterable (even if read_session is None). We should return early for the session-based test for ReadRowsIterable and call to_arrow() without arguments when use_session is False. Avoid using pytest.skip() here as it would mark the entire test as skipped; use a conditional return instead.
| 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) | |
| 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) |
References
- Use
pytest.skip()to skip an entire test. To skip only a portion of a test (e.g., a final verification step), use a conditionalreturninstead, aspytest.skip()would mark the entire test as skipped.
Problem
When reading from an empty BigQuery table using the Read API, the stream may return no messages. Previously, this could lead to an
AttributeErrorwhen trying to parse the schema from aNoneparser. While a fallback was added to return an empty table, this fallback uses an empty schema (zero columns), which can cause downstream consumers (like Vertex Ray) to fail if they expect specific columns to be present.Solution
reader.pyto clarify the flow betweenReadRowsStream.to_arrow()andReadRowsIterable.to_arrow(), and to inform users about the fallback behavior.test_reader_v1_arrow.pyto verify the behavior ofto_arrow()on empty streams, both with and without aread_sessionprovided.Notes to Reviewers
read_sessiontoReadRowsStream.to_arrow().Fixes internal bug #352600521