[Dataflow Streaming] [Multi Key] Flush streaming sinks at key boundaries and bundle completion by arunpandianp · Pull Request #39961 · apache/beam

arunpandianp · 2026-09-01T20:00:17Z

No description provided.

…ies and bundle completion When multi-key bundling is enabled in Dataflow Streaming Engine, outputs (productions) emitted during a key's processElement lifecycle must be attributed to that specific key's WorkItemCommitRequest, while outputs emitted during DoFn finishBundle must be attributed to the bundle-level MultiKeyWorkItemCommitRequest. This change: 1. Adds finishKey(@nullable Object key) to the SinkWriter interface and implements it across streaming sinks (WindmillSink, PubsubSink, PubsubDynamicSink) and non-streaming sinks. 2. In WindmillSink, PubsubSink, and PubsubDynamicSink: - When multi-key bundles are enabled: flushes to the active key's commit builder in finishKey; flushes remaining finishBundle outputs to bundle-level outputs (bundleOutputMessages / bundlePubsubMessages) in close(). - When multi-key bundles are disabled: finishKey does not flush; close() flushes all outputs directly into the single key's commit request. - Updates abort() to cleanly discard buffered messages and reset internal streams without flushing. 3. In StreamingModeExecutionContext and StreamingWorkScheduler, tracks bundleOutputMessages and bundlePubsubMessages and attaches them to MultiKeyWorkItemCommitRequest.

arunpandianp · 2026-09-03T12:27:34Z

R: @scwhittle

github-actions · 2026-09-03T12:28:54Z

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-09-07T08:32:06Z

    /** Adds a value to the sink. Returns the size in bytes of the data written. */
    public long add(ElemT value) throws IOException;

+    /** Called when all elements for a specific key have been processed. */


Should we note this is optional and only currently done for multi-key streaming execution?

This is called in streaming for all jobs. Batch does not call. Updated comment to say this is called only for streaming.

scwhittle · 2026-09-07T08:36:54Z

-      Windmill.PubSubMessageBundle pubsubMessages = outputBuilder.build();
-      if (pubsubMessages.getMessagesCount() > 0) {
-        context.getOutputBuilder().addPubsubMessages(pubsubMessages);
+      if (context.multiKeyBundleEnabled()) {


just pass multiKeyBundleEnabled() as param?

scwhittle · 2026-09-07T08:38:01Z

    @Override
    public void close() throws IOException {
-      outputBuilders.clear();
+      if (context.multiKeyBundleEnabled()) {


just pass multiKeyBundleEnabled() as param to flush

scwhittle · 2026-09-07T08:43:16Z

+    @Override
+    public void finishKey(@Nullable Object key) throws IOException {
+      if (context.multiKeyBundleEnabled()) {
+        flush(/* bundleLevel= */ false);


see above comment, is there a benefit to producing at the key level and not the bundle level? With bundle level we can reduce duplicating topic info and proto overhead

I think it is a good idea to merge the pubsub outputs. I'm thinking we can flush every 1MB or so to the current key. Do we need to worry about how it'll affect publishing with ordering keys (if/when we add support for that).

Added logic to merge outputs till they reach 1MB.

why not flush at bundle level? seems odd to attach outputs made from Key A to key B if limit is only exceeded when reaching key B.

Changed to flush at the bundle level.

scwhittle · 2026-09-07T08:44:59Z

-          context.getOutputBuilder().addOutputMessages(outputBuilder.build());
+          Windmill.OutputMessageBundle bundle = outputBuilder.build();
+          if (bundleLevel) {
+            context.addBundleOutputMessages(bundle);


same merging comments

In WindmillSink, attaching the outputs with the keys that produced them might be useful since doing so will preserve the key-key ordering as long as there are no buffered outputs. I'm thinking to keep it without cross key merging, till we figure out the story around ordering. wdyt?

scwhittle · 2026-09-15T08:28:13Z

    }

+    @Override
+    public void finishKey(@Nullable Object key) throws IOException {}


should we add a default no-op impl to the Sink base class? Since it's not called in batch anyway, forcing a lot of these batch sinks to implement it seems like unneeded clutter.

scwhittle · 2026-09-15T08:31:30Z

-      if (pubsubMessages.getMessagesCount() > 0) {
-        context.getOutputBuilder().addPubsubMessages(pubsubMessages);
+    public void finishKey(@Nullable Object key) throws IOException {
+      if (context.multiKeyBundleEnabled() && bufferedBytes >= MAX_PUBSUB_BUNDLE_BYTES) {


this will need to be per-key if we change windmill pubsub to support ordering keys. Maybe you could put a comment here and in the add method so we don't forget?

scwhittle · 2026-09-15T08:34:44Z

@@ -321,6 +323,8 @@ public void reset() {
    // these lists and maps are returned to callers after processing


see other comment, maybe we can avoid setting to emptyList and then reallocating if unused

scwhittle · 2026-09-15T08:40:24Z

what about changing context.flushState() to return ExecuteWorkResult and reset as needed?

Then we can get rid of the accessors just used for the flushing and we can be smarter about how we reset to avoid allocations.

good idea! done.

scwhittle · 2026-09-15T08:42:03Z

+
+    @Override
+    public void finishKey(@Nullable Object key) throws IOException {
+      if (context.multiKeyBundleEnabled() && bufferedBytes >= MAX_PUBSUB_BUNDLE_BYTES) {


ditto about ordered keys

scwhittle · 2026-09-15T08:44:15Z

+    @Override
+    public void finishKey(@Nullable Object key) throws IOException {
+      if (context.multiKeyBundleEnabled() && bufferedBytes >= MAX_PUBSUB_BUNDLE_BYTES) {
+        flush(/* bundleLevel= */ false);


why not flush at bundle level? It seems odd to attach output made from processing key A to key B later if the buffered size was not exceeded on A but later was on B.

why not flush at bundle level?

Wanted to spread the load across many keys. The bundle level productions eventually get attached to an arbitrary key in the backend and having lots of productions on a single key could run into limits or cause load skew.

Should we just change that implementation? If we do add key-ordering, putting these on a key may affect performance if we try to respect that they are attached to a particular key.

ack. Changed to flush at bundle level.

scwhittle · 2026-09-15T08:45:49Z

+    @Override
+    public void finishKey(@Nullable Object key) throws IOException {
+      if (context.multiKeyBundleEnabled()) {
+        flush(/* bundleLevel= */ false);


why not flush at bundle level? seems odd to attach outputs made from Key A to key B if limit is only exceeded when reaching key B.

codecov · 2026-09-16T07:46:38Z

Codecov Report

❌ Patch coverage is 87.39496% with 15 lines in your changes missing coverage. Please review.
✅ Project coverage is 55.98%. Comparing base (0b6057d) to head (36d3c71).
⚠️ Report is 142 commits behind head on master.

Files with missing lines	Patch %	Lines
...dataflow/worker/StreamingModeExecutionContext.java	79.36%	9 Missing and 4 partials ⚠️
...eam/runners/dataflow/worker/PubsubDynamicSink.java	94.44%	0 Missing and 1 partial ⚠️
...flow/worker/util/common/worker/WriteOperation.java	80.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #39961      +/-   ##
============================================
+ Coverage     55.06%   55.98%   +0.92%     
- Complexity     1722     2290     +568     
============================================
  Files          1075     1116      +41     
  Lines        171582   175894    +4312     
  Branches       1262     1486     +224     
============================================
+ Hits          94476    98471    +3995     
- Misses        74836    74911      +75     
- Partials       2270     2512     +242

Flag	Coverage Δ
java	`75.72% <87.39%> (+8.03%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

scwhittle · 2026-09-16T09:05:47Z

+    return this.bundlePubsubMessages;
+  }
+
+  private List<Windmill.WorkItemCommitRequest> getOrCreateWorkItemCommits() {


how about inlining this? it has a single usage and we can create a single-item list (like Immutable list which has specialized SingletonImmutableList) if not using multikeybundles and enforce only one key is present

could also inline the pubsub ones since they also just have single usage

scwhittle · 2026-09-16T09:16:45Z

+          if (bundleLevel) {
+            context.addBundleOutputMessages(bundle);
+          } else {
+            context.getOutputBuilder().addOutputMessages(bundle);


maybe we should change getOutputBuilder to be getKeyOutputBuilder? could do separately.

scwhittle · 2026-09-16T09:19:07Z

+          }
+        }
+      } finally {
+        outputBuilder = createOutputBuilder();


previously we cleared the builder (though see other comment that seems buggy).

However should we consider just clearing the messages from the bundle so we don't have to recreate builder, and it can cache message array etc?

scwhittle · 2026-09-16T09:23:06Z

-      if (pubsubMessages.getMessagesCount() > 0) {
-        context.getOutputBuilder().addPubsubMessages(pubsubMessages);
-      }
-      outputBuilder.clear();


how did this work previously? It seems if we ever closed and then reused this object, the outputBuilder would not have the topic etc set. Are these sinks not actually reused in streaming?

scwhittle · 2026-09-16T09:26:20Z

+      context.finishKey();
+      WindmillComputationKey computationKey = context.getComputationKey();
+      ExecuteWorkResult result = context.flushStateAndReset();
+      Windmill.WorkItemCommitRequest commitRequest = result.workItemCommits().get(0);


assert there is a single commit

scwhittle · 2026-09-16T09:26:46Z

+    ExecuteWorkResult result = executionContext.flushStateAndReset();

-    assertEquals(1234, executionContext.getOutputBuilder().getSourceBacklogBytes());
+    assertEquals(1234, result.workItemCommits().get(0).getSourceBacklogBytes());


assert there is a single commit

arunpandianp added 2 commits September 1, 2026 19:33

Merge remote-tracking branch 'beam/master' into flushWindmillSink

1027f56

github-actions Bot added runners dataflow labels Sep 1, 2026

arunpandianp added 2 commits September 3, 2026 12:06

improve tests

4cbbf84

add checks

9deeefc

arunpandianp marked this pull request as ready for review September 3, 2026 12:27

scwhittle requested changes Sep 7, 2026

View reviewed changes

arunpandianp force-pushed the flushWindmillSink branch from f8fd69b to 624b1ab Compare September 13, 2026 22:15

review comments

2b2be9e

arunpandianp force-pushed the flushWindmillSink branch from 624b1ab to 2b2be9e Compare September 13, 2026 22:34

arunpandianp requested a review from scwhittle September 13, 2026 22:37

spotless

1fec61c

scwhittle requested changes Sep 15, 2026

View reviewed changes

arunpandianp added 2 commits September 16, 2026 00:46

address comments

1ddbe0f

Refactor StreamingModeExecutionContext

a93c0e2

arunpandianp requested a review from scwhittle September 16, 2026 04:54

fix comment

36d3c71

scwhittle requested changes Sep 16, 2026

View reviewed changes

		@@ -321,6 +323,8 @@ public void reset() {
		// these lists and maps are returned to callers after processing

Uh oh!

Conversation

arunpandianp commented Sep 1, 2026

Uh oh!

arunpandianp commented Sep 3, 2026

Uh oh!

github-actions Bot commented Sep 3, 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!

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!

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!

codecov Bot commented Sep 16, 2026

Codecov Report

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