You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by jo...@apache.org on 2025/04/04 13:15:09 UTC
(beam) branch master updated: Auto-enable allowDuplicates for Redistribute with ALO streaming mode (#34490)
This is an automated email from the ASF dual-hosted git repository.
johncasey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new b111d730642 Auto-enable allowDuplicates for Redistribute with ALO streaming mode (#34490)
b111d730642 is described below
commit b111d7306420e9456bc578d4577db0d734072563
Author: Celeste Zeng <61...@users.noreply.github.com>
AuthorDate: Fri Apr 4 06:15:03 2025 -0700
Auto-enable allowDuplicates for Redistribute with ALO streaming mode (#34490)
* Extrac streaming mode from pipeline options and use it to turn on allowDuplicates for Redistribtue
* Fix format
* Add test.
* Fix test failure and formats.
---
.../beam/runners/dataflow/DataflowRunner.java | 5 ++-
.../dataflow/RedistributeByKeyOverrideFactory.java | 21 ++++++++--
.../beam/runners/dataflow/DataflowRunnerTest.java | 46 ++++++++++++++++++++++
3 files changed, 67 insertions(+), 5 deletions(-)
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java
index fd38a926faf..6ff804590a8 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java
+++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java
@@ -802,10 +802,13 @@ public class DataflowRunner extends PipelineRunner<DataflowPipelineJob> {
PTransformMatchers.classEqualTo(ParDo.SingleOutput.class),
new PrimitiveParDoSingleFactory()));
+ boolean usesAtLeastOnceStreamingMode =
+ options.getDataflowServiceOptions() != null
+ && options.getDataflowServiceOptions().contains("streaming_mode_at_least_once");
overridesBuilder.add(
PTransformOverride.of(
PTransformMatchers.classEqualTo(RedistributeByKey.class),
- new RedistributeByKeyOverrideFactory()));
+ new RedistributeByKeyOverrideFactory(usesAtLeastOnceStreamingMode)));
if (streaming) {
// For update compatibility, always use a Read for Create in streaming mode.
diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/RedistributeByKeyOverrideFactory.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/RedistributeByKeyOverrideFactory.java
index 9dc8daf8d43..509b35b1539 100644
--- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/RedistributeByKeyOverrideFactory.java
+++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/RedistributeByKeyOverrideFactory.java
@@ -47,6 +47,12 @@ class RedistributeByKeyOverrideFactory<K, V>
extends SingleInputOutputOverrideFactory<
PCollection<KV<K, V>>, PCollection<KV<K, V>>, RedistributeByKey<K, V>> {
+ private final boolean usesAtLeastOnceStreamingMode;
+
+ public RedistributeByKeyOverrideFactory(boolean usesAtLeastOnceStreamingMode) {
+ this.usesAtLeastOnceStreamingMode = usesAtLeastOnceStreamingMode;
+ }
+
@Override
public PTransformReplacement<PCollection<KV<K, V>>, PCollection<KV<K, V>>>
getReplacementTransform(
@@ -54,17 +60,24 @@ class RedistributeByKeyOverrideFactory<K, V>
transform) {
return PTransformOverrideFactory.PTransformReplacement.of(
PTransformReplacements.getSingletonMainInput(transform),
- new DataflowRedistributeByKey<>(transform.getTransform()));
+ new DataflowRedistributeByKey<>(transform.getTransform(), usesAtLeastOnceStreamingMode));
}
/** Specialized implementation of {@link RedistributeByKey} for Dataflow pipelines. */
- private static class DataflowRedistributeByKey<K, V>
+ public static class DataflowRedistributeByKey<K, V>
extends PTransform<PCollection<KV<K, V>>, PCollection<KV<K, V>>> {
private final RedistributeByKey<K, V> originalTransform;
+ private final boolean usesAtLeastOnceStreamingMode;
- private DataflowRedistributeByKey(RedistributeByKey<K, V> originalTransform) {
+ private DataflowRedistributeByKey(
+ RedistributeByKey<K, V> originalTransform, boolean usesAtLeastOnceStreamingMode) {
this.originalTransform = originalTransform;
+ this.usesAtLeastOnceStreamingMode = usesAtLeastOnceStreamingMode;
+ }
+
+ public boolean getAllowDuplicates() {
+ return this.usesAtLeastOnceStreamingMode || this.originalTransform.getAllowDuplicates();
}
@Override
@@ -84,7 +97,7 @@ class RedistributeByKeyOverrideFactory<K, V>
.apply("ReifyOriginalMetadata", Reify.windowsInValue());
PCollection<KV<K, Iterable<ValueInSingleWindow<V>>>> grouped;
- if (originalTransform.getAllowDuplicates()) {
+ if (getAllowDuplicates()) {
grouped = reified.apply(DataflowGroupByKey.createWithAllowDuplicates());
} else {
grouped = reified.apply(DataflowGroupByKey.create());
diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java
index 83cbfbafa79..c9bd50da0a5 100644
--- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java
+++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java
@@ -101,6 +101,9 @@ import org.apache.beam.runners.dataflow.util.PropertyNames;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.Pipeline.PipelineVisitor;
import org.apache.beam.sdk.coders.BigEndianIntegerCoder;
+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.coders.VoidCoder;
import org.apache.beam.sdk.extensions.gcp.auth.NoopCredentialFactory;
import org.apache.beam.sdk.extensions.gcp.auth.TestCredential;
@@ -144,6 +147,7 @@ import org.apache.beam.sdk.transforms.GroupIntoBatches;
import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.Redistribute;
import org.apache.beam.sdk.transforms.SerializableFunctions;
import org.apache.beam.sdk.transforms.SimpleFunction;
import org.apache.beam.sdk.transforms.resourcehints.ResourceHints;
@@ -2533,6 +2537,48 @@ public class DataflowRunnerTest implements Serializable {
assertTrue(sawPubsubOverride.get());
}
+ @Test
+ public void testEnableAllowDuplicatesForRedistributeWithALO() throws IOException {
+ DataflowPipelineOptions options = buildPipelineOptions();
+ options.setDataflowServiceOptions(ImmutableList.of("streaming_mode_at_least_once"));
+ Pipeline pipeline = Pipeline.create(options);
+
+ ImmutableList<KV<String, Integer>> abitraryKVs =
+ ImmutableList.of(
+ KV.of("k1", 3),
+ KV.of("k5", Integer.MAX_VALUE),
+ KV.of("k5", Integer.MIN_VALUE),
+ KV.of("k2", 66),
+ KV.of("k1", 4),
+ KV.of("k2", -33),
+ KV.of("k3", 0));
+ PCollection<KV<String, Integer>> input =
+ pipeline.apply(
+ Create.of(abitraryKVs).withCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())));
+ // The allowDuplicates for Redistribute is false by default.
+ PCollection<KV<String, Integer>> output = input.apply(Redistribute.byKey());
+ pipeline.run();
+
+ // The DataflowRedistributeByKey transform translated from Redistribute should have
+ // allowDuplicates set to true.
+ AtomicBoolean redistributeAllowDuplicates = new AtomicBoolean(false);
+ pipeline.traverseTopologically(
+ new PipelineVisitor.Defaults() {
+ @Override
+ public CompositeBehavior enterCompositeTransform(Node node) {
+ if (node.getTransform()
+ instanceof RedistributeByKeyOverrideFactory.DataflowRedistributeByKey) {
+ RedistributeByKeyOverrideFactory.DataflowRedistributeByKey<?, ?> redistribute =
+ (RedistributeByKeyOverrideFactory.DataflowRedistributeByKey<?, ?>)
+ node.getTransform();
+ redistributeAllowDuplicates.set(redistribute.getAllowDuplicates());
+ }
+ return CompositeBehavior.ENTER_TRANSFORM;
+ }
+ });
+ assertTrue(redistributeAllowDuplicates.get());
+ }
+
static class TestExpansionServiceClientFactory implements ExpansionServiceClientFactory {
ExpansionApi.ExpansionResponse response;