Preserve partitioning on temp FILE_LOADS tables#38833
Conversation
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Adds support for temporary-table BigQuery file loads to inherit partitioning settings from the final destination table when not explicitly provided, and validates this behavior with new unit tests.
Changes:
- Add helpers to detect/propagate
timePartitioningandrangePartitioninginto load job parameters for temporary table loads. - Update
TriggerLoadJobs.processto optionally fetch the destination table once to reuse schema and partitioning metadata. - Add tests verifying partitioning inheritance and that explicit partitioning parameters are not overridden.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| sdks/python/apache_beam/io/gcp/bigquery_file_loads.py | Propagates destination-table partitioning settings into temporary-table load job parameters and reuses fetched table metadata. |
| sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py | Adds unit tests for partitioning inheritance and explicit-parameter precedence. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| def _add_destination_partitioning_load_parameters( | ||
| additional_parameters, destination_table): | ||
| if not isinstance(destination_table, bigquery_tools.bigquery.Table): | ||
| return additional_parameters | ||
|
|
||
| additional_parameters = dict(additional_parameters) | ||
|
|
||
| if ('timePartitioning' not in additional_parameters and | ||
| getattr(destination_table, 'timePartitioning', None) is not None): | ||
| additional_parameters['timePartitioning'] = ( | ||
| destination_table.timePartitioning) | ||
|
|
||
| if ('rangePartitioning' not in additional_parameters and | ||
| getattr(destination_table, 'rangePartitioning', None) is not None): | ||
| additional_parameters['rangePartitioning'] = ( | ||
| destination_table.rangePartitioning) | ||
|
|
||
| return additional_parameters |
| destination_table = None | ||
| hashed_dest = bigquery_tools.get_hashable_destination(table_reference) | ||
| should_lookup_destination_table = ( | ||
| schema is None or | ||
| not _has_partitioning_load_parameters(additional_parameters)) | ||
| if should_lookup_destination_table: | ||
| try: | ||
| destination_table = self.bq_wrapper.get_table( | ||
| project_id=table_reference.projectId, | ||
| dataset_id=table_reference.datasetId, | ||
| table_id=table_reference.tableId) | ||
| except Exception as e: |
| load_call = dofn.bq_wrapper.perform_load_job.call_args.kwargs | ||
| self.assertEqual( | ||
| load_call['additional_load_parameters']['timePartitioning'], | ||
| destination_table.timePartitioning) | ||
| dofn.bq_wrapper.get_table.assert_called_once() |
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 improves the reliability of BigQuery file load operations by ensuring that partitioning settings are correctly propagated to temporary tables. By intelligently fetching destination table metadata and avoiding redundant lookups, the changes ensure that existing partitioning configurations are respected unless overridden by the user. This fix addresses issues where temporary table creation might otherwise lose critical partitioning constraints. 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 ensures that temporary table loads inherit time and range partitioning parameters from the destination BigQuery table. It introduces helper functions to check and apply these parameters, updates the TriggerLoadJobs process to fetch the destination table configuration, and adds corresponding unit tests. However, the current implementation introduces a performance regression by bypassing the schema cache and making redundant, synchronous get_table API calls on every bundle. It is recommended to cache the fetched destination table locally to avoid unnecessary API calls.
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.
| destination_table = None | ||
| hashed_dest = bigquery_tools.get_hashable_destination(table_reference) | ||
| should_lookup_destination_table = ( | ||
| schema is None or | ||
| not _has_partitioning_load_parameters(additional_parameters)) | ||
| if should_lookup_destination_table: | ||
| try: | ||
| destination_table = self.bq_wrapper.get_table( | ||
| project_id=table_reference.projectId, | ||
| dataset_id=table_reference.datasetId, | ||
| table_id=table_reference.tableId) | ||
| except Exception as e: | ||
| if schema is None: | ||
| _LOGGER.warning( | ||
| "Input schema is absent and could not fetch the final " | ||
| "destination table's schema [%s]. Creating temp table [%s] " | ||
| "will likely fail: %s", | ||
| hashed_dest, | ||
| job_name, | ||
| e) | ||
| destination_table = None |
There was a problem hiding this comment.
This implementation introduces a significant performance regression and redundant BigQuery API calls:
- Bypassing the Schema Cache: Since
should_lookup_destination_tableis evaluated before checkingself.schema_cache, ifschemaisNoneinitially,should_lookup_destination_tablewill always beTrue. This forces a synchronousget_tableAPI call on every single bundle/partition even if the schema is already cached inself.schema_cache. - Redundant Calls per Bundle: If
schemais provided but partitioning parameters are not (the default case),get_tableis called on every single bundle/partition without any caching.
We can resolve both issues by caching the fetched destination_table in a local cache (e.g., self._destination_table_cache) and only calling get_table if we actually need the schema (and it's not in self.schema_cache) or if we need the partitioning parameters.
hashed_dest = bigquery_tools.get_hashable_destination(table_reference)
if not hasattr(self, '_destination_table_cache'):
self._destination_table_cache = {}
destination_table = self._destination_table_cache.get(hashed_dest)
if destination_table is None:
need_schema = schema is None and hashed_dest not in self.schema_cache
need_partitioning = not _has_partitioning_load_parameters(additional_parameters)
if need_schema or need_partitioning:
try:
destination_table = self.bq_wrapper.get_table(
project_id=table_reference.projectId,
dataset_id=table_reference.datasetId,
table_id=table_reference.tableId)
self._destination_table_cache[hashed_dest] = destination_table
except Exception as e:
if schema is None and hashed_dest not in self.schema_cache:
_LOGGER.warning(
"Input schema is absent and could not fetch the final "
"destination table's schema [%s]. Creating temp table [%s] "
"will likely fail: %s",
hashed_dest,
job_name,
e)
destination_table = None
|
Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment |
Summary
FILE_LOADSwrites through temporary tablesFixes #38017.
Testing
python -m pytest apache_beam/io/gcp/bigquery_file_loads_test.py -k "temporary_table_load_inherits_destination_time_partitioning or temporary_table_load_inherits_destination_range_partitioning or temporary_table_load_keeps_explicit_partitioning_parameters" -qpython -m pytest apache_beam/io/gcp/bigquery_file_loads_test.py::TestBigQueryFileLoads::test_wait_for_load_job_completion -qpython -m py_compile apache_beam/io/gcp/bigquery_file_loads.py apache_beam/io/gcp/bigquery_file_loads_test.py