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 189482ed5f9d..943673dfe98c 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 @@ -50,6 +50,7 @@ public KafkaStreamsPipelineTranslator() { this( ImmutableMap.builder() .put(PTransformTranslation.IMPULSE_TRANSFORM_URN, new ImpulseTranslator()) + .put(PTransformTranslation.READ_TRANSFORM_URN, new ReadTranslator()) .put(PTransformTranslation.REDISTRIBUTE_ARBITRARILY_URN, new RedistributeTranslator()) .put(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN, new GroupByKeyTranslator()) .put(ExecutableStage.URN, new ExecutableStageTranslator()) diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java index 3d95eabafed6..0a045ff7395b 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KafkaStreamsTranslationContext.java @@ -19,6 +19,7 @@ import java.util.HashMap; import java.util.Map; +import java.util.regex.Pattern; import org.apache.beam.runners.fnexecution.provisioning.JobInfo; import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions; import org.apache.kafka.streams.Topology; @@ -34,6 +35,12 @@ public class KafkaStreamsTranslationContext { /** Prefix for the per-job bootstrap topic Impulse reads from. */ private static final String IMPULSE_BOOTSTRAP_TOPIC_PREFIX = "__beam_impulse_"; + /** Prefix for the per-transform bootstrap topic a primitive Read reads from. */ + private static final String READ_BOOTSTRAP_TOPIC_PREFIX = "__beam_read_"; + + /** Characters not legal in a Kafka topic name; a topic's legal set is {@code [a-zA-Z0-9._-]}. */ + private static final Pattern ILLEGAL_TOPIC_CHARS = Pattern.compile("[^a-zA-Z0-9._-]"); + private final JobInfo jobInfo; private final KafkaStreamsPipelineOptions pipelineOptions; private final Topology topology; @@ -100,4 +107,17 @@ public String getProcessorNameForPCollection(String pCollectionId) { public String getImpulseBootstrapTopic() { return IMPULSE_BOOTSTRAP_TOPIC_PREFIX + pipelineOptions.getApplicationId(); } + + /** + * Returns the dedicated bootstrap topic name a primitive Read reads from. Keyed by transform id + * (sanitized to Kafka's legal topic-name character set) so multiple Reads — and Impulse — never + * register the same topic on two source nodes, which Kafka Streams rejects. + */ + public String getReadBootstrapTopic(String transformId) { + String sanitizedTransformId = ILLEGAL_TOPIC_CHARS.matcher(transformId).replaceAll("_"); + return READ_BOOTSTRAP_TOPIC_PREFIX + + pipelineOptions.getApplicationId() + + "_" + + sanitizedTransformId; + } } 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 new file mode 100644 index 000000000000..604eca54014a --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadProcessor.java @@ -0,0 +1,207 @@ +/* + * 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.io.IOException; +import java.time.Duration; +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.io.BoundedSource; +import org.apache.beam.sdk.io.BoundedSource.BoundedReader; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.CoderUtils; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.kafka.streams.processor.Cancellable; +import org.apache.kafka.streams.processor.PunctuationType; +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.apache.kafka.streams.state.KeyValueStore; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Kafka Streams {@link Processor} implementing Beam's deprecated primitive {@code Read} + * (beam:transform:read:v1) over a {@link BoundedSource}. + * + *

For each task instance, reads the whole {@link BoundedSource} once and emits, in order: + * + *

    + *
  1. One {@link KStreamsPayload#data data} payload per source element, each wrapping a {@link + * WindowedValue} in the {@link org.apache.beam.sdk.transforms.windowing.GlobalWindow} at the + * element's own event time (from {@link BoundedReader#getCurrentTimestamp()}). + *
  2. A {@link KStreamsPayload#watermark watermark} payload at {@link + * BoundedWindow#TIMESTAMP_MAX_VALUE} telling downstream transforms the source is done. + *
+ * + *

Wire form. Unlike Impulse (whose element is already an opaque {@code byte[]}), a Read + * produces decoded Java objects. Downstream {@link ExecutableStageProcessor} feeds + * whatever it receives straight into the SDK harness, whose main-input receiver expects each + * element in the runner-side wire form — a raw object for a model coder, but a length-prefixed + * {@code byte[]} for a coder the runner does not know (e.g. {@code VarIntCoder}). Stage-to-stage + * edges already carry that wire form because harness outputs are decoded with the runner-side wire + * coder; this processor reproduces it for the source edge by transcoding each element through the + * SDK-side wire coder (encode) and back through the runner-side wire coder (decode). The two are + * byte-compatible by construction, so the transcode yields exactly the object the receiver expects, + * nesting and all. + * + *

This mirrors {@link ImpulseProcessor}: a persistent state store records whether the elements + * have already been emitted so task restarts do not duplicate them, while the terminal watermark is + * re-emitted on every restart so downstream watermark holds still release after recovery. The + * trigger is a wall-clock punctuator scheduled on {@link #init} so the processor fires even though + * its bootstrap source topic is empty. + * + *

The source is read in a single instance with no splitting — parallelism across the source's + * splits arrives with the topic-based shuffle work (#18479). Kafka Streams disallows negative + * record timestamps, so each forwarded {@link Record} carries the Unix epoch ({@code 0L}); the Beam + * event time lives inside the {@link WindowedValue}. + */ +class ReadProcessor implements Processor> { + + private static final Logger LOG = LoggerFactory.getLogger(ReadProcessor.class); + + /** Sole entry in the state store; the value tracks whether this processor has already emitted. */ + static final String FIRED_KEY = "fired"; + + /** How soon after {@link #init} the punctuator first fires. */ + private static final Duration PUNCTUATION_DELAY = Duration.ofMillis(50); + + private final BoundedSource source; + // Held as SerializablePipelineOptions (Beam's idiom for a reader holding options) so the + // processor's captured state is uniformly serializable alongside the BoundedSource and coders. + private final SerializablePipelineOptions options; + // Encodes a raw WindowedValue as the SDK harness would on the wire (length-prefixing the + // element coder if the runner does not know it); the runner-side coder then decodes it into the + // wire form the downstream stage's input receiver expects. See the class javadoc. + private final Coder> sdkWireCoder; + private final Coder> runnerWireCoder; + private final String stateStoreName; + private final String transformId; + + private @Nullable ProcessorContext> context; + private @Nullable KeyValueStore firedStore; + private @Nullable Cancellable scheduledPunctuator; + + ReadProcessor( + BoundedSource source, + SerializablePipelineOptions options, + Coder> sdkWireCoder, + Coder> runnerWireCoder, + String stateStoreName, + String transformId) { + this.source = source; + this.options = options; + this.sdkWireCoder = sdkWireCoder; + this.runnerWireCoder = runnerWireCoder; + this.stateStoreName = stateStoreName; + this.transformId = transformId; + } + + @Override + public void init(ProcessorContext> context) { + this.context = context; + this.firedStore = context.getStateStore(stateStoreName); + this.scheduledPunctuator = + context.schedule(PUNCTUATION_DELAY, PunctuationType.WALL_CLOCK_TIME, ts -> maybeFire()); + } + + @Override + public void process(Record record) { + // Records that happen to land on the bootstrap topic are not actual data; they just provide an + // extra opportunity to fire the read on restart. The state store still gates the emit. + maybeFire(); + } + + private void maybeFire() { + ProcessorContext> ctx = context; + KeyValueStore store = firedStore; + if (ctx == null || store == null) { + return; + } + if (Boolean.TRUE.equals(store.get(FIRED_KEY))) { + // Elements were already emitted in a previous task lifetime, but downstream watermark holds + // may still need to release after the restart — re-emit the terminal watermark and stop. + forwardWatermarkMax(ctx); + cancelPunctuator(); + return; + } + int count = readAndForward(ctx); + forwardWatermarkMax(ctx); + store.put(FIRED_KEY, Boolean.TRUE); + cancelPunctuator(); + LOG.debug("Read {} emitted {} elements and terminal watermark", transformId, count); + } + + /** Reads the whole bounded source and forwards each element, in wire form, as a data payload. */ + private int readAndForward(ProcessorContext> ctx) { + int count = 0; + try (BoundedReader reader = source.createReader(options.get())) { + for (boolean hasElement = reader.start(); hasElement; hasElement = reader.advance()) { + WindowedValue element = + WindowedValues.timestampedValueInGlobalWindow( + reader.getCurrent(), reader.getCurrentTimestamp()); + // The Read output PCollection is not keyed; use an empty byte[] as a placeholder key so + // downstream processors that adopt the byte[]-key convention see a consistent shape. + ctx.forward( + new Record>( + new byte[0], KStreamsPayload.data(toRunnerWire(element)), 0L)); + count++; + } + } catch (IOException e) { + throw new RuntimeException("Failed to read bounded source for transform " + transformId, e); + } + return count; + } + + /** Transcodes a raw element into the runner-side wire form the SDK harness input expects. */ + private WindowedValue toRunnerWire(WindowedValue element) { + try { + byte[] wireBytes = CoderUtils.encodeToByteArray(sdkWireCoder, element); + return CoderUtils.decodeFromByteArray(runnerWireCoder, wireBytes); + } catch (CoderException e) { + throw new RuntimeException( + "Failed to transcode a read element to wire form for transform " + transformId, e); + } + } + + /** + * 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. + */ + private static void forwardWatermarkMax(ProcessorContext> ctx) { + long maxMillis = BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(); + ctx.forward( + new Record>( + new byte[0], KStreamsPayload.watermark(maxMillis, 0, 1), 0L)); + } + + /** Cancels the wall-clock punctuator after the read has fired to stop periodic wakeups. */ + private void cancelPunctuator() { + Cancellable handle = scheduledPunctuator; + if (handle != null) { + handle.cancel(); + scheduledPunctuator = null; + } + } +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java new file mode 100644 index 000000000000..403499ca616c --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ReadTranslator.java @@ -0,0 +1,173 @@ +/* + * 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.io.IOException; +import org.apache.beam.model.pipeline.v1.RunnerApi; +import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload.WireCoderSetting; +import org.apache.beam.runners.core.construction.SerializablePipelineOptions; +import org.apache.beam.runners.fnexecution.wire.WireCoders; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.io.BoundedSource; +import org.apache.beam.sdk.util.construction.ReadTranslation; +import org.apache.beam.sdk.util.construction.RehydratedComponents; +import org.apache.beam.sdk.util.construction.graph.PipelineNode; +import org.apache.beam.sdk.util.construction.graph.PipelineNode.PCollectionNode; +import org.apache.beam.sdk.values.WindowedValue; +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; +import org.apache.kafka.streams.state.KeyValueBytesStoreSupplier; +import org.apache.kafka.streams.state.Stores; + +/** + * Translates the deprecated primitive {@code Read} URN ({@code beam:transform:read:v1}) over a + * {@link BoundedSource}. + * + *

The runner forces every {@code Read.Bounded} (including the one {@code Create} of two or more + * elements expands to) into this primitive read before translation — see {@code + * KafkaStreamsTestRunner.translate}, which applies {@code + * SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReads}. This deliberately avoids the + * default {@code BoundedSourceAsSDFWrapperFn} splittable-DoFn expansion, which the runner cannot + * execute yet (no SDF restriction protocol), as agreed with the mentor. + * + *

Adds the same three-node shape as {@link ImpulseTranslator}: + * + *

    + *
  • A {@code byte[]} source bound to a dedicated per-transform bootstrap topic (see {@link + * KafkaStreamsTranslationContext#getReadBootstrapTopic(String)}). Kafka Streams refuses to + * start a topology with no real source topic; records published to it are ignored by {@link + * ReadProcessor}. + *
  • The {@link ReadProcessor}, which reads the {@link BoundedSource} on a one-shot wall-clock + * punctuator and emits one data payload per element followed by a terminal watermark. + *
  • A per-processor persistent state store recording whether the read has already fired so task + * restarts do not duplicate elements. + *
+ * + *

The processor emits elements in the runner-side wire form the downstream stage's SDK harness + * expects, so it is handed the SDK-side and runner-side wire coders for the read's output + * PCollection (see {@link ReadProcessor} for why). Only {@link + * org.apache.beam.model.pipeline.v1.RunnerApi.IsBounded.Enum#BOUNDED bounded} sources are + * supported; {@link ReadTranslation#boundedSourceFromProto} rejects an unbounded payload. + */ +class ReadTranslator implements PTransformTranslator { + + static final String SOURCE_SUFFIX = "-source"; + static final String STATE_STORE_SUFFIX = "-state"; + + @Override + public void translate( + String transformId, RunnerApi.Pipeline pipeline, KafkaStreamsTranslationContext context) { + RunnerApi.PTransform transform = pipeline.getComponents().getTransformsOrThrow(transformId); + // Read produces exactly one output PCollection; downstream consumers are separate PTransforms + // whose inputs reference this PCollection id and are wired by their own translators. + String outputPCollectionId = Iterables.getOnlyElement(transform.getOutputsMap().values()); + addReadNodes( + transformId, + boundedSource(transform), + pipeline.getComponents(), + outputPCollectionId, + context); + } + + private static BoundedSource boundedSource(RunnerApi.PTransform transform) { + try { + RunnerApi.ReadPayload payload = + RunnerApi.ReadPayload.parseFrom(transform.getSpec().getPayload()); + return ReadTranslation.boundedSourceFromProto(payload); + } catch (IOException e) { + throw new RuntimeException( + "Failed to read the BoundedSource from transform " + transform.getUniqueName(), e); + } + } + + /** + * Adds the source, {@link ReadProcessor}, and state store for the read. The type variable {@code + * T} captures the {@link BoundedSource}'s element type so the processor and its wire coders are + * built consistently. + */ + private void addReadNodes( + String transformId, + BoundedSource source, + RunnerApi.Components components, + String outputPCollectionId, + KafkaStreamsTranslationContext context) { + PCollectionNode outputNode = + PipelineNode.pCollection( + outputPCollectionId, components.getPcollectionsOrThrow(outputPCollectionId)); + Coder> sdkWireCoder = sdkWireCoder(outputNode, components); + Coder> runnerWireCoder = runnerWireCoder(outputNode, components); + + Topology topology = context.getTopology(); + String sourceNodeName = transformId + SOURCE_SUFFIX; + String stateStoreName = transformId + STATE_STORE_SUFFIX; + String bootstrapTopic = context.getReadBootstrapTopic(transformId); + SerializablePipelineOptions options = + new SerializablePipelineOptions(context.getPipelineOptions()); + + topology.addSource( + sourceNodeName, + Serdes.ByteArray().deserializer(), + Serdes.ByteArray().deserializer(), + bootstrapTopic); + topology.addProcessor( + transformId, + () -> + new ReadProcessor<>( + source, options, sdkWireCoder, runnerWireCoder, stateStoreName, transformId), + sourceNodeName); + KeyValueBytesStoreSupplier storeSupplier = Stores.persistentKeyValueStore(stateStoreName); + topology.addStateStore( + Stores.keyValueStoreBuilder(storeSupplier, Serdes.String(), Serdes.Boolean()), transformId); + + context.registerPCollectionProducer(outputPCollectionId, transformId); + } + + /** The coder the SDK harness would use on the wire, keeping unknown element coders intact. */ + private static Coder> sdkWireCoder( + PCollectionNode outputNode, RunnerApi.Components components) { + try { + RunnerApi.Components.Builder builder = components.toBuilder(); + String coderId = + WireCoders.addSdkWireCoder(outputNode, builder, WireCoderSetting.getDefaultInstance()); + @SuppressWarnings("unchecked") + Coder> coder = + (Coder>) + RehydratedComponents.forComponents(builder.build()).getCoder(coderId); + return coder; + } catch (IOException e) { + throw new RuntimeException( + "Failed to build the SDK wire coder for PCollection " + outputNode.getId(), e); + } + } + + /** The coder the runner uses on the wire, replacing unknown element coders with byte arrays. */ + private static Coder> runnerWireCoder( + PCollectionNode outputNode, RunnerApi.Components components) { + try { + @SuppressWarnings("unchecked") + Coder> coder = + (Coder>) + (Coder) WireCoders.instantiateRunnerWireCoder(outputNode, components); + return coder; + } catch (IOException e) { + throw new RuntimeException( + "Failed to build the runner wire coder for PCollection " + outputNode.getId(), e); + } + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTestRunner.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTestRunner.java index 53b7714548a0..4d73eee9581a 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTestRunner.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTestRunner.java @@ -36,6 +36,7 @@ import org.apache.beam.sdk.util.construction.Environments; import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation; import org.apache.beam.sdk.util.construction.PipelineTranslation; +import org.apache.beam.sdk.util.construction.SplittableParDo; import org.apache.kafka.common.serialization.ByteArrayDeserializer; import org.apache.kafka.common.serialization.ByteArraySerializer; import org.apache.kafka.common.serialization.Serdes; @@ -88,6 +89,11 @@ public static PipelineOptions testOptions() { public static KafkaStreamsTranslationContext translate(Pipeline pipeline) { KafkaStreamsPipelineOptions options = pipeline.getOptions().as(KafkaStreamsPipelineOptions.class); + // Force every Read.Bounded (including the one Create of 2+ elements expands to) into the + // deprecated primitive Read the runner translates, instead of the default + // BoundedSourceAsSDFWrapperFn splittable-DoFn expansion the runner cannot execute yet. This is + // unconditional — it does not depend on the use_deprecated_read experiment (mentor's steer). + SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReads(pipeline); RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline); KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); JobInfo jobInfo = diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/CreateTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/CreateTest.java new file mode 100644 index 000000000000..e2fab243d4bf --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/CreateTest.java @@ -0,0 +1,73 @@ +/* + * 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.List; +import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.junit.Test; + +/** + * End-to-end test that {@code Create} of two or more elements runs on the Kafka Streams runner. + * + *

{@code Create.of(1, 2, 3)} expands to {@code Read.from(CreateSource)}, a bounded {@code + * Read.Bounded}. {@code KafkaStreamsTestRunner.translate} forces it into the deprecated primitive + * read the runner translates (see {@link ReadTranslator}), so the elements flow through {@link + * ReadProcessor} into a recording ParDo in the EMBEDDED harness. This is the multi-element Create + * case that previously failed because the default expansion is an unsupported splittable DoFn. + */ +public class CreateTest { + + /** Records every element the harness feeds it so the test can assert Create produced them. */ + 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 createOfMultipleElementsRunsThroughPrimitiveRead() { + try (SharedTestCollector collector = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + pipeline + .apply("create", Create.of(1, 2, 3)) + .apply("record", ParDo.of(new RecordingFn(collector))); + + KafkaStreamsTestRunner.run(pipeline); + + List recorded = collector.recorded(); + assertThat(recorded.size(), is(3)); + assertThat(recorded, hasItems(1, 2, 3)); + } + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ReadTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ReadTest.java new file mode 100644 index 000000000000..5b50aa272145 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ReadTest.java @@ -0,0 +1,75 @@ +/* + * 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.List; +import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.io.CountingSource; +import org.apache.beam.sdk.io.Read; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.junit.Test; + +/** + * End-to-end test for {@link ReadTranslator}: a primitive bounded {@code Read} over a {@link + * CountingSource} feeds a recording ParDo executed in the in-process (EMBEDDED) Java SDK harness. + * Verifies every source element reaches the DoFn. + * + *

{@code KafkaStreamsTestRunner.translate} forces the {@code Read.Bounded} into the deprecated + * primitive read the runner translates, rather than the default splittable-DoFn expansion — so this + * exercises the {@link ReadProcessor} + {@link ReadTranslator} path end to end. + */ +public class ReadTest { + + /** Records every element the harness feeds it so the test can assert the read produced them. */ + private static class RecordingFn extends DoFn { + private final SharedTestCollector collector; + + RecordingFn(SharedTestCollector collector) { + this.collector = collector; + } + + @ProcessElement + public void processElement(@Element Long input, OutputReceiver out) { + collector.record(input); + out.output(input); + } + } + + @Test + public void boundedReadEmitsEverySourceElement() { + try (SharedTestCollector collector = SharedTestCollector.create()) { + Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions()); + pipeline + .apply("read", Read.from(CountingSource.upTo(5))) + .apply("record", ParDo.of(new RecordingFn(collector))); + + KafkaStreamsTestRunner.run(pipeline); + + List recorded = collector.recorded(); + // CountingSource.upTo(5) yields 0..4 exactly once each. + assertThat(recorded.size(), is(5)); + assertThat(recorded, hasItems(0L, 1L, 2L, 3L, 4L)); + } + } +}