You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by ro...@apache.org on 2024/03/04 18:54:31 UTC
(beam) branch master updated: Initial prototype of using Beam transforms directly in a Flink pipeline. (#30332)
This is an automated email from the ASF dual-hosted git repository.
robertwb 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 9c3f209aa69 Initial prototype of using Beam transforms directly in a Flink pipeline. (#30332)
9c3f209aa69 is described below
commit 9c3f209aa69db0af7a972dea2f7fceef631bcd2a
Author: Robert Bradshaw <ro...@gmail.com>
AuthorDate: Mon Mar 4 10:54:25 2024 -0800
Initial prototype of using Beam transforms directly in a Flink pipeline. (#30332)
Handling DataStream, windowing, and more complex types will come in a future PR.
---
.../FlinkBatchPortablePipelineTranslator.java | 18 +-
.../beam/runners/flink/FlinkPipelineRunner.java | 28 +--
.../flink/FlinkPortablePipelineTranslator.java | 29 +++
.../flink/adapter/BeamAdapterCoderUtils.java | 83 +++++++
.../runners/flink/adapter/BeamAdapterUtils.java | 139 +++++++++++
.../flink/adapter/BeamFlinkDataSetAdapter.java | 265 +++++++++++++++++++++
.../beam/runners/flink/adapter/FlinkInput.java | 102 ++++++++
.../beam/runners/flink/adapter/FlinkOutput.java | 93 ++++++++
.../beam/runners/flink/adapter/package-info.java | 20 ++
.../flink/adapter/BeamFlinkDataSetAdapterTest.java | 135 +++++++++++
10 files changed, 883 insertions(+), 29 deletions(-)
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkBatchPortablePipelineTranslator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkBatchPortablePipelineTranslator.java
index 76a711a65cd..4d9a5a516c7 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkBatchPortablePipelineTranslator.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkBatchPortablePipelineTranslator.java
@@ -133,11 +133,24 @@ public class FlinkBatchPortablePipelineTranslator
ExecutionEnvironment executionEnvironment =
FlinkExecutionEnvironments.createBatchExecutionEnvironment(
pipelineOptions, filesToStage, confDir);
+ return createTranslationContext(jobInfo, pipelineOptions, executionEnvironment);
+ }
+
+ public static BatchTranslationContext createTranslationContext(
+ JobInfo jobInfo,
+ FlinkPipelineOptions pipelineOptions,
+ ExecutionEnvironment executionEnvironment) {
return new BatchTranslationContext(jobInfo, pipelineOptions, executionEnvironment);
}
/** Creates a batch translator. */
public static FlinkBatchPortablePipelineTranslator createTranslator() {
+ return createTranslator(ImmutableMap.of());
+ }
+
+ /** Creates a batch translator. */
+ public static FlinkBatchPortablePipelineTranslator createTranslator(
+ Map<String, PTransformTranslator> extraTranslations) {
ImmutableMap.Builder<String, PTransformTranslator> translatorMap = ImmutableMap.builder();
translatorMap.put(
PTransformTranslation.FLATTEN_TRANSFORM_URN,
@@ -153,6 +166,7 @@ public class FlinkBatchPortablePipelineTranslator
translatorMap.put(
PTransformTranslation.RESHUFFLE_URN,
FlinkBatchPortablePipelineTranslator::translateReshuffle);
+ translatorMap.putAll(extraTranslations);
return new FlinkBatchPortablePipelineTranslator(translatorMap.build());
}
@@ -225,7 +239,7 @@ public class FlinkBatchPortablePipelineTranslator
/** Transform translation interface. */
@FunctionalInterface
- private interface PTransformTranslator {
+ public interface PTransformTranslator {
/** Translate a PTransform into the given translation context. */
void translate(
PTransformNode transform, RunnerApi.Pipeline pipeline, BatchTranslationContext context);
@@ -233,7 +247,7 @@ public class FlinkBatchPortablePipelineTranslator
private final Map<String, PTransformTranslator> urnToTransformTranslator;
- private FlinkBatchPortablePipelineTranslator(
+ public FlinkBatchPortablePipelineTranslator(
Map<String, PTransformTranslator> urnToTransformTranslator) {
this.urnToTransformTranslator = urnToTransformTranslator;
}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineRunner.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineRunner.java
index eb21887a3e1..04b3b7e202d 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineRunner.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineRunner.java
@@ -37,13 +37,7 @@ import org.apache.beam.sdk.metrics.MetricsOptions;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.options.SdkHarnessOptions;
-import org.apache.beam.sdk.util.construction.PTransformTranslation;
import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation;
-import org.apache.beam.sdk.util.construction.graph.ExecutableStage;
-import org.apache.beam.sdk.util.construction.graph.GreedyPipelineFuser;
-import org.apache.beam.sdk.util.construction.graph.ProtoOverrides;
-import org.apache.beam.sdk.util.construction.graph.SplittableParDoExpander;
-import org.apache.beam.sdk.util.construction.graph.TrivialNativeTransformExpander;
import org.apache.beam.vendor.grpc.v1p60p1.com.google.protobuf.Struct;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
import org.apache.flink.api.common.JobExecutionResult;
@@ -105,30 +99,10 @@ public class FlinkPipelineRunner implements PortablePipelineRunner {
throws Exception {
LOG.info("Translating pipeline to Flink program.");
- // Expand any splittable ParDos within the graph to enable sizing and splitting of bundles.
- Pipeline pipelineWithSdfExpanded =
- ProtoOverrides.updateTransform(
- PTransformTranslation.PAR_DO_TRANSFORM_URN,
- pipeline,
- SplittableParDoExpander.createSizedReplacement());
-
- // Don't let the fuser fuse any subcomponents of native transforms.
- Pipeline trimmedPipeline =
- TrivialNativeTransformExpander.forKnownUrns(
- pipelineWithSdfExpanded, translator.knownUrns());
-
- // Fused pipeline proto.
- // TODO: Consider supporting partially-fused graphs.
- RunnerApi.Pipeline fusedPipeline =
- trimmedPipeline.getComponents().getTransformsMap().values().stream()
- .anyMatch(proto -> ExecutableStage.URN.equals(proto.getSpec().getUrn()))
- ? trimmedPipeline
- : GreedyPipelineFuser.fuse(trimmedPipeline).toPipeline();
-
FlinkPortablePipelineTranslator.Executor executor =
translator.translate(
translator.createTranslationContext(jobInfo, pipelineOptions, confDir, filesToStage),
- fusedPipeline);
+ translator.prepareForTranslation(pipeline));
final JobExecutionResult result = executor.execute(pipelineOptions.getJobName());
return createPortablePipelineResult(result, pipelineOptions);
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPortablePipelineTranslator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPortablePipelineTranslator.java
index de69959b09e..bbb9bd4999e 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPortablePipelineTranslator.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPortablePipelineTranslator.java
@@ -21,6 +21,12 @@ import java.util.List;
import java.util.Set;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
+import org.apache.beam.sdk.util.construction.PTransformTranslation;
+import org.apache.beam.sdk.util.construction.graph.ExecutableStage;
+import org.apache.beam.sdk.util.construction.graph.GreedyPipelineFuser;
+import org.apache.beam.sdk.util.construction.graph.ProtoOverrides;
+import org.apache.beam.sdk.util.construction.graph.SplittableParDoExpander;
+import org.apache.beam.sdk.util.construction.graph.TrivialNativeTransformExpander;
import org.apache.flink.api.common.JobExecutionResult;
import org.checkerframework.checker.nullness.qual.Nullable;
@@ -55,6 +61,29 @@ public interface FlinkPortablePipelineTranslator<
Set<String> knownUrns();
+ default RunnerApi.Pipeline prepareForTranslation(RunnerApi.Pipeline pipeline) {
+ // Expand any splittable ParDos within the graph to enable sizing and splitting of bundles.
+ RunnerApi.Pipeline pipelineWithSdfExpanded =
+ ProtoOverrides.updateTransform(
+ PTransformTranslation.PAR_DO_TRANSFORM_URN,
+ pipeline,
+ SplittableParDoExpander.createSizedReplacement());
+
+ // Don't let the fuser fuse any subcomponents of native transforms.
+ RunnerApi.Pipeline trimmedPipeline =
+ TrivialNativeTransformExpander.forKnownUrns(pipelineWithSdfExpanded, knownUrns());
+
+ // Fused pipeline proto.
+ // TODO: Consider supporting partially-fused graphs.
+ RunnerApi.Pipeline fusedPipeline =
+ trimmedPipeline.getComponents().getTransformsMap().values().stream()
+ .anyMatch(proto -> ExecutableStage.URN.equals(proto.getSpec().getUrn()))
+ ? trimmedPipeline
+ : GreedyPipelineFuser.fuse(trimmedPipeline).toPipeline();
+
+ return fusedPipeline;
+ }
+
/** Translates the given pipeline. */
Executor translate(T context, RunnerApi.Pipeline pipeline);
}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/adapter/BeamAdapterCoderUtils.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/adapter/BeamAdapterCoderUtils.java
new file mode 100644
index 00000000000..d20d5e2fc03
--- /dev/null
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/adapter/BeamAdapterCoderUtils.java
@@ -0,0 +1,83 @@
+/*
+ * 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.flink.adapter;
+
+import java.io.IOException;
+import java.util.Map;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.runners.flink.translation.types.CoderTypeInformation;
+import org.apache.beam.sdk.coders.CannotProvideCoderException;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.coders.IterableCoder;
+import org.apache.beam.sdk.coders.MapCoder;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.util.Preconditions;
+import org.apache.beam.sdk.util.construction.CoderTranslation;
+import org.apache.beam.sdk.util.construction.RehydratedComponents;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+
+class BeamAdapterCoderUtils {
+ static <T> Coder<T> typeInformationToCoder(
+ TypeInformation<T> typeInfo, CoderRegistry coderRegistry) {
+ Class<T> clazz = typeInfo.getTypeClass();
+ if (typeInfo instanceof CoderTypeInformation) {
+ return ((CoderTypeInformation) typeInfo).getCoder();
+ } else if (clazz.getTypeParameters().length == 0) {
+ try {
+ return coderRegistry.getCoder(clazz);
+ } catch (CannotProvideCoderException exn) {
+ throw new RuntimeException(exn);
+ }
+ } else if (Iterable.class.isAssignableFrom(clazz)) {
+ TypeInformation<?> elementType =
+ Preconditions.checkArgumentNotNull(typeInfo.getGenericParameters().get("T"));
+ return (Coder) IterableCoder.of(typeInformationToCoder(elementType, coderRegistry));
+ } else if (Map.class.isAssignableFrom(clazz)) {
+ TypeInformation<?> keyType =
+ Preconditions.checkArgumentNotNull(typeInfo.getGenericParameters().get("K"));
+ TypeInformation<?> valueType =
+ Preconditions.checkArgumentNotNull(typeInfo.getGenericParameters().get("V"));
+ return (Coder)
+ MapCoder.of(
+ typeInformationToCoder(keyType, coderRegistry),
+ typeInformationToCoder(valueType, coderRegistry));
+ } else {
+ throw new RuntimeException("Coder translation for " + typeInfo + " not yet supported.");
+ }
+ }
+
+ static <T> TypeInformation<T> coderToTypeInformation(Coder<T> coder, PipelineOptions options) {
+ // TODO(robertwb): Consider mapping some common types.
+ return new CoderTypeInformation<>(coder, options);
+ }
+
+ static <T> Coder<T> lookupCoder(RunnerApi.Pipeline p, String pCollectionId) {
+ try {
+ return (Coder<T>)
+ CoderTranslation.fromProto(
+ p.getComponents()
+ .getCodersOrThrow(
+ p.getComponents().getPcollectionsOrThrow(pCollectionId).getCoderId()),
+ RehydratedComponents.forComponents(p.getComponents()),
+ CoderTranslation.TranslationContext.DEFAULT);
+ } catch (IOException exn) {
+ throw new RuntimeException(exn);
+ }
+ }
+}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/adapter/BeamAdapterUtils.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/adapter/BeamAdapterUtils.java
new file mode 100644
index 00000000000..065796057c2
--- /dev/null
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/adapter/BeamAdapterUtils.java
@@ -0,0 +1,139 @@
+/*
+ * 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.flink.adapter;
+
+import java.util.Map;
+import java.util.function.BiFunction;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.PortablePipelineOptions;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.util.construction.Environments;
+import org.apache.beam.sdk.util.construction.PipelineTranslation;
+import org.apache.beam.sdk.util.construction.SdkComponents;
+import org.apache.beam.sdk.values.PBegin;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.PInput;
+import org.apache.beam.sdk.values.POutput;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.ExecutionEnvironment;
+
+class BeamAdapterUtils {
+ private BeamAdapterUtils() {}
+
+ interface PipelineFragmentTranslator<DataSetOrStreamT> {
+ Map<String, DataSetOrStreamT> translate(
+ Map<String, ? extends DataSetOrStreamT> inputs,
+ RunnerApi.Pipeline pipelineProto,
+ ExecutionEnvironment executionEnvironment);
+ }
+
+ @SuppressWarnings({"nullness", "rawtypes"})
+ static <DataSetOrStreamT, BeamInputT extends PInput, BeamOutputT extends POutput>
+ Map<String, DataSetOrStreamT> applyBeamPTransformInternal(
+ Map<String, ? extends DataSetOrStreamT> inputs,
+ BiFunction<Pipeline, Map<String, PCollection<?>>, BeamInputT> toBeamInput,
+ Function<BeamOutputT, Map<String, PCollection<?>>> fromBeamOutput,
+ PTransform<? super BeamInputT, BeamOutputT> transform,
+ ExecutionEnvironment executionEnvironment,
+ Function<DataSetOrStreamT, TypeInformation<?>> getTypeInformation,
+ PipelineOptions pipelineOptions,
+ CoderRegistry coderRegistry,
+ PipelineFragmentTranslator<DataSetOrStreamT> translator) {
+ Pipeline pipeline = Pipeline.create();
+
+ // Construct beam inputs corresponding to each Flink input.
+ Map<String, PCollection<?>> beamInputs =
+ // Copy as transformEntries lazy recomputes entries.
+ ImmutableMap.copyOf(
+ Maps.transformEntries(
+ inputs,
+ (key, flinkInput) ->
+ pipeline.apply(
+ new FlinkInput<>(
+ key,
+ BeamAdapterCoderUtils.typeInformationToCoder(
+ getTypeInformation.apply(flinkInput), coderRegistry)))));
+
+ // Actually apply the transform to create Beam outputs.
+ Map<String, PCollection<?>> beamOutputs =
+ fromBeamOutput.apply(applyTransform(toBeamInput.apply(pipeline, beamInputs), transform));
+
+ // This attaches PTransforms to each output which will be used to populate the Flink outputs
+ // during translation.
+ beamOutputs.entrySet().stream()
+ .forEach(
+ e -> {
+ ((PCollection<Object>) e.getValue()).apply(new FlinkOutput<Object>(e.getKey()));
+ });
+
+ // This "environment" executes the SDK harness in the parent worker process.
+ // TODO(robertwb): Support other modes.
+ // TODO(robertwb): In embedded mode, consider an optimized data (and state) channel rather than
+ // serializing everything over grpc protos.
+ pipelineOptions
+ .as(PortablePipelineOptions.class)
+ .setDefaultEnvironmentType(Environments.ENVIRONMENT_EMBEDDED);
+
+ // Extract the pipeline definition so that we can apply or Flink translation logic.
+ SdkComponents components = SdkComponents.create(pipelineOptions);
+ RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline, components);
+ return translator.translate(inputs, pipelineProto, executionEnvironment);
+ }
+
+ static Map<String, PCollection<?>> tupleToMap(PCollectionTuple tuple) {
+ return tuple.getAll().entrySet().stream()
+ .collect(Collectors.toMap(e -> e.getKey().getId(), e -> e.getValue()));
+ }
+
+ static PCollectionTuple mapToTuple(Pipeline p, Map<String, PCollection<?>> map) {
+ PCollectionTuple tuple = PCollectionTuple.empty(p);
+ for (Map.Entry<String, PCollection<?>> entry : map.entrySet()) {
+ tuple = tuple.and(entry.getKey(), entry.getValue());
+ }
+ return tuple;
+ }
+
+ /**
+ * This is required as there is no apply() method on the base PInput type due to the inability to
+ * declare type parameters as self types.
+ */
+ private static <BeamInputT extends PInput, BeamOutputT extends POutput>
+ BeamOutputT applyTransform(
+ BeamInputT beamInput, PTransform<? super BeamInputT, BeamOutputT> transform) {
+ if (beamInput instanceof PCollection) {
+ return (BeamOutputT) ((PCollection) beamInput).apply(transform);
+ } else if (beamInput instanceof PCollectionTuple) {
+ return (BeamOutputT) ((PCollectionTuple) beamInput).apply((PTransform) transform);
+ } else if (beamInput instanceof PBegin) {
+ return (BeamOutputT) ((PBegin) beamInput).apply((PTransform) transform);
+ } else {
+ // We should never get here as we control the creation of all Beam types above.
+ // If new types of transform inputs are supported, this enumeration may need to be updated.
+ throw new IllegalArgumentException(
+ "Unknown Beam input type " + beamInput.getClass().getTypeName() + " for " + beamInput);
+ }
+ }
+}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/adapter/BeamFlinkDataSetAdapter.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/adapter/BeamFlinkDataSetAdapter.java
new file mode 100644
index 00000000000..93b3a8492a6
--- /dev/null
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/adapter/BeamFlinkDataSetAdapter.java
@@ -0,0 +1,265 @@
+/*
+ * 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.flink.adapter;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.function.BiFunction;
+import java.util.function.Function;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.runners.flink.FlinkBatchPortablePipelineTranslator;
+import org.apache.beam.runners.flink.FlinkPipelineOptions;
+import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.beam.sdk.util.construction.PipelineOptionsTranslation;
+import org.apache.beam.sdk.util.construction.graph.PipelineNode;
+import org.apache.beam.sdk.values.PBegin;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.PInput;
+import org.apache.beam.sdk.values.POutput;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
+import org.apache.flink.api.java.DataSet;
+import org.apache.flink.api.java.ExecutionEnvironment;
+import org.apache.flink.api.java.operators.MapOperator;
+
+/** An adapter class that allows one to apply Apache Beam PTransforms directly to Flink DataSets. */
+public class BeamFlinkDataSetAdapter {
+ private final PipelineOptions pipelineOptions;
+ private final CoderRegistry coderRegistry = CoderRegistry.createDefault();
+
+ public BeamFlinkDataSetAdapter() {
+ this(PipelineOptionsFactory.create());
+ }
+
+ public BeamFlinkDataSetAdapter(PipelineOptions pipelineOptions) {
+ this.pipelineOptions = pipelineOptions;
+ }
+
+ @SuppressWarnings("nullness")
+ public <InputT, OutputT, CollectionT extends PCollection<? extends InputT>>
+ DataSet<OutputT> applyBeamPTransform(
+ DataSet<InputT> input, PTransform<CollectionT, PCollection<OutputT>> transform) {
+ return (DataSet)
+ this.<CollectionT, PCollection<OutputT>>applyBeamPTransformInternal(
+ ImmutableMap.of("input", input),
+ (pipeline, map) -> (CollectionT) map.get("input"),
+ (output) -> ImmutableMap.of("output", output),
+ transform,
+ input.getExecutionEnvironment())
+ .get("output");
+ }
+
+ @SuppressWarnings("nullness")
+ public <OutputT> DataSet<OutputT> applyBeamPTransform(
+ Map<String, ? extends DataSet<?>> inputs,
+ PTransform<PCollectionTuple, PCollection<OutputT>> transform) {
+ return (DataSet)
+ applyBeamPTransformInternal(
+ inputs,
+ BeamAdapterUtils::mapToTuple,
+ (output) -> ImmutableMap.of("output", output),
+ transform,
+ inputs.values().stream().findAny().get().getExecutionEnvironment())
+ .get("output");
+ }
+
+ @SuppressWarnings("nullness")
+ public <OutputT> DataSet<OutputT> applyBeamPTransform(
+ ExecutionEnvironment executionEnvironment,
+ PTransform<PBegin, PCollection<OutputT>> transform) {
+ return (DataSet)
+ applyBeamPTransformInternal(
+ ImmutableMap.<String, DataSet<?>>of(),
+ (pipeline, map) -> PBegin.in(pipeline),
+ (output) -> ImmutableMap.of("output", output),
+ transform,
+ executionEnvironment)
+ .get("output");
+ }
+
+ @SuppressWarnings("nullness")
+ public <InputT, CollectionT extends PCollection<? extends InputT>>
+ Map<String, DataSet<?>> applyMultiOutputBeamPTransform(
+ DataSet<InputT> input, PTransform<CollectionT, PCollectionTuple> transform) {
+ return applyBeamPTransformInternal(
+ ImmutableMap.of("input", input),
+ (pipeline, map) -> (CollectionT) map.get("input"),
+ BeamAdapterUtils::tupleToMap,
+ transform,
+ input.getExecutionEnvironment());
+ }
+
+ public Map<String, DataSet<?>> applyMultiOutputBeamPTransform(
+ Map<String, ? extends DataSet<?>> inputs,
+ PTransform<PCollectionTuple, PCollectionTuple> transform) {
+ return applyBeamPTransformInternal(
+ inputs,
+ BeamAdapterUtils::mapToTuple,
+ BeamAdapterUtils::tupleToMap,
+ transform,
+ inputs.values().stream().findAny().get().getExecutionEnvironment());
+ }
+
+ public Map<String, DataSet<?>> applyMultiOutputBeamPTransform(
+ ExecutionEnvironment executionEnvironment, PTransform<PBegin, PCollectionTuple> transform) {
+ return applyBeamPTransformInternal(
+ ImmutableMap.of(),
+ (pipeline, map) -> PBegin.in(pipeline),
+ BeamAdapterUtils::tupleToMap,
+ transform,
+ executionEnvironment);
+ }
+
+ @SuppressWarnings("nullness")
+ public <InputT, CollectionT extends PCollection<? extends InputT>>
+ void applyNoOutputBeamPTransform(
+ DataSet<InputT> input, PTransform<CollectionT, PDone> transform) {
+ applyBeamPTransformInternal(
+ ImmutableMap.of("input", input),
+ (pipeline, map) -> (CollectionT) map.get("input"),
+ pDone -> ImmutableMap.of(),
+ transform,
+ input.getExecutionEnvironment());
+ }
+
+ public void applyNoOutputBeamPTransform(
+ Map<String, ? extends DataSet<?>> inputs, PTransform<PCollectionTuple, PDone> transform) {
+ applyBeamPTransformInternal(
+ inputs,
+ BeamAdapterUtils::mapToTuple,
+ pDone -> ImmutableMap.of(),
+ transform,
+ inputs.values().stream().findAny().get().getExecutionEnvironment());
+ }
+
+ public void applyNoOutputBeamPTransform(
+ ExecutionEnvironment executionEnvironment, PTransform<PBegin, PDone> transform) {
+ applyBeamPTransformInternal(
+ ImmutableMap.of(),
+ (pipeline, map) -> PBegin.in(pipeline),
+ pDone -> ImmutableMap.of(),
+ transform,
+ executionEnvironment);
+ }
+
+ private <BeamInputT extends PInput, BeamOutputT extends POutput>
+ Map<String, DataSet<?>> applyBeamPTransformInternal(
+ Map<String, ? extends DataSet<?>> inputs,
+ BiFunction<Pipeline, Map<String, PCollection<?>>, BeamInputT> toBeamInput,
+ Function<BeamOutputT, Map<String, PCollection<?>>> fromBeamOutput,
+ PTransform<? super BeamInputT, BeamOutputT> transform,
+ ExecutionEnvironment executionEnvironment) {
+ return BeamAdapterUtils.<DataSet<?>, BeamInputT, BeamOutputT>applyBeamPTransformInternal(
+ inputs,
+ toBeamInput,
+ fromBeamOutput,
+ transform,
+ executionEnvironment,
+ dataSet -> dataSet.getType(),
+ pipelineOptions,
+ coderRegistry,
+ new BeamAdapterUtils.PipelineFragmentTranslator<DataSet<?>>() {
+ @Override
+ public Map<String, DataSet<?>> translate(
+ Map<String, ? extends DataSet<?>> inputs,
+ RunnerApi.Pipeline pipelineProto,
+ ExecutionEnvironment executionEnvironment) {
+ Map<String, DataSet<?>> outputs = new HashMap<>();
+ FlinkBatchPortablePipelineTranslator translator =
+ FlinkBatchPortablePipelineTranslator.createTranslator(
+ ImmutableMap.of(
+ FlinkInput.URN, flinkInputTranslator(inputs),
+ FlinkOutput.URN, flinkOutputTranslator(outputs)));
+ FlinkBatchPortablePipelineTranslator.BatchTranslationContext context =
+ FlinkBatchPortablePipelineTranslator.createTranslationContext(
+ JobInfo.create(
+ "unusedJobId",
+ "unusedJobName",
+ "unusedRetrievalToken",
+ PipelineOptionsTranslation.toProto(pipelineOptions)),
+ pipelineOptions.as(FlinkPipelineOptions.class),
+ executionEnvironment);
+ translator.translate(context, translator.prepareForTranslation(pipelineProto));
+ return outputs;
+ }
+ });
+ }
+
+ private <InputT> FlinkBatchPortablePipelineTranslator.PTransformTranslator flinkInputTranslator(
+ Map<String, ? extends DataSet<?>> inputMap) {
+ return (PipelineNode.PTransformNode t,
+ RunnerApi.Pipeline p,
+ FlinkBatchPortablePipelineTranslator.BatchTranslationContext context) -> {
+ // When we run into a FlinkInput operator, it "produces" the corresponding input as its
+ // "computed result."
+ String inputId = t.getTransform().getSpec().getPayload().toStringUtf8();
+ DataSet<InputT> flinkInput = (DataSet<InputT>) inputMap.get(inputId);
+ // To make the nullness checker happy...
+ if (flinkInput == null) {
+ throw new IllegalStateException("Missing input: " + inputId);
+ }
+ context.addDataSet(
+ Iterables.getOnlyElement(t.getTransform().getOutputsMap().values()),
+ // new MapOperator(...) rather than .map to manually designate the type information.
+ // Note that MapOperator is a subclass of DataSet.
+ new MapOperator<InputT, WindowedValue<InputT>>(
+ flinkInput,
+ BeamAdapterCoderUtils.coderToTypeInformation(
+ WindowedValue.getValueOnlyCoder(
+ BeamAdapterCoderUtils.typeInformationToCoder(
+ flinkInput.getType(), coderRegistry)),
+ pipelineOptions),
+ x -> WindowedValue.valueInGlobalWindow(x),
+ "AddGlobalWindows"));
+ };
+ }
+
+ private <InputT> FlinkBatchPortablePipelineTranslator.PTransformTranslator flinkOutputTranslator(
+ Map<String, DataSet<?>> outputMap) {
+ return (PipelineNode.PTransformNode t,
+ RunnerApi.Pipeline p,
+ FlinkBatchPortablePipelineTranslator.BatchTranslationContext context) -> {
+ DataSet<WindowedValue<InputT>> inputDataSet =
+ context.getDataSetOrThrow(
+ Iterables.getOnlyElement(t.getTransform().getInputsMap().values()));
+ // When we run into a FlinkOutput operator, we cache the computed PCollection to return to the
+ // user.
+ String outputId = t.getTransform().getSpec().getPayload().toStringUtf8();
+ Coder<InputT> outputCoder =
+ BeamAdapterCoderUtils.lookupCoder(
+ p, Iterables.getOnlyElement(t.getTransform().getInputsMap().values()));
+ // TODO(robertwb): Also handle or disable length prefix coding (for embedded mode at least).
+ outputMap.put(
+ outputId,
+ new MapOperator<WindowedValue<InputT>, InputT>(
+ inputDataSet,
+ BeamAdapterCoderUtils.coderToTypeInformation(outputCoder, pipelineOptions),
+ w -> w.getValue(),
+ "StripWindows"));
+ };
+ }
+}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/adapter/FlinkInput.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/adapter/FlinkInput.java
new file mode 100644
index 00000000000..0b76319e9ea
--- /dev/null
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/adapter/FlinkInput.java
@@ -0,0 +1,102 @@
+/*
+ * 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.flink.adapter;
+
+import com.google.auto.service.AutoService;
+import java.io.IOException;
+import java.util.Collections;
+import java.util.Map;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.runners.AppliedPTransform;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.util.construction.NativeTransforms;
+import org.apache.beam.sdk.util.construction.PTransformTranslation;
+import org.apache.beam.sdk.util.construction.SdkComponents;
+import org.apache.beam.sdk.util.construction.TransformPayloadTranslatorRegistrar;
+import org.apache.beam.sdk.values.PBegin;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.WindowingStrategy;
+import org.apache.beam.vendor.grpc.v1p60p1.com.google.protobuf.ByteString;
+
+/**
+ * A Beam PTransform used for marking an input that comes from Flink.
+ *
+ * @param <T>
+ */
+/*package private*/ class FlinkInput<T> extends PTransform<PBegin, PCollection<T>> {
+ public static final String URN = "beam:flink:internal:translation_input";
+
+ private final String identifier;
+
+ private final Coder<T> coder;
+
+ FlinkInput(String identifier, Coder<T> coder) {
+ this.identifier = identifier;
+ this.coder = coder;
+ }
+
+ @Override
+ public PCollection<T> expand(PBegin input) {
+ return PCollection.createPrimitiveOutputInternal(
+ input.getPipeline(),
+ WindowingStrategy.globalDefault(),
+ PCollection.IsBounded.BOUNDED,
+ coder);
+ }
+
+ // This Translator translates this kind of PTransform into a Beam proto representation.
+ @SuppressWarnings("nullness")
+ @AutoService(TransformPayloadTranslatorRegistrar.class)
+ public static class Translator
+ implements PTransformTranslation.TransformPayloadTranslator<FlinkInput<?>>,
+ TransformPayloadTranslatorRegistrar {
+ @Override
+ public String getUrn() {
+ return URN;
+ }
+
+ @Override
+ public RunnerApi.FunctionSpec translate(
+ AppliedPTransform<?, ?, FlinkInput<?>> application, SdkComponents components)
+ throws IOException {
+ return RunnerApi.FunctionSpec.newBuilder()
+ .setUrn(FlinkInput.URN)
+ .setPayload(ByteString.copyFromUtf8(application.getTransform().identifier))
+ .build();
+ }
+
+ @Override
+ @SuppressWarnings("rawtypes")
+ public Map<
+ ? extends Class<? extends PTransform>,
+ ? extends PTransformTranslation.TransformPayloadTranslator>
+ getTransformPayloadTranslators() {
+ return Collections.singletonMap(FlinkInput.class, this);
+ }
+ }
+
+ @AutoService(NativeTransforms.IsNativeTransform.class)
+ public static class FlinkInputOutputIsNativeTransform
+ implements NativeTransforms.IsNativeTransform {
+ @Override
+ public boolean test(RunnerApi.PTransform pTransform) {
+ return FlinkInput.URN.equals(PTransformTranslation.urnForTransformOrNull(pTransform));
+ }
+ }
+}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/adapter/FlinkOutput.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/adapter/FlinkOutput.java
new file mode 100644
index 00000000000..2cc25236c1c
--- /dev/null
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/adapter/FlinkOutput.java
@@ -0,0 +1,93 @@
+/*
+ * 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.flink.adapter;
+
+import com.google.auto.service.AutoService;
+import java.io.IOException;
+import java.util.Collections;
+import java.util.Map;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+import org.apache.beam.sdk.runners.AppliedPTransform;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.util.construction.NativeTransforms;
+import org.apache.beam.sdk.util.construction.PTransformTranslation;
+import org.apache.beam.sdk.util.construction.SdkComponents;
+import org.apache.beam.sdk.util.construction.TransformPayloadTranslatorRegistrar;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.vendor.grpc.v1p60p1.com.google.protobuf.ByteString;
+
+/**
+ * A Beam PTransform used for marking an output that should be passed back to Flink.
+ *
+ * @param <T>
+ */
+/*package private*/ class FlinkOutput<T> extends PTransform<PCollection<T>, PDone> {
+
+ public static final String URN = "beam:flink:internal:translation_output";
+
+ private final String identifier;
+
+ FlinkOutput(String identifier) {
+ this.identifier = identifier;
+ }
+
+ @Override
+ public PDone expand(PCollection<T> input) {
+ return PDone.in(input.getPipeline());
+ }
+
+ @SuppressWarnings("nullness")
+ @AutoService(TransformPayloadTranslatorRegistrar.class)
+ public static class Translator
+ implements PTransformTranslation.TransformPayloadTranslator<FlinkOutput<?>>,
+ TransformPayloadTranslatorRegistrar {
+ @Override
+ public String getUrn() {
+ return FlinkOutput.URN;
+ }
+
+ @Override
+ public RunnerApi.FunctionSpec translate(
+ AppliedPTransform<?, ?, FlinkOutput<?>> application, SdkComponents components)
+ throws IOException {
+ return RunnerApi.FunctionSpec.newBuilder()
+ .setUrn(FlinkOutput.URN)
+ .setPayload(ByteString.copyFromUtf8(application.getTransform().identifier))
+ .build();
+ }
+
+ @Override
+ @SuppressWarnings("rawtypes")
+ public Map<
+ ? extends Class<? extends PTransform>,
+ ? extends PTransformTranslation.TransformPayloadTranslator>
+ getTransformPayloadTranslators() {
+ return Collections.singletonMap(FlinkOutput.class, this);
+ }
+ }
+
+ @AutoService(NativeTransforms.IsNativeTransform.class)
+ public static class FlinkInputOutputIsNativeTransform
+ implements NativeTransforms.IsNativeTransform {
+ @Override
+ public boolean test(RunnerApi.PTransform pTransform) {
+ return FlinkOutput.URN.equals(PTransformTranslation.urnForTransformOrNull(pTransform));
+ }
+ }
+}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/adapter/package-info.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/adapter/package-info.java
new file mode 100644
index 00000000000..5cfba3d03d4
--- /dev/null
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/adapter/package-info.java
@@ -0,0 +1,20 @@
+/*
+ * 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.
+ */
+
+/** Adaptors for using Beam transforms in Apache Flink pipelines. */
+package org.apache.beam.runners.flink.adapter;
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/adapter/BeamFlinkDataSetAdapterTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/adapter/BeamFlinkDataSetAdapterTest.java
new file mode 100644
index 00000000000..b69106e44f3
--- /dev/null
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/adapter/BeamFlinkDataSetAdapterTest.java
@@ -0,0 +1,135 @@
+/*
+ * 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.flink.adapter;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.containsInAnyOrder;
+
+import java.util.Map;
+import org.apache.beam.sdk.transforms.Count;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.Flatten;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionList;
+import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import org.apache.flink.api.java.DataSet;
+import org.apache.flink.api.java.ExecutionEnvironment;
+import org.junit.Test;
+
+public class BeamFlinkDataSetAdapterTest {
+
+ private static PTransform<PCollection<? extends String>, PCollection<String>> withPrefix(
+ String prefix) {
+ return ParDo.of(
+ new DoFn<String, String>() {
+ @ProcessElement
+ public void processElement(@Element String word, OutputReceiver<String> out) {
+ out.output(prefix + word);
+ }
+ });
+ }
+
+ @Test
+ public void testApplySimpleTransform() throws Exception {
+ ExecutionEnvironment env = ExecutionEnvironment.createCollectionsEnvironment();
+
+ DataSet<String> input = env.fromCollection(ImmutableList.of("a", "b", "c"));
+ DataSet<String> result =
+ new BeamFlinkDataSetAdapter().applyBeamPTransform(input, withPrefix("x"));
+
+ assertThat(result.collect(), containsInAnyOrder("xa", "xb", "xc"));
+ }
+
+ @Test
+ public void testApplyCompositeTransform() throws Exception {
+ ExecutionEnvironment env = ExecutionEnvironment.createCollectionsEnvironment();
+
+ DataSet<String> input = env.fromCollection(ImmutableList.of("a", "b", "c"));
+ DataSet<String> result =
+ new BeamFlinkDataSetAdapter()
+ .applyBeamPTransform(
+ input,
+ new PTransform<PCollection<String>, PCollection<String>>() {
+ @Override
+ public PCollection<String> expand(PCollection<String> input) {
+ return input.apply(withPrefix("x")).apply(withPrefix("y"));
+ }
+ });
+
+ assertThat(result.collect(), containsInAnyOrder("yxa", "yxb", "yxc"));
+ }
+
+ @Test
+ public void testApplyMultiInputTransform() throws Exception {
+ ExecutionEnvironment env = ExecutionEnvironment.createCollectionsEnvironment();
+
+ DataSet<String> input1 = env.fromCollection(ImmutableList.of("a", "b", "c"));
+ DataSet<String> input2 = env.fromCollection(ImmutableList.of("d", "e", "f"));
+ DataSet<String> result =
+ new BeamFlinkDataSetAdapter()
+ .applyBeamPTransform(
+ ImmutableMap.of("x", input1, "y", input2),
+ new PTransform<PCollectionTuple, PCollection<String>>() {
+ @Override
+ public PCollection<String> expand(PCollectionTuple input) {
+ return PCollectionList.of(input.<String>get("x").apply(withPrefix("x")))
+ .and(input.<String>get("y").apply(withPrefix("y")))
+ .apply(Flatten.pCollections());
+ }
+ });
+
+ assertThat(result.collect(), containsInAnyOrder("xa", "xb", "xc", "yd", "ye", "yf"));
+ }
+
+ @Test
+ public void testApplyMultiOutputTransform() throws Exception {
+ ExecutionEnvironment env = ExecutionEnvironment.createCollectionsEnvironment();
+
+ DataSet<String> input = env.fromCollection(ImmutableList.of("a", "b", "c"));
+ Map<String, DataSet<?>> result =
+ new BeamFlinkDataSetAdapter()
+ .applyMultiOutputBeamPTransform(
+ input,
+ new PTransform<PCollection<String>, PCollectionTuple>() {
+ @Override
+ public PCollectionTuple expand(PCollection<String> input) {
+ return PCollectionTuple.of("x", input.apply(withPrefix("x")))
+ .and("y", input.apply(withPrefix("y")));
+ }
+ });
+
+ assertThat(result.get("x").collect(), containsInAnyOrder("xa", "xb", "xc"));
+ assertThat(result.get("y").collect(), containsInAnyOrder("ya", "yb", "yc"));
+ }
+
+ @Test
+ public void testApplyGroupingTransform() throws Exception {
+ ExecutionEnvironment env = ExecutionEnvironment.createCollectionsEnvironment();
+
+ DataSet<String> input = env.fromCollection(ImmutableList.of("a", "a", "b"));
+ DataSet<KV<String, Long>> result =
+ new BeamFlinkDataSetAdapter().applyBeamPTransform(input, Count.perElement());
+
+ assertThat(result.collect(), containsInAnyOrder(KV.of("a", 2L), KV.of("b", 1L)));
+ }
+}