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 |
|
Assigning reviewers: R: @shunping for label python. Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
|
Reminder, please take a look at this pr: @shunping |
|
Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment R: @tvalentyn for label python. Available commands:
|
|
Reminder, please take a look at this pr: @tvalentyn |
|
Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment R: @claudevdm for label python. Available commands:
|
|
/gemini review |
|
cc: @ahmedabu98 in case you looked at this issue already and would like to chime in. Thanks |
There was a problem hiding this comment.
Code Review
This pull request updates Apache Beam's BigQuery file loads to inherit partitioning parameters (time or range partitioning) from the destination table when creating temporary tables. It also introduces caching for the destination table to avoid redundant API calls. Feedback on the changes points out a potential performance bottleneck: if get_table raises an exception, the failure is not cached, which could lead to repeated API calls for subsequent partitions in the same bundle. It is recommended to cache None in self.destination_table_cache when an exception occurs.
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.
| except Exception as e: | ||
| if need_schema: | ||
| _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.
If get_table raises an exception (e.g., if the destination table does not exist yet, which is common when creating a new table), the exception is caught, but self.destination_table_cache is not updated. As a result, subsequent elements/partitions processed in the same bundle will repeatedly attempt to call get_table and fail, leading to a significant performance bottleneck and potential BigQuery API rate-limiting issues.
To prevent this, we should cache None in self.destination_table_cache when an exception occurs.
| except Exception as e: | |
| if need_schema: | |
| _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 | |
| except Exception as e: | |
| if need_schema: | |
| _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 | |
| self.destination_table_cache[hashed_dest] = None |
ahmedabu98
left a comment
ahmedabu98
left a comment
There was a problem hiding this comment.
I think this looks great.
Can we add an integration test to make sure it works against the real API?
| @@ -758,6 +808,11 @@ def process( | |||
| job_name, | |||
| e) | |||
There was a problem hiding this comment.
| elif destination_table is not None: | |
| schema = bigquery_tools.table_schema_to_dict( | |
| destination_table.schema) | |
| self.schema_cache[hashed_dest] = schema |
We can probable remove this try-catch? It was originally meant for the get_table() call
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