From a0adae53810427b1f35b11233ddc869ac7fc6297 Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Wed, 9 Sep 2026 04:33:33 +0500 Subject: [PATCH 1/2] Kafka Streams runner: add a flush marker payload variant First step toward bundles bounded by time (#39633). A bundle cannot be closed from a punctuator, because transactions are committed by the Kafka Streams runtime in the background and are not exposed, so instead a source will emit a marker that travels the topology as an ordinary record and a stage closes its bundle from process(). This adds the marker itself and nothing that emits or consumes one yet. It carries the producing partition and that transform's partition count, which is what will let it be addressed to a slice of the downstream partitions rather than broadcast: broadcasting would deliver one flush per upstream partition, so a downstream partition would see N times more flushes than the configured interval. --- .../main/proto/kafka_streams_payload.proto | 25 +++++++- .../streams/translation/FlushPayload.java | 46 ++++++++++++++ .../streams/translation/KStreamsPayload.java | 60 +++++++++++++++++-- .../translation/KStreamsPayloadSerde.java | 15 ++++- .../translation/KStreamsPayloadSerdeTest.java | 38 ++++++++++++ 5 files changed, 175 insertions(+), 9 deletions(-) create mode 100644 runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlushPayload.java diff --git a/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto b/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto index 4dadf2d14b70..15e6f4333151 100644 --- a/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto +++ b/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto @@ -50,9 +50,32 @@ message KafkaStreamsPayload { bytes value = 1; } - // Exactly one variant is set; the oneof case discriminates data vs watermark. + // A request to close the open bundle and flush its output. + // + // A bundle has to be bounded in time as well as in size, or on a sparse stream the elements + // already fed to it wait for the next watermark. Closing it from a wall-clock punctuator does + // not work: transactions are committed by the Kafka Streams runtime in the background, are + // agnostic to punctuations, and are deliberately not exposed, so a bundle cannot be aligned with + // one. Instead a source emits this marker on its own punctuator and it travels the topology as + // an ordinary record, so a stage closes its bundle from process() rather than beside it. + // + // The partition fields are what let the marker be targeted rather than broadcast. Broadcasting + // would deliver one flush per upstream partition, so a downstream partition would see N times + // more flushes than the configured interval. Instead the producing partition addresses a slice + // of the downstream partitions, and the slices tile the whole range, so each downstream + // partition receives exactly one flush per interval. + message FlushPayload { + // Which partition (physical instance) of the producing transform emitted this marker, in + // [0, total_partitions). + uint32 source_partition = 1; + // How many partitions (physical instances) the producing transform has in total. + uint32 total_partitions = 2; + } + + // Exactly one variant is set; the oneof case discriminates data vs watermark vs flush. oneof payload { WatermarkPayload watermark = 1; DataPayload data = 2; + FlushPayload flush = 3; } } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlushPayload.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlushPayload.java new file mode 100644 index 000000000000..b593a0538b4e --- /dev/null +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlushPayload.java @@ -0,0 +1,46 @@ +/* + * 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; + +/** + * The flush-only view of a {@link KStreamsPayload}, obtained via {@link KStreamsPayload#asFlush()}. + * As with {@link WatermarkPayload}, the accessors live here so they are only reachable once the + * caller has checked the kind and narrowed the payload. + * + *

A flush marker asks the stage that receives it to close its open bundle and flush the output, + * which is how a bundle is bounded in time. It arrives as an ordinary record, so the bundle is + * closed from {@code process()} rather than from a punctuator: transactions are committed by the + * Kafka Streams runtime in the background and are not exposed, so a bundle cannot be aligned with + * one, and trying to do it from a punctuator duplicated output against a real broker + * (https://github.com/apache/beam/issues/39633). + * + *

The partition fields exist so the marker can be targeted rather than broadcast. Broadcasting + * would give a downstream partition one flush per upstream partition, so N times more flushes than + * the interval asks for. Instead the producing partition addresses a slice of the downstream + * partitions and the slices tile the range, so each downstream partition gets exactly one flush per + * interval. Unlike a watermark, a flush needs no aggregation on arrival: there is nothing to hold + * and nothing to combine, because only one arrives. + */ +public interface FlushPayload { + + /** Which partition of the producing transform emitted this marker. */ + int getSourcePartition(); + + /** How many partitions the producing transform has in total. */ + int getTotalSourcePartitions(); +} diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java index 647c64953b06..a495edd5e308 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java @@ -24,10 +24,10 @@ import org.checkerframework.checker.nullness.qual.Nullable; /** - * Envelope for every record value passed between the runner's processors. It is either a {@link - * #isData() data} element wrapping a {@link WindowedValue}, or a {@link #isWatermark() watermark} - * report carrying an event time plus the partition fields the downstream {@link WatermarkManager} - * needs. + * Envelope for every record value passed between the runner's processors. It is a {@link #isData() + * data} element wrapping a {@link WindowedValue}, a {@link #isWatermark() watermark} report + * carrying an event time plus the partition fields the downstream {@link WatermarkManager} needs, + * or a {@link #isFlush() flush} marker asking the receiving stage to close its bundle. * *

One channel therefore carries both Beam data and the watermark coordination Kafka Streams has * no notion of. Across topic boundaries it is encoded by {@link KStreamsPayloadSerde}. @@ -38,7 +38,8 @@ public final class KStreamsPayload { private enum Kind { DATA, - WATERMARK + WATERMARK, + FLUSH } private final Kind kind; @@ -92,6 +93,25 @@ public static KStreamsPayload watermark( Kind.WATERMARK, null, watermarkMillis, transformId, sourcePartition, totalSourcePartitions); } + /** + * Returns a flush marker: a request to close the open bundle and flush its output, carrying the + * partition of the producing transform that emitted it and how many partitions that transform + * has. Those two fields are what let the marker be addressed to a slice of the downstream + * partitions rather than broadcast to all of them; see {@link FlushPayload}. + */ + public static KStreamsPayload flush(int sourcePartition, int totalSourcePartitions) { + Preconditions.checkArgument( + totalSourcePartitions > 0, + "totalSourcePartitions must be positive: %s", + totalSourcePartitions); + Preconditions.checkArgument( + sourcePartition >= 0 && sourcePartition < totalSourcePartitions, + "sourcePartition %s out of range for totalSourcePartitions %s", + sourcePartition, + totalSourcePartitions); + return new KStreamsPayload<>(Kind.FLUSH, null, 0L, "", sourcePartition, totalSourcePartitions); + } + public boolean isData() { return kind == Kind.DATA; } @@ -100,6 +120,10 @@ public boolean isWatermark() { return kind == Kind.WATERMARK; } + public boolean isFlush() { + return kind == Kind.FLUSH; + } + /** * Returns the wrapped data element. Caller must check {@link #isData()} first; calling this on a * watermark payload throws. @@ -121,6 +145,28 @@ public WatermarkPayload asWatermark() { return new WatermarkView(); } + /** + * Narrows this payload to its {@link FlushPayload} view. Caller must check {@link #isFlush()} + * first; calling this on any other payload throws. + */ + public FlushPayload asFlush() { + Preconditions.checkState(isFlush(), "Payload is not a flush marker: kind=%s", kind); + return new FlushView(); + } + + /** {@link FlushPayload} view backed by this payload's fields. */ + private final class FlushView implements FlushPayload { + @Override + public int getSourcePartition() { + return sourcePartition; + } + + @Override + public int getTotalSourcePartitions() { + return totalSourcePartitions; + } + } + /** {@link WatermarkPayload} view backed by this payload's fields. */ private final class WatermarkView implements WatermarkPayload { @Override @@ -172,6 +218,10 @@ public String toString() { MoreObjects.ToStringHelper helper = MoreObjects.toStringHelper(this).add("kind", kind); if (kind == Kind.DATA) { helper.add("data", data); + } else if (kind == Kind.FLUSH) { + helper + .add("sourcePartition", sourcePartition) + .add("totalSourcePartitions", totalSourcePartitions); } else { helper .add("watermarkMillis", watermarkMillis) diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java index 1363740d58bd..a7680ffa8b49 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java @@ -36,9 +36,9 @@ * *

The wire form is the {@link KafkaStreamsPayload} protobuf message — protobuf gives compatible * schema evolution and compact varint encoding. The data variant carries the {@link WindowedValue} - * encoded with the {@link Coder} supplied for the topic's PCollection; the watermark variant - * carries the coder-independent watermark report. A {@link KStreamsPayloadSerde} is therefore - * parameterized by the data {@link Coder} (different topics carry different element types). + * encoded with the {@link Coder} supplied for the topic's PCollection; the watermark and flush + * variants are coder-independent. A {@link KStreamsPayloadSerde} is therefore parameterized by the + * data {@link Coder} (different topics carry different element types). * *

The serde assumes non-null payloads: the topics it is used on (repartition and watermark * fan-out) are not log-compacted, so no tombstone (null-valued) records occur. @@ -77,6 +77,12 @@ public byte[] serialize(String topic, KStreamsPayload payload) { proto.setData( KafkaStreamsPayload.DataPayload.newBuilder() .setValue(ByteString.copyFrom(encoded.toByteArray()))); + } else if (payload.isFlush()) { + FlushPayload flush = payload.asFlush(); + proto.setFlush( + KafkaStreamsPayload.FlushPayload.newBuilder() + .setSourcePartition(flush.getSourcePartition()) + .setTotalPartitions(flush.getTotalSourcePartitions())); } else { WatermarkPayload watermark = payload.asWatermark(); proto.setWatermark( @@ -113,6 +119,9 @@ public KStreamsPayload deserialize(String topic, byte[] bytes) { watermark.getTransformId(), watermark.getSourcePartition(), watermark.getTotalPartitions()); + case FLUSH: + KafkaStreamsPayload.FlushPayload flush = proto.getFlush(); + return KStreamsPayload.flush(flush.getSourcePartition(), flush.getTotalPartitions()); case PAYLOAD_NOT_SET: default: throw new SerializationException( diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java index 95ce70b88578..b0383329db9e 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java @@ -47,6 +47,44 @@ private KStreamsPayload roundTrip(KStreamsPayload payload) { return deserializer.deserialize(TOPIC, serializer.serialize(TOPIC, payload)); } + @Test + public void roundTripsFlushPayload() { + KStreamsPayload payload = KStreamsPayload.flush(3, 8); + KStreamsPayload out = roundTrip(payload); + assertThat(out.isFlush(), is(true)); + assertThat(out.isData(), is(false)); + assertThat(out.isWatermark(), is(false)); + assertThat(out.asFlush().getSourcePartition(), is(3)); + assertThat(out.asFlush().getTotalSourcePartitions(), is(8)); + assertThat(out, is(payload)); + } + + @Test + public void aFlushPayloadSurvivesTheFirstAndLastPartition() { + // Partition 0 and the last partition are the boundary cases of the range check, and the last + // one also happens to be the only partition that emits a flush when fanning in. + assertThat(roundTrip(KStreamsPayload.flush(0, 1)).asFlush().getSourcePartition(), is(0)); + KStreamsPayload last = roundTrip(KStreamsPayload.flush(7, 8)); + assertThat(last.asFlush().getSourcePartition(), is(7)); + assertThat(last.asFlush().getTotalSourcePartitions(), is(8)); + } + + @Test + public void aFlushPayloadRejectsAPartitionOutsideItsRange() { + assertThrows(IllegalArgumentException.class, () -> KStreamsPayload.flush(8, 8)); + assertThrows(IllegalArgumentException.class, () -> KStreamsPayload.flush(-1, 8)); + assertThrows(IllegalArgumentException.class, () -> KStreamsPayload.flush(0, 0)); + } + + @Test + public void aFlushPayloadIsNotAWatermarkOrData() { + KStreamsPayload flush = KStreamsPayload.flush(0, 1); + assertThrows(IllegalStateException.class, flush::asWatermark); + assertThrows(IllegalStateException.class, flush::getData); + KStreamsPayload data = KStreamsPayload.data(WindowedValues.valueInGlobalWindow(1)); + assertThrows(IllegalStateException.class, data::asFlush); + } + @Test public void roundTripsDataPayload() { KStreamsPayload payload = KStreamsPayload.data(WindowedValues.valueInGlobalWindow(42)); From 60205a1a12ac0fcd4935d3a2d9aeeedf50ffd6bc Mon Sep 17 00:00:00 2001 From: junaiddshaukat Date: Fri, 11 Sep 2026 12:29:47 +0500 Subject: [PATCH 2/2] Carry target partitions in the flush marker, and shorten the comments The marker now carries the partitions it is addressed to rather than the producing partition and count. The producer knows all three numbers, so it can compute the targets itself, and the consumer never needed the source information. That also removes the empty case from the wire: a producer with nothing to address emits no marker at all. --- .../main/proto/kafka_streams_payload.proto | 25 ++------- .../streams/translation/FlushPayload.java | 31 ++++------- .../streams/translation/KStreamsPayload.java | 55 ++++++++++--------- .../translation/KStreamsPayloadSerde.java | 9 ++- .../translation/KStreamsPayloadSerdeTest.java | 26 +++------ 5 files changed, 59 insertions(+), 87 deletions(-) diff --git a/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto b/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto index 15e6f4333151..43428ffb7f52 100644 --- a/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto +++ b/runners/kafka-streams/proto/src/main/proto/kafka_streams_payload.proto @@ -50,26 +50,13 @@ message KafkaStreamsPayload { bytes value = 1; } - // A request to close the open bundle and flush its output. - // - // A bundle has to be bounded in time as well as in size, or on a sparse stream the elements - // already fed to it wait for the next watermark. Closing it from a wall-clock punctuator does - // not work: transactions are committed by the Kafka Streams runtime in the background, are - // agnostic to punctuations, and are deliberately not exposed, so a bundle cannot be aligned with - // one. Instead a source emits this marker on its own punctuator and it travels the topology as - // an ordinary record, so a stage closes its bundle from process() rather than beside it. - // - // The partition fields are what let the marker be targeted rather than broadcast. Broadcasting - // would deliver one flush per upstream partition, so a downstream partition would see N times - // more flushes than the configured interval. Instead the producing partition addresses a slice - // of the downstream partitions, and the slices tile the whole range, so each downstream - // partition receives exactly one flush per interval. + // A request to close the open bundle and flush its output. See + // https://github.com/apache/beam/issues/39633. message FlushPayload { - // Which partition (physical instance) of the producing transform emitted this marker, in - // [0, total_partitions). - uint32 source_partition = 1; - // How many partitions (physical instances) the producing transform has in total. - uint32 total_partitions = 2; + // Which partitions of the repartition topic this marker is for. The producer picks them so + // that each downstream partition gets exactly one flush per interval; broadcasting would give + // it one per upstream partition instead. + repeated uint32 target_partitions = 1; } // Exactly one variant is set; the oneof case discriminates data vs watermark vs flush. diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlushPayload.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlushPayload.java index b593a0538b4e..b00902cbdfdc 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlushPayload.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/FlushPayload.java @@ -17,30 +17,21 @@ */ package org.apache.beam.runners.kafka.streams.translation; +import java.util.Set; + /** * The flush-only view of a {@link KStreamsPayload}, obtained via {@link KStreamsPayload#asFlush()}. - * As with {@link WatermarkPayload}, the accessors live here so they are only reachable once the - * caller has checked the kind and narrowed the payload. - * - *

A flush marker asks the stage that receives it to close its open bundle and flush the output, - * which is how a bundle is bounded in time. It arrives as an ordinary record, so the bundle is - * closed from {@code process()} rather than from a punctuator: transactions are committed by the - * Kafka Streams runtime in the background and are not exposed, so a bundle cannot be aligned with - * one, and trying to do it from a punctuator duplicated output against a real broker - * (https://github.com/apache/beam/issues/39633). * - *

The partition fields exist so the marker can be targeted rather than broadcast. Broadcasting - * would give a downstream partition one flush per upstream partition, so N times more flushes than - * the interval asks for. Instead the producing partition addresses a slice of the downstream - * partitions and the slices tile the range, so each downstream partition gets exactly one flush per - * interval. Unlike a watermark, a flush needs no aggregation on arrival: there is nothing to hold - * and nothing to combine, because only one arrives. + *

A flush marker asks the stage that receives it to close its bundle, which is how a bundle is + * bounded in time. It arrives as a record so the bundle is closed from {@code process()} rather + * than from a punctuator; see https://github.com/apache/beam/issues/39633. */ public interface FlushPayload { - /** Which partition of the producing transform emitted this marker. */ - int getSourcePartition(); - - /** How many partitions the producing transform has in total. */ - int getTotalSourcePartitions(); + /** + * The repartition-topic partitions this marker is addressed to. Never empty: a producer with + * nothing to address emits no marker. The receiving stage ignores this; only the partitioner + * reads it. + */ + Set getTargetPartitions(); } diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java index a495edd5e308..e31e3719aaf3 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayload.java @@ -18,9 +18,11 @@ package org.apache.beam.runners.kafka.streams.translation; import java.util.Objects; +import java.util.Set; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; import org.checkerframework.checker.nullness.qual.Nullable; /** @@ -48,6 +50,7 @@ private enum Kind { private final String transformId; private final int sourcePartition; private final int totalSourcePartitions; + private final Set targetPartitions; private KStreamsPayload( Kind kind, @@ -55,18 +58,20 @@ private KStreamsPayload( long watermarkMillis, String transformId, int sourcePartition, - int totalSourcePartitions) { + int totalSourcePartitions, + Set targetPartitions) { this.kind = kind; this.data = data; this.watermarkMillis = watermarkMillis; this.transformId = transformId; this.sourcePartition = sourcePartition; this.totalSourcePartitions = totalSourcePartitions; + this.targetPartitions = targetPartitions; } /** Returns a data payload wrapping the given {@link WindowedValue}. */ public static KStreamsPayload data(WindowedValue value) { - return new KStreamsPayload<>(Kind.DATA, value, 0L, "", 0, 0); + return new KStreamsPayload<>(Kind.DATA, value, 0L, "", 0, 0, ImmutableSet.of()); } /** @@ -90,26 +95,24 @@ public static KStreamsPayload watermark( sourcePartition, totalSourcePartitions); return new KStreamsPayload<>( - Kind.WATERMARK, null, watermarkMillis, transformId, sourcePartition, totalSourcePartitions); + Kind.WATERMARK, + null, + watermarkMillis, + transformId, + sourcePartition, + totalSourcePartitions, + ImmutableSet.of()); } /** - * Returns a flush marker: a request to close the open bundle and flush its output, carrying the - * partition of the producing transform that emitted it and how many partitions that transform - * has. Those two fields are what let the marker be addressed to a slice of the downstream - * partitions rather than broadcast to all of them; see {@link FlushPayload}. + * Returns a flush marker addressed to the given repartition-topic partitions. A producer with no + * partitions to address emits no marker, so the set must not be empty. */ - public static KStreamsPayload flush(int sourcePartition, int totalSourcePartitions) { - Preconditions.checkArgument( - totalSourcePartitions > 0, - "totalSourcePartitions must be positive: %s", - totalSourcePartitions); + public static KStreamsPayload flush(Set targetPartitions) { Preconditions.checkArgument( - sourcePartition >= 0 && sourcePartition < totalSourcePartitions, - "sourcePartition %s out of range for totalSourcePartitions %s", - sourcePartition, - totalSourcePartitions); - return new KStreamsPayload<>(Kind.FLUSH, null, 0L, "", sourcePartition, totalSourcePartitions); + !targetPartitions.isEmpty(), "flush marker must target at least one partition"); + return new KStreamsPayload<>( + Kind.FLUSH, null, 0L, "", 0, 0, ImmutableSet.copyOf(targetPartitions)); } public boolean isData() { @@ -157,13 +160,8 @@ public FlushPayload asFlush() { /** {@link FlushPayload} view backed by this payload's fields. */ private final class FlushView implements FlushPayload { @Override - public int getSourcePartition() { - return sourcePartition; - } - - @Override - public int getTotalSourcePartitions() { - return totalSourcePartitions; + public Set getTargetPartitions() { + return targetPartitions; } } @@ -204,13 +202,20 @@ public boolean equals(@Nullable Object o) { && transformId.equals(that.transformId) && sourcePartition == that.sourcePartition && totalSourcePartitions == that.totalSourcePartitions + && targetPartitions.equals(that.targetPartitions) && Objects.equals(data, that.data); } @Override public int hashCode() { return Objects.hash( - kind, data, watermarkMillis, transformId, sourcePartition, totalSourcePartitions); + kind, + data, + watermarkMillis, + transformId, + sourcePartition, + totalSourcePartitions, + targetPartitions); } @Override diff --git a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java index a7680ffa8b49..1f39a9c424ce 100644 --- a/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java +++ b/runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerde.java @@ -24,6 +24,7 @@ import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.InvalidProtocolBufferException; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; import org.apache.kafka.common.errors.SerializationException; import org.apache.kafka.common.serialization.Deserializer; import org.apache.kafka.common.serialization.Serde; @@ -78,11 +79,9 @@ public byte[] serialize(String topic, KStreamsPayload payload) { KafkaStreamsPayload.DataPayload.newBuilder() .setValue(ByteString.copyFrom(encoded.toByteArray()))); } else if (payload.isFlush()) { - FlushPayload flush = payload.asFlush(); proto.setFlush( KafkaStreamsPayload.FlushPayload.newBuilder() - .setSourcePartition(flush.getSourcePartition()) - .setTotalPartitions(flush.getTotalSourcePartitions())); + .addAllTargetPartitions(payload.asFlush().getTargetPartitions())); } else { WatermarkPayload watermark = payload.asWatermark(); proto.setWatermark( @@ -120,8 +119,8 @@ public KStreamsPayload deserialize(String topic, byte[] bytes) { watermark.getSourcePartition(), watermark.getTotalPartitions()); case FLUSH: - KafkaStreamsPayload.FlushPayload flush = proto.getFlush(); - return KStreamsPayload.flush(flush.getSourcePartition(), flush.getTotalPartitions()); + return KStreamsPayload.flush( + ImmutableSet.copyOf(proto.getFlush().getTargetPartitionsList())); case PAYLOAD_NOT_SET: default: throw new SerializationException( diff --git a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java index b0383329db9e..f0cd19a781c9 100644 --- a/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java +++ b/runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/KStreamsPayloadSerdeTest.java @@ -27,6 +27,7 @@ import org.apache.beam.sdk.transforms.windowing.GlobalWindow; import org.apache.beam.sdk.values.WindowedValue; import org.apache.beam.sdk.values.WindowedValues; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; import org.apache.kafka.common.errors.SerializationException; import org.apache.kafka.common.serialization.Deserializer; import org.apache.kafka.common.serialization.Serializer; @@ -49,36 +50,25 @@ private KStreamsPayload roundTrip(KStreamsPayload payload) { @Test public void roundTripsFlushPayload() { - KStreamsPayload payload = KStreamsPayload.flush(3, 8); + KStreamsPayload payload = KStreamsPayload.flush(ImmutableSet.of(0, 3, 7)); KStreamsPayload out = roundTrip(payload); assertThat(out.isFlush(), is(true)); assertThat(out.isData(), is(false)); assertThat(out.isWatermark(), is(false)); - assertThat(out.asFlush().getSourcePartition(), is(3)); - assertThat(out.asFlush().getTotalSourcePartitions(), is(8)); + assertThat(out.asFlush().getTargetPartitions(), is(ImmutableSet.of(0, 3, 7))); assertThat(out, is(payload)); } @Test - public void aFlushPayloadSurvivesTheFirstAndLastPartition() { - // Partition 0 and the last partition are the boundary cases of the range check, and the last - // one also happens to be the only partition that emits a flush when fanning in. - assertThat(roundTrip(KStreamsPayload.flush(0, 1)).asFlush().getSourcePartition(), is(0)); - KStreamsPayload last = roundTrip(KStreamsPayload.flush(7, 8)); - assertThat(last.asFlush().getSourcePartition(), is(7)); - assertThat(last.asFlush().getTotalSourcePartitions(), is(8)); - } - - @Test - public void aFlushPayloadRejectsAPartitionOutsideItsRange() { - assertThrows(IllegalArgumentException.class, () -> KStreamsPayload.flush(8, 8)); - assertThrows(IllegalArgumentException.class, () -> KStreamsPayload.flush(-1, 8)); - assertThrows(IllegalArgumentException.class, () -> KStreamsPayload.flush(0, 0)); + public void aFlushPayloadMustTargetSomething() { + // A producer with nothing to address emits no marker at all, so an empty set is a bug rather + // than a case to encode. + assertThrows(IllegalArgumentException.class, () -> KStreamsPayload.flush(ImmutableSet.of())); } @Test public void aFlushPayloadIsNotAWatermarkOrData() { - KStreamsPayload flush = KStreamsPayload.flush(0, 1); + KStreamsPayload flush = KStreamsPayload.flush(ImmutableSet.of(0)); assertThrows(IllegalStateException.class, flush::asWatermark); assertThrows(IllegalStateException.class, flush::getData); KStreamsPayload data = KStreamsPayload.data(WindowedValues.valueInGlobalWindow(1));