Repository navigation
refactor(bigquery): move Arrow query result reading to _arrow_query - #18644
Conversation
There was a problem hiding this comment.
Code Review
This pull request refactors the Arrow query results downloading logic by extracting the _download_arrow_from_job_id method from google/cloud/bigquery/table.py into a new dedicated module google/cloud/bigquery/_arrow_query.py. The unit tests have been updated accordingly to mock pyarrow in the new module. A review comment suggests validating that the project, location, and job ID are all present before constructing the stream name to prevent invalid resource paths and cryptic API errors.
| project = row_iterator._project or ( | ||
| row_iterator.client.project if row_iterator.client else None | ||
| ) | ||
| location = row_iterator._location or ( | ||
| row_iterator.client.location if row_iterator.client else None | ||
| ) | ||
| stream_name = f"projects/{project}/locations/{location}/jobs/{row_iterator._job_id}/streams/_default" |
There was a problem hiding this comment.
If row_iterator.client is None (which can happen if bqstorage_client is provided externally) or if the project, location, or job ID is missing, constructing the stream_name with None values will result in an invalid resource path and cause a cryptic API error. We should validate that project, location, and job_id are all present before constructing the stream name.
project = row_iterator._project or (
row_iterator.client.project if row_iterator.client else None
)
location = row_iterator._location or (
row_iterator.client.location if row_iterator.client else None
)
job_id = row_iterator._job_id
if not project or not location or not job_id:
raise ValueError(
"Project, location, and job ID must all be present to stream Arrow results."
)
stream_name = f"projects/{project}/locations/{location}/jobs/{job_id}/streams/_default"There was a problem hiding this comment.
I'd rather keep this PR a pure move with no behavior changes. On this path the project, location and job ID come from the job reference that query_and_wait() sets on the RowIterator, so they're populated in practice. If we want stricter validation it should be its own change.
133e101 to
1851fa9
Compare
RowIterator._download_arrow_from_job_id() holds the whole ARROW read path inline in table.py, and upcoming Arrow work (serialization options, row limits, explicit destinations, empty results) would grow it further. _pyarrow_helpers isn't an option, because it's deprecated in favor of pandas-gbq. Move the method body unchanged into a new private module, _arrow_query.download_arrow_from_job(), with self renamed to row_iterator. The RowIterator method keeps its signature and docstring and delegates to it. Tests that mocked table.pyarrow for this path now mock _arrow_query.pyarrow. No behavior change.
1851fa9 to
db138d2
Compare
Moves the Arrow query results code out of
RowIterator._download_arrow_from_job_id()into a new private module,_arrow_query.py, with no behavior change. The rest of the Arrow changes build on it, and_pyarrow_helpers.pyis deprecated.RowIterator._download_arrow_from_job_id()keeps its signature and delegates to_arrow_query.download_arrow_from_job()._arrow_query.pyarrowinstead oftable.pyarrow.git diff --color-movedshows the moved block.