[GSoC 2026] Kafka Streams runner: run on a real broker, correctly across partitions

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

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();

runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsPipelineRunner.java

@@ -24,10 +24,10 @@
 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 KafkaStreamsPipelineRunner(KafkaStreamsPipelineOptions pipelineOptions) {
 
   @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 PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo)
     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) {

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;

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;
+  }
+}

runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessor.java

@@ -286,10 +286,11 @@ private void closeBundleAndFlush(Record<byte[], KStreamsPayload<?>> record) {
   }
 
   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<?>>(

runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageTranslator.java

@@ -74,6 +74,8 @@ public void translate(
     // 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 @@ public void translate(
             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);
     }
   }
 }

runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenProcessor.java

@@ -99,8 +99,7 @@ public void process(Record<byte[], KStreamsPayload<?>> record) {
     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(),

runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlattenTranslator.java

@@ -57,6 +57,9 @@ public void translate(
     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 @@ public void translate(
       String parentProcessor = context.getProcessorNameForPCollection(inputPCollectionId);
       parentProcessors.add(parentProcessor);
       upstreamTransformIds.add(parentProcessor);
+      partitionCount = Math.max(partitionCount, context.getPartitionCount(inputPCollectionId));
     }
 
     topology.addProcessor(
@@ -78,5 +82,6 @@ public void translate(
         parentProcessors.toArray(new String[0]));
 
     context.registerPCollectionProducer(outputPCollectionId, transformId);
+    context.registerPCollectionPartitionCount(outputPCollectionId, partitionCount);
   }
 }