You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@beam.apache.org by "boy397 (via GitHub)" <gi...@apache.org> on 2026/06/27 20:36:22 UTC
[PR] Implement GroupByKey translation for Kafka Streams runner (#18479) [beam]
boy397 opened a new pull request, #39136:
URL: https://github.com/apache/beam/pull/39136
- Registered GroupByKey in KafkaStreamsPipelineTranslator.
- Created GroupByKeyTranslator to repartition via internal topic and route into the GroupByKeyProcessor.
- Added GroupByKeyProcessor using Kafka Streams state stores to buffer the global window KV elements.
- Added GroupByKeyTranslatorTest to verify execution via TopologyTestDriver.
**Please** add a meaningful description for your change here
Description:
Implement PTransformTranslator.
Translate the beam:transform:group_by_key:v1 primitive.
Set up a repartitioning step using Kafka Streams by setting the record key to the extracted Beam GroupByKey key.
Route the data through an internal repartition topic using KStreamsPayloadSerde from #39051.
Attach a key-value state store (via Stores.keyValueStoreBuilder) to buffer the elements.
Implement a Kafka Streams Processor.
Buffer incoming WindowedValue<KV<K, V>> elements per key into the state store.
Listen for watermark advancements (using the runner's watermark management).
When the watermark reaches TIMESTAMP_MAX (the end of GlobalWindow), retrieve all buffered values for each key and emit KV<K, Iterable<V>> downstream.
Clear the state store for emitted keys.
Register GroupByKeyTranslator for beam:transform:group_by_key:v1 (known URNs) alongside other primitives like Redistribute.
Set up a TopologyTestDriver E2E test.
Pipeline structure: Impulse -> produce mock KVs -> GroupByKey -> assert grouped output matches the expected iterables.
Verify that elements are only emitted when the watermark reaches the end of the global window.
Wire in the first @ValidatesRunner test on the GroupByKey category for the global window.
------------------------
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
- [ ] Mention the appropriate issue in your description (for example: `addresses #123`), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment `fixes #<ISSUE NUMBER>` instead.
- [ ] Update `CHANGES.md` with noteworthy changes.
- [ ] If this contribution is large, please file an Apache [Individual Contributor License Agreement](https://www.apache.org/licenses/icla.pdf).
See the [Contributor Guide](https://beam.apache.org/contribute) for more tips on [how to make review process smoother](https://github.com/apache/beam/blob/master/CONTRIBUTING.md#make-the-reviewers-job-easier).
To check the build health, please visit [https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md](https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md)
GitHub Actions Tests Status (on master branch)
------------------------------------------------------------------------------------------------
[](https://github.com/apache/beam/actions?query=workflow%3A%22Build+python+source+distribution+and+wheels%22+branch%3Amaster+event%3Aschedule)
[](https://github.com/apache/beam/actions?query=workflow%3A%22Python+Tests%22+branch%3Amaster+event%3Aschedule)
[](https://github.com/apache/beam/actions?query=workflow%3A%22Java+Tests%22+branch%3Amaster+event%3Aschedule)
[](https://github.com/apache/beam/actions?query=workflow%3A%22Go+tests%22+branch%3Amaster+event%3Aschedule)
See [CI.md](https://github.com/apache/beam/blob/master/CI.md) for more information about GitHub Actions CI or the [workflows README](https://github.com/apache/beam/blob/master/.github/workflows/README.md) to see a list of phrases to trigger workflows.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement GroupByKey translation for Kafka Streams runner (#18479) [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on PR #39136:
URL: https://github.com/apache/beam/pull/39136#issuecomment-4821438940
## Summary of Changes
Hello, I'm Gemini Code Assist[^1]! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request introduces support for the GroupByKey primitive in the Apache Beam Kafka Streams runner. It establishes a robust mechanism for repartitioning data via internal Kafka topics and uses state stores to buffer windowed elements until the global window watermark is reached, ensuring correct grouping behavior in streaming pipelines.
### Highlights
* **GroupByKey Implementation**: Implemented the GroupByKey translation for the Kafka Streams runner, enabling grouping operations by buffering elements in a state store and emitting them upon watermark advancement.
* **Infrastructure Setup**: Added GroupByKeyTranslator and GroupByKeyProcessor to handle repartitioning and stateful buffering, and registered them in the KafkaStreamsPipelineTranslator.
* **Testing**: Added a comprehensive E2E test using TopologyTestDriver to verify the GroupByKey execution logic and watermark handling.
<details>
<summary><b>New Features</b></summary>
🧠You can now enable **Memory** (public preview) to help **Gemini Code Assist** learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. **Click [here](https://codeassist.google/code-review/login) to enable Memory in your admin console.**
</details>
<details>
<summary><b>Using Gemini Code Assist</b></summary>
<br>
The full guide for Gemini Code Assist can be found on our [documentation page](https://developers.google.com/gemini-code-assist/docs/review-github-code), here are some quick tips.
<b>Invoking Gemini</b>
You can request assistance from Gemini at any point by creating a comment using either `/gemini <command>` or `@gemini-code-assist <command>`. Below is a summary of the supported commands on the current page.
Feature | Command | Description
--- | --- | ---
Code Review | `/gemini review` | Performs a code review for the current pull request in its current state.
Pull Request Summary | `/gemini summary` | Provides a summary of the current pull request in its current state.
Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help | `/gemini help` | Displays a list of available commands.
<b>Customization</b>
To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a `.gemini/` folder in the base of the repository. Detailed instructions can be found [here](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github).
<b>Limitations & Feedback</b>
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with :thumbsup: and :thumbsdown: on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up [here](https://google.qualtrics.com/jfe/form/SV_2cyuGuTWsEw84yG).
</details>
[^1]: Review the [Privacy Notices](https://policies.google.com/privacy), [Generative AI Prohibited Use Policy](https://policies.google.com/terms/generative-ai/use-policy), [Terms of Service](https://policies.google.com/terms), and learn how to configure Gemini Code Assist in GitHub [here](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github). Gemini can make mistakes, so double check it and [use code with caution](https://support.google.com/legal/answer/13505487).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement GroupByKey translation for Kafka Streams runner (#18479) [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #39136:
URL: https://github.com/apache/beam/pull/39136#issuecomment-4914911497
Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment `assign to next reviewer`:
R: @Abacn added as fallback since no labels match configuration
Available commands:
- `stop reviewer notifications` - opt out of the automated review tooling
- `remind me after tests pass` - tag the comment author after tests pass
- `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement GroupByKey translation for Kafka Streams runner (#18479) [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #39136:
URL: https://github.com/apache/beam/pull/39136#issuecomment-4821887511
Assigning reviewers:
R: @Abacn added as fallback since no labels match configuration
Note: If you would like to opt out of this review, comment `assign to next reviewer`.
Available commands:
- `stop reviewer notifications` - opt out of the automated review tooling
- `remind me after tests pass` - tag the comment author after tests pass
- `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)
The PR bot will only process comments in the main thread (not review comments).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement GroupByKey translation for Kafka Streams runner (#18479) [beam]
Posted by "boy397 (via GitHub)" <gi...@apache.org>.
boy397 closed pull request #39136: Implement GroupByKey translation for Kafka Streams runner (#18479)
URL: https://github.com/apache/beam/pull/39136
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement GroupByKey translation for Kafka Streams runner (#18479) [beam]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #39136:
URL: https://github.com/apache/beam/pull/39136#discussion_r3486741923
##########
runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java:
##########
@@ -0,0 +1,133 @@
+/*
+ * 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 java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.ListCoder;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.WindowedValue;
+import org.apache.beam.sdk.values.WindowedValues;
+import org.apache.kafka.streams.processor.api.Processor;
+import org.apache.kafka.streams.processor.api.ProcessorContext;
+import org.apache.kafka.streams.processor.api.Record;
+import org.apache.kafka.streams.state.KeyValueIterator;
+import org.apache.kafka.streams.state.KeyValueStore;
+
+/** Processor for GroupByKey. */
+class GroupByKeyProcessor
+ implements Processor<
+ byte[],
+ KStreamsPayload<KV<Object, Object>>,
+ byte[],
+ KStreamsPayload<KV<Object, Iterable<Object>>>> {
+
+ private final String stateStoreName;
+ private final String transformId;
+ private final Coder<List<WindowedValue<KV<Object, Object>>>> listCoder;
+
+ private ProcessorContext<byte[], KStreamsPayload<KV<Object, Iterable<Object>>>> context;
+ private KeyValueStore<byte[], byte[]> stateStore;
+
+ GroupByKeyProcessor(
+ String stateStoreName,
+ String transformId,
+ Coder<WindowedValue<KV<Object, Object>>> inputCoder) {
+ this.stateStoreName = stateStoreName;
+ this.transformId = transformId;
+ this.listCoder = ListCoder.of(inputCoder);
+ }
Review Comment:
Refactor the processor to store the individual element coder (`inputCoder`) instead of wrapping it in a `ListCoder`. This enables a highly optimized append-only state storage pattern, avoiding the need to decode and re-encode the entire list of elements on every new record.
```suggestion
private final Coder<WindowedValue<KV<Object, Object>>> inputCoder;
private ProcessorContext<byte[], KStreamsPayload<KV<Object, Iterable<Object>>>> context;
private KeyValueStore<byte[], byte[]> stateStore;
GroupByKeyProcessor(
String stateStoreName,
String transformId,
Coder<WindowedValue<KV<Object, Object>>> inputCoder) {
this.stateStoreName = stateStoreName;
this.transformId = transformId;
this.inputCoder = inputCoder;
}
```
##########
runners/kafka-streams/src/main/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyProcessor.java:
##########
@@ -0,0 +1,133 @@
+/*
+ * 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 java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.ListCoder;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.WindowedValue;
+import org.apache.beam.sdk.values.WindowedValues;
+import org.apache.kafka.streams.processor.api.Processor;
+import org.apache.kafka.streams.processor.api.ProcessorContext;
+import org.apache.kafka.streams.processor.api.Record;
+import org.apache.kafka.streams.state.KeyValueIterator;
+import org.apache.kafka.streams.state.KeyValueStore;
+
+/** Processor for GroupByKey. */
+class GroupByKeyProcessor
+ implements Processor<
+ byte[],
+ KStreamsPayload<KV<Object, Object>>,
+ byte[],
+ KStreamsPayload<KV<Object, Iterable<Object>>>> {
+
+ private final String stateStoreName;
+ private final String transformId;
+ private final Coder<List<WindowedValue<KV<Object, Object>>>> listCoder;
+
+ private ProcessorContext<byte[], KStreamsPayload<KV<Object, Iterable<Object>>>> context;
+ private KeyValueStore<byte[], byte[]> stateStore;
+
+ GroupByKeyProcessor(
+ String stateStoreName,
+ String transformId,
+ Coder<WindowedValue<KV<Object, Object>>> inputCoder) {
+ this.stateStoreName = stateStoreName;
+ this.transformId = transformId;
+ this.listCoder = ListCoder.of(inputCoder);
+ }
+
+ @Override
+ public void init(
+ ProcessorContext<byte[], KStreamsPayload<KV<Object, Iterable<Object>>>> context) {
+ this.context = context;
+ this.stateStore = context.getStateStore(stateStoreName);
+ }
+
+ @Override
+ public void process(Record<byte[], KStreamsPayload<KV<Object, Object>>> record) {
+ KStreamsPayload<KV<Object, Object>> payload = record.value();
+
+ if (payload.isData()) {
+ byte[] keyBytes = record.key();
+ byte[] existingBytes = stateStore.get(keyBytes);
+ List<WindowedValue<KV<Object, Object>>> list;
+ if (existingBytes == null) {
+ list = new ArrayList<>();
+ } else {
+ try {
+ list = listCoder.decode(new ByteArrayInputStream(existingBytes));
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to decode buffered GroupByKey state", e);
+ }
+ }
+ list.add(payload.getData());
+ ByteArrayOutputStream os = new ByteArrayOutputStream();
+ try {
+ listCoder.encode(list, os);
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to encode buffered GroupByKey state", e);
+ }
+ stateStore.put(keyBytes, os.toByteArray());
+ } else {
+ WatermarkPayload watermark = payload.asWatermark();
+ if (watermark.getWatermarkMillis() == BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()) {
+ try (KeyValueIterator<byte[], byte[]> iterator = stateStore.all()) {
+ while (iterator.hasNext()) {
+ org.apache.kafka.streams.KeyValue<byte[], byte[]> kv = iterator.next();
+ List<WindowedValue<KV<Object, Object>>> buffered;
+ try {
+ buffered = listCoder.decode(new ByteArrayInputStream(kv.value));
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to decode buffered GroupByKey state on emit", e);
+ }
+ if (!buffered.isEmpty()) {
+ List<Object> values = new ArrayList<>();
+ for (WindowedValue<KV<Object, Object>> wv : buffered) {
+ values.add(wv.getValue().getValue());
+ }
+ Object key = buffered.get(0).getValue().getKey();
+ WindowedValue<KV<Object, Iterable<Object>>> outWv =
+ WindowedValues.valueInGlobalWindow(KV.of(key, values));
+ context.forward(
+ new Record<>(kv.key, KStreamsPayload.data(outWv), record.timestamp()));
+ }
+ }
+ }
+ // Since we fired everything for the global window, we can optionally clear the store here.
+ // But the pipeline is finishing.
+
+ // Forward the watermark downstream
+ context.forward(
+ new Record<>(
+ record.key(),
+ KStreamsPayload.watermark(
+ watermark.getWatermarkMillis(),
+ watermark.getSourcePartition(),
+ watermark.getTotalSourcePartitions()),
+ record.timestamp()));
+ }
+ }
+ }
Review Comment:
This refactoring addresses several critical issues in the `process` method:
1. **Performance ($O(N)$ vs $O(N^2)$)**: Instead of decoding and re-encoding the entire list of elements for every incoming record, we append the serialized bytes of the new element directly to the existing byte array in the state store. When emitting, we decode the elements sequentially from the stream.
2. **State Leak**: We collect the keys and delete them from the state store after forwarding the grouped elements, preventing persistent state leaks.
3. **Watermark Forwarding**: We forward all watermarks downstream (not just the terminal one) to allow downstream watermark progression and prevent pipeline stalls.
4. **Correct Output Timestamp**: We compute the minimum timestamp of the buffered elements (following `TimestampCombiner.EARLIEST`) and use `WindowedValue.timestampedValueInGlobalWindow` instead of defaulting to `TIMESTAMP_MIN_VALUE`, which would violate event-time semantics and cause downstream data loss.
```java
@Override
public void process(Record<byte[], KStreamsPayload<KV<Object, Object>>> record) {
KStreamsPayload<KV<Object, Object>> payload = record.value();
if (payload.isData()) {
byte[] keyBytes = record.key();
byte[] existingBytes = stateStore.get(keyBytes);
ByteArrayOutputStream os = new ByteArrayOutputStream();
if (existingBytes != null) {
try {
os.write(existingBytes);
} catch (IOException e) {
throw new RuntimeException("Failed to append to buffered GroupByKey state", e);
}
}
try {
inputCoder.encode(payload.getData(), os);
} catch (IOException e) {
throw new RuntimeException("Failed to encode buffered GroupByKey state", e);
}
stateStore.put(keyBytes, os.toByteArray());
} else {
WatermarkPayload watermark = payload.asWatermark();
if (watermark.getWatermarkMillis() == BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()) {
List<byte[]> keysToDelete = new ArrayList<>();
try (KeyValueIterator<byte[], byte[]> iterator = stateStore.all()) {
while (iterator.hasNext()) {
org.apache.kafka.streams.KeyValue<byte[], byte[]> kv = iterator.next();
keysToDelete.add(kv.key);
List<WindowedValue<KV<Object, Object>>> buffered = new ArrayList<>();
try (ByteArrayInputStream bis = new ByteArrayInputStream(kv.value)) {
while (bis.available() > 0) {
buffered.add(inputCoder.decode(bis));
}
} catch (IOException e) {
throw new RuntimeException("Failed to decode buffered GroupByKey state on emit", e);
}
if (!buffered.isEmpty()) {
org.joda.time.Instant minTimestamp = BoundedWindow.TIMESTAMP_MAX_VALUE;
List<Object> values = new ArrayList<>();
for (WindowedValue<KV<Object, Object>> wv : buffered) {
values.add(wv.getValue().getValue());
if (wv.getTimestamp().isBefore(minTimestamp)) {
minTimestamp = wv.getTimestamp();
}
}
Object key = buffered.get(0).getValue().getKey();
WindowedValue<KV<Object, Iterable<Object>>> outWv =
WindowedValue.timestampedValueInGlobalWindow(KV.of(key, values), minTimestamp);
context.forward(
new Record<>(kv.key, KStreamsPayload.data(outWv), record.timestamp()));
}
}
}
for (byte[] key : keysToDelete) {
stateStore.delete(key);
}
}
// Forward the watermark downstream
context.forward(
new Record<>(
record.key(),
KStreamsPayload.watermark(
watermark.getWatermarkMillis(),
watermark.getSourcePartition(),
watermark.getTotalSourcePartitions()),
record.timestamp()));
}
}
```
##########
runners/kafka-streams/src/test/java/org/apache/beam/runners/kafka/streams/translation/GroupByKeyTranslatorTest.java:
##########
@@ -0,0 +1,199 @@
+/*
+ * 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.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Properties;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.GroupByKey;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.util.construction.PipelineTranslation;
+import org.apache.beam.sdk.util.construction.RehydratedComponents;
+import org.apache.beam.sdk.values.KV;
+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.Iterables;
+import org.apache.kafka.common.serialization.Serdes;
+import org.apache.kafka.streams.StreamsConfig;
+import org.apache.kafka.streams.TestInputTopic;
+import org.apache.kafka.streams.Topology;
+import org.apache.kafka.streams.TopologyTestDriver;
+import org.apache.kafka.streams.processor.api.Processor;
+import org.apache.kafka.streams.processor.api.ProcessorContext;
+import org.apache.kafka.streams.processor.api.ProcessorSupplier;
+import org.apache.kafka.streams.processor.api.Record;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.junit.Test;
+
+/** Tests for {@link GroupByKeyTranslator}. */
+public class GroupByKeyTranslatorTest {
+
+ @Test
+ public void groupByKeyBuffersAndFiresAtTerminalWatermark() throws Exception {
+ KafkaStreamsTranslationContext context = KafkaStreamsPipelineTranslatorTest.newContext();
+
+ Pipeline sdkPipeline = Pipeline.create(PipelineOptionsFactory.create());
+ sdkPipeline
+ .apply("create", Create.empty(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())))
+ .apply("gbk", GroupByKey.create());
+
+ RunnerApi.Pipeline pipeline = PipelineTranslation.toProto(sdkPipeline);
+ String gbkTransformId = null;
+ for (String id : pipeline.getComponents().getTransformsMap().keySet()) {
+ if (pipeline
+ .getComponents()
+ .getTransformsOrThrow(id)
+ .getSpec()
+ .getUrn()
+ .equals("beam:transform:group_by_key:v1")) {
+ gbkTransformId = id;
+ }
+ }
+
+ RunnerApi.PTransform gbkTransform =
+ pipeline.getComponents().getTransformsOrThrow(gbkTransformId);
+ String inputPCollId = Iterables.getOnlyElement(gbkTransform.getInputsMap().values());
+ context.registerPCollectionProducer(inputPCollId, "mock-parent");
+
+ RehydratedComponents components = RehydratedComponents.forComponents(pipeline.getComponents());
+ Coder<?> rawInputCoder =
+ components.getCoder(
+ pipeline.getComponents().getPcollectionsOrThrow(inputPCollId).getCoderId());
+ @SuppressWarnings("unchecked")
+ Coder<WindowedValue<KV<String, String>>> inputCoder =
+ (Coder<WindowedValue<KV<String, String>>>) rawInputCoder;
+ KStreamsPayloadSerde<KV<String, String>> payloadSerde = new KStreamsPayloadSerde<>(inputCoder);
+
+ Topology topology = context.getTopology();
+ topology.addSource(
+ "mock-source",
+ Serdes.ByteArray().deserializer(),
+ payloadSerde.deserializer(),
+ "mock-topic");
+ // Just a pass-through to satisfy the parent reference
+ topology.addProcessor("mock-parent", PassThroughProcessor::new, "mock-source");
+
+ new GroupByKeyTranslator().translate(gbkTransformId, pipeline, context);
+
+ CapturingProcessor capture = new CapturingProcessor();
+ topology.addProcessor("capture", capture, gbkTransformId);
+
+ try (TopologyTestDriver driver = new TopologyTestDriver(topology, baseProps())) {
+ TestInputTopic<byte[], KStreamsPayload<KV<String, String>>> inputTopic =
+ driver.createInputTopic(
+ "mock-topic", Serdes.ByteArray().serializer(), payloadSerde.serializer());
+
+ // Send elements
+ inputTopic.pipeInput(
+ new byte[0], KStreamsPayload.data(WindowedValues.valueInGlobalWindow(KV.of("k1", "v1"))));
+ inputTopic.pipeInput(
+ new byte[0], KStreamsPayload.data(WindowedValues.valueInGlobalWindow(KV.of("k1", "v2"))));
+ inputTopic.pipeInput(
+ new byte[0], KStreamsPayload.data(WindowedValues.valueInGlobalWindow(KV.of("k2", "v3"))));
+
+ // No output yet
+ assertThat(capture.received.size(), is(0));
+
+ // Send terminal watermark
+ inputTopic.pipeInput(
+ new byte[0],
+ KStreamsPayload.watermark(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(), 0, 1));
+ }
+
+ // Now it should have fired. It fires 1 for each key, plus the watermark
+ assertThat(
+ capture.received.size(),
+ is(3)); // 2 data + 1 watermark (wait, maybe order is deterministic?)
+
+ int dataCount = 0;
+ int watermarkCount = 0;
+ for (KStreamsPayload<?> payload : capture.received) {
+ if (payload.isData()) {
+ dataCount++;
+ KV<?, ?> kv = (KV<?, ?>) payload.getData().getValue();
+ Iterable<?> iter = (Iterable<?>) kv.getValue();
+ if ("k1".equals(kv.getKey())) {
+ assertThat(Iterables.size(iter), is(2));
+ } else if ("k2".equals(kv.getKey())) {
+ assertThat(Iterables.size(iter), is(1));
+ }
+ } else {
+ watermarkCount++;
+ assertThat(
+ payload.asWatermark().getWatermarkMillis(),
+ is(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()));
+ }
+ }
+ assertThat(dataCount, is(2));
+ assertThat(watermarkCount, is(1));
Review Comment:
Update the test to verify that non-terminal watermarks are correctly forwarded downstream without triggering early emission, and that the terminal watermark is also forwarded alongside the emitted data.
```java
// No output yet
assertThat(capture.received.size(), is(0));
// Send non-terminal watermark
inputTopic.pipeInput(
new byte[0],
KStreamsPayload.watermark(1000L, 0, 1));
// Watermark should be forwarded, but no data emitted yet
assertThat(capture.received.size(), is(1));
assertThat(capture.received.get(0).isData(), is(false));
assertThat(capture.received.get(0).asWatermark().getWatermarkMillis(), is(1000L));
// Send terminal watermark
inputTopic.pipeInput(
new byte[0],
KStreamsPayload.watermark(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(), 0, 1));
}
// Now it should have fired. It fires 1 for each key, plus the watermarks
assertThat(
capture.received.size(),
is(4)); // 2 data + 2 watermarks
int dataCount = 0;
int watermarkCount = 0;
for (KStreamsPayload<?> payload : capture.received) {
if (payload.isData()) {
dataCount++;
KV<?, ?> kv = (KV<?, ?>) payload.getData().getValue();
Iterable<?> iter = (Iterable<?>) kv.getValue();
if ("k1".equals(kv.getKey())) {
assertThat(Iterables.size(iter), is(2));
} else if ("k2".equals(kv.getKey())) {
assertThat(Iterables.size(iter), is(1));
}
} else {
watermarkCount++;
long wm = payload.asWatermark().getWatermarkMillis();
assertThat(
wm == 1000L || wm == BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis(),
is(true));
}
}
assertThat(dataCount, is(2));
assertThat(watermarkCount, is(2));
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Implement GroupByKey translation for Kafka Streams runner (#18479) [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #39136:
URL: https://github.com/apache/beam/pull/39136#issuecomment-4886071598
Reminder, please take a look at this pr: @Abacn
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org