junaiddshaukat commented on code in PR #39410: URL: https://github.com/apache/beam/pull/39410#discussion_r3629603262
########## runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/MultiOutputStageTest.java: ########## @@ -0,0 +1,85 @@ +/* + * 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; + +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Count; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.TupleTagList; +import org.junit.Rule; +import org.junit.Test; + +/** + * Verifies that a fused executable stage with more than one output routes each output to its own + * downstream — including the case the mentor asked about, where each output is followed by a + * separate GroupByKey (so each output ends up in its own repartition topic). + * + * <p>A DoFn splits its input to a main and an additional (side) output; each output is grouped with + * {@code Count.perElement} (a GroupByKey under the hood), and a {@link PAssert} on each confirms + * the two outputs carried the right elements to the right GroupByKey. + */ +public class MultiOutputStageTest { + + private static final TupleTag<Integer> EVENS = new TupleTag<Integer>() {}; + private static final TupleTag<Integer> ODDS = new TupleTag<Integer>() {}; + + /** Routes even elements to the main output and odd elements to the side output. */ + private static class SplitByParityFn extends DoFn<Integer, Integer> { + @ProcessElement + public void processElement(@Element Integer input, MultiOutputReceiver out) { + if (input % 2 == 0) { + out.get(EVENS).output(input); + } else { + out.get(ODDS).output(input); + } + } + } + + private static PipelineOptions options() { + PipelineOptions options = PipelineOptionsFactory.create(); + options.setRunner(TestKafkaStreamsRunner.class); + return options; + } + + @Rule public final transient TestPipeline pipeline = TestPipeline.fromOptions(options()); + + @Test + public void eachStageOutputFeedsItsOwnGroupByKey() { + PCollectionTuple split = + pipeline + .apply("create", Create.of(1, 2, 3, 4, 5, 6)) + .apply( + "split", + ParDo.of(new SplitByParityFn()).withOutputTags(EVENS, TupleTagList.of(ODDS))); + + PAssert.that(split.get(EVENS).apply("countEvens", Count.perElement())) Review Comment: I made the odds branch assert directly with no Count, and kept the evens branch through Count/GBK so we still cover an output followed by a repartition. ########## runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/StageOutputProcessor.java: ########## @@ -0,0 +1,88 @@ +/* + * 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 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; + +/** + * One output port of a multi-output {@link ExecutableStageProcessor}: a relay node that stands in + * as the producer of one of the stage's output PCollections. + * + * <p>A Kafka Streams processor forwards a record either to all of its children or to one named + * child, but downstream transforms are wired to a <em>producer node</em> by PCollection id, and + * that node's identity has to be known when the stage is translated — before the downstream + * transforms are. So each output PCollection of a multi-output stage gets its own relay: the stage + * routes each harness output to the matching relay by name, and downstream transforms wire to the + * relay. (A single-output stage needs none of this and forwards directly.) + * + * <p>The relay forwards data records unchanged, and re-stamps the watermark with its own transform + * id so a downstream watermark aggregator, which expects to hear from this relay, sees a consistent + * source. The stage is a single instance and its watermark is already the monotonic aggregate over + * its input, so the relay only relabels it — there is exactly one upstream partition to relay, no + * aggregation to do. + */ +class StageOutputProcessor + implements Processor<byte[], KStreamsPayload<?>, byte[], KStreamsPayload<?>> { + + private static final Logger LOG = LoggerFactory.getLogger(StageOutputProcessor.class); + + private final String transformId; + private @Nullable ProcessorContext<byte[], KStreamsPayload<?>> context; + + StageOutputProcessor(String transformId) { + this.transformId = transformId; + } + + @Override + public void init(ProcessorContext<byte[], KStreamsPayload<?>> context) { + this.context = context; + } + + @Override + public void process(Record<byte[], KStreamsPayload<?>> record) { + KStreamsPayload<?> payload = record.value(); + ProcessorContext<byte[], KStreamsPayload<?>> ctx = context; + if (ctx == null) { + throw new IllegalStateException("StageOutputProcessor used before init()"); + } + if (payload == null) { + LOG.warn( + "Stage output {} dropping record with null payload (external write or tombstone)", + transformId); + return; + } + if (!payload.isWatermark()) { + // Data for this output: forward unchanged. + ctx.forward(record); + return; + } + // Re-stamp the stage's watermark with this output port's own id. Single upstream instance, so + // its value is already monotonic; just relabel it as source 0 of 1. + WatermarkPayload report = payload.asWatermark(); + ctx.forward( + new Record<byte[], KStreamsPayload<?>>( + record.key(), + KStreamsPayload.watermark(report.getWatermarkMillis(), transformId, 0, 1), Review Comment: Good catch, you're right. Hardcoding (0 of 1) drops the partition info the moment the stage runs in more than one instance, and downstream needs it to know how many instances there are. Fixed: the relay now keeps the reporting stage instance's source_partition / total_partitions and only relabels the transform id. Added StageOutputProcessorTest that feeds reports from several stage partitions and asserts they stay distinct downstream (via MockProcessorContext, since TopologyTestDriver is single-instance so I can't spin up parallel instances there). -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
