Fixes metadata propagation from doFn output to next parDo by introducing new interface and plumbing propagation

[stankiewicz]

When data is outputted with processContext, outputReceiver or builder, it's going through sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnOutputReceivers.java which was still using context.outputWindowedValue(element, ts, windows, pane ...) interface which we decided to abandon as adding new metadata requires changing this interface and we should use builder and windowedValue interface.

• adds new outputWindowedValue(WV) public interface with implementations
• add tests to cover metadata propagation across parDos in single and mulitple stages for direct runner
• plumbs metadata propagation, where it was lost previously.

---

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, 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 where metadata associated with WindowedValue objects was not fully propagated between DoFn outputs and subsequent ParDo operations. The existing output mechanism in DoFnOutputReceivers was deconstructing WindowedValue into individual components, leading to the loss of richer metadata like CausedByDrain, record IDs, and offsets. To resolve this, a new outputWindowedValue(WindowedValue) interface has been introduced, allowing DoFns to output the complete WindowedValue object directly. This change ensures that all relevant metadata is preserved and correctly plumbed through the Beam pipeline, improving data fidelity and enabling more robust processing, particularly in scenarios involving draining and record tracking.

Highlights

• New outputWindowedValue(WindowedValue) Interface: Introduced new outputWindowedValue(WindowedValue) methods in DoFn.Context and its implementations across various runners and testers to allow direct output of WindowedValue objects.
• Refactored Output Mechanism: Refactored DoFnOutputReceivers to utilize the new outputWindowedValue(WindowedValue) interface, ensuring complete metadata propagation by avoiding deconstruction of WindowedValue components.
• Enhanced Metadata Handling: Enhanced WindowedValues and ValueInSingleWindow to correctly handle and encode CausedByDrain, record IDs, and record offsets, ensuring these critical metadata fields are preserved throughout the pipeline.
• New Metadata Propagation Tests: Added new integration tests to validate the correct propagation of metadata, specifically CausedByDrain, across ParDo and Redistribute transforms.

🧠 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

• runners/core-java/src/main/java/org/apache/beam/runners/core/OutputAndTimeBoundedSplittableProcessElementInvoker.java
  • Added new outputWindowedValue methods to support direct WindowedValue output.
• runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java
  • Implemented new outputWindowedValue methods in DoFnProcessContext and DoFnStartBundleContext to handle WindowedValue objects.
• sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFn.java
  • Imported WindowedValue and added abstract outputWindowedValue methods to DoFn.Context for direct WindowedValue output.
• sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnOutputReceivers.java
  • Updated the output method to use the newly introduced outputWindowedValue(WindowedValue) interface, simplifying metadata handling.
• sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/DoFnTester.java
  • Added new outputWindowedValue methods to correctly process and store WindowedValue objects during testing.
• sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java
  • Modified the processElement method to propagate the CausedByDrain metadata using the output builder.
• sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Reify.java
  • Updated the processElement method to accept ProcessContext and capture currentRecordId, currentRecordOffset, and causedByDrain for ValueInSingleWindow.
• sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/SplittableParDoNaiveBounded.java
  • Imported WindowedValue, updated OutputBuilder to use the new outputWindowedValue interface, and added new outputWindowedValue methods to ProcessContext and StartBundleContext.
• sdks/java/core/src/main/java/org/apache/beam/sdk/values/ValueInSingleWindow.java
  • Enhanced ElementMetadata encoding to include CausedByDrain status.
• sdks/java/core/src/main/java/org/apache/beam/sdk/values/WindowedValues.java
  • Modified builder and of methods to ensure recordOffset, recordId, and causedByDrain are correctly propagated when creating WindowedValue instances.
• sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/MetadataPropagationTest.java
  • Added a new test file containing CausedByDrainSettingDoFn and CausedByDrainExtractingDoFn to verify metadata propagation across transforms.
• sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
  • Implemented new outputWindowedValue methods in various contexts to directly handle WindowedValue objects and removed a redundant causedByDrain override.

Activity

• The pull request was created by stankiewicz.
• The author provided a detailed description of the problem and the proposed solution, including a checklist for contribution.

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. ↩

[stankiewicz]

@kennknowles I'm really curious if we are able to avoid that. You've spent quite a lot of time on outputbuilder work to avoid changing public interface and this method is exposed in ProcessContext so any user can use it if they want.

[kennknowles]

Ah, yes. The old interface between the SDK and runner v1 is that the runner v1 provides a ProcessContext that will route the outputs to the correct place. When we added OutputReceiver we just made it a wrapper instead of porting runner v1 to the new design fully.

I think it is still OK to make the change you have here. The OutputBuilder work was also to make sure that the metadata was propagated, so it is still necessary. If a user choose to make a new WindowedValue from scratch and output it... well, it is not ideal but it could be worse. TBH I wish to deprecate the use of Contexts by users (ever since https://s.apache.org/a-new-dofn) since it is a hopeless interface in some ways.

I think the only alternative is to adjust runner v1 (and other runners) to implement OutputReceiver directly as though ProcessContext (and friends) do not exist. Never invested in this because of runner v2 being the focus plus it didn't seem too bad.

[stankiewicz]

few tests to fix as it looks like this approached skipped some timestamp validation.

[stankiewicz]

/gemini review

[stankiewicz]

Run Java_PVR_Prism_Loopback PreCommit

[gemini-code-assist]

Code Review

This pull request introduces a new outputWindowedValue(WindowedValue) interface to better propagate metadata, replacing the older interface that took exploded parameters. The changes are well-implemented and plumbed through various runners and components. The new tests for metadata propagation are a great addition.

I've found a few areas for improvement:

• There are a couple of instances of code duplication in new methods that could be refactored for better maintainability.
• I noticed a potential correctness issue in SimpleDoFnRunner where a timestamp check is missing in newly added methods.
• The new test class MetadataPropagationTest has a test isolation issue due to a shared static flag, which could lead to flaky tests.

Overall, this is a solid contribution that improves the SDK's metadata handling capabilities. Addressing the points above will make it even better.

[stankiewicz]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a new outputWindowedValue(WindowedValue) interface to simplify metadata propagation, deprecating the older interface that required deconstructing the WindowedValue. The changes are plumbed through various DoFn runners and contexts, and new tests are added to verify the propagation of CausedByDrain metadata. The changes are well-implemented and consistent. I have a couple of minor suggestions for the new test file to improve maintainability.

[stankiewicz]

Run Java_PVR_Prism_Loopback PreCommit

[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

[kennknowles]

This should be getting the record metadata through magic parametrs, not process context

[stankiewicz]

I don't have yet currentRecordId currentRecordOffset exposed in doFn. Let me prepare follow up PR to make it work.

[stankiewicz]

Run Java_PVR_Prism_Loopback PreCommit

[stankiewicz]

Run Java_PVR_Prism_Loopback PreCommit

[github-actions]

Assigning reviewers:

R: @chamikaramj 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).