[GSoC 2026] Kafka Streams runner: bound a bundle by element count

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

@@ -49,7 +49,12 @@ public interface KafkaStreamsPipelineOptions extends PortablePipelineOptions {
 
   void setMaxBundleSize(int maxBundleSize);
 
-  @Description("Soft cap on bundle wall-clock duration in milliseconds.")
+  @Description(
+      "Intended cap on how long a bundle may stay open, in milliseconds. NOT APPLIED YET: closing a"
+          + " bundle from a wall-clock punctuator made a pipeline with two chained GroupByKeys"
+          + " across several partitions emit its groups repeatedly against a real broker, so only"
+          + " the element-count bound is enforced for now. See"
+          + " https://github.com/apache/beam/issues/18479.")
   @Default.Integer(1000)
   int getMaxBundleTimeMs();
 

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

@@ -65,6 +65,18 @@
  * across the upstream transform's partitions actually advances. Until every partition has reported,
  * the watermark is held and nothing is forwarded — but data is still processed in the meantime.
  *
+ * <p>A bundle is also bounded in size, by {@code --maxBundleSize}, and closed once that many
+ * elements have been fed to it. Without the bound a bundle stays open until the next watermark,
+ * which on a stream that produces steadily lets it grow without limit. The bound is checked as
+ * elements arrive. A time bound ({@code --maxBundleTimeMs}) is not applied yet — see the option's
+ * own documentation.
+ *
+ * <p>Closing a bundle asks Kafka Streams to commit, so the elements a bundle consumed and the
+ * records it produced are committed together and a restart replays either all of the bundle or none
+ * of it. Note that this aligns commits <em>to</em> bundle boundaries but does not stop Kafka
+ * Streams from committing on its own interval part-way through a bundle; closing the bundle first
+ * from a pre-commit hook would be needed to rule that out entirely.
+ *
  * <p>This is the Kafka Streams analogue of Flink's {@code ExecutableStageDoFnOperator} and Spark's
  * {@code SparkExecutableStageFunction}. State, timers, and side inputs are out of scope for this
  * first version: the stage is executed with {@link StateRequestHandler#unsupported()} and no timer
@@ -107,6 +119,15 @@ class ExecutableStageProcessor
   private @Nullable StageBundleFactory stageBundleFactory;
   private @Nullable RemoteBundle currentBundle;
 
+  /** Bound on how many elements may be fed to one bundle. */
+  private final int maxBundleSize;
+
+  /** Elements fed to the open bundle, for the size bound above. */
+  private int elementsInBundle;
+
+  /** The key of the last record processed, used when a close has no record to take one from. */
+  private byte[] lastKey = new byte[0];
+
   /**
    * @param transformId this stage's own transform id, stamped on the watermarks it emits
    * @param upstreamTransformIds the transform ids feeding this stage (known from the pipeline
@@ -120,13 +141,15 @@ class ExecutableStageProcessor
       String transformId,
       Set<String> upstreamTransformIds,
       MetricsContainerImpl metricsContainer,
-      Map<String, String> outputChildByPCollectionId) {
+      Map<String, String> outputChildByPCollectionId,
+      int maxBundleSize) {
     this.stagePayload = stagePayload;
     this.jobInfo = jobInfo;
     this.transformId = transformId;
     this.watermarkAggregator = new WatermarkAggregator(upstreamTransformIds);
     this.metricsContainer = metricsContainer;
     this.outputChildByPCollectionId = ImmutableMap.copyOf(outputChildByPCollectionId);
+    this.maxBundleSize = maxBundleSize;
   }
 
   /** A harness output element together with the id of the output PCollection it belongs to. */
@@ -182,12 +205,20 @@ public void process(Record<byte[], KStreamsPayload<?>> record) {
       }
       return;
     }
+    byte[] key = record.key();
+    if (key != null) {
+      lastKey = key;
+    }
     try {
       ensureBundleOpen();
       mainInputReceiver().accept(payload.getData());
+      elementsInBundle++;
     } catch (Exception e) {
       throw new RuntimeException("Failed to process element through SDK harness", e);
     }
+    if (elementsInBundle >= maxBundleSize) {
+      closeBundleAndFlush(record);
+    }
   }
 
   private void ensureBundleOpen() throws Exception {
@@ -241,6 +272,7 @@ public void onCompleted(ProcessBundleResponse response) {
     currentBundle =
         factory.getBundle(
             outputReceiverFactory, StateRequestHandler.unsupported(), progressHandler);
+    elementsInBundle = 0;
   }
 
   private FnDataReceiver<WindowedValue<?>> mainInputReceiver() {
@@ -253,6 +285,18 @@ private FnDataReceiver<WindowedValue<?>> mainInputReceiver() {
   }
 
   private void closeBundleAndFlush(Record<byte[], KStreamsPayload<?>> record) {
+    byte[] key = record.key();
+    closeBundleAndFlush(key == null ? lastKey : key, record.timestamp());
+  }
+
+  /**
+   * Finishes the open bundle, forwards everything it produced, and asks Kafka Streams to commit.
+   *
+   * <p>The commit request is what ties a bundle to a transaction: the elements the bundle consumed
+   * and the records it produced are then committed together, so a restart either replays the whole
+   * bundle or none of it.
+   */
+  private void closeBundleAndFlush(byte[] key, long timestamp) {
     RemoteBundle bundle = currentBundle;
     if (bundle == null) {
       return;
@@ -265,6 +309,7 @@ private void closeBundleAndFlush(Record<byte[], KStreamsPayload<?>> record) {
       throw new RuntimeException("Failed to close SDK harness bundle", e);
     } finally {
       currentBundle = null;
+      elementsInBundle = 0;
     }
     ProcessorContext<byte[], KStreamsPayload<?>> ctx = checkInitialized(context);
     // The harness has finished the bundle (close() returned) so no further enqueues happen.
@@ -275,14 +320,15 @@ private void closeBundleAndFlush(Record<byte[], KStreamsPayload<?>> record) {
     while ((output = pendingOutputs.poll()) != null) {
       Record<byte[], KStreamsPayload<?>> outputRecord =
           new Record<byte[], KStreamsPayload<?>>(
-              record.key(), KStreamsPayload.data(output.value), record.timestamp());
+              key, KStreamsPayload.data(output.value), timestamp);
       String childNode = outputChildByPCollectionId.get(output.pCollectionId);
       if (childNode == null) {
         ctx.forward(outputRecord);
       } else {
         ctx.forward(outputRecord, childNode);
       }
     }
+    ctx.commit();
   }
 
   private void forwardWatermark(Record<byte[], KStreamsPayload<?>> record, long watermarkMillis) {

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

@@ -108,7 +108,8 @@ public void translate(
                 transformId,
                 ImmutableSet.of(parentProcessor),
                 context.getMetricsContainerStepMap().getContainer(transformId),
-                outputChildByPCollectionId),
+                outputChildByPCollectionId,
+                context.getPipelineOptions().getMaxBundleSize()),
         parentProcessor);
 
     if (multiOutput) {

runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/BundleBoundaryTest.java

@@ -0,0 +1,118 @@
+/*
+ * 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 static org.hamcrest.Matchers.greaterThanOrEqualTo;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions;
+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.Before;
+import org.junit.Test;
+
+/**
+ * Tests that a bundle is closed once it reaches {@code --maxBundleSize}, rather than staying open
+ * until the next watermark.
+ *
+ * <p>The bound is observable through the DoFn's own lifecycle: {@code @FinishBundle} runs once per
+ * bundle the SDK harness processes, so feeding a known number of elements with a known bound tells
+ * us how many bundles the stage actually opened.
+ *
+ * <p>The elements have to reach the stage as separate records for the bound to see them, since it
+ * counts what is fed to the stage rather than what the user's code emits inside it. A DoFn that
+ * fans one element out into many would be fused into the same stage and still be a single input, so
+ * these pipelines read the elements from a {@link Create} instead — the runner translates that to a
+ * primitive Read, which forwards one record per element.
+ */
+public class BundleBoundaryTest {
+
+  private static final int ELEMENTS = 50;
+  private static final int MAX_BUNDLE_SIZE = 10;
+
+  /** Records how many times {@code @FinishBundle} fired, i.e. how many bundles were processed. */
+  private static final List<String> FINISHED_BUNDLES =
+      Collections.synchronizedList(new ArrayList<>());
+
+  @Before
+  public void resetCounters() {
+    FINISHED_BUNDLES.clear();
+  }
+
+  /** Counts the bundles it is asked to process. */
+  private static class CountBundlesFn extends DoFn<Integer, Integer> {
+    @ProcessElement
+    public void processElement(@Element Integer element, OutputReceiver<Integer> out) {
+      out.output(element);
+    }
+
+    @FinishBundle
+    public void finishBundle() {
+      FINISHED_BUNDLES.add("bundle");
+    }
+  }
+
+  private static List<Integer> elements() {
+    List<Integer> elements = new ArrayList<>();
+    for (int i = 0; i < ELEMENTS; i++) {
+      elements.add(i);
+    }
+    return elements;
+  }
+
+  private static Pipeline buildPipeline(KafkaStreamsPipelineOptions options) {
+    Pipeline pipeline = Pipeline.create(options);
+    pipeline
+        .apply("read", Create.of(elements()))
+        .apply("countBundles", ParDo.of(new CountBundlesFn()));
+    return pipeline;
+  }
+
+  private static Pipeline pipelineWithBundleSize(int maxBundleSize) {
+    KafkaStreamsPipelineOptions options =
+        KafkaStreamsTestRunner.testOptions().as(KafkaStreamsPipelineOptions.class);
+    options.setMaxBundleSize(maxBundleSize);
+    return buildPipeline(options);
+  }
+
+  @Test
+  public void aBundleIsClosedOnceItReachesTheSizeBound() {
+    KafkaStreamsTestRunner.run(pipelineWithBundleSize(MAX_BUNDLE_SIZE));
+
+    // 50 elements bounded at 10 cannot have gone through in fewer than 5 bundles. Without the
+    // bound the whole run is one bundle, so this is what tells the two apart. The count is a lower
+    // bound rather than exact: a watermark arriving mid-bundle also closes one.
+    assertThat(FINISHED_BUNDLES.size(), is(greaterThanOrEqualTo(ELEMENTS / MAX_BUNDLE_SIZE)));
+  }
+
+  @Test
+  public void aBoundLargerThanTheInputLeavesASingleBundle() {
+    // The control: with a bound nothing reaches, the stage keeps one bundle open until the
+    // terminal watermark closes it.
+    KafkaStreamsTestRunner.run(pipelineWithBundleSize(ELEMENTS * 10));
+
+    assertThat(FINISHED_BUNDLES.size(), is(1));
+  }
+}

runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ExecutableStageProcessorWatermarkTest.java

@@ -62,7 +62,9 @@ private static ExecutableStageProcessor newProcessor() {
         ImmutableSet.of(UPSTREAM_ID),
         new MetricsContainerImpl(STAGE_ID),
         // Single-output: no per-output routing (this test drives the watermark path directly).
-        ImmutableMap.of());
+        ImmutableMap.of(),
+        // The bundle size bound is irrelevant to the watermark path this test drives.
+        1000);
   }
 
   /** A report from the upstream transform's given partition. */

runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/MetricsAcrossBundlesTest.java

@@ -0,0 +1,81 @@
+/*
+ * 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 org.apache.beam.runners.kafka.streams.KafkaStreamsPipelineOptions;
+import org.apache.beam.runners.kafka.streams.KafkaStreamsTestRunner;
+import org.apache.beam.sdk.Pipeline;
+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.MetricResults;
+import org.apache.beam.sdk.metrics.Metrics;
+import org.apache.beam.sdk.metrics.MetricsFilter;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
+import org.junit.Test;
+
+/**
+ * Checks that a user counter stays correct when the bundle size bound makes a stage run many small
+ * bundles instead of one large one.
+ *
+ * <p>The runner folds the metrics the SDK harness reports into the job's step map as each bundle
+ * completes, and those updates add rather than replace. That is right only if each report covers
+ * its own bundle, so splitting the same input across more bundles must not change the total.
+ */
+public class MetricsAcrossBundlesTest {
+  private static class CountingFn extends DoFn<Integer, Integer> {
+    private final Counter counter = Metrics.counter("probe", "elements");
+
+    @ProcessElement
+    public void processElement(@Element Integer in, OutputReceiver<Integer> out) {
+      counter.inc();
+      out.output(in);
+    }
+  }
+
+  private static long run(int maxBundleSize) {
+    KafkaStreamsPipelineOptions options =
+        KafkaStreamsTestRunner.testOptions().as(KafkaStreamsPipelineOptions.class);
+    options.setMaxBundleSize(maxBundleSize);
+    Pipeline p = Pipeline.create(options);
+    p.apply(Create.of(1, 2, 3, 4, 5, 6)).apply(ParDo.of(new CountingFn()));
+    MetricResults metrics = KafkaStreamsTestRunner.run(p);
+    MetricQueryResults q =
+        metrics.queryMetrics(
+            MetricsFilter.builder()
+                .addNameFilter(MetricNameFilter.named("probe", "elements"))
+                .build());
+    return Iterables.getOnlyElement(q.getCounters()).getAttempted();
+  }
+
+  @Test
+  public void oneBundle() {
+    assertThat(run(1000), is(6L));
+  }
+
+  @Test
+  public void manyBundles() {
+    assertThat(run(1), is(6L));
+  }
+}