This is an automated email from the ASF dual-hosted git repository.
scwhittle pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 0d374bf001e [Dataflow Streaming] [Multikey] Route only finalize
commits via multi key path (#40495)
0d374bf001e is described below
commit 0d374bf001e2f2727c124267e8da1b6babc23099
Author: Arun Pandian <[email protected]>
AuthorDate: Mon Oct 12 02:58:03 2026 -0700
[Dataflow Streaming] [Multikey] Route only finalize commits via multi key
path (#40495)
---
.../work/processing/StreamingWorkScheduler.java | 8 +-
.../worker/StreamingDataflowWorkerTest.java | 106 ++++++++++++++++++++-
2 files changed, 111 insertions(+), 3 deletions(-)
diff --git
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java
index aff42b8468f..9e6a88c62cd 100644
---
a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java
+++
b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java
@@ -360,8 +360,12 @@ public class StreamingWorkScheduler {
Windmill.WorkItemCommitRequest.Builder outputBuilder =
initializeOutputBuilder(workItem.getKey(), workItem);
outputBuilder.setSourceStateUpdates(Windmill.SourceState.newBuilder().setOnlyFinalize(true));
- work.setState(Work.State.COMMIT_QUEUED);
- work.queueCommit(outputBuilder.build(), computationState);
+ commitWorkBatch(
+ computationState,
+ ImmutableList.of(work),
+ ImmutableList.of(outputBuilder.build()),
+ ImmutableList.of(),
+ ImmutableList.of());
}
private StageInfo getStageInfo(ComputationState computationState) {
diff --git
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java
index 6fc19c30324..b679ca40a95 100644
---
a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java
+++
b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java
@@ -3638,7 +3638,111 @@ public class StreamingDataflowWorkerTest {
commit = result.get(3L);
assertThat(
- commit,
+ removeDynamicFields(commit),
+ equalTo(
+ parseCommitRequest(
+ "key: \"0000000000000001\" "
+ + "sharding_key: 1 "
+ + "work_token: 3 "
+ + "cache_token: 3 "
+ + "source_state_updates {"
+ + " only_finalize: true"
+ + "} ")
+ .build()));
+
+ assertThat(finalizeTracker, contains(0));
+ worker.stop();
+ }
+
+ @Test
+ public void testUnboundedSourcesDrain_multiKeyBundleEnabled() throws
Exception {
+ if (!streamingEngine) {
+ return;
+ }
+ List<Integer> finalizeTracker = Lists.newArrayList();
+ TestCountingSource.setFinalizeTracker(finalizeTracker);
+ server.clearCommitsReceived();
+
+ StreamingDataflowWorker worker =
+ makeWorker(
+
defaultWorkerParams("--experiments=unstable_enable_multi_key_bundle")
+ .setInstructions(makeUnboundedSourcePipeline())
+ .publishCounters()
+ .build());
+ worker.start();
+
+ server
+ .whenGetWorkCalled()
+ .thenReturn(
+ buildInput(
+ "work {"
+ + " computation_id: \"computation\""
+ + " input_data_watermark: 0"
+ + " work {"
+ + " key: \"0000000000000001\""
+ + " sharding_key: 1"
+ + " work_token: 2"
+ + " cache_token: 3"
+ + " key_group { high: 0 low: 1 }"
+ + " }"
+ + "}",
+ null));
+
+ Map<Long, Windmill.WorkItemCommitRequest> result =
server.waitForAndGetCommits(1);
+ Windmill.WorkItemCommitRequest commit = result.get(2L);
+ assertEquals(1, commit.getSourceStateUpdates().getFinalizeIdsCount());
+ UnsignedLong finalizeId =
+
UnsignedLong.fromLongBits(commit.getSourceStateUpdates().getFinalizeIds(0));
+
+ // Test drain work item with only_finalize.
+ server
+ .whenGetWorkCalled()
+ .thenReturn(
+ buildInput(
+ "work {"
+ + " computation_id: \"computation\""
+ + " input_data_watermark: 0"
+ + " work {"
+ + " key: \"0000000000000001\""
+ + " sharding_key: 1"
+ + " work_token: 3"
+ + " cache_token: 3"
+ + " key_group { high: 0 low: 1 }"
+ + " source_state {"
+ + " only_finalize: true"
+ + " finalize_ids: "
+ + finalizeId
+ + " }"
+ + " }"
+ + "}",
+ null));
+
+ result = server.waitForAndGetCommits(1);
+ commit = result.get(3L);
+
+ assertThat(
+ removeDynamicFields(commit),
+ equalTo(
+ parseCommitRequest(
+ "key: \"0000000000000001\" "
+ + "sharding_key: 1 "
+ + "work_token: 3 "
+ + "cache_token: 3 "
+ + "source_state_updates {"
+ + " only_finalize: true"
+ + "} ")
+ .build()));
+
+ List<Windmill.MultiKeyWorkItemCommitRequest> multiKeyCommits =
+ server.getMultiKeyCommitsReceived();
+ assertEquals(2, multiKeyCommits.size());
+ Windmill.MultiKeyWorkItemCommitRequest finalizeMultiKeyCommit =
multiKeyCommits.get(1);
+ assertEquals(
+ Windmill.Uint128Proto.newBuilder().setHigh(0).setLow(1).build(),
+ finalizeMultiKeyCommit.getKeyGroup());
+ assertEquals(1, finalizeMultiKeyCommit.getRequestsCount());
+ assertThat(
+ removeDynamicFields(finalizeMultiKeyCommit.getRequests(0)),
equalTo(
parseCommitRequest(
"key: \"0000000000000001\" "
