This is an automated email from the ASF dual-hosted git repository.
je-ik pushed a commit to branch feat/18479-kafka-streams-runner-skeleton
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to
refs/heads/feat/18479-kafka-streams-runner-skeleton by this push:
new 27198768f89 [GSoC 2026] Kafka Streams runner: run on a real broker,
correctly across partitions (#39546)
27198768f89 is described below
commit 27198768f89f5b177e65e9c00c32dbba2d94cfb1
Author: M Junaid Shaukat <[email protected]>
AuthorDate: Sat Aug 1 14:31:20 2026 +0500
[GSoC 2026] Kafka Streams runner: run on a real broker, correctly across
partitions (#39546)
* [GSoC 2026] Kafka Streams runner: run against a real broker
Adds what the runner needs to execute on an actual Kafka cluster, and an
integration test that runs a pipeline through the production
KafkaStreamsPipelineRunner against a Kafka container. Everything until now
ran through TopologyTestDriver, which fakes the topics and never builds a
KafkaStreams application, so the production path had not been executed.
---
runners/kafka-streams/build.gradle | 23 ++
.../kafka/streams/KafkaStreamsPipelineOptions.java | 17 ++
.../kafka/streams/KafkaStreamsPipelineRunner.java | 38 ++-
.../KafkaStreamsPortablePipelineResult.java | 4 +
.../kafka/streams/KafkaStreamsTopicManager.java | 171 ++++++++++++
.../translation/ExecutableStageProcessor.java | 9 +-
.../translation/ExecutableStageTranslator.java | 4 +
.../streams/translation/FlattenProcessor.java | 3 +-
.../streams/translation/FlattenTranslator.java | 5 +
.../streams/translation/GroupByKeyTranslator.java | 12 +-
.../KafkaStreamsTranslationContext.java | 30 +++
.../translation/RedistributeTranslator.java | 3 +
.../streams/translation/ShuffleByKeyProcessor.java | 37 ++-
.../translation/WindowedGroupByKeyProcessor.java | 3 +-
.../kafka/streams/KafkaStreamsRunnerBrokerIT.java | 286 +++++++++++++++++++++
.../translation/ShuffleByKeyProcessorTest.java | 123 +++++++++
.../StandardWindowFnTranslationTest.java | 152 +++++++++++
17 files changed, 900 insertions(+), 20 deletions(-)
diff --git a/runners/kafka-streams/build.gradle
b/runners/kafka-streams/build.gradle
index bdf7e3be058..d1326c079e5 100644
--- a/runners/kafka-streams/build.gradle
+++ b/runners/kafka-streams/build.gradle
@@ -17,6 +17,7 @@
*/
import groovy.json.JsonOutput
+import java.time.Duration
plugins { id 'org.apache.beam.module' }
@@ -74,6 +75,7 @@ dependencies {
testImplementation library.java.junit
testImplementation library.java.mockito_core
testImplementation "org.apache.kafka:kafka-streams-test-utils:$kafka_version"
+ testImplementation library.java.testcontainers_kafka
// Beam's @ValidatesRunner suite: the test classes come from the SDK core
test jar; the runner
// (TestKafkaStreamsRunner) and its TopologyTestDriver harness come from
this module's test
@@ -85,6 +87,27 @@ dependencies {
}
+// The broker integration test drives the production runner against a real
Kafka in Docker, so it
+// is not part of the default build. Run it with
:runners:kafka-streams:brokerIntegrationTest.
+test {
+ filter {
+ excludeTestsMatching 'org.apache.beam.runners.kafka.streams.*IT'
+ }
+}
+
+tasks.register("brokerIntegrationTest", Test) {
+ group = "Verification"
+ description = "Runs the Kafka Streams runner against a real broker (requires
Docker)."
+ outputs.upToDateWhen { false }
+ testClassesDirs = sourceSets.test.output.classesDirs
+ classpath = sourceSets.test.runtimeClasspath
+ filter {
+ includeTestsMatching 'org.apache.beam.runners.kafka.streams.*IT'
+ }
+ // A container start plus a streaming run is well past the default per-test
expectations.
+ timeout = Duration.ofMinutes(15)
+}
+
// Known-failing @ValidatesRunner tests, excluded until the feature they need
lands.
def sickbayTests = [
// Merging (session) windows are not supported yet: ReduceFnRunner drives
them through a merging
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java
index 2fa992e66e7..a5a8bb9328b 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineOptions.java
@@ -55,6 +55,23 @@ public interface KafkaStreamsPipelineOptions extends
PortablePipelineOptions {
void setMaxBundleTimeMs(int maxBundleTimeMs);
+ @Description(
+ "How many partitions the runner gives the internal topics it creates to
shuffle a pipeline"
+ + " through, which is the parallelism the shuffled parts of that
pipeline can reach. A"
+ + " GroupByKey runs one task per partition of its repartition topic,
so this is the"
+ + " number of instances its state and its downstream stages are
spread over. Must be at"
+ + " least 1.")
+ @Default.Integer(1)
+ int getInternalParallelism();
+
+ void setInternalParallelism(int internalParallelism);
+
+ @Description("Replication factor for the internal topics the runner creates
for a pipeline.")
+ @Default.Short(1)
+ short getTopicReplicationFactor();
+
+ void setTopicReplicationFactor(short topicReplicationFactor);
+
@Description("Directory where Kafka Streams stores local state.")
@Default.InstanceFactory(StateDirDefaultFactory.class)
String getStateDir();
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunner.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunner.java
index cdb59f67cce..96b4b2cd7f9 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunner.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunner.java
@@ -24,10 +24,10 @@ import
org.apache.beam.runners.jobsubmission.PortablePipelineResult;
import org.apache.beam.runners.jobsubmission.PortablePipelineRunner;
import
org.apache.beam.runners.kafka.streams.translation.KafkaStreamsPipelineTranslator;
import
org.apache.beam.runners.kafka.streams.translation.KafkaStreamsTranslationContext;
-import org.apache.beam.sdk.options.PipelineOptionsValidator;
import org.apache.kafka.streams.KafkaStreams;
import org.apache.kafka.streams.StreamsConfig;
import org.apache.kafka.streams.Topology;
+import org.checkerframework.checker.nullness.qual.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -44,9 +44,22 @@ public class KafkaStreamsPipelineRunner implements
PortablePipelineRunner {
@Override
public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo
jobInfo) {
- // Surface a clear error if a required option (e.g. applicationId) is
missing instead of
- // letting Properties.put fail with a raw NullPointerException further
down.
- PipelineOptionsValidator.validate(KafkaStreamsPipelineOptions.class,
pipelineOptions);
+ // Surface a clear error if an option this runner needs is missing,
instead of letting
+ // Properties.put fail with a raw NullPointerException further down. Only
the options that are
+ // meaningful here are checked, rather than validating the whole
interface: this runs on the job
+ // server, executing a pipeline that has already been submitted, so the
client-side options
+ // PortablePipelineOptions marks required — jobEndpoint above all — do not
apply. Flink's
+ // equivalent PortablePipelineRunner does not validate here either.
+ checkRequiredOption("applicationId", pipelineOptions.getApplicationId());
+ checkRequiredOption("bootstrapServers",
pipelineOptions.getBootstrapServers());
+ // A topic cannot have fewer than one partition, and the value is also the
number of watermark
+ // reports a shuffle's consumer waits for, so a non-positive value would
leave it waiting
+ // forever rather than failing.
+ if (pipelineOptions.getInternalParallelism() < 1) {
+ throw new IllegalArgumentException(
+ "--internalParallelism must be at least 1, but was "
+ + pipelineOptions.getInternalParallelism());
+ }
KafkaStreamsPipelineTranslator translator = new
KafkaStreamsPipelineTranslator();
KafkaStreamsTranslationContext context =
@@ -55,15 +68,28 @@ public class KafkaStreamsPipelineRunner implements
PortablePipelineRunner {
translator.translate(context, prepared);
Topology topology = context.getTopology();
+ // The runner names its own bootstrap and repartition topics, which Kafka
Streams treats as
+ // user topics and will not create; it refuses to start if a source topic
is missing.
+ KafkaStreamsTopicManager.createMissingTopics(topology, pipelineOptions);
LOG.info(
"Translated pipeline {} into Kafka Streams topology:\n{}",
jobInfo.jobId(),
topology.describe());
KafkaStreams kafkaStreams = new KafkaStreams(topology,
streamsConfig(jobInfo));
+ // Build the result before starting: it registers a state listener, and
Kafka Streams only
+ // accepts one while the application is still in the CREATED state.
+ KafkaStreamsPortablePipelineResult result =
+ new KafkaStreamsPortablePipelineResult(kafkaStreams,
context.getMetricsContainerStepMap());
kafkaStreams.start();
- return new KafkaStreamsPortablePipelineResult(
- kafkaStreams, context.getMetricsContainerStepMap());
+ return result;
+ }
+
+ private static void checkRequiredOption(String name, @Nullable String value)
{
+ if (value == null || value.isEmpty()) {
+ throw new IllegalArgumentException(
+ "Missing required pipeline option --" + name + " for the Kafka
Streams runner");
+ }
}
private Properties streamsConfig(JobInfo jobInfo) {
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResult.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResult.java
index 972afa15c35..0d508b8189f 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResult.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPortablePipelineResult.java
@@ -48,6 +48,10 @@ class KafkaStreamsPortablePipelineResult implements
PortablePipelineResult {
private final CountDownLatch terminated = new CountDownLatch(1);
private volatile boolean cancelled = false;
+ /**
+ * Must be constructed before {@link KafkaStreams#start()} is called: it
registers a state
+ * listener, and Kafka Streams rejects one once the application has left the
CREATED state.
+ */
KafkaStreamsPortablePipelineResult(
KafkaStreams kafkaStreams, MetricsContainerStepMap
metricsContainerStepMap) {
this.kafkaStreams = kafkaStreams;
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTopicManager.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTopicManager.java
new file mode 100644
index 00000000000..0e5f46f53ea
--- /dev/null
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTopicManager.java
@@ -0,0 +1,171 @@
+/*
+ * 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 java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.AdminClientConfig;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.common.errors.TopicExistsException;
+import org.apache.kafka.streams.Topology;
+import org.apache.kafka.streams.TopologyDescription;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Creates the topics a translated pipeline needs before the Kafka Streams
application starts.
+ *
+ * <p>The runner shuffles data through topics it names itself: a bootstrap
topic per Impulse and per
+ * primitive Read, and a repartition topic per GroupByKey. Kafka Streams does
create the internal
+ * topics it manages on its own, but these are declared with explicit names
through {@code
+ * addSource} and {@code addSink}, so to Kafka Streams they are ordinary user
topics — it will not
+ * create them, and refuses to start with {@code MissingSourceTopicException}
if a source topic is
+ * absent. Relying on the broker's {@code auto.create.topics.enable} is not an
option either: it is
+ * off on many clusters, and a topic auto-created on first fetch gets the
broker's default partition
+ * count rather than the pipeline's.
+ *
+ * <p>Only topics carrying one of the runner's own prefixes are created. Any
other topic in the
+ * topology belongs to the user (a source or sink they named), and creating
those implicitly would
+ * hide a misconfiguration behind an empty topic.
+ */
+class KafkaStreamsTopicManager {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(KafkaStreamsTopicManager.class);
+
+ /**
+ * Prefixes of the bootstrap topics, which must have exactly one partition.
+ *
+ * <p>An Impulse or a primitive Read emits its elements once per task, gated
by a state store that
+ * is itself per task. Kafka Streams creates one task per partition of the
source topic, so a
+ * bootstrap topic with several partitions would make the same Impulse fire
once per partition and
+ * the same source be read once per partition.
+ */
+ private static final List<String> SINGLE_PARTITION_TOPIC_PREFIXES =
+ java.util.Arrays.asList("__beam_impulse_", "__beam_read_");
+
+ /**
+ * Prefixes of the topics whose partition count sets the pipeline's
parallelism — the repartition
+ * topic a GroupByKey shuffles through.
+ */
+ private static final List<String> PARTITIONED_TOPIC_PREFIXES =
+ java.util.Arrays.asList("__beam_gbk_");
+
+ private KafkaStreamsTopicManager() {}
+
+ /**
+ * Creates any runner-owned topic in {@code topology} that does not exist
yet.
+ *
+ * <p>Safe to run concurrently with another instance of the same job: a
topic that appears between
+ * the existence check and the create request surfaces as {@link
TopicExistsException}, which is
+ * treated as success.
+ */
+ static void createMissingTopics(Topology topology,
KafkaStreamsPipelineOptions options) {
+ Set<String> runnerTopics = runnerOwnedTopics(topology);
+ if (runnerTopics.isEmpty()) {
+ return;
+ }
+ Properties adminConfig = new Properties();
+ adminConfig.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG,
options.getBootstrapServers());
+ try (Admin admin = Admin.create(adminConfig)) {
+ Set<String> existing = admin.listTopics().names().get();
+ List<NewTopic> toCreate = new ArrayList<>();
+ for (String topic : runnerTopics) {
+ if (!existing.contains(topic)) {
+ toCreate.add(
+ new NewTopic(
+ topic, partitionsFor(topic, options),
options.getTopicReplicationFactor()));
+ }
+ }
+ if (toCreate.isEmpty()) {
+ return;
+ }
+ LOG.info("Creating {} runner-owned topic(s): {}", toCreate.size(),
toCreate);
+ createAll(admin, toCreate);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while creating the pipeline's
Kafka topics", e);
+ } catch (ExecutionException e) {
+ throw new RuntimeException("Failed to create the pipeline's Kafka
topics", e);
+ }
+ }
+
+ private static void createAll(Admin admin, Collection<NewTopic> topics)
+ throws InterruptedException, ExecutionException {
+ try {
+ admin.createTopics(topics).all().get();
+ } catch (ExecutionException e) {
+ // Another instance of the same application may have created them first,
which is fine.
+ if (!(e.getCause() instanceof TopicExistsException)) {
+ throw e;
+ }
+ LOG.debug("Some topics already existed; another instance created them
first", e);
+ }
+ }
+
+ /** The topics in the topology that the runner named, and so is responsible
for creating. */
+ private static Set<String> runnerOwnedTopics(Topology topology) {
+ Set<String> topics = new HashSet<>();
+ for (TopologyDescription.Subtopology subtopology :
topology.describe().subtopologies()) {
+ for (TopologyDescription.Node node : subtopology.nodes()) {
+ if (node instanceof TopologyDescription.Source) {
+ Set<String> sourceTopics = ((TopologyDescription.Source)
node).topicSet();
+ if (sourceTopics != null) {
+ topics.addAll(sourceTopics);
+ }
+ } else if (node instanceof TopologyDescription.Sink) {
+ String topic = ((TopologyDescription.Sink) node).topic();
+ if (topic != null) {
+ topics.add(topic);
+ }
+ }
+ }
+ }
+ topics.removeIf(topic -> !isRunnerOwned(topic));
+ return topics;
+ }
+
+ /**
+ * The partition count a runner-owned topic is created with: one for a
bootstrap topic, and the
+ * configured parallelism for a shuffle topic.
+ */
+ private static int partitionsFor(String topic, KafkaStreamsPipelineOptions
options) {
+ return hasAnyPrefix(topic, SINGLE_PARTITION_TOPIC_PREFIXES)
+ ? 1
+ : options.getInternalParallelism();
+ }
+
+ private static boolean isRunnerOwned(String topic) {
+ return hasAnyPrefix(topic, SINGLE_PARTITION_TOPIC_PREFIXES)
+ || hasAnyPrefix(topic, PARTITIONED_TOPIC_PREFIXES);
+ }
+
+ private static boolean hasAnyPrefix(String topic, List<String> prefixes) {
+ for (String prefix : prefixes) {
+ if (topic.startsWith(prefix)) {
+ return true;
+ }
+ }
+ return false;
+ }
+}
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 15265073dd2..63fefe3e15c 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
@@ -286,10 +286,11 @@ class ExecutableStageProcessor
}
private void forwardWatermark(Record<byte[], KStreamsPayload<?>> record,
long watermarkMillis) {
- // 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).
+ // Labelled as the only source a consumer will see. Forwarding here is
in-process, to the
+ // stage's
+ // fused children, so exactly one instance of this stage reaches each of
them. Where the output
+ // instead crosses a shuffle, ShuffleByKeyProcessor relabels the report
with the real partition
+ // identity, because the broadcast then delivers every instance's report
to every consumer.
ProcessorContext<byte[], KStreamsPayload<?>> ctx =
checkInitialized(context);
ctx.forward(
new Record<byte[], KStreamsPayload<?>>(
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 71e24f32c1f..caefa6534fa 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
@@ -74,6 +74,8 @@ class ExecutableStageTranslator implements
PTransformTranslator {
// is unambiguous even before we add side-input support.
String inputPCollectionId = stagePayload.getInput();
String parentProcessor =
context.getProcessorNameForPCollection(inputPCollectionId);
+ // A fused stage runs wherever its input runs: same task, so same
partition identity.
+ int partitionCount = context.getPartitionCount(inputPCollectionId);
// A multi-output stage (a DoFn with side outputs, or a Read whose SDF
wrapper produces several
// outputs) needs each output routed to the right downstream. Since
downstream transforms are
@@ -117,9 +119,11 @@ class ExecutableStageTranslator implements
PTransformTranslator {
topology.addProcessor(
relayName, () -> new StageOutputProcessor(relayName),
transformId);
context.registerPCollectionProducer(outputPCollectionId,
relayName);
+ context.registerPCollectionPartitionCount(outputPCollectionId,
partitionCount);
});
} else if (!outputPCollectionIds.isEmpty()) {
context.registerPCollectionProducer(outputPCollectionIds.get(0),
transformId);
+ context.registerPCollectionPartitionCount(outputPCollectionIds.get(0),
partitionCount);
}
}
}
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
index d09fed8185a..e37da677448 100644
---
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
@@ -99,8 +99,7 @@ class FlattenProcessor
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).
+ // Labelled as the only source a consumer will see; a shuffle downstream
relabels it.
ctx.forward(
new Record<byte[], KStreamsPayload<?>>(
record.key(),
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
index 3ac4c1304da..a5c8ce05bae 100644
---
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
@@ -57,6 +57,9 @@ class FlattenTranslator implements PTransformTranslator {
Set<String> seenInputs = new HashSet<>();
List<String> parentProcessors = new ArrayList<>();
Set<String> upstreamTransformIds = new HashSet<>();
+ // Kafka Streams puts a processor and the parents it is wired to in one
subtopology, so the
+ // inputs are co-partitioned and this Flatten runs at their partition
count.
+ int partitionCount = 1;
for (String inputPCollectionId : transform.getInputsMap().values()) {
if (!seenInputs.add(inputPCollectionId)) {
throw new UnsupportedOperationException(
@@ -70,6 +73,7 @@ class FlattenTranslator implements PTransformTranslator {
String parentProcessor =
context.getProcessorNameForPCollection(inputPCollectionId);
parentProcessors.add(parentProcessor);
upstreamTransformIds.add(parentProcessor);
+ partitionCount = Math.max(partitionCount,
context.getPartitionCount(inputPCollectionId));
}
topology.addProcessor(
@@ -78,5 +82,6 @@ class FlattenTranslator implements PTransformTranslator {
parentProcessors.toArray(new String[0]));
context.registerPCollectionProducer(outputPCollectionId, transformId);
+ context.registerPCollectionPartitionCount(outputPCollectionId,
partitionCount);
}
}
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 c5327e28e06..c460436eedf 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
@@ -95,6 +95,9 @@ class GroupByKeyTranslator implements PTransformTranslator {
hydrateWindowingStrategy(pipeline, inputPCollectionId);
String parentProcessor =
context.getProcessorNameForPCollection(inputPCollectionId);
+ // The shuffle is what changes the parallelism: everything from the
repartition topic onwards
+ // runs one task per partition of it.
+ int partitionCount = context.getPipelineOptions().getInternalParallelism();
String shuffleName = transformId + SHUFFLE_SUFFIX;
String sinkName = transformId + SINK_SUFFIX;
@@ -110,7 +113,13 @@ class GroupByKeyTranslator implements PTransformTranslator
{
Topology topology = context.getTopology();
// Re-key data records by the encoded Beam key; pass watermark reports
through.
- topology.addProcessor(shuffleName, () -> new
ShuffleByKeyProcessor(keyCoder), parentProcessor);
+ // The shuffle runs in the upstream transform's task, so it relabels each
report with that
+ // transform's instance identity before the sink broadcasts it to every
partition.
+ int upstreamPartitionCount = context.getPartitionCount(inputPCollectionId);
+ topology.addProcessor(
+ shuffleName,
+ () -> new ShuffleByKeyProcessor(keyCoder, upstreamPartitionCount),
+ parentProcessor);
// Shuffle through the repartition topic: data partitioned by key,
watermark broadcast.
topology.addSink(
@@ -168,6 +177,7 @@ class GroupByKeyTranslator implements PTransformTranslator {
transformId);
context.registerPCollectionProducer(outputPCollectionId, transformId);
+ context.registerPCollectionPartitionCount(outputPCollectionId,
partitionCount);
}
/** Hydrates the input PCollection's windowing strategy from the pipeline
proto. */
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 ec1b3f26ade..d0316961566 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
@@ -46,6 +46,12 @@ public class KafkaStreamsTranslationContext {
private final KafkaStreamsPipelineOptions pipelineOptions;
private final Topology topology;
private final Map<String, String> pCollectionIdToProcessorName;
+
+ /**
+ * How many partitions the transform producing each PCollection runs across.
A PCollection that
+ * has not been registered is produced by a single instance; only a shuffle
raises the count.
+ */
+ private final Map<String, Integer> pCollectionIdToPartitionCount = new
HashMap<>();
// Accumulates the Beam metrics reported by the SDK harness, one container
per executable stage.
// Processors update it as bundles complete (in-JVM reference sharing); the
pipeline result
// exposes it as MetricResults. Sharing one container across a stage's
parallel tasks is safe and
@@ -113,6 +119,30 @@ public class KafkaStreamsTranslationContext {
}
}
+ /**
+ * Records how many partitions the transform producing {@code pCollectionId}
runs across.
+ *
+ * <p>This is the {@code totalSourcePartitions} its watermark reports carry,
and what a downstream
+ * {@link WatermarkAggregator} waits to hear from before it lets the
watermark advance. It changes
+ * only at a shuffle: everything fused downstream of one runs at the shuffle
topic's partition
+ * count, and everything else runs as a single instance.
+ */
+ public void registerPCollectionPartitionCount(String pCollectionId, int
partitionCount) {
+ pCollectionIdToPartitionCount.put(pCollectionId, partitionCount);
+ }
+
+ /**
+ * How many partitions the transform producing {@code pCollectionId} runs
across; one unless a
+ * shuffle upstream raised it.
+ *
+ * <p>Always at least one: an unregistered PCollection is produced by a
single instance, and the
+ * only value ever registered is {@code --internalParallelism}, which the
runner rejects below one
+ * before translating.
+ */
+ public int getPartitionCount(String pCollectionId) {
+ return pCollectionIdToPartitionCount.getOrDefault(pCollectionId, 1);
+ }
+
/** Returns the processor node name producing the given PCollection. */
public String getProcessorNameForPCollection(String pCollectionId) {
String name = pCollectionIdToProcessorName.get(pCollectionId);
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/RedistributeTranslator.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/RedistributeTranslator.java
index 72c47db8d4b..a44740887b3 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/RedistributeTranslator.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/RedistributeTranslator.java
@@ -50,5 +50,8 @@ class RedistributeTranslator implements PTransformTranslator {
// Passthrough: downstream lookups for the output PCollection resolve to
the producer of the
// input PCollection. No KS Processor / state store / source is added.
context.registerPCollectionProducer(outputPCollectionId, parentProcessor);
+ // A pass-through: the output is produced by the same processor, so same
partition identity.
+ context.registerPCollectionPartitionCount(
+ outputPCollectionId, context.getPartitionCount(inputPCollectionId));
}
}
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 774d36db0f2..79595cba7c9 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
@@ -34,9 +34,16 @@ import org.slf4j.LoggerFactory;
* <p>This is not GroupByKey-specific: any transform that needs the values of
a key co-located on
* one partition uses it — GroupByKey today, and stateful ParDo later. For a
data record it sets the
* Kafka record key to the encoded Beam key (taken from the {@code KV}), so
the downstream
- * repartition sink co-locates every value of a key. Watermark reports are
forwarded unchanged — the
- * {@link GroupByKeyBroadcastPartitioner} fans them out to all partitions so
every downstream task
- * can fire.
+ * repartition sink co-locates every value of a key.
+ *
+ * <p>Watermark reports are relabelled here with the reporting instance's real
partition identity.
+ * This is the point at which a report stops being delivered in-process and
starts crossing a topic:
+ * upstream of it a transform forwards to its fused children, which see
exactly one instance of it,
+ * so the report names a single source. The {@link
GroupByKeyBroadcastPartitioner} on the sink below
+ * fans each report out to <em>every</em> partition, so a downstream task
instead sees a report from
+ * every instance of the upstream transform, and has to be able to tell them
apart to know when it
+ * has heard from all of them. The transform id is left alone, so the report
still names the
+ * transform that produced it.
*/
class ShuffleByKeyProcessor
implements Processor<byte[], KStreamsPayload<?>, byte[],
KStreamsPayload<?>> {
@@ -44,15 +51,25 @@ class ShuffleByKeyProcessor
private static final Logger LOG =
LoggerFactory.getLogger(ShuffleByKeyProcessor.class);
private final Coder<Object> keyCoder;
+
+ /** How many instances the transform being shuffled runs as, and which one
this is. */
+ private final int upstreamPartitionCount;
+
+ private int upstreamPartition;
+
private @Nullable ProcessorContext<byte[], KStreamsPayload<?>> context;
- ShuffleByKeyProcessor(Coder<Object> keyCoder) {
+ ShuffleByKeyProcessor(Coder<Object> keyCoder, int upstreamPartitionCount) {
this.keyCoder = keyCoder;
+ this.upstreamPartitionCount = upstreamPartitionCount;
}
@Override
public void init(ProcessorContext<byte[], KStreamsPayload<?>> context) {
this.context = context;
+ // This processor runs in the upstream transform's task, so the task's
partition is the
+ // identity of the instance whose reports it is forwarding.
+ this.upstreamPartition = context.taskId().partition();
}
@Override
@@ -82,8 +99,16 @@ class ShuffleByKeyProcessor
}
ctx.forward(record.withKey(encodedKey));
} else {
- // Watermark report: forward as-is; the sink's partitioner broadcasts it
to all partitions.
- ctx.forward(record);
+ WatermarkPayload report = payload.asWatermark();
+ ctx.forward(
+ new Record<byte[], KStreamsPayload<?>>(
+ record.key(),
+ KStreamsPayload.watermark(
+ report.getWatermarkMillis(),
+ report.getTransformId(),
+ upstreamPartition,
+ upstreamPartitionCount),
+ record.timestamp()));
}
}
diff --git
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java
index 1b0ffae23fd..d00b01a3753 100644
---
a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java
+++
b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/WindowedGroupByKeyProcessor.java
@@ -284,7 +284,8 @@ class WindowedGroupByKeyProcessor<K, V, W extends
BoundedWindow>
private void forwardWatermark(Record<byte[], KStreamsPayload<?>> trigger,
long watermarkMillis) {
ProcessorContext<byte[], KStreamsPayload<?>> ctx =
checkInitialized(context);
- // Stamped with this transform's own id; GroupByKey is a single instance
for now (0 of 1).
+ // Labelled as the only source a consumer will see; see
ExecutableStageProcessor for why an
+ // in-process edge reports a single source and a shuffle relabels.
ctx.forward(
new Record<byte[], KStreamsPayload<?>>(
trigger.key(),
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java
new file mode 100644
index 00000000000..bb82a403e55
--- /dev/null
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java
@@ -0,0 +1,286 @@
+/*
+ * 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 static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+import java.nio.file.Files;
+import java.util.UUID;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.PipelineResult;
+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.metrics.Counter;
+import org.apache.beam.sdk.metrics.MetricNameFilter;
+import org.apache.beam.sdk.metrics.MetricQueryResults;
+import org.apache.beam.sdk.metrics.MetricResult;
+import org.apache.beam.sdk.metrics.Metrics;
+import org.apache.beam.sdk.metrics.MetricsFilter;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.options.PortablePipelineOptions;
+import org.apache.beam.sdk.testing.CrashingRunner;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.GroupByKey;
+import org.apache.beam.sdk.transforms.Impulse;
+import org.apache.beam.sdk.transforms.ParDo;
+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.beam.sdk.values.KV;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
+import org.joda.time.Duration;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.testcontainers.kafka.KafkaContainer;
+import org.testcontainers.utility.DockerImageName;
+
+/**
+ * Runs a pipeline through the production {@link KafkaStreamsPipelineRunner}
against a real Kafka
+ * broker, rather than through the {@code TopologyTestDriver} the rest of the
suite uses.
+ *
+ * <p>The test driver stands in for a broker well enough for translation and
windowing logic, but it
+ * runs one instance in one thread and fakes the topics. Everything that only
exists on a real
+ * cluster is untested by it: the runner creating its own bootstrap and
repartition topics, records
+ * actually round-tripping through a repartition topic, exactly-once
processing, the state stores'
+ * changelog, and the Kafka Streams application lifecycle. This test covers
that path.
+ *
+ * <p>It needs Docker and so is not part of the default build; the {@code
brokerIntegrationTest}
+ * Gradle task runs it.
+ */
+@RunWith(JUnit4.class)
+public class KafkaStreamsRunnerBrokerIT {
+
+ private static final String NAMESPACE = "brokerIT";
+ private static final String GROUPS_COUNTER = "groups";
+
+ /** How long to wait for the streaming application to work through the
pipeline. */
+ private static final Duration TIMEOUT = Duration.standardMinutes(2);
+
+ private static KafkaContainer kafka;
+
+ @BeforeClass
+ public static void startBroker() {
+ // The official Apache Kafka image. 4.0.0 rather than the 3.9.0 the
runner's client is built
+ // against: Testcontainers' KafkaContainer cannot bring up the 3.9.0 image
(it exits during
+ // startup), and a client talking to a newer broker is the compatibility
direction Kafka
+ // supports anyway.
+ kafka = new KafkaContainer(DockerImageName.parse("apache/kafka:4.0.0"));
+ kafka.start();
+ }
+
+ @AfterClass
+ public static void stopBroker() {
+ if (kafka != null) {
+ kafka.stop();
+ }
+ }
+
+ /** Emits a fixed set of keyed elements, one per key group. */
+ private static class EmitKvsFn extends DoFn<byte[], KV<String, Integer>> {
+ @ProcessElement
+ public void processElement(OutputReceiver<KV<String, Integer>> out) {
+ out.output(KV.of("a", 1));
+ out.output(KV.of("a", 2));
+ out.output(KV.of("b", 3));
+ }
+ }
+
+ /** Counts the groups that come out of the GroupByKey. */
+ private static class CountGroupsFn extends DoFn<KV<String,
Iterable<Integer>>, Void> {
+ private final Counter groups = Metrics.counter(NAMESPACE, GROUPS_COUNTER);
+
+ @ProcessElement
+ public void processElement() {
+ groups.inc();
+ }
+ }
+
+ private KafkaStreamsPipelineOptions options() {
+ return options(1);
+ }
+
+ private KafkaStreamsPipelineOptions options(int topicPartitions) {
+ KafkaStreamsPipelineOptions options =
+ PipelineOptionsFactory.create().as(KafkaStreamsPipelineOptions.class);
+ options.setRunner(CrashingRunner.class);
+ options.setBootstrapServers(kafka.getBootstrapServers());
+ options.setApplicationId("ks-broker-it-" + UUID.randomUUID());
+ options.setInternalParallelism(topicPartitions);
+ options
+ .as(PortablePipelineOptions.class)
+ .setDefaultEnvironmentType(Environments.ENVIRONMENT_EMBEDDED);
+ try {
+
options.setStateDir(Files.createTempDirectory("ks-broker-it").toString());
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ return options;
+ }
+
+ @Test
+ public void groupByKeyRunsThroughARealBrokerAndReportsMetrics() throws
Exception {
+ KafkaStreamsPipelineOptions options = options();
+ Pipeline pipeline = Pipeline.create(options);
+ pipeline
+ .apply(Impulse.create())
+ .apply("emit", ParDo.of(new EmitKvsFn()))
+ .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))
+ .apply(GroupByKey.create())
+ .apply("countGroups", ParDo.of(new CountGroupsFn()));
+
+ // The same conversion the test runner does: translate Read-based sources
as the primitive Read
+ // the runner supports rather than the splittable-DoFn expansion.
+ SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReads(pipeline);
+ RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline);
+ JobInfo jobInfo =
+ JobInfo.create(
+ options.getApplicationId(),
+ options.getJobName(),
+ "",
+ PipelineOptionsTranslation.toProto(options));
+
+ PipelineResult result = new
KafkaStreamsPipelineRunner(options).run(pipelineProto, jobInfo);
+ try {
+ // Two keys in, so two groups out once the elements have travelled
through the repartition
+ // topic and the watermark has closed the global window.
+ assertThat(awaitCounter(result, 2L), is(2L));
+ } finally {
+ result.cancel();
+ }
+ }
+
+ /**
+ * Collapses every group onto one key, so the next GroupByKey has to shuffle
across partitions.
+ */
+ private static class ToSingleKeyFn
+ extends DoFn<KV<String, Iterable<Integer>>, KV<String, Integer>> {
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, Iterable<Integer>> group,
OutputReceiver<KV<String, Integer>> out) {
+ int sum = 0;
+ for (int value : group.getValue()) {
+ sum += value;
+ }
+ out.output(KV.of("all", sum));
+ }
+ }
+
+ /** Builds the two-GroupByKey pipeline used by the chained tests. */
+ private static void buildChainedPipeline(Pipeline pipeline) {
+ pipeline
+ .apply(Impulse.create())
+ .apply("emit", ParDo.of(new EmitKvsFn()))
+ .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))
+ .apply("groupPerKey", GroupByKey.create())
+ .apply("toSingleKey", ParDo.of(new ToSingleKeyFn()))
+ .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))
+ .apply("groupAll", GroupByKey.create())
+ .apply("countGroups", ParDo.of(new CountGroupsFn()));
+ }
+
+ private static PipelineResult runPipeline(
+ Pipeline pipeline, KafkaStreamsPipelineOptions options) {
+ SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReads(pipeline);
+ RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline);
+ JobInfo jobInfo =
+ JobInfo.create(
+ options.getApplicationId(),
+ options.getJobName(),
+ "",
+ PipelineOptionsTranslation.toProto(options));
+ return new KafkaStreamsPipelineRunner(options).run(pipelineProto, jobInfo);
+ }
+
+ @Test
+ public void chainedGroupByKeysAreCorrectOnOnePartition() throws Exception {
+ // The control for the partitioned case below: the same shape, one
partition throughout.
+ KafkaStreamsPipelineOptions options = options(1);
+ Pipeline pipeline = Pipeline.create(options);
+ buildChainedPipeline(pipeline);
+
+ PipelineResult result = runPipeline(pipeline, options);
+ try {
+ assertThat(awaitCounter(result, 1L), is(1L));
+ } finally {
+ result.cancel();
+ }
+ }
+
+ @Test
+ public void chainedGroupByKeysAreCorrectAcrossPartitions() throws Exception {
+ // Two GroupByKeys with a partitioned shuffle between them, which is what
makes each task's
+ // watermark identity matter. The second GroupByKey aggregates the reports
of every task of the
+ // first, so those tasks have to report under their own partition: if each
claimed to be the
+ // only partition, the second would advance its watermark on the first
report it saw and fire
+ // before the remaining partitions had contributed their groups.
+ KafkaStreamsPipelineOptions options = options(4);
+ Pipeline pipeline = Pipeline.create(options);
+ buildChainedPipeline(pipeline);
+
+ PipelineResult result = runPipeline(pipeline, options);
+ try {
+ // Everything collapses onto one key, so the second GroupByKey emits
exactly one group — and
+ // only once every partition of the first has contributed to it.
+ assertThat(awaitCounter(result, 1L), is(1L));
+ // A premature firing would show up as a second group, so give one a
chance to appear.
+ Thread.sleep(5_000L);
+ assertThat(counterValue(result), is(1L));
+ } finally {
+ result.cancel();
+ }
+ }
+
+ /**
+ * Polls the pipeline's metrics until the counter reaches {@code expected}
or the timeout hits.
+ */
+ private static long awaitCounter(PipelineResult result, long expected)
throws Exception {
+ long deadline = System.currentTimeMillis() + TIMEOUT.getMillis();
+ long value = 0;
+ while (System.currentTimeMillis() < deadline) {
+ value = counterValue(result);
+ if (value >= expected) {
+ return value;
+ }
+ Thread.sleep(500L);
+ }
+ return value;
+ }
+
+ private static long counterValue(PipelineResult result) {
+ MetricQueryResults query =
+ result
+ .metrics()
+ .queryMetrics(
+ MetricsFilter.builder()
+ .addNameFilter(MetricNameFilter.named(NAMESPACE,
GROUPS_COUNTER))
+ .build());
+ if (Iterables.isEmpty(query.getCounters())) {
+ return 0L;
+ }
+ MetricResult<Long> counter = Iterables.getOnlyElement(query.getCounters());
+ return counter.getAttempted();
+ }
+}
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessorTest.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessorTest.java
new file mode 100644
index 00000000000..38669138075
--- /dev/null
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessorTest.java
@@ -0,0 +1,123 @@
+/*
+ * 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.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+import java.util.Properties;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.values.WindowedValues;
+import org.apache.kafka.streams.processor.TaskId;
+import org.apache.kafka.streams.processor.api.MockProcessorContext;
+import org.apache.kafka.streams.processor.api.Record;
+import org.junit.Test;
+
+/**
+ * Tests how {@link ShuffleByKeyProcessor} restamps a watermark report as it
is about to cross a
+ * repartition topic.
+ *
+ * <p>Upstream of the shuffle a transform forwards its watermark in process,
to its fused children,
+ * which see exactly one instance of it — so the report names a single source.
The sink below the
+ * shuffle broadcasts each report to every partition, so a downstream task
sees a report from every
+ * instance of the upstream transform and has to tell them apart to know when
it has heard from all
+ * of them. The shuffle is where that identity is attached.
+ */
+public class ShuffleByKeyProcessorTest {
+
+ private static final String UPSTREAM_ID = "upstream";
+
+ @SuppressWarnings("unchecked")
+ private static ShuffleByKeyProcessor processorFor(int taskPartition, int
upstreamPartitions) {
+ ShuffleByKeyProcessor processor =
+ new ShuffleByKeyProcessor(
+ (org.apache.beam.sdk.coders.Coder<Object>)
+ (org.apache.beam.sdk.coders.Coder<?>) StringUtf8Coder.of(),
+ upstreamPartitions);
+ MockProcessorContext<byte[], KStreamsPayload<?>> ctx =
+ new MockProcessorContext<>(new Properties(), new TaskId(0,
taskPartition), null);
+ processor.init(ctx);
+ lastContext = ctx;
+ return processor;
+ }
+
+ private static MockProcessorContext<byte[], KStreamsPayload<?>> lastContext;
+
+ private static Record<byte[], KStreamsPayload<?>> watermark(long millis) {
+ // As forwarded in process by the upstream transform: a single source,
since a fused child sees
+ // exactly one instance of it.
+ return new Record<>(new byte[0], KStreamsPayload.watermark(millis,
UPSTREAM_ID, 0, 1), 0L);
+ }
+
+ @Test
+ public void restampsTheWatermarkWithTheUpstreamInstanceIdentity() {
+ // Instance 2 of a 4-instance upstream transform.
+ ShuffleByKeyProcessor processor = processorFor(2, 4);
+
+ processor.process(watermark(500L));
+
+ assertThat(lastContext.forwarded().size(), is(1));
+ WatermarkPayload out =
lastContext.forwarded().get(0).record().value().asWatermark();
+ assertThat(out.getWatermarkMillis(), is(500L));
+ // The transform id still names the producer, so a downstream aggregator
matches it to the
+ // upstream it expects; the partition identity is what it counts.
+ assertThat(out.getTransformId(), is(UPSTREAM_ID));
+ assertThat(out.getSourcePartition(), is(2));
+ assertThat(out.getTotalSourcePartitions(), is(4));
+ }
+
+ @Test
+ public void distinctUpstreamInstancesRestampDistinctly() {
+ processorFor(0, 4).process(watermark(100L));
+ WatermarkPayload first =
lastContext.forwarded().get(0).record().value().asWatermark();
+ processorFor(3, 4).process(watermark(100L));
+ WatermarkPayload second =
lastContext.forwarded().get(0).record().value().asWatermark();
+
+ // Two instances of the same transform must be distinguishable downstream,
or a consumer would
+ // treat one report as if every instance had already reported.
+ assertThat(first.getSourcePartition(), is(0));
+ assertThat(second.getSourcePartition(), is(3));
+ }
+
+ @Test
+ public void anUnpartitionedUpstreamStillReportsASingleSource() {
+ ShuffleByKeyProcessor processor = processorFor(0, 1);
+
+ processor.process(watermark(700L));
+
+ WatermarkPayload out =
lastContext.forwarded().get(0).record().value().asWatermark();
+ assertThat(out.getSourcePartition(), is(0));
+ assertThat(out.getTotalSourcePartitions(), is(1));
+ }
+
+ @Test
+ public void dataIsRekeyedByTheBeamKeyAndNotRestamped() {
+ ShuffleByKeyProcessor processor = processorFor(1, 4);
+
+ processor.process(
+ new Record<>(
+ new byte[0],
+ KStreamsPayload.data(
+ WindowedValues.valueInGlobalWindow(
+ org.apache.beam.sdk.values.KV.of("key", "value"))),
+ 0L));
+
+ assertThat(lastContext.forwarded().size(), is(1));
+ assertThat(lastContext.forwarded().get(0).record().value().isData(),
is(true));
+ }
+}
diff --git
a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StandardWindowFnTranslationTest.java
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StandardWindowFnTranslationTest.java
new file mode 100644
index 00000000000..912d51d6b5d
--- /dev/null
+++
b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StandardWindowFnTranslationTest.java
@@ -0,0 +1,152 @@
+/*
+ * 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.instanceOf;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.CoreMatchers.not;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+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.DoFn;
+import org.apache.beam.sdk.transforms.GroupByKey;
+import org.apache.beam.sdk.transforms.Impulse;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.windowing.FixedWindows;
+import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
+import org.apache.beam.sdk.transforms.windowing.SlidingWindows;
+import org.apache.beam.sdk.transforms.windowing.Window;
+import org.apache.beam.sdk.transforms.windowing.WindowFn;
+import org.apache.beam.sdk.util.construction.PipelineTranslation;
+import org.apache.beam.sdk.util.construction.RehydratedComponents;
+import org.apache.beam.sdk.util.construction.WindowingStrategyTranslation;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.WindowingStrategy;
+import org.joda.time.Duration;
+import org.junit.Test;
+
+/**
+ * Checks that the runner reconstructs the standard WindowFns from the
language-neutral windowing
+ * strategy in the pipeline proto.
+ *
+ * <p>The runner executes GroupAlsoByWindow itself (see {@link
WindowedGroupByKeyProcessor}), so it
+ * has to rebuild the WindowFn from the proto rather than call into the SDK.
Beam gives the standard
+ * WindowFns a URN and a parameter payload — {@code
beam:window_fn:fixed_windows:v1} and friends —
+ * and those are what {@link org.apache.beam.runners.core.ReduceFnRunner}
interprets directly. Every
+ * SDK emits the same URNs for them, so a pipeline built in another language
that uses fixed,
+ * sliding, session or global windows produces a strategy this runner can
rebuild; these tests pin
+ * that down.
+ *
+ * <p>What this does <em>not</em> cover is a WindowFn the user wrote
themselves. That cannot be
+ * interpreted runner-side at all: it is opaque to the runner and would have
to be executed through
+ * the SDK harness that owns it. The runner does not support that today —
{@code
+ * WindowingStrategyTranslation.windowFnFromProto} rejects an unrecognised URN
— and it is the case
+ * that would really exercise cross-language windowing.
+ */
+public class StandardWindowFnTranslationTest {
+
+ private static final Duration WINDOW_SIZE = Duration.millis(10);
+
+ private static class EmitKvFn extends DoFn<byte[], KV<String, Integer>> {
+ @ProcessElement
+ public void processElement(OutputReceiver<KV<String, Integer>> out) {
+ out.output(KV.of("a", 1));
+ }
+ }
+
+ /** A windowed GroupByKey pipeline, as a proto — the form the runner is
handed. */
+ private static RunnerApi.Pipeline windowedPipelineProto(WindowFn<Object, ?>
windowFn) {
+ Pipeline pipeline = Pipeline.create(KafkaStreamsTestRunner.testOptions());
+ pipeline
+ .apply(Impulse.create())
+ .apply(ParDo.of(new EmitKvFn()))
+ .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))
+ .apply(Window.into(windowFn))
+ .apply(GroupByKey.create());
+ return PipelineTranslation.toProto(pipeline);
+ }
+
+ /** The windowing strategy of the GroupByKey's input, which is what the
translator reads. */
+ private static RunnerApi.WindowingStrategy
nonGlobalStrategy(RunnerApi.Pipeline proto) {
+ for (RunnerApi.WindowingStrategy strategy :
+ proto.getComponents().getWindowingStrategiesMap().values()) {
+ if (!strategy
+ .getWindowFn()
+ .getUrn()
+ .equals(WindowingStrategyTranslation.GLOBAL_WINDOWS_URN)) {
+ return strategy;
+ }
+ }
+ throw new AssertionError("pipeline has no non-global windowing strategy");
+ }
+
+ @Test
+ public void fixedWindowsTravelAsTheStandardUrn() {
+ RunnerApi.WindowingStrategy strategy =
+ nonGlobalStrategy(windowedPipelineProto(FixedWindows.of(WINDOW_SIZE)));
+
+ // The same URN and payload any SDK emits for fixed windows.
+ assertThat(strategy.getWindowFn().getUrn(),
is(WindowingStrategyTranslation.FIXED_WINDOWS_URN));
+ }
+
+ @Test
+ public void slidingWindowsTravelAsTheStandardUrn() {
+ RunnerApi.WindowingStrategy strategy =
+ nonGlobalStrategy(
+
windowedPipelineProto(SlidingWindows.of(WINDOW_SIZE).every(Duration.millis(5))));
+
+ assertThat(
+ strategy.getWindowFn().getUrn(),
is(WindowingStrategyTranslation.SLIDING_WINDOWS_URN));
+ }
+
+ @Test
+ public void aStandardWindowFnNeedsNoJavaSerialization() {
+ RunnerApi.Pipeline proto =
windowedPipelineProto(FixedWindows.of(WINDOW_SIZE));
+
+ // Java serialization is the fallback for a WindowFn with no standard URN.
A strategy that fell
+ // back to it here would only be reconstructable by a Java runner reading
a Java pipeline.
+ for (RunnerApi.WindowingStrategy strategy :
+ proto.getComponents().getWindowingStrategiesMap().values()) {
+ assertThat(
+ strategy.getWindowFn().getUrn(),
+ is(not(WindowingStrategyTranslation.SERIALIZED_JAVA_WINDOWFN_URN)));
+ }
+ }
+
+ @Test
+ public void hydratingTheStandardUrnRebuildsTheWindowFnTheRunnerRunsWith()
throws Exception {
+ RunnerApi.Pipeline proto =
windowedPipelineProto(FixedWindows.of(WINDOW_SIZE));
+ RunnerApi.WindowingStrategy strategy = nonGlobalStrategy(proto);
+
+ // The translator's own path: rebuild the strategy from the proto alone.
+ WindowingStrategy<?, ?> hydrated =
+ WindowingStrategyTranslation.fromProto(
+ strategy,
RehydratedComponents.forComponents(proto.getComponents()));
+
+ assertThat(hydrated.getWindowFn(), instanceOf(FixedWindows.class));
+ assertThat(((FixedWindows) hydrated.getWindowFn()).getSize(),
is(WINDOW_SIZE));
+ // The window coder the runner encodes state and timers with comes from
this WindowFn, so it is
+ // the standard interval-window coder rather than anything SDK-specific.
+ assertThat(hydrated.getWindowFn().windowCoder(),
is(IntervalWindow.getCoder()));
+ }
+}
