gemini-code-assist[bot] commented on PR #37535: URL: https://github.com/apache/beam/pull/37535#issuecomment-3859652701
## Summary of Changes Hello @stankiewicz, I'm Gemini Code Assist[^1]! 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](https://codeassist.google/code-review/login) to enable Memory in your admin console.** <details> <summary><b>Changelog</b></summary> * **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. </details> <details> <summary><b>Activity</b></summary> * 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. </details> <details> <summary><b>Using Gemini Code Assist</b></summary> <br> The full guide for Gemini Code Assist can be found on our [documentation page](https://developers.google.com/gemini-code-assist/docs/review-github-code), here are some quick tips. <b>Invoking Gemini</b> 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. <b>Customization</b> 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](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github). <b>Limitations & Feedback</b> Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with :thumbsup: and :thumbsdown: 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](https://google.qualtrics.com/jfe/form/SV_2cyuGuTWsEw84yG). <b>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](https://cloud.google.com/products/gemini/code-assist).</b> </details> [^1]: Review the [Privacy Notices](https://policies.google.com/privacy), [Generative AI Prohibited Use Policy](https://policies.google.com/terms/generative-ai/use-policy), [Terms of Service](https://policies.google.com/terms), and learn how to configure Gemini Code Assist in GitHub [here](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github). Gemini can make mistakes, so double check it and [use code with caution](https://support.google.com/legal/answer/13505487). -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
