Refactor: separate SplittableTruncateSizedRestrictions by kennknowles · Pull Request #35021 · apache/beam

kennknowles · 2025-05-21T16:23:27Z

This follows #34919, splis out calls to @TruncateRestriction. This is the first Fn API transform that deals with responding to splits so it is a bit more complex than the prior refactors that only touched @GetSize and @GetInitialRestriction and @SplitRestriction that do not respond to split requests.

As I understand it, this responds to split requests by delegating them downstream, as it is expected to be fused with final element+restriction processing.

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.

kennknowles · 2025-06-16T22:22:19Z

R: @scwhittle

github-actions · 2025-06-16T22:23:20Z

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

kennknowles · 2025-06-16T22:23:43Z

I was spending my time on #34902 but honestly these refactors make that a lot more manageable. So I put an hour into this one today to fix the last remaining harness test failure. CC @stankiewicz @talatuyarer

kennknowles · 2025-06-17T20:09:08Z

Please do review. All the failing tests pass locally and I believe are flaking. I will keep kicking them and trying to repro locally but so far non are actually broken.

kennknowles · 2025-06-17T21:39:30Z

Green! Huzzah!

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java

.../src/main/java/org/apache/beam/fn/harness/SplittableTruncateSizedRestrictionsDoFnRunner.java

scwhittle · 2025-06-20T13:19:58Z

.../src/main/java/org/apache/beam/fn/harness/SplittableTruncateSizedRestrictionsDoFnRunner.java

+    // Note that the assumption here is the fullInputCoder of the Truncate transform should be the
+    // the same as the SDF/Process transform.
+    Coder<WindowedValue<?>> fullInputCoder =
+        (Coder<WindowedValue<?>>) (Object) WindowedValues.getFullCoder(inputCoder, windowCoder);


can this be enforced earlier during constuction?

Hmm. I don't really understand this comment TBH.

I thought having the coders line up is an invariant on the ProcessBundleDescriptor. I had actually thought that it would be up to the descriptor to say what coder to use, and that in the Fn API the fact that the data place is using a WindowedValueCoder is explicit in the instruction graph. But I do see that we pull the coder from the proto and it is not a WindowedValueCoder, so I guess we have left that implicit? That actually seems like a problem we may have to fix...

For now I actually think just removing this comment is fair, no? This is just the same process that is used in all the variants of the FnApiDoFnRunner. I think the TODO would be "just use the coder the runner tells you to"? Like... we want to be able to tell the harness to use coders that may or may not be windowed value coders? As in, the harness would be a UDF server that just does as it is told with as few assumptions as possible.

I was wondering if this cast could ever fail, and if it could should we be creating the fullInputCoder during the construction of this class so we fail earlier during setup, since I believe we have inputCoder and windowCoder available then.

The cast can't fail because it is actually just forgetting some actual type parameters for type compatibility with other code. It converts FullWindowedValueCoder<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>> to Coder<WindowedValue<?>>.

Because the erased type is Coder in all cases it literally cannot fail.

It can be incorrect, but we don't have a way to catch it until it goes weirdly wrong later (as with all such "unchecked" casts)

The wacky cast "through" Object avoids raw types, which would require suppression. This used to be a cast through the raw type Coder.

But it can, indeed, be constructed earlier. I've moved it to the constructor.

.../src/main/java/org/apache/beam/fn/harness/SplittableTruncateSizedRestrictionsDoFnRunner.java

scwhittle

Just some minor comments

scwhittle · 2025-06-25T13:25:40Z

sdks/java/core/src/main/java/org/apache/beam/sdk/util/Holder.java

+ * check values for null to determine presence/absence. Instead you can store a {@code @Nullable
+ * Holder<T>}.
+ */
+public class Holder<T> {


mark Internal?

Technically, all of sdk/util is "internal". Not sure whether that or an annotation will stop a user from complaining if we break it though :-). I added the annotation.

scwhittle · 2025-06-25T13:26:38Z

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/ProgressUtils.java

+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+
+/** Miscellaneous methods for working with progress. */
+public abstract class ProgressUtils {


Technically, users should never have a compile-time dep on org.apache.beam.fn.harness. The SDK harness is a runtime-only non-user-facing library. I added the annotation anyhow, though I worry that suggests other classes in it are not internal..

scwhittle · 2025-06-25T13:30:44Z

.../src/main/java/org/apache/beam/fn/harness/SplittableTruncateSizedRestrictionsDoFnRunner.java

+    // Note that the assumption here is the fullInputCoder of the Truncate transform should be the
+    // the same as the SDF/Process transform.
+    Coder<WindowedValue<?>> fullInputCoder =
+        (Coder<WindowedValue<?>>) (Object) WindowedValues.getFullCoder(inputCoder, windowCoder);


I was wondering if this cast could ever fail, and if it could should we be creating the fullInputCoder during the construction of this class so we fail earlier during setup, since I believe we have inputCoder and windowCoder available then.

kennknowles

Thanks for the thorough review! I will get the tests green, run our internal tests, then merge.

kennknowles · 2025-07-01T19:20:25Z

.../src/main/java/org/apache/beam/fn/harness/SplittableTruncateSizedRestrictionsDoFnRunner.java

+    // Note that the assumption here is the fullInputCoder of the Truncate transform should be the
+    // the same as the SDF/Process transform.
+    Coder<WindowedValue<?>> fullInputCoder =
+        (Coder<WindowedValue<?>>) (Object) WindowedValues.getFullCoder(inputCoder, windowCoder);


The cast can't fail because it is actually just forgetting some actual type parameters for type compatibility with other code. It converts FullWindowedValueCoder<KV<KV<InputT, KV<RestrictionT, WatermarkEstimatorStateT>>, Double>> to Coder<WindowedValue<?>>.

Because the erased type is Coder in all cases it literally cannot fail.

It can be incorrect, but we don't have a way to catch it until it goes weirdly wrong later (as with all such "unchecked" casts)

The wacky cast "through" Object avoids raw types, which would require suppression. This used to be a cast through the raw type Coder.

But it can, indeed, be constructed earlier. I've moved it to the constructor.

kennknowles · 2025-07-01T19:23:05Z

sdks/java/core/src/main/java/org/apache/beam/sdk/util/Holder.java

+ * check values for null to determine presence/absence. Instead you can store a {@code @Nullable
+ * Holder<T>}.
+ */
+public class Holder<T> {


Technically, all of sdk/util is "internal". Not sure whether that or an annotation will stop a user from complaining if we break it though :-). I added the annotation.

kennknowles · 2025-07-01T19:24:19Z

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/ProgressUtils.java

+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+
+/** Miscellaneous methods for working with progress. */
+public abstract class ProgressUtils {


Technically, users should never have a compile-time dep on org.apache.beam.fn.harness. The SDK harness is a runtime-only non-user-facing library. I added the annotation anyhow, though I worry that suggests other classes in it are not internal..

Abacn · 2025-07-24T14:52:09Z

This PR breaks PubsubIO Load Test, specifically a use case of UnboundedSyntheticSource stopped producing data, see #35194 (comment) for details.

github-actions bot added the java label May 21, 2025

kennknowles force-pushed the FnApiDoFnRunner-TruncateRestriction branch 5 times, most recently from b5b739b to 5178d01 Compare June 16, 2025 22:21

kennknowles marked this pull request as ready for review June 16, 2025 22:22

github-actions bot added runners local labels Jun 17, 2025

kennknowles force-pushed the FnApiDoFnRunner-TruncateRestriction branch 2 times, most recently from 300dc36 to bc12475 Compare June 17, 2025 17:00

github-actions bot removed runners local labels Jun 17, 2025

kennknowles force-pushed the FnApiDoFnRunner-TruncateRestriction branch from bc12475 to 89fbea4 Compare June 17, 2025 18:11

scwhittle requested changes Jun 20, 2025

View reviewed changes

kennknowles force-pushed the FnApiDoFnRunner-TruncateRestriction branch from 89fbea4 to f577796 Compare June 20, 2025 18:20

scwhittle approved these changes Jun 25, 2025

View reviewed changes

Refactor: separate SplittableTruncateSizedRestrictions

5b4c05e

kennknowles force-pushed the FnApiDoFnRunner-TruncateRestriction branch from f577796 to 5b4c05e Compare July 1, 2025 19:38

kennknowles commented Jul 1, 2025

View reviewed changes

kennknowles merged commit 8eba198 into apache:master Jul 18, 2025
25 of 27 checks passed

kennknowles deleted the FnApiDoFnRunner-TruncateRestriction branch July 18, 2025 21:51

Abacn mentioned this pull request Jul 24, 2025

The LoadTests Java PubsubIO job is flaky #35194

Closed

Abacn mentioned this pull request Jul 25, 2025

Revert "Refactor: separate SplittableTruncateSizedRestrictions" #35700

Merged

3 tasks

Conversation

kennknowles commented May 21, 2025

GitHub Actions Tests Status (on master branch)

Uh oh!

kennknowles commented Jun 16, 2025

Uh oh!

github-actions bot commented Jun 16, 2025

Uh oh!

kennknowles commented Jun 16, 2025

Uh oh!

kennknowles commented Jun 17, 2025

Uh oh!

kennknowles commented Jun 17, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

scwhittle left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kennknowles left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Abacn commented Jul 24, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants