[Dataflow Streaming] Commit size validation for multi key commits by arunpandianp · Pull Request #39473 · apache/beam

arunpandianp · 2026-07-24T07:56:54Z

Please add a meaningful description for your change here

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 · 2026-07-24T07:56:57Z

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

…_failed

This reverts commit 3670a03.

This reverts commit 2c12786.

…_failed

arunpandianp · 2026-08-14T02:43:41Z

R: @scwhittle

github-actions · 2026-08-14T02:44:48Z

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

scwhittle · 2026-08-14T11:08:52Z

      }
+
+      // MultiKeyBatchingDisabled items should not be in keyGroupWorkList
+      checkState(!Node.isMultiKeyBatchingDisabled(firstNode.task));


do we need Node. ?

Yes, because the code here is not inside Node. static import fails saying the method is private.

scwhittle · 2026-08-14T11:18:53Z

      return;
    }

+    // Look at budgetHandle instead of executedWorks because when intermediate work items are


Comment was a little unclear at first. Maybe something like the following?

// If this is a multi-key work item, then we need to retry all of the individual work items without merging so that we can identify large commits to truncate. We determine the work items that were part of the bundle by looking at the budgethandle instead of executedWorks because validateCommitRequestSize is called when transitioning and the handle has been updated but executedWorks has not.

scwhittle · 2026-08-14T11:22:33Z

    if (elapsedNanos >= multiKeyBundleOptions.maxKeyGroupBatchTimeNanos()) {
      return true;
    }
    return getBytesSinked() >= multiKeyBundleOptions.maxKeyGroupBatchSinkBytes();


bytessinked currently is just the output messages. There is a todo to measure state. I think that we should update the flushing of key state to add in the state bytes so that we can stop merging keys if each key is modifying state and it adds up. Then we will be less likely to need to truncate since the sinkbytes limit is lower than the max limit.

It is a good idea to stop merging when we hit a limit including state updates. We can accumulate the individual commit size from flushStateInternal and compare that against a threshold. Will send a separate change for that.

We flush key state per key and truncation is only based on individual key commit size, so I'm not sure if counting the state updates will reduce truncations.

Maybe my terminology was confusing, but I was thinking of the case where each key is flushing to state that we might build up too large a commit to apply. Then we have to retry all of the bundles locally without merging instead of just stopping earlier. If we were counting the state bytes we could just stop before we built up such a large commit and wouldn't have to retry.

…_failed

…dation

…ation

scwhittle · 2026-08-17T11:13:47Z

    Node(Runnable task) {
      this.task = task;
-      if (task instanceof QueuedWork) {
+      if (task instanceof QueuedWork && !isMultiKeyBatchingDisabled(task)) {


could just be
if (!isMultiKeyBatchingDisabled(task)

scwhittle · 2026-08-17T11:15:59Z

          computationState, handle.getWorkBatch(), computationId, systemName, work, t);
    } finally {
-      List<Work> processedWorkBatch = workBatch != null ? workBatch : ImmutableList.of(work);
+      List<Work> processedWorkBatch = workBatch != null ? workBatch : handle.getWorkBatch();


is it always correct to use the handle batch? should we just use that and possibly remove from ExecuteWorkResult as well?

good idea, done and removed executedWorks. I added handle.getWorkBatch after adding executedWorks and didn't think about it to replace executedWorks

scwhittle · 2026-08-17T11:18:49Z

              .build());
    }

+    Windmill.MultiKeyWorkItemCommitRequest multiKeyCommitRequest = multiKeyBuilder.build();


does this possibly throw and thus is done earlier now? If so add a comment, if not maybe revert.

scwhittle · 2026-08-17T11:23:08Z

    if (elapsedNanos >= multiKeyBundleOptions.maxKeyGroupBatchTimeNanos()) {
      return true;
    }
    return getBytesSinked() >= multiKeyBundleOptions.maxKeyGroupBatchSinkBytes();


Maybe my terminology was confusing, but I was thinking of the case where each key is flushing to state that we might build up too large a commit to apply. Then we have to retry all of the bundles locally without merging instead of just stopping earlier. If we were counting the state bytes we could just stop before we built up such a large commit and wouldn't have to retry.

scwhittle · 2026-08-17T11:24:07Z

+    StreamingDataflowWorker worker =
+        makeWorker(
+            defaultWorkerParams(
+                    "--experiments=unstable_enable_multi_key_bundle,windmill_max_key_group_batch_time_ms=5000",


larger timeout just so it isn't possibly flaky?

scwhittle · 2026-08-17T11:26:15Z

+    List<ParallelInstruction> instructions =
+        Arrays.asList(
+            makeSourceInstruction(kvCoder),
+            makeDoFnInstruction(new FixedSizeCommitFn(500), 0, kvCoder),


maybe have a dofn that you can see that the work is re-executed?

done. Updated tests to use a DoFn writing to state instead of using productions. Added a todo to add tests using productions. Need to send a WindmillSink change to flush productions in finishKey.

scwhittle · 2026-08-17T11:31:04Z

+    Windmill.Uint128Proto keyGroup =
+        Windmill.Uint128Proto.newBuilder().setHigh(1).setLow(2).build();
+
+    Work work1 =


I only see one work in this test so it doesn't seem like it is necessarily stopping because the batching is disabled.

Updated and renamed the test.

scwhittle · 2026-08-18T20:42:51Z

does this need to use the handle works as well? Can we get rid of executedWorks?

Good idea, removed executedWorks

arunpandianp added 2 commits June 11, 2026 13:27

Drop failed work in BoundedQueueExecutor::pollWork

2c12786

address comment

3670a03

github-actions Bot added runners dataflow labels Jul 24, 2026

arunpandianp added 13 commits August 5, 2026 22:57

Merge remote-tracking branch 'beam/master' into multikey_queue_filter…

546739c

…_failed

Revert "address comment"

2ddc10e

This reverts commit 3670a03.

Revert "Drop failed work in BoundedQueueExecutor::pollWork"

fc15302

This reverts commit 2c12786.

[Dataflow Streaming] Remove finalizeCommits from processWork

7938f7f

Merge branch 'fixfinalizer' into multikey_queue_filter_failed

ebfe88a

Plumb ComputationState to ProcessingContext

fce54b0

Drop failed workitems during pollwork

840ecda

Improve tests

bbda033

Merge remote-tracking branch 'beam/master' into multikey_queue_filter…

bc9b4d0

…_failed

address comments

88b7b18

Merge remote-tracking branch 'beam/master' into multikey_queue_filter…

cce0bb3

…_failed

address comments

a02260e

address comments

b3d191a

arunpandianp force-pushed the multikey_commit_validation branch 2 times, most recently from 130dafe to 8553e14 Compare August 14, 2026 00:26

[Dataflow Streaming] Commit size validation for multi key commits

49db96a

arunpandianp force-pushed the multikey_commit_validation branch from 8553e14 to 49db96a Compare August 14, 2026 00:50

arunpandianp changed the title ~~[wip] Multikey commit validation~~ [Dataflow Streaming] Commit size validation for multi key commits Aug 14, 2026

arunpandianp marked this pull request as ready for review August 14, 2026 02:43

scwhittle requested changes Aug 14, 2026

View reviewed changes

arunpandianp added 3 commits August 14, 2026 22:06

Merge remote-tracking branch 'beam/master' into multikey_queue_filter…

1c1c988

…_failed

Merge branch 'multikey_queue_filter_failed' into multikey_commit_vali…

444eb1c

…dation

address comments

96c3b0c

arunpandianp added 3 commits August 15, 2026 00:12

fix test

cc27c02

Merge remote-tracking branch 'beam/master' into multikey_commit_valid…

ebf8b34

…ation

fix merge

dd8267a

scwhittle requested changes Aug 17, 2026

View reviewed changes

arunpandianp added 3 commits August 18, 2026 16:51

address comments

0505956

address comments

d4ae796

remove unrelated diff

61a4e0f

scwhittle requested changes Aug 18, 2026

View reviewed changes

address comments

48b6e82

Uh oh!

Conversation

arunpandianp commented Jul 24, 2026

GitHub Actions Tests Status (on master branch)

Uh oh!

gemini-code-assist Bot commented Jul 24, 2026

Uh oh!

arunpandianp commented Aug 14, 2026

Uh oh!

github-actions Bot commented Aug 14, 2026

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!

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!

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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants