Fix temp table cleanup for continuous read in BigQueryIO

[stankiewicz]

[BigQuery] Fix temporary table leakage during continuous dynamic reads

🚨 Problem

When using BigQueryIO for continuous dynamic reads in unbounded streaming pipelines, temporary tables created to hold query results were not being reliably cleaned up. The original implementation either deleted the temporary table too early (before parallel workers fully consumed the streams) or lacked a mechanism to track stream completions across parallel workers altogether. As a result, unbounded streaming jobs over time would leak BigQuery storage resources, leaving orphaned temporary datasets and tables.

Global Window-based cleanup is insufficient for this unbounded streaming scenario, requiring a more granular stream-tracking mechanism.

🛠️ Solution

This PR implements a robust, state-based cleanup mechanism that accurately tracks when parallel workers finish reading from BigQuery Storage API streams, deleting the temporary datasets and tables only when they are truly no longer needed.

Key Changes:

1. Stateful Cleanup Tracking (CleanupTempTableDoFn):
   • Introduced a stateful DoFn that uses ValueState to track the total number of streams created for a given query job, alongside a counter for how many streams have successfully completed.
   • When the completed stream count equals the total expected streams, the DoFn safely drops the temporary BigQuery table (and the temporary dataset if it was created by Beam).
2. Stream Initialization & Side Outputs (CreateBoundedSourceForTable):
   • Augmented the source creation step to emit a side output CleanupOperationMessage.initialize() containing metadata (CleanupInfo with projectId, datasetId, tableId, and totalStreams).
3. Completion Signaling (ReadDynamicStreamSource):
   • Wrapped the underlying read operations. As each parallel stream is fully consumed, it emits a CleanupOperationMessage.streamComplete() signal to the cleanup DoFn.
   • Added a .withFromQuery() context toggle to BigQueryStorageStreamSource so consumers know if the stream is bound to a temporary query table that requires tracking.
4. Serialization Safety:
   • Replaced the non-serializable Google API TableReference inside CleanupInfo with primitive String fields to prevent NotSerializableExceptions during state persistence and shuffling.

🧪 Testing

• Added CleanupTempTableDoFnTest: Built a dedicated unit test utilizing FakeDatasetService and FakeBigQueryServices. It validates that:
  • Sequential or out-of-order stream completion signals properly aggregate in the stateful DoFn.
  • The mock BigQuery API deleteTable / deleteDataset methods are invoked exactly once after the final stream completes.
  • Attempting to access the tables post-cleanup correctly results in a 404 Not Found exception.
• Verified that spotlessApply and compileJava pass locally without format or syntax violations.

---

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

• Mention the appropriate issue in your description (for example: 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, comment fixes #<ISSUE NUMBER> instead.
• Update CHANGES.md with noteworthy changes.
• If this contribution is large, please file an Apache Individual Contributor License Agreement.

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.

[gemini-code-assist]

Summary of Changes

Hello @stankiewicz, I'm Gemini Code Assist¹! 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 addresses a critical issue in BigQueryIO where temporary tables created for continuous dynamic reads in unbounded streaming pipelines were not being reliably cleaned up, leading to resource leakage. The solution implements a robust, state-based mechanism to accurately track the completion of parallel BigQuery Storage API streams. By introducing new DoFns and state management, the system now ensures that temporary resources are only deleted when they are truly no longer needed, thereby preventing orphaned BigQuery storage resources and improving pipeline efficiency.

Highlights

• Stateful Cleanup Tracking: Introduced a new stateful CleanupTempTableDoFn that uses ValueState to track the total number of streams created for a query job and the count of completed streams. This DoFn safely deletes temporary BigQuery tables and datasets only when all associated streams have finished.
• Stream Initialization & Side Outputs: The CreateBoundedSourceForTable DoFn was augmented to emit a side output containing CleanupOperationMessage.initialize() with metadata (CleanupInfo) including project, dataset, table IDs, and the total number of streams for cleanup tracking.
• Completion Signaling: A new ReadDynamicStreamSource wraps the underlying read operations. As each parallel stream is fully consumed, it emits a CleanupOperationMessage.streamComplete() signal to the cleanup DoFn. A .withFromQuery() context toggle was added to BigQueryStorageStreamSource to identify streams requiring tracking.
• Serialization Safety: Replaced the non-serializable Google API TableReference within CleanupInfo with primitive String fields (projectId, datasetId, tableId) to prevent NotSerializableException during state persistence and data shuffling.
• Dedicated Unit Testing: Added CleanupTempTableDoFnTest to validate the stateful DoFn's behavior, ensuring proper aggregation of stream completion signals and correct, single invocation of table/dataset deletion.

🧠 New Feature in Public Preview: You can now enable Memory 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.

Changelog

• sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
  • Added imports for firstNonNull, state management classes (StateSpec, StateSpecs, ValueState), DoFn.StateId, Flatten, and PCollectionList.
  • Modified CreateBoundedSourceForTable to accept a cleanupInfoTag and use MultiOutputReceiver to output both stream sources and cleanup initialization messages.
  • Updated the expand method in DynamicRead to integrate the new cleanup mechanism, including flattening cleanup messages and applying CleanupTempTableDoFn.
  • Introduced CleanupInfo class to store serializable table reference details (project, dataset, table IDs as strings) and total stream count.
  • Introduced CleanupOperationMessage class to differentiate between stream initialization and completion signals.
  • Added CleanupTempTableDoFn, a stateful DoFn responsible for tracking stream completions and initiating the deletion of temporary BigQuery tables and datasets.
  • Created ReadDynamicStreamSource to wrap the core read operation and emit stream completion signals for temporary query tables.
• sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryStorageStreamSource.java
  • Added a fromQuery boolean field to the BigQueryStorageStreamSource constructor and class to indicate if the source originates from a temporary query table.
  • Introduced a withFromQuery() method to create a new BigQueryStorageStreamSource instance with the fromQuery flag set to true.
  • Modified fromExisting methods to correctly propagate the fromQuery flag.
  • Added a getFromQuery() getter method.
• sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/CleanupTempTableDoFnTest.java
  • Added a new unit test file CleanupTempTableDoFnTest.
  • The test uses FakeDatasetService and FakeBigQueryServices to simulate BigQuery interactions.
  • It verifies that CleanupTempTableDoFn correctly processes initialization and stream completion messages.
  • The test asserts that temporary tables and datasets are created and subsequently deleted only after all expected stream completion signals are received.

Activity

• The pull request introduces a new stateful cleanup mechanism to prevent temporary BigQuery table leakage in continuous dynamic reads.
• New classes and DoFns (CleanupInfo, CleanupOperationMessage, CleanupTempTableDoFn, ReadDynamicStreamSource) have been added to manage the lifecycle of temporary resources.
• Existing DoFns (CreateBoundedSourceForTable) and classes (BigQueryStorageStreamSource) have been modified to support the new cleanup signaling.
• A dedicated unit test (CleanupTempTableDoFnTest) has been added to thoroughly validate the new cleanup logic.
• The author has verified that spotlessApply and compileJava pass locally.

Using Gemini Code Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize 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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. 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.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[github-actions]

Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment assign set of reviewers

[github-actions]

Assigning reviewers:

R: @Abacn for label java.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).