From 0f8f1378dca97ce4bbe7948220ad27494414887f Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Wed, 29 Jul 2026 14:47:41 +0500 Subject: [PATCH 1/3] [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. KafkaStreamsTopicManager creates the topics the runner names for itself - a bootstrap topic per Impulse and primitive Read, and a repartition topic per GroupByKey - with an AdminClient before the application starts. Kafka Streams treats them as user topics because they are declared with explicit names, so it does not create them and refuses to start when a source topic is missing; the broker's auto-creation is not a substitute, since it is often disabled and gives the topic the broker's default partition count. Only topics carrying one of the runner's prefixes are created, so a topic the user named is never created implicitly. Adds topicPartitions and topicReplicationFactor options. Fixes two bugs the test driver could not catch. Validation demanded jobEndpoint, which KafkaStreamsPipelineOptions inherits as required from PortablePipelineOptions but which is a client-side option: this code runs on the job server for an already-submitted pipeline, so it rejected every valid job. Only the options the runner reads are checked now, as Flink's equivalent PortablePipelineRunner does. Separately, the result object registers a KafkaStreams state listener in its constructor but was built after start(), and Kafka Streams only accepts a listener in the CREATED state, so every real startup threw and the listener behind waitUntilFinish and cancel was never installed; the result is now built before the application starts. The integration test needs Docker, so the default test task excludes it and a brokerIntegrationTest task runs it. Also adds PortableWindowingStrategyTest, checking that the windowing the runner reconstructs comes from the standard beam:window_fn URNs every SDK emits rather than anything Java-specific. --- runners/kafka-streams/build.gradle | 23 ++ .../streams/KafkaStreamsPipelineOptions.java | 15 ++ .../streams/KafkaStreamsPipelineRunner.java | 30 ++- .../KafkaStreamsPortablePipelineResult.java | 4 + .../streams/KafkaStreamsTopicManager.java | 142 +++++++++++++ .../streams/KafkaStreamsRunnerBrokerIT.java | 196 ++++++++++++++++++ .../PortableWindowingStrategyTest.java | 155 ++++++++++++++ 7 files changed, 559 insertions(+), 6 deletions(-) create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTopicManager.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/PortableWindowingStrategyTest.java diff --git a/runners/kafka-streams/build.gradle b/runners/kafka-streams/build.gradle index bdf7e3be0585..d1326c079e5d 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 2fa992e66e7e..eb43351a3422 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,21 @@ public interface KafkaStreamsPipelineOptions extends PortablePipelineOptions { void setMaxBundleTimeMs(int maxBundleTimeMs); + @Description( + "Number of partitions for the topics the runner creates for a pipeline (the bootstrap topic" + + " of each Impulse and Read, and the repartition topic of each GroupByKey). This is the" + + " parallelism the shuffled parts of the pipeline can reach.") + @Default.Integer(1) + int getTopicPartitions(); + + void setTopicPartitions(int topicPartitions); + + @Description("Replication factor for the 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 cdb59f67cce9..189d448ec185 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.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,14 @@ 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()); KafkaStreamsPipelineTranslator translator = new KafkaStreamsPipelineTranslator(); KafkaStreamsTranslationContext context = @@ -55,15 +60,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) { 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 972afa15c358..0d508b8189fd 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 000000000000..8c37accb6330 --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/KafkaStreamsTopicManager.java @@ -0,0 +1,142 @@ +/* + * 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. + * + *

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. + * + *

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 topics the runner owns and may create. */ + private static final List RUNNER_TOPIC_PREFIXES = + java.util.Arrays.asList("__beam_impulse_", "__beam_read_", "__beam_gbk_"); + + private KafkaStreamsTopicManager() {} + + /** + * Creates any runner-owned topic in {@code topology} that does not exist yet. + * + *

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 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 existing = admin.listTopics().names().get(); + List toCreate = new ArrayList<>(); + for (String topic : runnerTopics) { + if (!existing.contains(topic)) { + toCreate.add( + new NewTopic( + topic, options.getTopicPartitions(), 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 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 runnerOwnedTopics(Topology topology) { + Set topics = new HashSet<>(); + for (TopologyDescription.Subtopology subtopology : topology.describe().subtopologies()) { + for (TopologyDescription.Node node : subtopology.nodes()) { + if (node instanceof TopologyDescription.Source) { + Set 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; + } + + private static boolean isRunnerOwned(String topic) { + for (String prefix : RUNNER_TOPIC_PREFIXES) { + if (topic.startsWith(prefix)) { + return true; + } + } + return false; + } +} 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 000000000000..0bdff36a200e --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/KafkaStreamsRunnerBrokerIT.java @@ -0,0 +1,196 @@ +/* + * 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.containers.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. + * + *

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. + * + *

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() { + kafka = new KafkaContainer(DockerImageName.parse("confluentinc/cp-kafka:7.6.1")); + 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> { + @ProcessElement + public void processElement(OutputReceiver> 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>, Void> { + private final Counter groups = Metrics.counter(NAMESPACE, GROUPS_COUNTER); + + @ProcessElement + public void processElement() { + groups.inc(); + } + } + + private KafkaStreamsPipelineOptions options() { + KafkaStreamsPipelineOptions options = + PipelineOptionsFactory.create().as(KafkaStreamsPipelineOptions.class); + options.setRunner(CrashingRunner.class); + options.setBootstrapServers(kafka.getBootstrapServers()); + options.setApplicationId("ks-broker-it-" + UUID.randomUUID()); + 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(); + } + } + + /** + * 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 counter = Iterables.getOnlyElement(query.getCounters()); + return counter.getAttempted(); + } +} diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/PortableWindowingStrategyTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/PortableWindowingStrategyTest.java new file mode 100644 index 000000000000..724da90a4461 --- /dev/null +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/PortableWindowingStrategyTest.java @@ -0,0 +1,155 @@ +/* + * 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.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's windowing is driven by the language-neutral windowing strategy in the + * pipeline proto, not by anything specific to the Java SDK. + * + *

This matters because the runner executes GroupAlsoByWindow itself (see {@link + * WindowedGroupByKeyProcessor}), so it has to reconstruct the WindowFn from the proto. Beam encodes + * the standard WindowFns as a URN plus a parameter payload — {@code + * beam:window_fn:fixed_windows:v1} and friends — and only falls back to a serialized Java object + * for a custom WindowFn. An SDK in another language emits exactly those same standard URNs, so if + * the runner works from the URN form it works for any SDK, and if it had come to depend on the + * Java-serialized form it would only ever have worked for Java. + * + *

These tests assert the proto a windowed pipeline produces carries the standard URN, and that + * hydrating it back — which is what the translator does — yields the right WindowFn. They do not + * replace running a pipeline from another SDK end to end, which additionally exercises the coders + * and the SDK harness; that needs the job server against a real broker. + */ +public class PortableWindowingStrategyTest { + + private static final Duration WINDOW_SIZE = Duration.millis(10); + + private static class EmitKvFn extends DoFn> { + @ProcessElement + public void processElement(OutputReceiver> 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 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 fixedWindowsTravelAsTheStandardUrnNotSerializedJava() { + RunnerApi.WindowingStrategy strategy = + nonGlobalStrategy(windowedPipelineProto(FixedWindows.of(WINDOW_SIZE))); + + // The same bytes any SDK would send for FixedWindows. + assertThat(strategy.getWindowFn().getUrn(), is(WindowingStrategyTranslation.FIXED_WINDOWS_URN)); + assertThat( + strategy.getWindowFn().getUrn(), + is(not(WindowingStrategyTranslation.SERIALIZED_JAVA_WINDOWFN_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 noWindowingStrategyInAWindowedPipelineNeedsJavaSerialization() { + RunnerApi.Pipeline proto = windowedPipelineProto(FixedWindows.of(WINDOW_SIZE)); + + // If any strategy needed the Java-serialized form, the runner could not reconstruct it for a + // pipeline submitted from another SDK. + 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(org.apache.beam.sdk.transforms.windowing.IntervalWindow.getCoder())); + } +} From 5a9c0e48f7f367de221c96aa486e692d5e4c4264 Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Fri, 31 Jul 2026 02:10:12 +0500 Subject: [PATCH 2/3] Make the watermark correct when a pipeline runs across partitions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every transform stamped its watermark reports as partition 0 of 1. That is right for a report handed to a fused child, which runs in the same task and sees exactly one instance of its parent, but wrong for a report crossing a repartition topic: the sink's partitioner broadcasts each report to every partition, so a downstream task sees one from every instance of the upstream transform. With all of them claiming to be the only partition it would treat the first report as if every instance had reported, advance its watermark, and fire before the other partitions had delivered their data. Identity is now attached where the report stops being delivered in process and starts crossing a topic: ShuffleByKeyProcessor restamps it. That processor runs in the upstream transform's own task, so taskId().partition() is exactly the instance the report is for, and the transform id is left alone so a downstream aggregator still matches the report to the upstream it expects. Transforms go on reporting a single source when forwarding in process. The partition count a transform runs at is threaded through translation: one everywhere, raised only by a GroupByKey's shuffle, and inherited by everything fused downstream of it. Also pins the bootstrap topics to a single partition. An Impulse or a Read emits once per task, gated by a state store that is itself per task, so a multi-partition bootstrap topic would fire the same Impulse once per partition; only the repartition topic of a GroupByKey carries the pipeline's parallelism. Adds a broker test running two chained GroupByKeys across four partitions, which is the shape this affects — the second GroupByKey aggregates the reports of every task of the first, and a single-GroupByKey pipeline passes either way — plus a one-partition control and unit tests for the restamping. --- .../streams/KafkaStreamsTopicManager.java | 37 +++++- .../translation/ExecutableStageProcessor.java | 8 +- .../ExecutableStageTranslator.java | 4 + .../streams/translation/FlattenProcessor.java | 3 +- .../translation/FlattenTranslator.java | 5 + .../translation/GroupByKeyTranslator.java | 12 +- .../KafkaStreamsTranslationContext.java | 26 ++++ .../translation/RedistributeTranslator.java | 3 + .../translation/ShuffleByKeyProcessor.java | 37 +++++- .../WindowedGroupByKeyProcessor.java | 3 +- .../streams/KafkaStreamsRunnerBrokerIT.java | 86 ++++++++++++ .../ShuffleByKeyProcessorTest.java | 123 ++++++++++++++++++ 12 files changed, 328 insertions(+), 19 deletions(-) create mode 100644 runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/ShuffleByKeyProcessorTest.java 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 index 8c37accb6330..a2e6d2b7abd0 100644 --- 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 @@ -53,9 +53,23 @@ class KafkaStreamsTopicManager { private static final Logger LOG = LoggerFactory.getLogger(KafkaStreamsTopicManager.class); - /** Prefixes of the topics the runner owns and may create. */ - private static final List RUNNER_TOPIC_PREFIXES = - java.util.Arrays.asList("__beam_impulse_", "__beam_read_", "__beam_gbk_"); + /** + * Prefixes of the bootstrap topics, which must have exactly one partition. + * + *

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 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 PARTITIONED_TOPIC_PREFIXES = + java.util.Arrays.asList("__beam_gbk_"); private KafkaStreamsTopicManager() {} @@ -80,7 +94,7 @@ static void createMissingTopics(Topology topology, KafkaStreamsPipelineOptions o if (!existing.contains(topic)) { toCreate.add( new NewTopic( - topic, options.getTopicPartitions(), options.getTopicReplicationFactor())); + topic, partitionsFor(topic, options), options.getTopicReplicationFactor())); } } if (toCreate.isEmpty()) { @@ -131,8 +145,21 @@ private static Set runnerOwnedTopics(Topology topology) { 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.getTopicPartitions(); + } + private static boolean isRunnerOwned(String topic) { - for (String prefix : RUNNER_TOPIC_PREFIXES) { + return hasAnyPrefix(topic, SINGLE_PARTITION_TOPIC_PREFIXES) + || hasAnyPrefix(topic, PARTITIONED_TOPIC_PREFIXES); + } + + private static boolean hasAnyPrefix(String topic, List prefixes) { + for (String prefix : prefixes) { if (topic.startsWith(prefix)) { return true; } 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 15265073dd21..936107d1c231 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,10 @@ private void closeBundleAndFlush(Record> record) { } private void forwardWatermark(Record> 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). + // Stamped 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 restamps the report with the real partition + // identity, because the broadcast then delivers every instance's report to every consumer. ProcessorContext> ctx = checkInitialized(context); ctx.forward( new Record>( 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 71e24f32c1f5..caefa6534fad 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 @@ 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); } } } 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 d09fed8185a1..9af705c5af97 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 @@ public void process(Record> 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). + // Stamped as the only source a consumer will see; a shuffle downstream restamps. ctx.forward( new Record>( 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 3ac4c1304daa..a5c8ce05baeb 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 @@ public void translate( Set seenInputs = new HashSet<>(); List parentProcessors = new ArrayList<>(); Set 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); } } 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 c5327e28e069..75785cc25553 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 @@ public void translate( 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().getTopicPartitions(); String shuffleName = transformId + SHUFFLE_SUFFIX; String sinkName = transformId + SINK_SUFFIX; @@ -110,7 +113,13 @@ public void translate( 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 restamps 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 @@ public void translate( 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 ec1b3f26aded..34daf41b591a 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 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 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,26 @@ public void registerPCollectionProducer(String pCollectionId, String processorNa } } + /** + * Records how many partitions the transform producing {@code pCollectionId} runs across. + * + *

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. + */ + 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 72c47db8d4b1..a44740887b3e 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 @@ public void translate( // 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 774d36db0f2e..a07e7856cc85 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 @@ *

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. + * + *

Watermark reports are restamped 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 every 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> { @@ -44,15 +51,25 @@ class ShuffleByKeyProcessor private static final Logger LOG = LoggerFactory.getLogger(ShuffleByKeyProcessor.class); private final Coder 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> context; - ShuffleByKeyProcessor(Coder keyCoder) { + ShuffleByKeyProcessor(Coder keyCoder, int upstreamPartitionCount) { this.keyCoder = keyCoder; + this.upstreamPartitionCount = upstreamPartitionCount; } @Override public void init(ProcessorContext> 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 @@ public void process(Record> record) { } 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>( + 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 1b0ffae23fde..f00abdcb635e 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 @@ private void forwardData( private void forwardWatermark(Record> trigger, long watermarkMillis) { ProcessorContext> ctx = checkInitialized(context); - // Stamped with this transform's own id; GroupByKey is a single instance for now (0 of 1). + // Stamped as the only source a consumer will see; see ExecutableStageProcessor for why an + // in-process edge reports a single source and a shuffle restamps. ctx.forward( new Record>( 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 index 0bdff36a200e..9e686d8d07e9 100644 --- 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 @@ -115,11 +115,16 @@ public void processElement() { } 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.setTopicPartitions(topicPartitions); options .as(PortablePipelineOptions.class) .setDefaultEnvironmentType(Environments.ENVIRONMENT_EMBEDDED); @@ -163,6 +168,87 @@ public void groupByKeyRunsThroughARealBrokerAndReportsMetrics() throws Exception } } + /** + * Collapses every group onto one key, so the next GroupByKey has to shuffle across partitions. + */ + private static class ToSingleKeyFn + extends DoFn>, KV> { + @ProcessElement + public void processElement( + @Element KV> group, OutputReceiver> 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. */ 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 000000000000..38669138075c --- /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. + * + *

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) + (org.apache.beam.sdk.coders.Coder) StringUtf8Coder.of(), + upstreamPartitions); + MockProcessorContext> ctx = + new MockProcessorContext<>(new Properties(), new TaskId(0, taskPartition), null); + processor.init(ctx); + lastContext = ctx; + return processor; + } + + private static MockProcessorContext> lastContext; + + private static Record> 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)); + } +} From 7294bff8ba1ce4dba2a8afa8959a7c603ddeafa6 Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Fri, 31 Jul 2026 14:11:29 +0500 Subject: [PATCH 3/3] Make the watermark correct when a pipeline runs across partitions Every transform reported its watermark as partition 0 of 1. That is right for a report handed to a fused child, which runs in the same task and sees exactly one instance of its parent, but wrong for a report crossing a repartition topic: the sink's partitioner broadcasts each report to every partition, so a downstream task sees one from every instance of the upstream transform. With all of them claiming to be the only partition it would treat the first report as if every instance had reported, advance its watermark, and fire before the other partitions had delivered their data. The identity is now attached where the report stops being delivered in process and starts crossing a topic: ShuffleByKeyProcessor relabels it. That processor runs in the upstream transform's own task, so taskId().partition() is exactly the instance the report is for, and the transform id is left alone so a downstream aggregator still matches the report to the upstream it expects. The partition count a transform runs at is threaded through translation: one everywhere, raised only by a GroupByKey's shuffle, and inherited by everything fused downstream of it. Also pins the bootstrap topics to a single partition. An Impulse or a Read emits once per task, gated by a state store that is itself per task, so a multi-partition bootstrap topic would fire the same Impulse once per partition; only the repartition topic of a GroupByKey carries the pipeline's parallelism. Review changes: renames --topicPartitions to --internalParallelism, which says where it applies; rejects a value below one, since it is also the number of watermark reports a shuffle's consumer waits for and a non-positive value would leave it waiting rather than failing; drops the word stamp for label, which does not read like timestamp; and moves the broker test to the official apache/kafka image. Adds a broker test running two chained GroupByKeys across four partitions, which is the shape this affects, plus a one-partition control and unit tests for the relabelling. Renames PortableWindowingStrategyTest to StandardWindowFnTranslationTest and corrects what it claims: the standard WindowFns are interpreted runner-side by ReduceFnRunner, so the test pins down that they are rebuilt from the language-neutral URNs, and a WindowFn the user wrote themselves would instead have to run through the SDK harness, which the runner does not support. --- .../streams/KafkaStreamsPipelineOptions.java | 14 +++--- .../streams/KafkaStreamsPipelineRunner.java | 8 ++++ .../streams/KafkaStreamsTopicManager.java | 4 +- .../translation/ExecutableStageProcessor.java | 5 +- .../streams/translation/FlattenProcessor.java | 2 +- .../translation/GroupByKeyTranslator.java | 4 +- .../KafkaStreamsTranslationContext.java | 4 ++ .../translation/ShuffleByKeyProcessor.java | 2 +- .../WindowedGroupByKeyProcessor.java | 4 +- .../streams/KafkaStreamsRunnerBrokerIT.java | 10 ++-- ...a => StandardWindowFnTranslationTest.java} | 47 +++++++++---------- 11 files changed, 61 insertions(+), 43 deletions(-) rename runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/{PortableWindowingStrategyTest.java => StandardWindowFnTranslationTest.java} (74%) 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 eb43351a3422..a5a8bb9328b1 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 @@ -56,15 +56,17 @@ public interface KafkaStreamsPipelineOptions extends PortablePipelineOptions { void setMaxBundleTimeMs(int maxBundleTimeMs); @Description( - "Number of partitions for the topics the runner creates for a pipeline (the bootstrap topic" - + " of each Impulse and Read, and the repartition topic of each GroupByKey). This is the" - + " parallelism the shuffled parts of the pipeline can reach.") + "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 getTopicPartitions(); + int getInternalParallelism(); - void setTopicPartitions(int topicPartitions); + void setInternalParallelism(int internalParallelism); - @Description("Replication factor for the topics the runner creates for a pipeline.") + @Description("Replication factor for the internal topics the runner creates for a pipeline.") @Default.Short(1) short getTopicReplicationFactor(); 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 189d448ec185..96b4b2cd7f92 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 @@ -52,6 +52,14 @@ public PortablePipelineResult run(RunnerApi.Pipeline pipeline, JobInfo jobInfo) // 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 = 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 index a2e6d2b7abd0..0e5f46f53eac 100644 --- 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 @@ -150,7 +150,9 @@ private static Set runnerOwnedTopics(Topology topology) { * configured parallelism for a shuffle topic. */ private static int partitionsFor(String topic, KafkaStreamsPipelineOptions options) { - return hasAnyPrefix(topic, SINGLE_PARTITION_TOPIC_PREFIXES) ? 1 : options.getTopicPartitions(); + return hasAnyPrefix(topic, SINGLE_PARTITION_TOPIC_PREFIXES) + ? 1 + : options.getInternalParallelism(); } private static boolean isRunnerOwned(String topic) { 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 936107d1c231..63fefe3e15c4 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,9 +286,10 @@ private void closeBundleAndFlush(Record> record) { } private void forwardWatermark(Record> record, long watermarkMillis) { - // Stamped as the only source a consumer will see. Forwarding here is in-process, to the stage's + // 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 restamps the report with the real partition + // 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> ctx = checkInitialized(context); ctx.forward( 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 9af705c5af97..e37da6774489 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,7 +99,7 @@ public void process(Record> record) { Instant advanced = watermarkAggregator.advance(); if (advanced.isAfter(lastForwardedWatermark)) { lastForwardedWatermark = advanced; - // Stamped as the only source a consumer will see; a shuffle downstream restamps. + // Labelled as the only source a consumer will see; a shuffle downstream relabels it. ctx.forward( new Record>( record.key(), 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 75785cc25553..c460436eedf8 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 @@ -97,7 +97,7 @@ public void translate( 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().getTopicPartitions(); + int partitionCount = context.getPipelineOptions().getInternalParallelism(); String shuffleName = transformId + SHUFFLE_SUFFIX; String sinkName = transformId + SINK_SUFFIX; @@ -113,7 +113,7 @@ public void translate( Topology topology = context.getTopology(); // Re-key data records by the encoded Beam key; pass watermark reports through. - // The shuffle runs in the upstream transform's task, so it restamps each report with that + // 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( 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 34daf41b591a..d03169615665 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 @@ -134,6 +134,10 @@ public void registerPCollectionPartitionCount(String pCollectionId, int partitio /** * How many partitions the transform producing {@code pCollectionId} runs across; one unless a * shuffle upstream raised it. + * + *

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); 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 a07e7856cc85..79595cba7c96 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 @@ -36,7 +36,7 @@ * 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 restamped here with the reporting instance's real partition identity. + *

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 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 f00abdcb635e..d00b01a3753d 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,8 +284,8 @@ private void forwardData( private void forwardWatermark(Record> trigger, long watermarkMillis) { ProcessorContext> ctx = checkInitialized(context); - // Stamped as the only source a consumer will see; see ExecutableStageProcessor for why an - // in-process edge reports a single source and a shuffle restamps. + // 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>( 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 index 9e686d8d07e9..bb82a403e557 100644 --- 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 @@ -54,7 +54,7 @@ import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; -import org.testcontainers.containers.KafkaContainer; +import org.testcontainers.kafka.KafkaContainer; import org.testcontainers.utility.DockerImageName; /** @@ -83,7 +83,11 @@ public class KafkaStreamsRunnerBrokerIT { @BeforeClass public static void startBroker() { - kafka = new KafkaContainer(DockerImageName.parse("confluentinc/cp-kafka:7.6.1")); + // 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(); } @@ -124,7 +128,7 @@ private KafkaStreamsPipelineOptions options(int topicPartitions) { options.setRunner(CrashingRunner.class); options.setBootstrapServers(kafka.getBootstrapServers()); options.setApplicationId("ks-broker-it-" + UUID.randomUUID()); - options.setTopicPartitions(topicPartitions); + options.setInternalParallelism(topicPartitions); options .as(PortablePipelineOptions.class) .setDefaultEnvironmentType(Environments.ENVIRONMENT_EMBEDDED); diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/PortableWindowingStrategyTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StandardWindowFnTranslationTest.java similarity index 74% rename from runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/PortableWindowingStrategyTest.java rename to runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StandardWindowFnTranslationTest.java index 724da90a4461..912d51d6b5d1 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/PortableWindowingStrategyTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/StandardWindowFnTranslationTest.java @@ -33,6 +33,7 @@ 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; @@ -45,23 +46,24 @@ import org.junit.Test; /** - * Checks that the runner's windowing is driven by the language-neutral windowing strategy in the - * pipeline proto, not by anything specific to the Java SDK. + * Checks that the runner reconstructs the standard WindowFns from the language-neutral windowing + * strategy in the pipeline proto. * - *

This matters because the runner executes GroupAlsoByWindow itself (see {@link - * WindowedGroupByKeyProcessor}), so it has to reconstruct the WindowFn from the proto. Beam encodes - * the standard WindowFns as a URN plus a parameter payload — {@code - * beam:window_fn:fixed_windows:v1} and friends — and only falls back to a serialized Java object - * for a custom WindowFn. An SDK in another language emits exactly those same standard URNs, so if - * the runner works from the URN form it works for any SDK, and if it had come to depend on the - * Java-serialized form it would only ever have worked for Java. + *

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. * - *

These tests assert the proto a windowed pipeline produces carries the standard URN, and that - * hydrating it back — which is what the translator does — yields the right WindowFn. They do not - * replace running a pipeline from another SDK end to end, which additionally exercises the coders - * and the SDK harness; that needs the job server against a real broker. + *

What this does not 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 PortableWindowingStrategyTest { +public class StandardWindowFnTranslationTest { private static final Duration WINDOW_SIZE = Duration.millis(10); @@ -99,15 +101,12 @@ private static RunnerApi.WindowingStrategy nonGlobalStrategy(RunnerApi.Pipeline } @Test - public void fixedWindowsTravelAsTheStandardUrnNotSerializedJava() { + public void fixedWindowsTravelAsTheStandardUrn() { RunnerApi.WindowingStrategy strategy = nonGlobalStrategy(windowedPipelineProto(FixedWindows.of(WINDOW_SIZE))); - // The same bytes any SDK would send for FixedWindows. + // The same URN and payload any SDK emits for fixed windows. assertThat(strategy.getWindowFn().getUrn(), is(WindowingStrategyTranslation.FIXED_WINDOWS_URN)); - assertThat( - strategy.getWindowFn().getUrn(), - is(not(WindowingStrategyTranslation.SERIALIZED_JAVA_WINDOWFN_URN))); } @Test @@ -121,11 +120,11 @@ public void slidingWindowsTravelAsTheStandardUrn() { } @Test - public void noWindowingStrategyInAWindowedPipelineNeedsJavaSerialization() { + public void aStandardWindowFnNeedsNoJavaSerialization() { RunnerApi.Pipeline proto = windowedPipelineProto(FixedWindows.of(WINDOW_SIZE)); - // If any strategy needed the Java-serialized form, the runner could not reconstruct it for a - // pipeline submitted from another SDK. + // 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( @@ -148,8 +147,6 @@ public void hydratingTheStandardUrnRebuildsTheWindowFnTheRunnerRunsWith() throws 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(org.apache.beam.sdk.transforms.windowing.IntervalWindow.getCoder())); + assertThat(hydrated.getWindowFn().windowCoder(), is(IntervalWindow.getCoder())); } }