Skip to content

refactor(bigquery): move Arrow query result reading to _arrow_query - #18644

Merged
ohmayr merged 1 commit into
mainfrom
feat/bigquery-arrow-query-module
Oct 15, 2026
Merged

ohmayr merged 1 commit into
mainfrom
feat/bigquery-arrow-query-module

Conversation

@ohmayr

@ohmayr ohmayr commented Oct 13, 2026 •

Copy link
Copy Markdown
Contributor

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.py is deprecated.

  • RowIterator._download_arrow_from_job_id() keeps its signature and delegates to _arrow_query.download_arrow_from_job().
  • Tests patch _arrow_query.pyarrow instead of table.pyarrow.

git diff --color-moved shows the moved block.

@ohmayr
ohmayr requested review from a team as code owners October 13, 2026 07:59
@ohmayr
ohmayr requested review from sycai and removed request for a team October 13, 2026 07:59

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +99 to +105
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"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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"

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@ohmayr
ohmayr added this pull request to stack #18648 October 13, 2026 08:17
@ohmayr
ohmayr force-pushed the feat/bigquery-arrow-query-module branch from 133e101 to 1851fa9 Compare October 13, 2026 10:02
@ohmayr
ohmayr requested a review from tswast October 13, 2026 20:40
Base automatically changed from feat/bigquery-storage-arrow-read-rows to main October 15, 2026 02:29
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.
@ohmayr
ohmayr force-pushed the feat/bigquery-arrow-query-module branch from 1851fa9 to db138d2 Compare October 15, 2026 02:29
@ohmayr
ohmayr merged commit 09a3823 into main Oct 15, 2026
53 checks passed
@ohmayr
ohmayr deleted the feat/bigquery-arrow-query-module branch October 15, 2026 03:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants