Add query_output_schema to ReadFromBigQuery for BEAM_ROW + query support#39160
Add query_output_schema to ReadFromBigQuery for BEAM_ROW + query support#39160nikitagrover19 wants to merge 2 commits into
Conversation
Schema cannot be auto-derived from a table when a query is used, so this adds an explicit query_output_schema param for that case. Fixes apache#36988
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request enables the use of BEAM_ROW output with BigQuery queries by allowing users to explicitly provide the output schema. Previously, this was blocked because schemas could not be auto-derived from queries. The changes introduce a new parameter to the ReadFromBigQuery transform and the corresponding YAML interface, ensuring that schema-based pipelines can now function correctly with query-based inputs. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for specifying a query_output_schema when reading from BigQuery using a query with the BEAM_ROW output type, allowing schemas to be defined when they cannot be auto-derived from a table. The changes span both the standard Python SDK and the YAML-based IO transforms, along with accompanying unit tests. The review feedback highlights a critical runtime issue where the schema needs to be normalized to a TableSchema object using bigquery_tools.get_dict_table_schema to prevent an AttributeError. Additionally, the reviewer suggests raising an error if a schema is incorrectly provided for table-based reads in YAML, and adding corresponding tests for both of these cases.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if self._kwargs.get('query', None) is not None: | ||
| return output_pcollection | bigquery_schema_tools.convert_to_usertype( | ||
| self.query_output_schema, self._kwargs.get('selected_fields', None)) |
There was a problem hiding this comment.
The query_output_schema parameter can be passed as a dictionary or a JSON string (as documented in the docstring and used in YAML). However, bigquery_schema_tools.convert_to_usertype expects a TableSchema object. Passing a dictionary or string directly will result in an AttributeError at runtime (e.g., 'dict' object has no attribute 'fields'). Use bigquery_tools.get_dict_table_schema to normalize the schema before passing it to convert_to_usertype.
| if self._kwargs.get('query', None) is not None: | |
| return output_pcollection | bigquery_schema_tools.convert_to_usertype( | |
| self.query_output_schema, self._kwargs.get('selected_fields', None)) | |
| if self._kwargs.get('query', None) is not None: | |
| user_schema = bigquery_tools.get_dict_table_schema(self.query_output_schema) | |
| return output_pcollection | bigquery_schema_tools.convert_to_usertype( | |
| user_schema, self._kwargs.get('selected_fields', None)) |
| mock_get_table.assert_not_called() | ||
| mock_convert.assert_called_once_with(schema, None) |
There was a problem hiding this comment.
Since query_output_schema is now normalized to a TableSchema object using bigquery_tools.get_dict_table_schema, the mocked convert_to_usertype will be called with the converted TableSchema object rather than the raw dictionary. Update the assertion to reflect this conversion.
| mock_get_table.assert_not_called() | |
| mock_convert.assert_called_once_with(schema, None) | |
| mock_get_table.assert_not_called() | |
| mock_convert.assert_called_once_with( | |
| bigquery_tools.get_dict_table_schema(schema), None) |
| if query is None: | ||
| assert table is not None | ||
| else: |
There was a problem hiding this comment.
If a user provides a schema parameter for a table-based read (where query is None), it will be silently ignored because ReadFromBigQuery only uses query_output_schema when a query is specified. To prevent silent failures and improve usability, raise a ValueError if schema is provided for a table-based read.
if query is None:
assert table is not None
if schema is not None:
raise ValueError(
"The 'schema' parameter is only supported when reading with a 'query'. "
"For table-based reads, the schema is automatically derived.")
else:| self.assertIsNone(call_kwargs.get('query_output_schema')) | ||
|
|
There was a problem hiding this comment.
Add a unit test to verify that providing a schema for a table-based read raises a ValueError as expected.
| self.assertIsNone(call_kwargs.get('query_output_schema')) | |
| self.assertIsNone(call_kwargs.get('query_output_schema')) | |
| def test_table_with_schema_raises(self): | |
| from apache_beam.yaml.yaml_io import read_from_bigquery | |
| with self.assertRaisesRegex(ValueError, 'only supported when reading with a'): | |
| read_from_bigquery(table='project:dataset.table', schema={'fields': []}) |
|
Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment |
1 participant
Schema cannot be auto-derived from a table when a query is used, so this adds an explicit query_output_schema param for that case.
Fixes #36988
Problem
BEAM_ROW output was disallowed with query= since schema can't be derived from a table when there's no table. This blocked the YAML transform from using BEAM_ROW with queries.
Change
Adds an optional query_output_schema param to ReadFromBigQuery, used directly as the row schema when set. YAML's read_from_bigquery() exposes this as schema, required when query is set. Backward compatible, table-based reads unaffected.
Testing
Unit tests cover constructor validation, schema pass-through, and that get_table() is skipped when query_output_schema is supplied. Not covered: correctness of convert_to_usertype output against a real query execution - would need a live BQ connection or a heavier fake.
Known testing gap
These tests verify the schema is passed through and used (i.e., convert_to_usertype is called with the right arguments, and get_table is correctly skipped). They do not verify that convert_to_usertype produces correct Beam Rows at runtime from a real query execution, that would require either a live BigQuery connection or a more substantial fake of the job-execution lifecycle. Open to adding that if reviewers feel it's necessary for merge.
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, commentfixes #<ISSUE NUMBER>instead.CHANGES.mdwith noteworthy changes.See the Contributor Guide for more tips on how to make review process smoother.
To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md
GitHub Actions Tests Status (on master branch)
See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.