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 87e3de44397b..4dadf2d14b70 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 00f85032d041..38d3601e814f 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 @@ * ProcessorContext#forward} must only be called from the processing thread, so outputs are never * forwarded directly from a harness callback. * - *

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. + *

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. * *

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> 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 upstreamTransformIds) { this.stagePayload = stagePayload; this.jobInfo = jobInfo; + this.transformId = transformId; + this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds); } @Override @@ -116,18 +131,21 @@ private void ensureStageBundleFactory() { @Override public void process(Record> 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 @@ private void closeBundleAndFlush(Record> record) { } private void forwardWatermark(Record> 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> ctx = checkInitialized(context); ctx.forward( new Record>( - 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 dc56d57f57c2..a2e6ed837c7d 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 @@ 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 @@ public void translate( 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 000000000000..d09fed8185a1 --- /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. + * + *

Data records are forwarded straight through unchanged — the merge of the N parents' + * data streams is the flatten. + * + *

Watermark 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 + * every input branch has reported, so a downstream GroupByKey does not fire before all + * flattened branches are drained. + * + *

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> { + + 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> 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 upstreamTransformIds) { + this.transformId = transformId; + this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds); + } + + @Override + public void init(ProcessorContext> context) { + this.context = context; + } + + @Override + public void process(Record> 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> 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>( + record.key(), + KStreamsPayload.watermark(advanced.getMillis(), transformId, 0, 1), + record.timestamp())); + } + } + + private static ProcessorContext> checkInitialized( + @Nullable ProcessorContext> 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 000000000000..8f6ce2b546dd --- /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. + * + *

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. + * + *

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 seenInputs = new HashSet<>(); + List parentProcessors = new ArrayList<>(); + Set 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 a6c0cb8c4061..3e82935b807d 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 @@ 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.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). * *

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>} 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>} and the buffer cleared, then the watermark is + * forwarded downstream. * *

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 @@ class GroupByKeyProcessor implements Processor, 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 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> context; private @Nullable KeyValueStore 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 keyCoder, Coder<@Nullable Object> valueCoder) { + String stateStoreName, + String transformId, + Set upstreamTransformIds, + Coder 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 @@ public void init(ProcessorContext> context) { @Override public void process(Record> 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 @@ public void process(Record> record) { 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 @@ private void fireAll(Record> trigger) { private void forwardWatermark(Record> trigger, long watermarkMillis) { ProcessorContext> 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>( - 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 d7c4a309d9c3..9e23dbb5cfb0 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.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 @@ public void translate( payloadSerde.deserializer(), repartitionTopic); - // Buffer values per key and fire KV> at the terminal watermark. + // Buffer values per key and fire KV> 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 675bee4d8591..bac91978a298 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 @@ private void maybeFire() { } /** - * Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to downstream processors. - * - *

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> ctx) { + private void forwardWatermarkMax(ProcessorContext> ctx) { long maxMillis = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(); ctx.forward( new Record>( - 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 93b40346761a..c165f0e875d4 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 @@ private enum Kind { private final Kind kind; private final @Nullable WindowedValue data; private final long watermarkMillis; + private final String transformId; private final int sourcePartition; private final int totalSourcePartitions; @@ -64,27 +65,33 @@ private KStreamsPayload( Kind kind, @Nullable WindowedValue 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 KStreamsPayload data(WindowedValue 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 KStreamsPayload 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 static KStreamsPayload watermark( 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 long getWatermarkMillis() { return watermarkMillis; } + @Override + public String getTransformId() { + return transformId; + } + @Override public int getSourcePartition() { return sourcePartition; @@ -156,6 +168,7 @@ public boolean equals(@Nullable Object o) { 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 boolean equals(@Nullable Object o) { @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 String toString() { } 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 912eb5fb6047..1363740d58bd 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 byte[] serialize(String topic, KStreamsPayload payload) { proto.setWatermark( KafkaStreamsPayload.WatermarkPayload.newBuilder() .setMillis(watermark.getWatermarkMillis()) + .setTransformId(watermark.getTransformId()) .setSourcePartition(watermark.getSourcePartition()) .setTotalPartitions(watermark.getTotalSourcePartitions())); } @@ -109,6 +110,7 @@ public KStreamsPayload deserialize(String topic, byte[] bytes) { 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 943673dfe98c..a9e26ecd6412 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 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 604eca54014a..eb20a8b83588 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 @@ private WindowedValue toRunnerWire(WindowedValue element) { } /** - * Forwards a terminal {@code TIMESTAMP_MAX_VALUE} watermark payload to downstream processors. - * - *

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> ctx) { + private void forwardWatermarkMax(ProcessorContext> ctx) { long maxMillis = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(); ctx.forward( new Record>( - 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 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 3184d838e09d..774d36db0f2e 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.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 @@ class ShuffleByKeyProcessor implements Processor, byte[], KStreamsPayload> { + private static final Logger LOG = LoggerFactory.getLogger(ShuffleByKeyProcessor.class); + private final Coder keyCoder; private @Nullable ProcessorContext> context; @@ -55,6 +59,12 @@ public void init(ProcessorContext> context) { public void process(Record> record) { ProcessorContext> 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 000000000000..c9af03df26a9 --- /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. + * + *

A watermark report carries three orthogonal pieces of information (see {@link + * WatermarkPayload}): which transform produced it, which partition (physical + * instance) of that transform it is for, and how many partitions 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): + * + *

    + *
  • 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). + *
  • 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. + *
  • The aggregate input watermark is the {@code min()} across the upstream transforms' + * watermarks, defined only once every expected upstream transform is ready; until + * then {@link #advance()} returns {@link BoundedWindow#TIMESTAMP_MIN_VALUE} and the caller + * emits nothing. + *
+ * + *

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 expectedUpstreamTransformIds; + + /** Per-upstream-transform partition tracking. */ + private final Map managerByTransformId = new HashMap<>(); + + WatermarkAggregator(Set 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 194be701ecec..bac5314a6da7 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 @@ * {@link KStreamsPayload#isWatermark()} and narrowed the payload, so there is no kind check to do * on each accessor. * - *

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. + *

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 fc5797a12c26..74169acb81cb 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.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. * *

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 @@ */ 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 @@ private static ExecutableStageProcessor newProcessor() { "", 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> 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 void singleSourcePartitionForwardsImmediatelyStampedAsItsOwnSource() { 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 000000000000..a5af0cdc9223 --- /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. + * + *

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 { + @ProcessElement + public void processElement(@Element Integer input, OutputReceiver 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 { + private final SharedTestCollector collector; + + RecordingFn(SharedTestCollector collector) { + this.collector = collector; + } + + @ProcessElement + public void processElement(@Element Integer input, OutputReceiver out) { + collector.record(input); + out.output(input); + } + } + + @Test + public void flattenUnionsEveryBranch() { + try (SharedTestCollector collector = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + PCollection a = + pipeline.apply("createA", Create.of(1, 2)).apply("idA", ParDo.of(new IdentityFn())); + PCollection 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 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 left = SharedTestCollector.create(); + SharedTestCollector right = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + PCollection input1 = + pipeline.apply("c1", Create.of(1, 2)).apply("id1", ParDo.of(new IdentityFn())); + PCollection input2 = + pipeline.apply("c2", Create.of(3, 4)).apply("id2", ParDo.of(new IdentityFn())); + PCollection 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> { + @ProcessElement + public void processElement(@Element Integer input, OutputReceiver> out) { + out.output(KV.of("k", input)); + } + } + + /** Records each grouped result as {@code "key=[sorted values]"}. */ + private static class RecordGroupFn extends DoFn>, Void> { + private final SharedTestCollector collector; + + RecordGroupFn(SharedTestCollector collector) { + this.collector = collector; + } + + @ProcessElement + public void processElement(@Element KV> group) { + List 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 collector = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + PCollection> a = + pipeline + .apply("createA", Create.of(1, 2)) + .apply("kvA", ParDo.of(new ToKvFn())) + .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())); + PCollection> 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 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 collector = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + PCollection 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 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 19346f54763e..95ce70b88578 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 void roundTripsDataPayload() { @Test public void roundTripsWatermarkPayload() { - KStreamsPayload payload = KStreamsPayload.watermark(12345L, 2, 4); + KStreamsPayload payload = KStreamsPayload.watermark(12345L, "transform-a", 2, 4); KStreamsPayload 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 void roundTripsWatermarkPayload() { @Test public void roundTripsTerminalMaxWatermark() { KStreamsPayload 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 void roundTripsTerminalMaxWatermark() { public void roundTripsNegativeWatermark() { // Beam event times can be negative; sint64 must round-trip them losslessly. KStreamsPayload 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 000000000000..88d541bdc0c2 --- /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())); + } +}