This is an automated email from the ASF dual-hosted git repository.
je-ik pushed a commit to branch feat/18479-kafka-streams-runner-skeleton
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to
refs/heads/feat/18479-kafka-streams-runner-skeleton by this push:
new ff323ebd6db [GSoC 2026] Kafka Streams runner #39273: Kafka Streams
runner: Flatten support
ff323ebd6db is described below
commit ff323ebd6dbcb47fa1dd11c300aad721527bd0b8
Author: M Junaid Shaukat <[email protected]>
AuthorDate: Wed Jul 15 15:19:13 2026 +0500
[GSoC 2026] Kafka Streams runner #39273: Kafka Streams runner: Flatten
support
---
.../src/main/proto/kafka_streams_payload.proto | 14 +-
.../translation/ExecutableStageProcessor.java | 66 ++++---
.../translation/ExecutableStageTranslator.java | 8 +-
.../streams/translation/FlattenProcessor.java | 119 ++++++++++++
.../streams/translation/FlattenTranslator.java | 82 ++++++++
.../streams/translation/GroupByKeyProcessor.java | 56 ++++--
.../streams/translation/GroupByKeyTranslator.java | 13 +-
.../streams/translation/ImpulseProcessor.java | 13 +-
.../kafka/streams/translation/KStreamsPayload.java | 27 ++-
.../streams/translation/KStreamsPayloadSerde.java | 2 +
.../KafkaStreamsPipelineTranslator.java | 1 +
.../kafka/streams/translation/ReadProcessor.java | 13 +-
.../streams/translation/ShuffleByKeyProcessor.java | 10 +
.../streams/translation/WatermarkAggregator.java | 107 +++++++++++
.../streams/translation/WatermarkPayload.java | 18 +-
.../ExecutableStageProcessorWatermarkTest.java | 21 +-
.../kafka/streams/translation/FlattenTest.java | 212 +++++++++++++++++++++
.../translation/KStreamsPayloadSerdeTest.java | 7 +-
.../translation/WatermarkAggregatorTest.java | 160 ++++++++++++++++
19 files changed, 873 insertions(+), 76 deletions(-)
diff --git
a/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto
b/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto
index 87e3de44397..4dadf2d14b7 100644
--- a/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto
+++ b/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto
@@ -27,16 +27,22 @@ option java_outer_classname = "KafkaStreamsPayloadProtos";
// topic boundary (e.g. the GroupByKey repartition topic and the watermark
fan-out). Protobuf is
// used for compatible schema evolution and compact varint encoding.
message KafkaStreamsPayload {
- // A watermark report: the watermark plus the in-band coordination fields
the downstream
- // WatermarkManager needs.
+ // A watermark report: the watermark plus the in-band coordination fields a
downstream
+ // watermark aggregator needs to reconstruct its input watermark.
message WatermarkPayload {
// Event-time watermark in milliseconds. Signed (sint64, zigzag-encoded)
because Beam event
// times can be negative, e.g. BoundedWindow.TIMESTAMP_MIN_VALUE.
sint64 millis = 1;
- // The source partition this report is for.
+ // Which partition (physical instance) of the producing transform this
report is for, in
+ // [0, total_partitions).
uint32 source_partition = 2;
- // The total number of source partitions feeding the downstream stage.
+ // How many partitions (physical instances) the producing transform has in
total.
uint32 total_partitions = 3;
+ // Globally unique id of the transform that produced this report. A
producer stamps its own id
+ // without regard to who consumes the report; a consumer with several
upstream transforms
+ // (e.g. Flatten) aggregates per producing transform, holding its output
watermark until every
+ // partition of every upstream transform it expects has reported.
+ string transform_id = 4;
}
// A data element: the Beam WindowedValue encoded with the PCollection's
windowed-value coder.
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java
index 00f85032d04..38d3601e814 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java
@@ -18,6 +18,7 @@
package org.apache.beam.runners.kafka.streams.translation;
import java.util.Queue;
+import java.util.Set;
import java.util.concurrent.ConcurrentLinkedQueue;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.runners.fnexecution.control.BundleProgressHandler;
@@ -51,12 +52,12 @@ import org.slf4j.LoggerFactory;
* ProcessorContext#forward} must only be called from the processing thread,
so outputs are never
* forwarded directly from a harness callback.
*
- * <p>A {@link KStreamsPayload#isWatermark() watermark} payload is a
per-source-partition report and
- * marks a bundle boundary: the open bundle (if any) is closed (flushing
outputs), the report is fed
- * to the {@link WatermarkManager}, and the stage's output watermark is
forwarded downstream only
- * when the {@code min()} across its source partitions actually advances.
Until every source
- * partition has reported, the watermark is held and nothing is forwarded —
but data is still
- * processed in the meantime.
+ * <p>A {@link KStreamsPayload#isWatermark() watermark} payload is a report
from one partition of
+ * one upstream transform and marks a bundle boundary: the open bundle (if
any) is closed (flushing
+ * outputs), the report is fed to the {@link WatermarkAggregator}, and the
stage's output watermark
+ * is forwarded downstream — stamped with this stage's own transform id — only
when the aggregate
+ * across the upstream transform's partitions actually advances. Until every
partition has reported,
+ * the watermark is held and nothing is forwarded — but data is still
processed in the meantime.
*
* <p>This is the Kafka Streams analogue of Flink's {@code
ExecutableStageDoFnOperator} and Spark's
* {@code SparkExecutableStageFunction}. State, timers, and side inputs are
out of scope for this
@@ -70,6 +71,9 @@ class ExecutableStageProcessor
private final RunnerApi.ExecutableStagePayload stagePayload;
private final JobInfo jobInfo;
+ // This stage's own transform id, stamped on every watermark it forwards so
downstream watermark
+ // aggregators know which transform the report came from — regardless of who
consumes it.
+ private final String transformId;
// pendingOutputs is enqueued by SDK harness threads (inside the
OutputReceiverFactory callback)
// and drained by the Kafka Streams processing thread on bundle close; needs
to be thread-safe.
@@ -79,9 +83,9 @@ class ExecutableStageProcessor
// only safe because the Impulse output coder happens to be ByteArrayCoder.
private final Queue<WindowedValue<?>> pendingOutputs = new
ConcurrentLinkedQueue<>();
- // Computes this stage's output watermark as min() over its source
partitions' reported
- // watermarks, holding until every source partition has reported (see
WatermarkManager).
- private final WatermarkManager watermarkManager = new WatermarkManager();
+ // Computes this stage's input watermark from its upstream transform's
reports, holding until
+ // every partition of the upstream transform has reported (see
WatermarkAggregator).
+ private final WatermarkAggregator watermarkAggregator;
// The last watermark actually forwarded downstream, so we only forward when
it advances.
private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE;
@@ -90,9 +94,20 @@ class ExecutableStageProcessor
private @Nullable StageBundleFactory stageBundleFactory;
private @Nullable RemoteBundle currentBundle;
- ExecutableStageProcessor(RunnerApi.ExecutableStagePayload stagePayload,
JobInfo jobInfo) {
+ /**
+ * @param transformId this stage's own transform id, stamped on the
watermarks it emits
+ * @param upstreamTransformIds the transform ids feeding this stage (known
from the pipeline
+ * graph), whose reports the {@link WatermarkAggregator} waits for
+ */
+ ExecutableStageProcessor(
+ RunnerApi.ExecutableStagePayload stagePayload,
+ JobInfo jobInfo,
+ String transformId,
+ Set<String> upstreamTransformIds) {
this.stagePayload = stagePayload;
this.jobInfo = jobInfo;
+ this.transformId = transformId;
+ this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds);
}
@Override
@@ -116,18 +131,21 @@ class ExecutableStageProcessor
@Override
public void process(Record<byte[], KStreamsPayload<?>> record) {
KStreamsPayload<?> payload = record.value();
+ if (payload == null) {
+ // A topic feeding the runner can always be written to from outside (or
carry a tombstone),
+ // so recover from the obvious error instead of crashing the task: warn
and drop.
+ LOG.warn(
+ "Stage {} dropping record with null payload (external write or
tombstone)", transformId);
+ return;
+ }
if (payload.isWatermark()) {
// Emit any buffered outputs before the watermark. Data is processed
regardless of watermark
// readiness; only the watermark itself is held until every source
partition has reported.
closeBundleAndFlush(record);
- // Feed the report into the WatermarkManager and forward the stage's
output watermark only
- // when min() across the source partitions actually advances, not on
every received watermark.
- WatermarkPayload report = payload.asWatermark();
- watermarkManager.observe(
- report.getSourcePartition(),
- new Instant(report.getWatermarkMillis()),
- report.getTotalSourcePartitions());
- Instant advanced = watermarkManager.advance();
+ // Feed the report into the aggregator and forward the stage's output
watermark only when the
+ // aggregate across the upstream transform's partitions actually
advances.
+ watermarkAggregator.observe(payload.asWatermark());
+ Instant advanced = watermarkAggregator.advance();
if (advanced.isAfter(lastForwardedWatermark)) {
lastForwardedWatermark = advanced;
forwardWatermark(record, advanced.getMillis());
@@ -203,14 +221,16 @@ class ExecutableStageProcessor
}
private void forwardWatermark(Record<byte[], KStreamsPayload<?>> record,
long watermarkMillis) {
- // This stage is a single instance for now, so it forwards its watermark
as the only source
- // partition (0 of 1). Fanning the watermark out to every downstream
partition — and producing
- // it atomically with the offset commit so it is durable — lands with the
topic-based shuffle
- // work, when there are real source partitions to track (#18479).
+ // Stamped with this stage's own transform id; this stage is a single
instance for now, so the
+ // report is for its only partition (0 of 1). Fanning the watermark out to
every downstream
+ // partition — and producing it atomically with the offset commit so it is
durable — lands with
+ // the topic-based shuffle work (#18479).
ProcessorContext<byte[], KStreamsPayload<?>> ctx =
checkInitialized(context);
ctx.forward(
new Record<byte[], KStreamsPayload<?>>(
- record.key(), KStreamsPayload.watermark(watermarkMillis, 0, 1),
record.timestamp()));
+ record.key(),
+ KStreamsPayload.watermark(watermarkMillis, transformId, 0, 1),
+ record.timestamp()));
}
@Override
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java
index dc56d57f57c..a2e6ed837c7 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java
@@ -19,6 +19,7 @@ package org.apache.beam.runners.kafka.streams.translation;
import java.io.IOException;
import org.apache.beam.model.pipeline.v1.RunnerApi;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
import org.apache.kafka.streams.Topology;
@@ -83,9 +84,14 @@ class ExecutableStageTranslator implements
PTransformTranslator {
String parentProcessor =
context.getProcessorNameForPCollection(inputPCollectionId);
Topology topology = context.getTopology();
+ // The stage stamps its own transform id on the watermarks it emits, and
aggregates its input
+ // watermark from the reports of its single upstream transform (the
producer of its input
+ // PCollection, whose node name is the upstream transform id).
topology.addProcessor(
transformId,
- () -> new ExecutableStageProcessor(stagePayload, context.getJobInfo()),
+ () ->
+ new ExecutableStageProcessor(
+ stagePayload, context.getJobInfo(), transformId,
ImmutableSet.of(parentProcessor)),
parentProcessor);
if (!transform.getOutputsMap().isEmpty()) {
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java
new file mode 100644
index 00000000000..d09fed8185a
--- /dev/null
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java
@@ -0,0 +1,119 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.kafka.streams.translation;
+
+import java.util.Set;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.kafka.streams.processor.api.Processor;
+import org.apache.kafka.streams.processor.api.ProcessorContext;
+import org.apache.kafka.streams.processor.api.Record;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Kafka Streams {@link Processor} implementing Beam's {@code Flatten}
primitive ({@code
+ * beam:transform:flatten:v1}): the union of N input PCollections into one
output PCollection.
+ *
+ * <p><b>Data</b> records are forwarded straight through unchanged — the merge
of the N parents'
+ * data streams <em>is</em> the flatten.
+ *
+ * <p><b>Watermark</b> reports are where Flatten does real work, and it owns
its output watermark
+ * the same way GroupByKey does: it runs a {@link WatermarkAggregator} over
its inputs, forwards its
+ * own watermark only when the {@code min()} across them advances, and stamps
that as a single
+ * source ({@code 0 of 1}) to its downstream. This holds the output watermark
back until
+ * <em>every</em> input branch has reported, so a downstream GroupByKey does
not fire before all
+ * flattened branches are drained.
+ *
+ * <p>The {@link WatermarkAggregator} tells the input branches apart by the
transform id each
+ * branch's producer stamps on its watermark (Kafka Streams does not tell a
processor which parent
+ * forwarded a record). Each producer stamps its own identity regardless of
who consumes it, so a
+ * PCollection feeding several Flattens reports one identity and every Flatten
still waits only for
+ * the upstream transforms it expects — the set handed to it at construction
from the pipeline
+ * graph.
+ */
+class FlattenProcessor
+ implements Processor<byte[], KStreamsPayload<?>, byte[],
KStreamsPayload<?>> {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(FlattenProcessor.class);
+
+ // This transform's own id, stamped on every watermark it forwards
downstream.
+ private final String transformId;
+ // Computes the output watermark as min() over the upstream transforms'
reports, holding until
+ // every partition of every expected upstream transform has reported (see
WatermarkAggregator).
+ private final WatermarkAggregator watermarkAggregator;
+ // The last watermark actually forwarded downstream, so we only forward when
it advances.
+ private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE;
+
+ private @Nullable ProcessorContext<byte[], KStreamsPayload<?>> context;
+
+ /**
+ * @param transformId this Flatten's own transform id, stamped on the
watermarks it emits
+ * @param upstreamTransformIds the producers of this Flatten's input
PCollections (known from the
+ * pipeline graph), whose reports the {@link WatermarkAggregator} waits
for
+ */
+ FlattenProcessor(String transformId, Set<String> upstreamTransformIds) {
+ this.transformId = transformId;
+ this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds);
+ }
+
+ @Override
+ public void init(ProcessorContext<byte[], KStreamsPayload<?>> context) {
+ this.context = context;
+ }
+
+ @Override
+ public void process(Record<byte[], KStreamsPayload<?>> record) {
+ KStreamsPayload<?> payload = record.value();
+ if (payload == null) {
+ // A topic feeding the runner can always be written to from outside (or
carry a tombstone),
+ // so recover from the obvious error instead of crashing the task: warn
and drop.
+ LOG.warn(
+ "Flatten {} dropping record with null payload (external write or
tombstone)",
+ transformId);
+ return;
+ }
+ ProcessorContext<byte[], KStreamsPayload<?>> ctx =
checkInitialized(context);
+ if (!payload.isWatermark()) {
+ // Data: the union of the parents' data streams is the flatten — forward
unchanged.
+ ctx.forward(record);
+ return;
+ }
+ watermarkAggregator.observe(payload.asWatermark());
+ Instant advanced = watermarkAggregator.advance();
+ if (advanced.isAfter(lastForwardedWatermark)) {
+ lastForwardedWatermark = advanced;
+ // Stamped with this Flatten's own transform id; Flatten is a single
instance for now, so the
+ // report is for its only partition (0 of 1).
+ ctx.forward(
+ new Record<byte[], KStreamsPayload<?>>(
+ record.key(),
+ KStreamsPayload.watermark(advanced.getMillis(), transformId, 0,
1),
+ record.timestamp()));
+ }
+ }
+
+ private static ProcessorContext<byte[], KStreamsPayload<?>> checkInitialized(
+ @Nullable ProcessorContext<byte[], KStreamsPayload<?>> context) {
+ if (context == null) {
+ throw new IllegalStateException("FlattenProcessor used before init()");
+ }
+ return context;
+ }
+}
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java
new file mode 100644
index 00000000000..8f6ce2b546d
--- /dev/null
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.kafka.streams.translation;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
+import org.apache.kafka.streams.Topology;
+
+/**
+ * Translates Beam's {@code Flatten} primitive ({@code
beam:transform:flatten:v1}): the union of N
+ * input PCollections into one output PCollection.
+ *
+ * <p>Wires a single {@link FlattenProcessor} node to the producer of every
input PCollection (Kafka
+ * Streams lets a processor have many parents), so the parents' data streams
merge into it, and
+ * registers it as the producer of the flattened output so downstream
translators wire to it. The
+ * processor forwards data through and owns its output watermark via a {@link
WatermarkAggregator},
+ * which is handed the producers of the input PCollections — the upstream
transform ids whose
+ * watermark reports the Flatten must hear from. Producers stamp their own
transform id on the
+ * reports they emit, without regard to who consumes them, so an input shared
with another Flatten
+ * needs no special handling.
+ *
+ * <p>A user-written self-flatten never reaches this translator: the fuser
folds the Flatten into
+ * the consuming SDK-harness stage, which performs the duplication itself. The
duplicate-input check
+ * below is defensive — if a runner-executed Flatten ever did receive the same
PCollection twice,
+ * Kafka Streams could not wire the same parent to a child twice and the
duplicate copy would be
+ * silently dropped, so failing fast is safer.
+ */
+class FlattenTranslator implements PTransformTranslator {
+
+ @Override
+ public void translate(
+ String transformId, RunnerApi.Pipeline pipeline,
KafkaStreamsTranslationContext context) {
+ RunnerApi.PTransform transform =
pipeline.getComponents().getTransformsOrThrow(transformId);
+ // Flatten produces exactly one output PCollection, fed by all of its
input PCollections.
+ String outputPCollectionId =
Iterables.getOnlyElement(transform.getOutputsMap().values());
+
+ Set<String> seenInputs = new HashSet<>();
+ List<String> parentProcessors = new ArrayList<>();
+ Set<String> upstreamTransformIds = new HashSet<>();
+ for (String inputPCollectionId : transform.getInputsMap().values()) {
+ if (!seenInputs.add(inputPCollectionId)) {
+ throw new UnsupportedOperationException(
+ "Flatten "
+ + transform.getUniqueName()
+ + " has PCollection "
+ + inputPCollectionId
+ + " as an input more than once; a self-flatten is not yet
supported by the Kafka"
+ + " Streams runner.");
+ }
+ String parentProcessor =
context.getProcessorNameForPCollection(inputPCollectionId);
+ parentProcessors.add(parentProcessor);
+ upstreamTransformIds.add(parentProcessor);
+ }
+
+ Topology topology = context.getTopology();
+ topology.addProcessor(
+ transformId,
+ () -> new FlattenProcessor(transformId, upstreamTransformIds),
+ parentProcessors.toArray(new String[0]));
+
+ context.registerPCollectionProducer(outputPCollectionId, transformId);
+ }
+}
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java
index a6c0cb8c406..3e82935b807 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java
@@ -19,6 +19,7 @@ package org.apache.beam.runners.kafka.streams.translation;
import java.util.ArrayList;
import java.util.List;
+import java.util.Set;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderException;
import org.apache.beam.sdk.coders.IterableCoder;
@@ -35,16 +36,18 @@ import org.apache.kafka.streams.state.KeyValueIterator;
import org.apache.kafka.streams.state.KeyValueStore;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.Instant;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
* Executes a {@code GroupByKey} (GlobalWindow, default trigger, no allowed
lateness).
*
* <p>Records arrive on the repartition topic keyed by the encoded Beam key,
so every value of a key
* is co-located here. Each value is appended to a per-key buffer in a Kafka
Streams state store.
- * Watermark reports are fed to a {@link WatermarkManager}; when the input
watermark reaches {@link
- * BoundedWindow#TIMESTAMP_MAX_VALUE} (the end of the global window) every
buffered key is emitted
- * once as {@code KV<K, Iterable<V>>} and the buffer cleared, then the
watermark is forwarded
- * downstream.
+ * Watermark reports are fed to a {@link WatermarkAggregator}; when the input
watermark reaches
+ * {@link BoundedWindow#TIMESTAMP_MAX_VALUE} (the end of the global window)
every buffered key is
+ * emitted once as {@code KV<K, Iterable<V>>} and the buffer cleared, then the
watermark is
+ * forwarded downstream.
*
* <p>Buffering whole value lists and re-encoding on each append is O(n^2) per
key; fine for this
* first GroupByKey, and replaced when this moves to runner-core {@code
GroupAlsoByWindow}.
@@ -52,11 +55,18 @@ import org.joda.time.Instant;
class GroupByKeyProcessor
implements Processor<byte[], KStreamsPayload<?>, byte[],
KStreamsPayload<?>> {
+ private static final Logger LOG =
LoggerFactory.getLogger(GroupByKeyProcessor.class);
+
private final String stateStoreName;
+ // This transform's own id, stamped on every watermark it forwards
downstream.
+ private final String transformId;
private final Coder<Object> keyCoder;
private final IterableCoder<@Nullable Object> bufferCoder;
- private final WatermarkManager watermarkManager = new WatermarkManager();
+ // Aggregates the input watermark from the upstream transform's reports,
which arrive through the
+ // repartition topic with the upstream producer's transform id intact (the
shuffle forwards
+ // watermark payloads unchanged).
+ private final WatermarkAggregator watermarkAggregator;
private Instant lastForwardedWatermark = BoundedWindow.TIMESTAMP_MIN_VALUE;
// The global window fires exactly once, when the watermark first reaches
its end. Later watermark
// reports (e.g. the same terminal watermark broadcast across repartition
partitions) must not
@@ -69,9 +79,20 @@ class GroupByKeyProcessor
private @Nullable ProcessorContext<byte[], KStreamsPayload<?>> context;
private @Nullable KeyValueStore<byte[], byte[]> store;
+ /**
+ * @param transformId this transform's own id, stamped on the watermarks it
emits
+ * @param upstreamTransformIds the transform ids feeding this GroupByKey
(known from the pipeline
+ * graph), whose reports the {@link WatermarkAggregator} waits for
+ */
GroupByKeyProcessor(
- String stateStoreName, Coder<Object> keyCoder, Coder<@Nullable Object>
valueCoder) {
+ String stateStoreName,
+ String transformId,
+ Set<String> upstreamTransformIds,
+ Coder<Object> keyCoder,
+ Coder<@Nullable Object> valueCoder) {
this.stateStoreName = stateStoreName;
+ this.transformId = transformId;
+ this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds);
this.keyCoder = keyCoder;
this.bufferCoder = IterableCoder.of(valueCoder);
}
@@ -85,6 +106,14 @@ class GroupByKeyProcessor
@Override
public void process(Record<byte[], KStreamsPayload<?>> record) {
KStreamsPayload<?> payload = record.value();
+ if (payload == null) {
+ // The repartition topic can be written to from outside the runner (or
carry a tombstone),
+ // so recover from the obvious error instead of crashing the task: warn
and drop.
+ LOG.warn(
+ "GroupByKey {} dropping record with null payload (external write or
tombstone)",
+ transformId);
+ return;
+ }
if (payload.isData()) {
byte[] encodedKey = record.key();
Object element = payload.getData().getValue();
@@ -94,12 +123,8 @@ class GroupByKeyProcessor
appendValue(encodedKey, element);
return;
}
- WatermarkPayload report = payload.asWatermark();
- watermarkManager.observe(
- report.getSourcePartition(),
- new Instant(report.getWatermarkMillis()),
- report.getTotalSourcePartitions());
- Instant advanced = watermarkManager.advance();
+ watermarkAggregator.observe(payload.asWatermark());
+ Instant advanced = watermarkAggregator.advance();
if (!fired && !advanced.isBefore(BoundedWindow.TIMESTAMP_MAX_VALUE)) {
fireAll(record);
fired = true;
@@ -153,10 +178,13 @@ class GroupByKeyProcessor
private void forwardWatermark(Record<byte[], KStreamsPayload<?>> trigger,
long watermarkMillis) {
ProcessorContext<byte[], KStreamsPayload<?>> ctx =
checkInitialized(context);
- // GroupByKey is a single logical source for the next stage; report it as
partition 0 of 1.
+ // Stamped with this transform's own id; GroupByKey is a single instance
for now, so the report
+ // is for its only partition (0 of 1).
ctx.forward(
new Record<byte[], KStreamsPayload<?>>(
- trigger.key(), KStreamsPayload.watermark(watermarkMillis, 0, 1),
trigger.timestamp()));
+ trigger.key(),
+ KStreamsPayload.watermark(watermarkMillis, transformId, 0, 1),
+ trigger.timestamp()));
}
private byte[] encodeBuffer(List<@Nullable Object> values) {
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java
index d7c4a309d9c..9e23dbb5cfb 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslator.java
@@ -24,6 +24,7 @@ import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.WindowedValues;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
import org.apache.kafka.common.serialization.Serdes;
import org.apache.kafka.streams.Topology;
@@ -110,10 +111,18 @@ class GroupByKeyTranslator implements
PTransformTranslator {
payloadSerde.deserializer(),
repartitionTopic);
- // Buffer values per key and fire KV<K, Iterable<V>> at the terminal
watermark.
+ // Buffer values per key and fire KV<K, Iterable<V>> at the terminal
watermark. Watermark
+ // reports cross the repartition topic unchanged, so they still carry the
id of the transform
+ // that produced this GroupByKey's input — the parent the shuffle is
attached to.
topology.addProcessor(
transformId,
- () -> new GroupByKeyProcessor(stateStoreName, keyCoder, valueCoder),
+ () ->
+ new GroupByKeyProcessor(
+ stateStoreName,
+ transformId,
+ ImmutableSet.of(parentProcessor),
+ keyCoder,
+ valueCoder),
sourceName);
topology.addStateStore(
Stores.keyValueStoreBuilder(
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java
index 675bee4d859..bac91978a29 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ImpulseProcessor.java
@@ -122,17 +122,16 @@ class ImpulseProcessor implements Processor<byte[],
byte[], byte[], KStreamsPayl
}
/**
- * Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to
downstream processors.
- *
- * <p>Impulse is a single-instance source, so the report is stamped as the
only source partition:
- * {@code sourcePartition=0} of {@code totalSourcePartitions=1}. Real
per-partition identities
- * arrive once the topology gains topic-based shuffle.
+ * Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to
downstream processors,
+ * stamped with this transform's id. Impulse is a single-instance source, so
the report is for its
+ * only partition: {@code sourcePartition=0} of {@code
totalSourcePartitions=1}. Real
+ * per-partition identities arrive once the topology gains topic-based
shuffle.
*/
- private static void forwardWatermarkMax(ProcessorContext<byte[],
KStreamsPayload<byte[]>> ctx) {
+ private void forwardWatermarkMax(ProcessorContext<byte[],
KStreamsPayload<byte[]>> ctx) {
long maxMillis = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis();
ctx.forward(
new Record<byte[], KStreamsPayload<byte[]>>(
- new byte[0], KStreamsPayload.watermark(maxMillis, 0, 1), 0L));
+ new byte[0], KStreamsPayload.watermark(maxMillis, transformId, 0,
1), 0L));
}
/** Cancels the wall-clock punctuator after the impulse has fired to stop
periodic wakeups. */
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java
index 93b40346761..c165f0e875d 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java
@@ -57,6 +57,7 @@ public final class KStreamsPayload<T> {
private final Kind kind;
private final @Nullable WindowedValue<T> data;
private final long watermarkMillis;
+ private final String transformId;
private final int sourcePartition;
private final int totalSourcePartitions;
@@ -64,27 +65,33 @@ public final class KStreamsPayload<T> {
Kind kind,
@Nullable WindowedValue<T> data,
long watermarkMillis,
+ String transformId,
int sourcePartition,
int totalSourcePartitions) {
this.kind = kind;
this.data = data;
this.watermarkMillis = watermarkMillis;
+ this.transformId = transformId;
this.sourcePartition = sourcePartition;
this.totalSourcePartitions = totalSourcePartitions;
}
/** Returns a data payload wrapping the given {@link WindowedValue}. */
public static <T> KStreamsPayload<T> data(WindowedValue<T> value) {
- return new KStreamsPayload<>(Kind.DATA, value, 0L, 0, 0);
+ return new KStreamsPayload<>(Kind.DATA, value, 0L, "", 0, 0);
}
/**
* Returns a watermark report payload: the event-time milliseconds together
with the in-band
- * coordination fields the downstream stage's {@link WatermarkManager} needs
— which source
- * partition this report is for and how many source partitions feed the
stage in total.
+ * coordination fields a downstream watermark aggregator needs — which
transform produced the
+ * report ({@code transformId}, stamped by the producer without regard to
who consumes it), which
+ * of that transform's partitions this report is for, and how many
partitions that transform has
+ * in total.
*/
public static <T> KStreamsPayload<T> watermark(
- long watermarkMillis, int sourcePartition, int totalSourcePartitions) {
+ long watermarkMillis, String transformId, int sourcePartition, int
totalSourcePartitions) {
+ Preconditions.checkArgument(
+ transformId != null && !transformId.isEmpty(), "transformId must be
non-empty");
Preconditions.checkArgument(
totalSourcePartitions > 0,
"totalSourcePartitions must be positive: %s",
@@ -95,7 +102,7 @@ public final class KStreamsPayload<T> {
sourcePartition,
totalSourcePartitions);
return new KStreamsPayload<>(
- Kind.WATERMARK, null, watermarkMillis, sourcePartition,
totalSourcePartitions);
+ Kind.WATERMARK, null, watermarkMillis, transformId, sourcePartition,
totalSourcePartitions);
}
public boolean isData() {
@@ -134,6 +141,11 @@ public final class KStreamsPayload<T> {
return watermarkMillis;
}
+ @Override
+ public String getTransformId() {
+ return transformId;
+ }
+
@Override
public int getSourcePartition() {
return sourcePartition;
@@ -156,6 +168,7 @@ public final class KStreamsPayload<T> {
KStreamsPayload<?> that = (KStreamsPayload<?>) o;
return kind == that.kind
&& watermarkMillis == that.watermarkMillis
+ && transformId.equals(that.transformId)
&& sourcePartition == that.sourcePartition
&& totalSourcePartitions == that.totalSourcePartitions
&& Objects.equals(data, that.data);
@@ -163,7 +176,8 @@ public final class KStreamsPayload<T> {
@Override
public int hashCode() {
- return Objects.hash(kind, data, watermarkMillis, sourcePartition,
totalSourcePartitions);
+ return Objects.hash(
+ kind, data, watermarkMillis, transformId, sourcePartition,
totalSourcePartitions);
}
@Override
@@ -174,6 +188,7 @@ public final class KStreamsPayload<T> {
} else {
helper
.add("watermarkMillis", watermarkMillis)
+ .add("transformId", transformId)
.add("sourcePartition", sourcePartition)
.add("totalSourcePartitions", totalSourcePartitions);
}
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java
index 912eb5fb604..1363740d58b 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java
@@ -82,6 +82,7 @@ public final class KStreamsPayloadSerde<T> implements
Serde<KStreamsPayload<T>>
proto.setWatermark(
KafkaStreamsPayload.WatermarkPayload.newBuilder()
.setMillis(watermark.getWatermarkMillis())
+ .setTransformId(watermark.getTransformId())
.setSourcePartition(watermark.getSourcePartition())
.setTotalPartitions(watermark.getTotalSourcePartitions()));
}
@@ -109,6 +110,7 @@ public final class KStreamsPayloadSerde<T> implements
Serde<KStreamsPayload<T>>
KafkaStreamsPayload.WatermarkPayload watermark =
proto.getWatermark();
return KStreamsPayload.watermark(
watermark.getMillis(),
+ watermark.getTransformId(),
watermark.getSourcePartition(),
watermark.getTotalPartitions());
case PAYLOAD_NOT_SET:
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java
index 943673dfe98..a9e26ecd641 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsPipelineTranslator.java
@@ -52,6 +52,7 @@ public class KafkaStreamsPipelineTranslator {
.put(PTransformTranslation.IMPULSE_TRANSFORM_URN, new
ImpulseTranslator())
.put(PTransformTranslation.READ_TRANSFORM_URN, new
ReadTranslator())
.put(PTransformTranslation.REDISTRIBUTE_ARBITRARILY_URN, new
RedistributeTranslator())
+ .put(PTransformTranslation.FLATTEN_TRANSFORM_URN, new
FlattenTranslator())
.put(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN, new
GroupByKeyTranslator())
.put(ExecutableStage.URN, new ExecutableStageTranslator())
.build());
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java
index 604eca54014..eb20a8b8358 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java
@@ -183,17 +183,16 @@ class ReadProcessor<T> implements Processor<byte[],
byte[], byte[], KStreamsPayl
}
/**
- * Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to
downstream processors.
- *
- * <p>Read is a single-instance source, so the report is stamped as the only
source partition:
- * {@code sourcePartition=0} of {@code totalSourcePartitions=1}. Real
per-partition identities
- * arrive once the topology gains topic-based shuffle.
+ * Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to
downstream processors,
+ * stamped with this transform's id. Read is a single-instance source, so
the report is for its
+ * only partition: {@code sourcePartition=0} of {@code
totalSourcePartitions=1}. Real
+ * per-partition identities arrive once the topology gains topic-based
shuffle.
*/
- private static void forwardWatermarkMax(ProcessorContext<byte[],
KStreamsPayload<?>> ctx) {
+ private void forwardWatermarkMax(ProcessorContext<byte[],
KStreamsPayload<?>> ctx) {
long maxMillis = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis();
ctx.forward(
new Record<byte[], KStreamsPayload<?>>(
- new byte[0], KStreamsPayload.<Object>watermark(maxMillis, 0, 1),
0L));
+ new byte[0], KStreamsPayload.<Object>watermark(maxMillis,
transformId, 0, 1), 0L));
}
/** Cancels the wall-clock punctuator after the read has fired to stop
periodic wakeups. */
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java
index 3184d838e09..774d36db0f2 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessor.java
@@ -25,6 +25,8 @@ import org.apache.kafka.streams.processor.api.Processor;
import org.apache.kafka.streams.processor.api.ProcessorContext;
import org.apache.kafka.streams.processor.api.Record;
import org.checkerframework.checker.nullness.qual.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
/**
* Re-keys a {@code KV}-valued stream by the Beam key so Kafka Streams
shuffles by it.
@@ -39,6 +41,8 @@ import org.checkerframework.checker.nullness.qual.Nullable;
class ShuffleByKeyProcessor
implements Processor<byte[], KStreamsPayload<?>, byte[],
KStreamsPayload<?>> {
+ private static final Logger LOG =
LoggerFactory.getLogger(ShuffleByKeyProcessor.class);
+
private final Coder<Object> keyCoder;
private @Nullable ProcessorContext<byte[], KStreamsPayload<?>> context;
@@ -55,6 +59,12 @@ class ShuffleByKeyProcessor
public void process(Record<byte[], KStreamsPayload<?>> record) {
ProcessorContext<byte[], KStreamsPayload<?>> ctx =
checkInitialized(context);
KStreamsPayload<?> payload = record.value();
+ if (payload == null) {
+ // A topic feeding the runner can always be written to from outside (or
carry a tombstone),
+ // so recover from the obvious error instead of crashing the task: warn
and drop.
+ LOG.warn("Shuffle dropping record with null payload (external write or
tombstone)");
+ return;
+ }
if (payload.isData()) {
Object element = payload.getData().getValue();
if (element == null) {
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregator.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregator.java
new file mode 100644
index 00000000000..c9af03df26a
--- /dev/null
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregator.java
@@ -0,0 +1,107 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.kafka.streams.translation;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
+import org.joda.time.Instant;
+
+/**
+ * Computes a transform's input watermark from the watermark reports of its
upstream transforms.
+ *
+ * <p>A watermark report carries three orthogonal pieces of information (see
{@link
+ * WatermarkPayload}): <i>which transform</i> produced it, <i>which
partition</i> (physical
+ * instance) of that transform it is for, and <i>how many partitions</i> that
transform has. A
+ * producer stamps its own identity without regard to who consumes the report.
This aggregator is
+ * the consuming side, used by every transform that aggregates a watermark —
ExecutableStage,
+ * GroupByKey, Flatten (and CombinePerKey later):
+ *
+ * <ul>
+ * <li>It is constructed with the set of upstream transform ids the consumer
expects, known from
+ * the pipeline graph at translation time (a single-input transform
passes its one parent; a
+ * Flatten passes the producers of all of its input PCollections).
+ * <li>Per upstream transform it tracks partitions with a dedicated {@link
WatermarkManager},
+ * which holds until every partition of that transform has reported and
keeps each partition
+ * monotonic.
+ * <li>The aggregate input watermark is the {@code min()} across the
upstream transforms'
+ * watermarks, defined only once <em>every</em> expected upstream
transform is ready; until
+ * then {@link #advance()} returns {@link
BoundedWindow#TIMESTAMP_MIN_VALUE} and the caller
+ * emits nothing.
+ * </ul>
+ *
+ * <p>Not thread-safe; the caller (a single Kafka Streams processor thread)
serializes access.
+ */
+final class WatermarkAggregator {
+
+ /** Upstream transform ids this consumer must hear from, fixed by the
pipeline graph. */
+ private final Set<String> expectedUpstreamTransformIds;
+
+ /** Per-upstream-transform partition tracking. */
+ private final Map<String, WatermarkManager> managerByTransformId = new
HashMap<>();
+
+ WatermarkAggregator(Set<String> expectedUpstreamTransformIds) {
+ Preconditions.checkArgument(
+ !expectedUpstreamTransformIds.isEmpty(), "expectedUpstreamTransformIds
must not be empty");
+ this.expectedUpstreamTransformIds =
ImmutableSet.copyOf(expectedUpstreamTransformIds);
+ }
+
+ /**
+ * Records one upstream watermark report. A report from a transform this
consumer does not expect
+ * indicates a translation wiring bug and fails fast.
+ */
+ void observe(WatermarkPayload report) {
+ String transformId = report.getTransformId();
+ if (!expectedUpstreamTransformIds.contains(transformId)) {
+ throw new IllegalStateException(
+ "Received a watermark report from unexpected transform "
+ + transformId
+ + "; expected one of "
+ + expectedUpstreamTransformIds);
+ }
+ managerByTransformId
+ .computeIfAbsent(transformId, id -> new WatermarkManager())
+ .observe(
+ report.getSourcePartition(),
+ new Instant(report.getWatermarkMillis()),
+ report.getTotalSourcePartitions());
+ }
+
+ /**
+ * Returns the aggregate input watermark: {@code min()} across all expected
upstream transforms,
+ * or {@link BoundedWindow#TIMESTAMP_MIN_VALUE} while any upstream transform
has not yet fully
+ * reported (the hold).
+ */
+ Instant advance() {
+ if (managerByTransformId.size() < expectedUpstreamTransformIds.size()) {
+ return BoundedWindow.TIMESTAMP_MIN_VALUE;
+ }
+ Instant min = BoundedWindow.TIMESTAMP_MAX_VALUE;
+ for (WatermarkManager manager : managerByTransformId.values()) {
+ // A not-yet-ready manager advances to TIMESTAMP_MIN_VALUE, which
correctly holds the min.
+ Instant watermark = manager.advance();
+ if (watermark.isBefore(min)) {
+ min = watermark;
+ }
+ }
+ return min;
+ }
+}
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java
index 194be701ece..bac5314a6da 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WatermarkPayload.java
@@ -24,18 +24,26 @@ package org.apache.beam.runners.kafka.streams.translation;
* {@link KStreamsPayload#isWatermark()} and narrowed the payload, so there is
no kind check to do
* on each accessor.
*
- * <p>A watermark report is the in-band coordination message a downstream
stage's {@link
- * WatermarkManager} consumes: the watermark value plus which source partition
reported it and how
- * many source partitions feed the stage in total.
+ * <p>A watermark report is the in-band coordination message a downstream
watermark aggregator
+ * consumes: the watermark value, which transform produced it, which of that
transform's partitions
+ * reported it, and how many partitions that transform has in total. The
producer stamps its own
+ * identity without regard to who consumes the report; a consumer with several
upstream transforms
+ * (e.g. Flatten) aggregates per producing transform.
*/
public interface WatermarkPayload {
/** The reported watermark, in event-time milliseconds. */
long getWatermarkMillis();
- /** The source partition this report is for, in {@code [0,
getTotalSourcePartitions())}. */
+ /** Globally unique id of the transform that produced this report. */
+ String getTransformId();
+
+ /**
+ * Which partition (physical instance) of the producing transform this
report is for, in {@code
+ * [0, getTotalSourcePartitions())}.
+ */
int getSourcePartition();
- /** The total number of source partitions feeding the downstream stage. */
+ /** How many partitions (physical instances) the producing transform has in
total. */
int getTotalSourcePartitions();
}
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java
index fc5797a12c2..74169acb81c 100644
---
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java
@@ -25,13 +25,14 @@ import
org.apache.beam.runners.fnexecution.provisioning.JobInfo;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
import org.apache.kafka.streams.processor.api.MockProcessorContext;
import org.apache.kafka.streams.processor.api.Record;
import org.junit.Test;
/**
* Tests the watermark wiring of {@link ExecutableStageProcessor}: how it
feeds incoming watermark
- * reports to the {@link WatermarkManager} and forwards the stage's output
watermark.
+ * reports to the {@link WatermarkAggregator} and forwards the stage's output
watermark.
*
* <p>Only the watermark path is exercised, so the SDK harness is never
started (it is created
* lazily on the first data element). A {@link MockProcessorContext} captures
what the processor
@@ -39,6 +40,12 @@ import org.junit.Test;
*/
public class ExecutableStageProcessorWatermarkTest {
+ /** The stage's own transform id, expected on every watermark it forwards. */
+ private static final String STAGE_ID = "stage";
+
+ /** The single upstream transform whose reports the stage aggregates. */
+ private static final String UPSTREAM_ID = "upstream";
+
private static ExecutableStageProcessor newProcessor() {
JobInfo jobInfo =
JobInfo.create(
@@ -47,13 +54,17 @@ public class ExecutableStageProcessorWatermarkTest {
"",
PipelineOptionsTranslation.toProto(PipelineOptionsFactory.create()));
return new ExecutableStageProcessor(
- RunnerApi.ExecutableStagePayload.getDefaultInstance(), jobInfo);
+ RunnerApi.ExecutableStagePayload.getDefaultInstance(),
+ jobInfo,
+ STAGE_ID,
+ ImmutableSet.of(UPSTREAM_ID));
}
+ /** A report from the upstream transform's given partition. */
private static Record<byte[], KStreamsPayload<?>> watermark(
long millis, int sourcePartition, int totalSourcePartitions) {
KStreamsPayload<?> payload =
- KStreamsPayload.watermark(millis, sourcePartition,
totalSourcePartitions);
+ KStreamsPayload.watermark(millis, UPSTREAM_ID, sourcePartition,
totalSourcePartitions);
return new Record<>(new byte[0], payload, 0L);
}
@@ -75,7 +86,9 @@ public class ExecutableStageProcessorWatermarkTest {
assertThat(out.isWatermark(), is(true));
WatermarkPayload report = out.asWatermark();
assertThat(report.getWatermarkMillis(), is(100L));
- // The stage forwards as its own single source (0 of 1), not the
upstream's identity.
+ // The stage forwards under its own identity — its transform id and its
own single partition
+ // (0 of 1) — not the upstream's.
+ assertThat(report.getTransformId(), is(STAGE_ID));
assertThat(report.getSourcePartition(), is(0));
assertThat(report.getTotalSourcePartitions(), is(1));
}
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java
new file mode 100644
index 00000000000..a5af0cdc922
--- /dev/null
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/FlattenTest.java
@@ -0,0 +1,212 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.kafka.streams.translation;
+
+import static org.hamcrest.CoreMatchers.hasItems;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.Flatten;
+import org.apache.beam.sdk.transforms.GroupByKey;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionList;
+import org.junit.Test;
+
+/**
+ * End-to-end test for {@link FlattenTranslator}: two branches are flattened
into one PCollection
+ * and a recording ParDo sees every element from both.
+ *
+ * <p>Each branch is a {@code Create -> identity ParDo}, so its producer
feeding the Flatten is an
+ * {@link ExecutableStageProcessor} — the same shape PAssert's {@code
GroupGlobally} produces. This
+ * exercises the per-producing-transform watermark aggregation: each branch's
producer stamps its
+ * own transform id on its watermark, and the Flatten holds its output
watermark until every
+ * upstream transform it expects has reported. Without that, the Flatten would
release its watermark
+ * after the first branch drained and the downstream stage's bundle would
close early, dropping the
+ * second branch's elements.
+ */
+public class FlattenTest {
+
+ private static class IdentityFn extends DoFn<Integer, Integer> {
+ @ProcessElement
+ public void processElement(@Element Integer input, OutputReceiver<Integer>
out) {
+ out.output(input);
+ }
+ }
+
+ /** Records every element the harness feeds it so the test can assert the
flatten's union. */
+ private static class RecordingFn extends DoFn<Integer, Integer> {
+ private final SharedTestCollector<Integer> collector;
+
+ RecordingFn(SharedTestCollector<Integer> collector) {
+ this.collector = collector;
+ }
+
+ @ProcessElement
+ public void processElement(@Element Integer input, OutputReceiver<Integer>
out) {
+ collector.record(input);
+ out.output(input);
+ }
+ }
+
+ @Test
+ public void flattenUnionsEveryBranch() {
+ try (SharedTestCollector<Integer> collector =
SharedTestCollector.create()) {
+ Pipeline pipeline =
Pipeline.create(KafkaStreamsTestRunner.testOptions());
+ PCollection<Integer> a =
+ pipeline.apply("createA", Create.of(1, 2)).apply("idA", ParDo.of(new
IdentityFn()));
+ PCollection<Integer> b =
+ pipeline.apply("createB", Create.of(3, 4)).apply("idB", ParDo.of(new
IdentityFn()));
+ PCollectionList.of(a)
+ .and(b)
+ .apply("flatten", Flatten.pCollections())
+ .apply("record", ParDo.of(new RecordingFn(collector)));
+
+ KafkaStreamsTestRunner.run(pipeline);
+
+ List<Integer> recorded = collector.recorded();
+ assertThat(recorded.size(), is(4));
+ assertThat(recorded, hasItems(1, 2, 3, 4));
+ }
+ }
+
+ @Test
+ public void pCollectionFeedingTwoFlattensIsSupported() {
+ // input2 feeds both flattens, so its producer's watermark report is
consumed by two different
+ // aggregators. The producer stamps its own transform id once, and each
flatten holds until its
+ // own two branches drain. Verify both flattens produce the right union.
+ try (SharedTestCollector<Integer> left = SharedTestCollector.create();
+ SharedTestCollector<Integer> right = SharedTestCollector.create()) {
+ Pipeline pipeline =
Pipeline.create(KafkaStreamsTestRunner.testOptions());
+ PCollection<Integer> input1 =
+ pipeline.apply("c1", Create.of(1, 2)).apply("id1", ParDo.of(new
IdentityFn()));
+ PCollection<Integer> input2 =
+ pipeline.apply("c2", Create.of(3, 4)).apply("id2", ParDo.of(new
IdentityFn()));
+ PCollection<Integer> input3 =
+ pipeline.apply("c3", Create.of(5, 6)).apply("id3", ParDo.of(new
IdentityFn()));
+ PCollectionList.of(input1)
+ .and(input2)
+ .apply("l1", Flatten.pCollections())
+ .apply("recordL1", ParDo.of(new RecordingFn(left)));
+ PCollectionList.of(input2)
+ .and(input3)
+ .apply("l2", Flatten.pCollections())
+ .apply("recordL2", ParDo.of(new RecordingFn(right)));
+
+ KafkaStreamsTestRunner.run(pipeline);
+
+ assertThat(left.recorded().size(), is(4));
+ assertThat(left.recorded(), hasItems(1, 2, 3, 4));
+ assertThat(right.recorded().size(), is(4));
+ assertThat(right.recorded(), hasItems(3, 4, 5, 6));
+ }
+ }
+
+ /** Maps each int to {@code KV("k", int)} so a downstream GroupByKey groups
all branches. */
+ private static class ToKvFn extends DoFn<Integer, KV<String, Integer>> {
+ @ProcessElement
+ public void processElement(@Element Integer input,
OutputReceiver<KV<String, Integer>> out) {
+ out.output(KV.of("k", input));
+ }
+ }
+
+ /** Records each grouped result as {@code "key=[sorted values]"}. */
+ private static class RecordGroupFn extends DoFn<KV<String,
Iterable<Integer>>, Void> {
+ private final SharedTestCollector<String> collector;
+
+ RecordGroupFn(SharedTestCollector<String> collector) {
+ this.collector = collector;
+ }
+
+ @ProcessElement
+ public void processElement(@Element KV<String, Iterable<Integer>> group) {
+ List<Integer> values = new ArrayList<>();
+ group.getValue().forEach(values::add);
+ Collections.sort(values);
+ collector.record(group.getKey() + "=" + values);
+ }
+ }
+
+ @Test
+ public void watermarkPropagatesThroughFlattenAndFiresDownstreamGroupByKey() {
+ // GroupByKey fires exactly once, when its input watermark reaches the end
of the global
+ // window, and the Flatten forwards its watermark only after every branch
has drained. Both
+ // branches share the key, so a single group holding the elements of both
branches proves the
+ // watermark propagated through the Flatten at the right time — a
premature release would fire
+ // a partial group instead.
+ try (SharedTestCollector<String> collector = SharedTestCollector.create())
{
+ Pipeline pipeline =
Pipeline.create(KafkaStreamsTestRunner.testOptions());
+ PCollection<KV<String, Integer>> a =
+ pipeline
+ .apply("createA", Create.of(1, 2))
+ .apply("kvA", ParDo.of(new ToKvFn()))
+ .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()));
+ PCollection<KV<String, Integer>> b =
+ pipeline
+ .apply("createB", Create.of(3, 4))
+ .apply("kvB", ParDo.of(new ToKvFn()))
+ .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()));
+ PCollectionList.of(a)
+ .and(b)
+ .apply("flatten", Flatten.pCollections())
+ .apply("gbk", GroupByKey.create())
+ .apply("record", ParDo.of(new RecordGroupFn(collector)));
+
+ KafkaStreamsTestRunner.run(pipeline);
+
+ List<String> groups = collector.recorded();
+ assertThat(groups.size(), is(1));
+ assertThat(groups, hasItems("k=[1, 2, 3, 4]"));
+ }
+ }
+
+ @Test
+ public void flattenOfOneBranchTwiceDuplicatesEveryElement() {
+ // A self-flatten is a bag union with itself: every element must appear
twice. The fuser folds
+ // the Flatten into the SDK-harness stage (it never reaches the runner's
Flatten translator as a
+ // duplicate-input node), so the duplication happens in the harness.
+ try (SharedTestCollector<Integer> collector =
SharedTestCollector.create()) {
+ Pipeline pipeline =
Pipeline.create(KafkaStreamsTestRunner.testOptions());
+ PCollection<Integer> branch =
+ pipeline.apply("create", Create.of(1, 2)).apply("id", ParDo.of(new
IdentityFn()));
+ PCollectionList.of(branch)
+ .and(branch)
+ .apply("flatten", Flatten.pCollections())
+ .apply("record", ParDo.of(new RecordingFn(collector)));
+
+ KafkaStreamsTestRunner.run(pipeline);
+
+ List<Integer> recorded = collector.recorded();
+ assertThat(recorded.size(), is(4));
+ assertThat(recorded, hasItems(1, 2));
+ assertThat(recorded.stream().filter(v -> v == 1).count(), is(2L));
+ assertThat(recorded.stream().filter(v -> v == 2).count(), is(2L));
+ }
+ }
+}
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java
index 19346f54763..95ce70b8857 100644
---
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java
@@ -58,10 +58,11 @@ public class KStreamsPayloadSerdeTest {
@Test
public void roundTripsWatermarkPayload() {
- KStreamsPayload<Integer> payload = KStreamsPayload.watermark(12345L, 2, 4);
+ KStreamsPayload<Integer> payload = KStreamsPayload.watermark(12345L,
"transform-a", 2, 4);
KStreamsPayload<Integer> out = roundTrip(payload);
assertThat(out.isWatermark(), is(true));
assertThat(out.asWatermark().getWatermarkMillis(), is(12345L));
+ assertThat(out.asWatermark().getTransformId(), is("transform-a"));
assertThat(out.asWatermark().getSourcePartition(), is(2));
assertThat(out.asWatermark().getTotalSourcePartitions(), is(4));
assertThat(out, is(payload));
@@ -70,7 +71,7 @@ public class KStreamsPayloadSerdeTest {
@Test
public void roundTripsTerminalMaxWatermark() {
KStreamsPayload<Integer> payload =
-
KStreamsPayload.watermark(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(), 0, 1);
+
KStreamsPayload.watermark(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(), "t",
0, 1);
assertThat(
roundTrip(payload).asWatermark().getWatermarkMillis(),
is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()));
@@ -80,7 +81,7 @@ public class KStreamsPayloadSerdeTest {
public void roundTripsNegativeWatermark() {
// Beam event times can be negative; sint64 must round-trip them
losslessly.
KStreamsPayload<Integer> payload =
-
KStreamsPayload.watermark(BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis(), 0, 1);
+
KStreamsPayload.watermark(BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis(), "t",
0, 1);
assertThat(
roundTrip(payload).asWatermark().getWatermarkMillis(),
is(BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis()));
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregatorTest.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregatorTest.java
new file mode 100644
index 00000000000..88d541bdc0c
--- /dev/null
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/WatermarkAggregatorTest.java
@@ -0,0 +1,160 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.kafka.streams.translation;
+
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.Assert.assertThrows;
+
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
+import org.joda.time.Instant;
+import org.junit.Test;
+
+/**
+ * Unit tests for {@link WatermarkAggregator}: aggregation of watermark
reports across multiple
+ * upstream transforms, each with its own partition set. Per-partition
behaviour within one upstream
+ * transform (monotonicity, repartition reset) is covered in depth by {@code
WatermarkManagerTest};
+ * here it is exercised through the aggregate.
+ */
+public class WatermarkAggregatorTest {
+
+ private static Instant ts(long millis) {
+ return new Instant(millis);
+ }
+
+ /** A report from the given transform's partition {@code sourcePartition} of
{@code total}. */
+ private static WatermarkPayload report(
+ String transformId, long millis, int sourcePartition, int total) {
+ return KStreamsPayload.watermark(millis, transformId, sourcePartition,
total).asWatermark();
+ }
+
+ @Test
+ public void holdsBeforeAnyReport() {
+ WatermarkAggregator aggregator = new
WatermarkAggregator(ImmutableSet.of("a"));
+ assertThat(aggregator.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE));
+ }
+
+ @Test
+ public void singleUpstreamSinglePartitionAdvances() {
+ WatermarkAggregator aggregator = new
WatermarkAggregator(ImmutableSet.of("a"));
+ aggregator.observe(report("a", 100L, 0, 1));
+ assertThat(aggregator.advance(), is(ts(100L)));
+ }
+
+ @Test
+ public void holdsUntilEveryExpectedUpstreamReports() {
+ WatermarkAggregator aggregator = new
WatermarkAggregator(ImmutableSet.of("a", "b"));
+ aggregator.observe(report("a", 100L, 0, 1));
+ // Only one of the two expected upstream transforms has reported — still
holding.
+ assertThat(aggregator.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE));
+
+ aggregator.observe(report("b", 200L, 0, 1));
+ assertThat(aggregator.advance(), is(ts(100L)));
+ }
+
+ @Test
+ public void holdsUntilEveryPartitionOfEachUpstreamReports() {
+ WatermarkAggregator aggregator = new
WatermarkAggregator(ImmutableSet.of("a", "b"));
+ aggregator.observe(report("a", 100L, 0, 1));
+ aggregator.observe(report("b", 200L, 0, 2));
+ // Upstream "b" has reported only one of its two partitions — still
holding.
+ assertThat(aggregator.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE));
+
+ aggregator.observe(report("b", 300L, 1, 2));
+ // Ready: min(100, min(200, 300)) = 100.
+ assertThat(aggregator.advance(), is(ts(100L)));
+ }
+
+ @Test
+ public void aggregateIsMinAcrossUpstreams() {
+ WatermarkAggregator aggregator = new
WatermarkAggregator(ImmutableSet.of("a", "b", "c"));
+ aggregator.observe(report("a", 300L, 0, 1));
+ aggregator.observe(report("b", 100L, 0, 1));
+ aggregator.observe(report("c", 500L, 0, 1));
+ assertThat(aggregator.advance(), is(ts(100L)));
+
+ // The slowest upstream advances; the aggregate follows the new min.
+ aggregator.observe(report("b", 400L, 0, 1));
+ assertThat(aggregator.advance(), is(ts(300L)));
+ }
+
+ @Test
+ public void perUpstreamWatermarkIsMonotonic() {
+ WatermarkAggregator aggregator = new
WatermarkAggregator(ImmutableSet.of("a"));
+ aggregator.observe(report("a", 200L, 0, 1));
+ assertThat(aggregator.advance(), is(ts(200L)));
+
+ // A late lower report from the same upstream partition is ignored.
+ aggregator.observe(report("a", 100L, 0, 1));
+ assertThat(aggregator.advance(), is(ts(200L)));
+ }
+
+ @Test
+ public void duplicateReportDoesNotAdvance() {
+ WatermarkAggregator aggregator = new
WatermarkAggregator(ImmutableSet.of("a"));
+ aggregator.observe(report("a", 100L, 0, 1));
+ assertThat(aggregator.advance(), is(ts(100L)));
+
+ // The same report again (e.g. a broadcast duplicate) leaves the aggregate
unchanged.
+ aggregator.observe(report("a", 100L, 0, 1));
+ assertThat(aggregator.advance(), is(ts(100L)));
+ }
+
+ @Test
+ public void terminalMaxWatermarkAggregates() {
+ long max = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis();
+ WatermarkAggregator aggregator = new
WatermarkAggregator(ImmutableSet.of("a", "b"));
+ aggregator.observe(report("a", max, 0, 1));
+ assertThat(aggregator.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE));
+
+ aggregator.observe(report("b", max, 0, 1));
+ assertThat(aggregator.advance(), is(BoundedWindow.TIMESTAMP_MAX_VALUE));
+ }
+
+ @Test
+ public void repartitionOfOneUpstreamReopensTheHold() {
+ WatermarkAggregator aggregator = new
WatermarkAggregator(ImmutableSet.of("a", "b"));
+ aggregator.observe(report("a", 100L, 0, 1));
+ aggregator.observe(report("b", 200L, 0, 1));
+ assertThat(aggregator.advance(), is(ts(100L)));
+
+ // Upstream "b" changes its partition count (repartition): its
per-partition state resets and
+ // the aggregate holds again until b's new full partition set has reported.
+ aggregator.observe(report("b", 250L, 0, 2));
+ assertThat(aggregator.advance(), is(BoundedWindow.TIMESTAMP_MIN_VALUE));
+
+ aggregator.observe(report("b", 300L, 1, 2));
+ assertThat(aggregator.advance(), is(ts(100L)));
+ }
+
+ @Test
+ public void reportFromUnexpectedTransformFailsFast() {
+ WatermarkAggregator aggregator = new
WatermarkAggregator(ImmutableSet.of("a"));
+ IllegalStateException thrown =
+ assertThrows(
+ IllegalStateException.class, () ->
aggregator.observe(report("intruder", 100L, 0, 1)));
+ assertThat(thrown.getMessage(), containsString("intruder"));
+ }
+
+ @Test
+ public void emptyExpectedUpstreamsIsRejected() {
+ assertThrows(IllegalArgumentException.class, () -> new
WatermarkAggregator(ImmutableSet.of()));
+ }
+}
