[IcebergIO] Improve TableCache#38882
Conversation
|
@dejii any chance you can take a look? |
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 centralizes and optimizes the TableCache mechanism within the Iceberg IO connector. By transitioning to a process-wide cache keyed by catalog configuration and table identifiers, the changes significantly reduce overhead associated with redundant table loading and metadata refreshing. These improvements enhance robustness for both source and sink operations and simplify the codebase by offloading cache management from individual IO components. 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 refactors the Iceberg table caching mechanism by introducing a process-wide, thread-safe TableCache keyed by both IcebergCatalogConfig and TableIdentifier. This replaces various local caches and simplifies several DoFns across the SDK. The review feedback correctly identifies a critical bug in WritePartitionedRowsToFiles.java where storing specId and partitionFieldMap as single instance variables in WriteDoFn can cause state collisions when processing multiple tables in the same bundle; these should be keyed by TableIdentifier instead.
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.
| private int specId = Integer.MIN_VALUE; | ||
| private Map<String, PartitionField> partitionFieldMap = Maps.newHashMap(); |
There was a problem hiding this comment.
Storing specId and partitionFieldMap as single instance variables in WriteDoFn is incorrect because a single WriteDoFn instance can process elements for multiple different tables (destinations) within the same bundle. Since different tables can have the same specId (e.g., 0), this can lead to using the partition field map of one table for another table, causing incorrect partition paths or failures. We should key these by TableIdentifier instead.
| private int specId = Integer.MIN_VALUE; | |
| private Map<String, PartitionField> partitionFieldMap = Maps.newHashMap(); | |
| private final Map<TableIdentifier, Integer> specIds = Maps.newHashMap(); | |
| private final Map<TableIdentifier, Map<String, PartitionField>> partitionFieldMaps = Maps.newHashMap(); |
| Table table = lastRefreshedTable.table; | ||
| partitionPath = getPartitionDataPath(partitionPath, lastRefreshedTable.partitionFieldMap); | ||
| Table table = getOrCreateTable(destination, dataSchema); | ||
| partitionPath = getPartitionDataPath(partitionPath, getPartitionFieldMap(table)); |
There was a problem hiding this comment.
Update the call to getPartitionFieldMap to pass the table identifier so that the correct partition field map is retrieved for the specific table.
| partitionPath = getPartitionDataPath(partitionPath, getPartitionFieldMap(table)); | |
| partitionPath = getPartitionDataPath(partitionPath, getPartitionFieldMap(destination.getTableIdentifier(), table)); |
| private Map<String, PartitionField> getPartitionFieldMap(Table table) { | ||
| if (table.spec().specId() == this.specId) { | ||
| return partitionFieldMap; | ||
| } | ||
|
|
||
| /** | ||
| * Refreshes the table metadata if it is considered stale (older than 2 minutes). | ||
| * | ||
| * <p>This method first performs a non-synchronized check on the table's freshness. This | ||
| * provides a lock-free fast path that avoids synchronization overhead in the common case | ||
| * where the table does not need to be refreshed. If the table might be stale, it then enters | ||
| * a synchronized block to ensure that only one thread performs the refresh operation. | ||
| */ | ||
| void refreshIfStale() { | ||
| // Fast path: Avoid entering the synchronized block if the table is not stale. | ||
| if (lastRefreshTime.isAfter(Instant.now().minus(STALENESS_THRESHOLD))) { | ||
| return; | ||
| } | ||
| synchronized (this) { | ||
| if (lastRefreshTime.isBefore(Instant.now().minus(STALENESS_THRESHOLD))) { | ||
| table.refresh(); | ||
| lastRefreshTime = Instant.now(); | ||
| if (table.spec().specId() != this.specId) { | ||
| partitionFieldMap = Maps.newHashMap(); | ||
| for (PartitionField partitionField : table.spec().fields()) { | ||
| partitionFieldMap.put(partitionField.name(), partitionField); | ||
| } | ||
| this.specId = table.spec().specId(); | ||
| } | ||
| } | ||
| } | ||
| Map<String, PartitionField> partitionFieldMap = Maps.newHashMap(); | ||
| for (PartitionField partitionField : table.spec().fields()) { | ||
| partitionFieldMap.put(partitionField.name(), partitionField); | ||
| } | ||
| this.specId = table.spec().specId(); | ||
| this.partitionFieldMap = partitionFieldMap; | ||
| return partitionFieldMap; | ||
| } |
There was a problem hiding this comment.
Update getPartitionFieldMap to accept the TableIdentifier and look up/store the partition field map and spec ID per table identifier.
| private Map<String, PartitionField> getPartitionFieldMap(Table table) { | |
| if (table.spec().specId() == this.specId) { | |
| return partitionFieldMap; | |
| } | |
| /** | |
| * Refreshes the table metadata if it is considered stale (older than 2 minutes). | |
| * | |
| * <p>This method first performs a non-synchronized check on the table's freshness. This | |
| * provides a lock-free fast path that avoids synchronization overhead in the common case | |
| * where the table does not need to be refreshed. If the table might be stale, it then enters | |
| * a synchronized block to ensure that only one thread performs the refresh operation. | |
| */ | |
| void refreshIfStale() { | |
| // Fast path: Avoid entering the synchronized block if the table is not stale. | |
| if (lastRefreshTime.isAfter(Instant.now().minus(STALENESS_THRESHOLD))) { | |
| return; | |
| } | |
| synchronized (this) { | |
| if (lastRefreshTime.isBefore(Instant.now().minus(STALENESS_THRESHOLD))) { | |
| table.refresh(); | |
| lastRefreshTime = Instant.now(); | |
| if (table.spec().specId() != this.specId) { | |
| partitionFieldMap = Maps.newHashMap(); | |
| for (PartitionField partitionField : table.spec().fields()) { | |
| partitionFieldMap.put(partitionField.name(), partitionField); | |
| } | |
| this.specId = table.spec().specId(); | |
| } | |
| } | |
| } | |
| Map<String, PartitionField> partitionFieldMap = Maps.newHashMap(); | |
| for (PartitionField partitionField : table.spec().fields()) { | |
| partitionFieldMap.put(partitionField.name(), partitionField); | |
| } | |
| this.specId = table.spec().specId(); | |
| this.partitionFieldMap = partitionFieldMap; | |
| return partitionFieldMap; | |
| } | |
| private Map<String, PartitionField> getPartitionFieldMap(TableIdentifier tableId, Table table) { | |
| int currentSpecId = table.spec().specId(); | |
| if (currentSpecId == specIds.getOrDefault(tableId, Integer.MIN_VALUE)) { | |
| return partitionFieldMaps.get(tableId); | |
| } | |
| Map<String, PartitionField> fieldMap = Maps.newHashMap(); | |
| for (PartitionField partitionField : table.spec().fields()) { | |
| fieldMap.put(partitionField.name(), partitionField); | |
| } | |
| specIds.put(tableId, currentSpecId); | |
| partitionFieldMaps.put(tableId, fieldMap); | |
| return fieldMap; | |
| } |
|
Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment |
1 participant
Makes
TableCachemore robust so it can support both source and sink use cases.Previously, source and sink paths used separate table caches, which led to unnecessary
loadTableandrefreshcalls. This changes Iceberg IO to use one process-wide table cache keyed byIcebergCatalogConfigand table identifier. Callers no longer need to load/create the underlyingCatalogjust to check for a cached table. The catalog is only loaded on cache miss.Supports normal cache lookup, force refresh, and refresh if stale. Also allows sinks to pass loaders to create the table if needed.
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.