This is an automated email from the ASF dual-hosted git repository. guoweijie pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/flink.git
commit ceafa5a5705b6547eb8b55449de2c912a9a494f6 Author: Weijie Guo <[email protected]> AuthorDate: Fri Mar 22 17:13:22 2024 +0800 [FLINK-34548][API] Implement datastream --- .../datastream/impl/stream/AbstractDataStream.java | 83 +++++ .../impl/stream/BroadcastStreamImpl.java | 127 +++++++ .../datastream/impl/stream/GlobalStreamImpl.java | 177 ++++++++++ .../impl/stream/KeyedPartitionStreamImpl.java | 389 +++++++++++++++++++++ .../impl/stream/NonKeyedPartitionStreamImpl.java | 195 +++++++++++ .../flink/datastream/impl/utils/StreamUtils.java | 231 ++++++++++++ .../impl/stream/BroadcastStreamImplTest.java | 86 +++++ .../impl/stream/GlobalStreamImplTest.java | 75 ++++ .../impl/stream/KeyedPartitionStreamImplTest.java | 152 ++++++++ .../stream/NonKeyedPartitionStreamImplTest.java | 118 +++++++ .../datastream/impl/stream/StreamTestUtils.java | 119 +++++++ .../datastream/impl/utils/StreamUtilsTest.java | 197 +++++++++++ 12 files changed, 1949 insertions(+) diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/AbstractDataStream.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/AbstractDataStream.java new file mode 100644 index 00000000000..6c067d6602d --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/AbstractDataStream.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.flink.datastream.impl.stream; + +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.dag.Transformation; +import org.apache.flink.datastream.api.stream.DataStream; +import org.apache.flink.datastream.impl.ExecutionEnvironmentImpl; +import org.apache.flink.streaming.api.transformations.SideOutputTransformation; +import org.apache.flink.util.OutputTag; +import org.apache.flink.util.Preconditions; + +import java.util.HashMap; +import java.util.Map; + +/** + * Base class for all streams. + * + * <p>Note: This is only used for internal implementation. It must not leak to user face api. + */ +public abstract class AbstractDataStream<T> implements DataStream { + protected final ExecutionEnvironmentImpl environment; + + protected final Transformation<T> transformation; + + /** + * We keep track of the side outputs that were already requested and their types. With this, we + * can catch the case when a side output with a matching id is requested for a different type + * because this would lead to problems at runtime. + */ + protected final Map<OutputTag<?>, TypeInformation<?>> requestedSideOutputs = new HashMap<>(); + + public AbstractDataStream( + ExecutionEnvironmentImpl environment, Transformation<T> transformation) { + this.environment = + Preconditions.checkNotNull(environment, "Execution Environment must not be null."); + this.transformation = + Preconditions.checkNotNull( + transformation, "Stream Transformation must not be null."); + } + + public TypeInformation<T> getType() { + return transformation.getOutputType(); + } + + /** This is only used for internal implementation. It must not leak to user face api. */ + public Transformation<T> getTransformation() { + return transformation; + } + + public ExecutionEnvironmentImpl getEnvironment() { + return environment; + } + + public <X> Transformation<X> getSideOutputTransform(OutputTag<X> outputTag) { + TypeInformation<?> type = requestedSideOutputs.get(outputTag); + if (type != null && !type.equals(outputTag.getTypeInfo())) { + throw new UnsupportedOperationException( + "A side output with a matching id was " + + "already requested with a different type. This is not allowed, side output " + + "ids need to be unique."); + } + requestedSideOutputs.put(outputTag, outputTag.getTypeInfo()); + + return new SideOutputTransformation<>(getTransformation(), outputTag); + } +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/BroadcastStreamImpl.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/BroadcastStreamImpl.java new file mode 100644 index 00000000000..cf80e92f85e --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/BroadcastStreamImpl.java @@ -0,0 +1,127 @@ +/* + * 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.flink.datastream.impl.stream; + +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.dag.Transformation; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.api.java.typeutils.TypeExtractor; +import org.apache.flink.datastream.api.function.TwoInputBroadcastStreamProcessFunction; +import org.apache.flink.datastream.api.stream.BroadcastStream; +import org.apache.flink.datastream.api.stream.KeyedPartitionStream; +import org.apache.flink.datastream.api.stream.NonKeyedPartitionStream; +import org.apache.flink.datastream.impl.ExecutionEnvironmentImpl; +import org.apache.flink.datastream.impl.operators.KeyedTwoInputBroadcastProcessOperator; +import org.apache.flink.datastream.impl.operators.TwoInputBroadcastProcessOperator; +import org.apache.flink.datastream.impl.utils.StreamUtils; +import org.apache.flink.streaming.api.transformations.PartitionTransformation; +import org.apache.flink.streaming.runtime.partitioner.BroadcastPartitioner; + +/** The implementation of {@link BroadcastStream}. */ +public class BroadcastStreamImpl<T> extends AbstractDataStream<T> implements BroadcastStream<T> { + public BroadcastStreamImpl( + ExecutionEnvironmentImpl environment, Transformation<T> transformation) { + this( + environment, + new PartitionTransformation<>(transformation, new BroadcastPartitioner<>())); + } + + private BroadcastStreamImpl( + ExecutionEnvironmentImpl environment, PartitionTransformation<T> transformation) { + super(environment, transformation); + } + + @Override + public <K, T_OTHER, OUT> NonKeyedPartitionStream<OUT> connectAndProcess( + KeyedPartitionStream<K, T_OTHER> other, + TwoInputBroadcastStreamProcessFunction<T_OTHER, T, OUT> processFunction) { + TypeInformation<OUT> outTypeInfo = + StreamUtils.getOutputTypeForTwoInputBroadcastProcessFunction( + processFunction, + ((KeyedPartitionStreamImpl<K, T_OTHER>) other).getType(), + getType()); + KeyedTwoInputBroadcastProcessOperator<K, T_OTHER, T, OUT> processOperator = + new KeyedTwoInputBroadcastProcessOperator<>(processFunction); + Transformation<OUT> outTransformation = + StreamUtils.getTwoInputTransformation( + "Broadcast-Keyed-TwoInput-Process", + (KeyedPartitionStreamImpl<K, T_OTHER>) other, + // we should always take the broadcast input as second input. + this, + outTypeInfo, + processOperator); + environment.addOperator(outTransformation); + return new NonKeyedPartitionStreamImpl<>(environment, outTransformation); + } + + @Override + public <T_OTHER, OUT> NonKeyedPartitionStream<OUT> connectAndProcess( + NonKeyedPartitionStream<T_OTHER> other, + TwoInputBroadcastStreamProcessFunction<T_OTHER, T, OUT> processFunction) { + TypeInformation<OUT> outTypeInfo = + StreamUtils.getOutputTypeForTwoInputBroadcastProcessFunction( + processFunction, + ((NonKeyedPartitionStreamImpl<T_OTHER>) other).getType(), + getType()); + TwoInputBroadcastProcessOperator<T_OTHER, T, OUT> processOperator = + new TwoInputBroadcastProcessOperator<>(processFunction); + Transformation<OUT> outTransformation = + StreamUtils.getTwoInputTransformation( + "Broadcast-TwoInput-Process", + (NonKeyedPartitionStreamImpl<T_OTHER>) other, + // we should always take the broadcast input as second input. + this, + outTypeInfo, + processOperator); + environment.addOperator(outTransformation); + return new NonKeyedPartitionStreamImpl<>(environment, outTransformation); + } + + @Override + public <K, T_OTHER, OUT> KeyedPartitionStream<K, OUT> connectAndProcess( + KeyedPartitionStream<K, T_OTHER> other, + TwoInputBroadcastStreamProcessFunction<T_OTHER, T, OUT> processFunction, + KeySelector<OUT, K> newKeySelector) { + TypeInformation<OUT> outTypeInfo = + StreamUtils.getOutputTypeForTwoInputBroadcastProcessFunction( + processFunction, + ((KeyedPartitionStreamImpl<K, T_OTHER>) other).getType(), + getType()); + KeyedTwoInputBroadcastProcessOperator<K, T_OTHER, T, OUT> processOperator = + new KeyedTwoInputBroadcastProcessOperator<>(processFunction); + Transformation<OUT> outTransformation = + StreamUtils.getTwoInputTransformation( + "Broadcast-Keyed-TwoInput-Process", + (KeyedPartitionStreamImpl<K, T_OTHER>) other, + // we should always take the broadcast input as second input. + this, + outTypeInfo, + processOperator); + + NonKeyedPartitionStreamImpl<OUT> outputStream = + new NonKeyedPartitionStreamImpl<>(environment, outTransformation); + environment.addOperator(outTransformation); + // Construct a keyed stream directly without partitionTransformation to avoid shuffle. + return new KeyedPartitionStreamImpl<>( + outputStream, + outTransformation, + newKeySelector, + TypeExtractor.getKeySelectorTypes(newKeySelector, outputStream.getType())); + } +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/GlobalStreamImpl.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/GlobalStreamImpl.java new file mode 100644 index 00000000000..f899f04e3d5 --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/GlobalStreamImpl.java @@ -0,0 +1,177 @@ +/* + * 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.flink.datastream.impl.stream; + +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.dag.Transformation; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.api.java.tuple.Tuple2; +import org.apache.flink.datastream.api.function.OneInputStreamProcessFunction; +import org.apache.flink.datastream.api.function.TwoInputNonBroadcastStreamProcessFunction; +import org.apache.flink.datastream.api.function.TwoOutputStreamProcessFunction; +import org.apache.flink.datastream.api.stream.BroadcastStream; +import org.apache.flink.datastream.api.stream.GlobalStream; +import org.apache.flink.datastream.api.stream.KeyedPartitionStream; +import org.apache.flink.datastream.api.stream.NonKeyedPartitionStream; +import org.apache.flink.datastream.impl.ExecutionEnvironmentImpl; +import org.apache.flink.datastream.impl.operators.ProcessOperator; +import org.apache.flink.datastream.impl.operators.TwoInputNonBroadcastProcessOperator; +import org.apache.flink.datastream.impl.operators.TwoOutputProcessOperator; +import org.apache.flink.datastream.impl.utils.StreamUtils; +import org.apache.flink.streaming.api.operators.OneInputStreamOperator; +import org.apache.flink.streaming.api.operators.SimpleUdfStreamOperatorFactory; +import org.apache.flink.streaming.api.transformations.OneInputTransformation; +import org.apache.flink.streaming.api.transformations.PartitionTransformation; +import org.apache.flink.streaming.runtime.partitioner.ShufflePartitioner; +import org.apache.flink.util.OutputTag; + +/** The implementation of {@link GlobalStream}. */ +public class GlobalStreamImpl<T> extends AbstractDataStream<T> implements GlobalStream<T> { + public GlobalStreamImpl( + ExecutionEnvironmentImpl environment, Transformation<T> transformation) { + super(environment, transformation); + } + + @Override + public <OUT> GlobalStream<OUT> process(OneInputStreamProcessFunction<T, OUT> processFunction) { + TypeInformation<OUT> outType = + StreamUtils.getOutputTypeForOneInputProcessFunction(processFunction, getType()); + ProcessOperator<T, OUT> operator = new ProcessOperator<>(processFunction); + return transform("Global Process", outType, operator); + } + + @Override + public <OUT1, OUT2> TwoGlobalStreams<OUT1, OUT2> process( + TwoOutputStreamProcessFunction<T, OUT1, OUT2> processFunction) { + Tuple2<TypeInformation<OUT1>, TypeInformation<OUT2>> twoOutputType = + StreamUtils.getOutputTypesForTwoOutputProcessFunction(processFunction, getType()); + TypeInformation<OUT1> firstOutputType = twoOutputType.f0; + TypeInformation<OUT2> secondOutputType = twoOutputType.f1; + OutputTag<OUT2> secondOutputTag = new OutputTag<OUT2>("Second-Output", secondOutputType); + + TwoOutputProcessOperator<T, OUT1, OUT2> operator = + new TwoOutputProcessOperator<>(processFunction, secondOutputTag); + GlobalStreamImpl<OUT1> firstStream = + transform("Two-Output-Operator", firstOutputType, operator); + GlobalStreamImpl<OUT2> secondStream = + new GlobalStreamImpl<>( + environment, firstStream.getSideOutputTransform(secondOutputTag)); + return TwoGlobalStreamsImpl.of(firstStream, secondStream); + } + + @Override + public <T_OTHER, OUT> GlobalStream<OUT> connectAndProcess( + GlobalStream<T_OTHER> other, + TwoInputNonBroadcastStreamProcessFunction<T, T_OTHER, OUT> processFunction) { + TypeInformation<OUT> outTypeInfo = + StreamUtils.getOutputTypeForTwoInputNonBroadcastProcessFunction( + processFunction, getType(), ((GlobalStreamImpl<T_OTHER>) other).getType()); + TwoInputNonBroadcastProcessOperator<T, T_OTHER, OUT> processOperator = + new TwoInputNonBroadcastProcessOperator<>(processFunction); + Transformation<OUT> outTransformation = + StreamUtils.getTwoInputTransformation( + "Global-Global-TwoInput-Process", + this, + (GlobalStreamImpl<T_OTHER>) other, + outTypeInfo, + processOperator); + // Operator parallelism should always be 1 for global stream. + // parallelismConfigured should be true to avoid overwritten by AdaptiveBatchScheduler. + outTransformation.setParallelism(1, true); + environment.addOperator(outTransformation); + return new GlobalStreamImpl<>(environment, outTransformation); + } + + // TODO add toSink method. + + // --------------------- + // Partitioning + // --------------------- + + @Override + public <K> KeyedPartitionStream<K, T> keyBy(KeySelector<T, K> keySelector) { + return new KeyedPartitionStreamImpl<>(this, keySelector); + } + + @Override + public NonKeyedPartitionStream<T> shuffle() { + return new NonKeyedPartitionStreamImpl<>( + environment, + new PartitionTransformation<>(getTransformation(), new ShufflePartitioner<>())); + } + + @Override + public BroadcastStream<T> broadcast() { + return new BroadcastStreamImpl<>(environment, getTransformation()); + } + + private <R> GlobalStreamImpl<R> transform( + String operatorName, + TypeInformation<R> outputTypeInfo, + OneInputStreamOperator<T, R> operator) { + // read the output type of the input Transform to coax out errors about MissingTypeInfo + transformation.getOutputType(); + + OneInputTransformation<T, R> resultTransform = + new OneInputTransformation<>( + this.transformation, + operatorName, + SimpleUdfStreamOperatorFactory.of(operator), + outputTypeInfo, + // Operator parallelism should always be 1 for global stream. + 1, + // parallelismConfigured should be true to avoid overwritten by + // AdaptiveBatchScheduler. + true); + + GlobalStreamImpl<R> returnStream = new GlobalStreamImpl<>(environment, resultTransform); + + environment.addOperator(resultTransform); + + return returnStream; + } + + private static class TwoGlobalStreamsImpl<OUT1, OUT2> implements TwoGlobalStreams<OUT1, OUT2> { + + private final GlobalStream<OUT1> firstStream; + + private final GlobalStream<OUT2> secondStream; + + public static <OUT1, OUT2> TwoGlobalStreamsImpl<OUT1, OUT2> of( + GlobalStreamImpl<OUT1> firstStream, GlobalStreamImpl<OUT2> secondStream) { + return new TwoGlobalStreamsImpl<>(firstStream, secondStream); + } + + private TwoGlobalStreamsImpl( + GlobalStreamImpl<OUT1> firstStream, GlobalStreamImpl<OUT2> secondStream) { + this.firstStream = firstStream; + this.secondStream = secondStream; + } + + @Override + public GlobalStream<OUT1> getFirst() { + return firstStream; + } + + @Override + public GlobalStream<OUT2> getSecond() { + return secondStream; + } + } +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/KeyedPartitionStreamImpl.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/KeyedPartitionStreamImpl.java new file mode 100644 index 00000000000..c53c28a85e1 --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/KeyedPartitionStreamImpl.java @@ -0,0 +1,389 @@ +/* + * 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.flink.datastream.impl.stream; + +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.dag.Transformation; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.api.java.tuple.Tuple2; +import org.apache.flink.api.java.typeutils.TypeExtractor; +import org.apache.flink.datastream.api.function.OneInputStreamProcessFunction; +import org.apache.flink.datastream.api.function.TwoInputBroadcastStreamProcessFunction; +import org.apache.flink.datastream.api.function.TwoInputNonBroadcastStreamProcessFunction; +import org.apache.flink.datastream.api.function.TwoOutputStreamProcessFunction; +import org.apache.flink.datastream.api.stream.BroadcastStream; +import org.apache.flink.datastream.api.stream.GlobalStream; +import org.apache.flink.datastream.api.stream.KeyedPartitionStream; +import org.apache.flink.datastream.api.stream.NonKeyedPartitionStream; +import org.apache.flink.datastream.api.stream.NonKeyedPartitionStream.TwoNonKeyedPartitionStreams; +import org.apache.flink.datastream.impl.operators.KeyedProcessOperator; +import org.apache.flink.datastream.impl.operators.KeyedTwoInputBroadcastProcessOperator; +import org.apache.flink.datastream.impl.operators.KeyedTwoInputNonBroadcastProcessOperator; +import org.apache.flink.datastream.impl.operators.KeyedTwoOutputProcessOperator; +import org.apache.flink.datastream.impl.stream.NonKeyedPartitionStreamImpl.TwoNonKeyedPartitionStreamsImpl; +import org.apache.flink.datastream.impl.utils.StreamUtils; +import org.apache.flink.streaming.api.graph.StreamGraphGenerator; +import org.apache.flink.streaming.api.transformations.PartitionTransformation; +import org.apache.flink.streaming.runtime.partitioner.GlobalPartitioner; +import org.apache.flink.streaming.runtime.partitioner.KeyGroupStreamPartitioner; +import org.apache.flink.streaming.runtime.partitioner.ShufflePartitioner; +import org.apache.flink.util.OutputTag; + +import static org.apache.flink.util.Preconditions.checkNotNull; + +/** The implementation of {@link KeyedPartitionStream}. */ +public class KeyedPartitionStreamImpl<K, V> extends AbstractDataStream<V> + implements KeyedPartitionStream<K, V> { + + /** + * The key selector that can get the key by which the stream if partitioned from the elements. + */ + private final KeySelector<V, K> keySelector; + + /** The type of the key by which the stream is partitioned. */ + private final TypeInformation<K> keyType; + + public KeyedPartitionStreamImpl( + AbstractDataStream<V> dataStream, KeySelector<V, K> keySelector) { + this( + dataStream, + keySelector, + TypeExtractor.getKeySelectorTypes(keySelector, dataStream.getType())); + } + + public KeyedPartitionStreamImpl( + AbstractDataStream<V> dataStream, + KeySelector<V, K> keySelector, + TypeInformation<K> keyType) { + this( + dataStream, + new PartitionTransformation<>( + dataStream.getTransformation(), + new KeyGroupStreamPartitioner<>( + keySelector, + StreamGraphGenerator.DEFAULT_LOWER_BOUND_MAX_PARALLELISM)), + keySelector, + keyType); + } + + /** + * This can construct a keyed stream directly without partitionTransformation to avoid shuffle. + */ + public KeyedPartitionStreamImpl( + AbstractDataStream<V> dataStream, + Transformation<V> partitionTransformation, + KeySelector<V, K> keySelector, + TypeInformation<K> keyType) { + super(dataStream.getEnvironment(), partitionTransformation); + this.keySelector = keySelector; + this.keyType = keyType; + } + + @Override + public <OUT> NonKeyedPartitionStream<OUT> process( + OneInputStreamProcessFunction<V, OUT> processFunction) { + TypeInformation<OUT> outType; + outType = StreamUtils.getOutputTypeForOneInputProcessFunction(processFunction, getType()); + + KeyedProcessOperator<K, V, OUT> operator = new KeyedProcessOperator<>(processFunction); + Transformation<OUT> transform = + StreamUtils.getOneInputKeyedTransformation( + "KeyedProcess", this, outType, operator, keySelector, keyType); + environment.addOperator(transform); + return new NonKeyedPartitionStreamImpl<>(environment, transform); + } + + @Override + public <OUT> KeyedPartitionStream<K, OUT> process( + OneInputStreamProcessFunction<V, OUT> processFunction, + KeySelector<OUT, K> newKeySelector) { + TypeInformation<OUT> outType = + StreamUtils.getOutputTypeForOneInputProcessFunction(processFunction, getType()); + KeyedProcessOperator<K, V, OUT> operator = + new KeyedProcessOperator<>(processFunction, checkNotNull(newKeySelector)); + Transformation<OUT> transform = + StreamUtils.getOneInputKeyedTransformation( + "KeyedProcess", this, outType, operator, keySelector, keyType); + NonKeyedPartitionStreamImpl<OUT> outputStream = + new NonKeyedPartitionStreamImpl<>(environment, transform); + environment.addOperator(transform); + // Construct a keyed stream directly without partitionTransformation to avoid shuffle. + return new KeyedPartitionStreamImpl<>( + outputStream, + transform, + newKeySelector, + TypeExtractor.getKeySelectorTypes(newKeySelector, outputStream.getType())); + } + + @Override + public <OUT1, OUT2> TwoKeyedPartitionStreams<K, OUT1, OUT2> process( + TwoOutputStreamProcessFunction<V, OUT1, OUT2> processFunction, + KeySelector<OUT1, K> keySelector1, + KeySelector<OUT2, K> keySelector2) { + Tuple2<TypeInformation<OUT1>, TypeInformation<OUT2>> twoOutputType = + StreamUtils.getOutputTypesForTwoOutputProcessFunction(processFunction, getType()); + TypeInformation<OUT1> firstOutputType = twoOutputType.f0; + TypeInformation<OUT2> secondOutputType = twoOutputType.f1; + OutputTag<OUT2> secondOutputTag = new OutputTag<>("Second-Output", secondOutputType); + + KeyedTwoOutputProcessOperator<K, V, OUT1, OUT2> operator = + new KeyedTwoOutputProcessOperator<>( + processFunction, secondOutputTag, keySelector1, keySelector2); + Transformation<OUT1> mainOutputTransform = + StreamUtils.getOneInputKeyedTransformation( + "Two-Output-Process", + this, + firstOutputType, + operator, + keySelector, + keyType); + NonKeyedPartitionStreamImpl<OUT1> nonKeyedMainOutputStream = + new NonKeyedPartitionStreamImpl<>(environment, mainOutputTransform); + Transformation<OUT2> sideOutputTransform = + nonKeyedMainOutputStream.getSideOutputTransform(secondOutputTag); + NonKeyedPartitionStreamImpl<OUT2> nonKeyedSideStream = + new NonKeyedPartitionStreamImpl<>(environment, sideOutputTransform); + + // Construct a keyed stream directly without partitionTransformation to avoid shuffle. + KeyedPartitionStreamImpl<K, OUT1> keyedMainOutputStream = + new KeyedPartitionStreamImpl<>( + nonKeyedMainOutputStream, + mainOutputTransform, + keySelector1, + TypeExtractor.getKeySelectorTypes( + keySelector1, nonKeyedMainOutputStream.getType())); + KeyedPartitionStreamImpl<K, OUT2> keyedSideOutputStream = + new KeyedPartitionStreamImpl<>( + nonKeyedSideStream, + sideOutputTransform, + keySelector2, + TypeExtractor.getKeySelectorTypes( + keySelector2, nonKeyedSideStream.getType())); + environment.addOperator(mainOutputTransform); + return TwoKeyedPartitionStreamsImpl.of(keyedMainOutputStream, keyedSideOutputStream); + } + + @Override + public <OUT1, OUT2> TwoNonKeyedPartitionStreams<OUT1, OUT2> process( + TwoOutputStreamProcessFunction<V, OUT1, OUT2> processFunction) { + Tuple2<TypeInformation<OUT1>, TypeInformation<OUT2>> twoOutputType = + StreamUtils.getOutputTypesForTwoOutputProcessFunction(processFunction, getType()); + TypeInformation<OUT1> firstOutputType = twoOutputType.f0; + TypeInformation<OUT2> secondOutputType = twoOutputType.f1; + OutputTag<OUT2> secondOutputTag = new OutputTag<>("Second-Output", secondOutputType); + + KeyedTwoOutputProcessOperator<K, V, OUT1, OUT2> operator = + new KeyedTwoOutputProcessOperator<>(processFunction, secondOutputTag); + Transformation<OUT1> firstTransformation = + StreamUtils.getOneInputKeyedTransformation( + "Two-Output-Process", + this, + firstOutputType, + operator, + keySelector, + keyType); + NonKeyedPartitionStreamImpl<OUT1> firstStream = + new NonKeyedPartitionStreamImpl<>(environment, firstTransformation); + NonKeyedPartitionStreamImpl<OUT2> secondStream = + new NonKeyedPartitionStreamImpl<>( + environment, firstStream.getSideOutputTransform(secondOutputTag)); + environment.addOperator(firstTransformation); + return TwoNonKeyedPartitionStreamsImpl.of(firstStream, secondStream); + } + + @Override + public <T_OTHER, OUT> NonKeyedPartitionStream<OUT> connectAndProcess( + KeyedPartitionStream<K, T_OTHER> other, + TwoInputNonBroadcastStreamProcessFunction<V, T_OTHER, OUT> processFunction) { + TypeInformation<OUT> outTypeInfo = + StreamUtils.getOutputTypeForTwoInputNonBroadcastProcessFunction( + processFunction, + getType(), + ((KeyedPartitionStreamImpl<K, T_OTHER>) other).getType()); + + KeyedTwoInputNonBroadcastProcessOperator<K, V, T_OTHER, OUT> processOperator = + new KeyedTwoInputNonBroadcastProcessOperator<>(processFunction); + Transformation<OUT> outTransformation = + StreamUtils.getTwoInputTransformation( + "Keyed-TwoInput-Process", + this, + (KeyedPartitionStreamImpl<K, T_OTHER>) other, + outTypeInfo, + processOperator); + environment.addOperator(outTransformation); + return new NonKeyedPartitionStreamImpl<>(environment, outTransformation); + } + + @Override + public <T_OTHER, OUT> KeyedPartitionStream<K, OUT> connectAndProcess( + KeyedPartitionStream<K, T_OTHER> other, + TwoInputNonBroadcastStreamProcessFunction<V, T_OTHER, OUT> processFunction, + KeySelector<OUT, K> newKeySelector) { + TypeInformation<OUT> outTypeInfo = + StreamUtils.getOutputTypeForTwoInputNonBroadcastProcessFunction( + processFunction, + getType(), + ((KeyedPartitionStreamImpl<K, T_OTHER>) other).getType()); + + KeyedTwoInputNonBroadcastProcessOperator<K, V, T_OTHER, OUT> processOperator = + new KeyedTwoInputNonBroadcastProcessOperator<>(processFunction, newKeySelector); + Transformation<OUT> outTransformation = + StreamUtils.getTwoInputTransformation( + "Keyed-TwoInput-Process", + this, + (KeyedPartitionStreamImpl<K, T_OTHER>) other, + outTypeInfo, + processOperator); + NonKeyedPartitionStreamImpl<OUT> nonKeyedOutputStream = + new NonKeyedPartitionStreamImpl<>(environment, outTransformation); + environment.addOperator(outTransformation); + // Construct a keyed stream directly without partitionTransformation to avoid shuffle. + return new KeyedPartitionStreamImpl<>( + nonKeyedOutputStream, + outTransformation, + newKeySelector, + TypeExtractor.getKeySelectorTypes(newKeySelector, nonKeyedOutputStream.getType())); + } + + @Override + public <T_OTHER, OUT> NonKeyedPartitionStream<OUT> connectAndProcess( + BroadcastStream<T_OTHER> other, + TwoInputBroadcastStreamProcessFunction<V, T_OTHER, OUT> processFunction) { + TypeInformation<OUT> outTypeInfo = + StreamUtils.getOutputTypeForTwoInputBroadcastProcessFunction( + processFunction, + getType(), + ((BroadcastStreamImpl<T_OTHER>) other).getType()); + KeyedTwoInputBroadcastProcessOperator<K, V, T_OTHER, OUT> processOperator = + new KeyedTwoInputBroadcastProcessOperator<>(processFunction); + Transformation<OUT> outTransformation = + StreamUtils.getTwoInputTransformation( + "Keyed-TwoInput-Broadcast-Process", + this, + // we should always take the broadcast input as second input. + (BroadcastStreamImpl<T_OTHER>) other, + outTypeInfo, + processOperator); + environment.addOperator(outTransformation); + return new NonKeyedPartitionStreamImpl<>(environment, outTransformation); + } + + @Override + public <T_OTHER, OUT> KeyedPartitionStream<K, OUT> connectAndProcess( + BroadcastStream<T_OTHER> other, + TwoInputBroadcastStreamProcessFunction<V, T_OTHER, OUT> processFunction, + KeySelector<OUT, K> newKeySelector) { + TypeInformation<OUT> outTypeInfo = + StreamUtils.getOutputTypeForTwoInputBroadcastProcessFunction( + processFunction, + getType(), + ((BroadcastStreamImpl<T_OTHER>) other).getType()); + KeyedTwoInputBroadcastProcessOperator<K, V, T_OTHER, OUT> processOperator = + new KeyedTwoInputBroadcastProcessOperator<>( + processFunction, checkNotNull(newKeySelector)); + Transformation<OUT> outTransformation = + StreamUtils.getTwoInputTransformation( + "Keyed-TwoInput-Broadcast-Process", + this, + // we should always take the broadcast input as second input. + (BroadcastStreamImpl<T_OTHER>) other, + outTypeInfo, + processOperator); + + NonKeyedPartitionStreamImpl<OUT> outputStream = + new NonKeyedPartitionStreamImpl<>(environment, outTransformation); + environment.addOperator(outTransformation); + // Construct a keyed stream directly without partitionTransformation to avoid shuffle. + return new KeyedPartitionStreamImpl<>( + outputStream, + outTransformation, + newKeySelector, + TypeExtractor.getKeySelectorTypes(newKeySelector, outputStream.getType())); + } + + public TypeInformation<K> getKeyType() { + return keyType; + } + + public KeySelector<V, K> getKeySelector() { + return keySelector; + } + + // TODO add toSink method. + + // --------------------- + // Partitioning + // --------------------- + + @Override + public GlobalStream<V> global() { + return new GlobalStreamImpl<>( + environment, + new PartitionTransformation<>(transformation, new GlobalPartitioner<>())); + } + + @Override + public <NEW_KEY> KeyedPartitionStream<NEW_KEY, V> keyBy(KeySelector<V, NEW_KEY> keySelector) { + // Create a new keyed stream with different key selector. + return new KeyedPartitionStreamImpl<>(this, keySelector); + } + + @Override + public NonKeyedPartitionStream<V> shuffle() { + return new NonKeyedPartitionStreamImpl<>( + environment, + new PartitionTransformation<>(getTransformation(), new ShufflePartitioner<>())); + } + + @Override + public BroadcastStream<V> broadcast() { + return new BroadcastStreamImpl<>(environment, getTransformation()); + } + + static class TwoKeyedPartitionStreamsImpl<K, OUT1, OUT2> + implements TwoKeyedPartitionStreams<K, OUT1, OUT2> { + + private final KeyedPartitionStream<K, OUT1> firstStream; + + private final KeyedPartitionStream<K, OUT2> secondStream; + + public static <K, OUT1, OUT2> TwoKeyedPartitionStreamsImpl<K, OUT1, OUT2> of( + KeyedPartitionStreamImpl<K, OUT1> firstStream, + KeyedPartitionStreamImpl<K, OUT2> secondStream) { + return new TwoKeyedPartitionStreamsImpl<>(firstStream, secondStream); + } + + private TwoKeyedPartitionStreamsImpl( + KeyedPartitionStreamImpl<K, OUT1> firstStream, + KeyedPartitionStreamImpl<K, OUT2> secondStream) { + this.firstStream = firstStream; + this.secondStream = secondStream; + } + + @Override + public KeyedPartitionStream<K, OUT1> getFirst() { + return firstStream; + } + + @Override + public KeyedPartitionStream<K, OUT2> getSecond() { + return secondStream; + } + } +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/NonKeyedPartitionStreamImpl.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/NonKeyedPartitionStreamImpl.java new file mode 100644 index 00000000000..2bac6c0cd1e --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/stream/NonKeyedPartitionStreamImpl.java @@ -0,0 +1,195 @@ +/* + * 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.flink.datastream.impl.stream; + +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.dag.Transformation; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.api.java.tuple.Tuple2; +import org.apache.flink.datastream.api.function.OneInputStreamProcessFunction; +import org.apache.flink.datastream.api.function.TwoInputBroadcastStreamProcessFunction; +import org.apache.flink.datastream.api.function.TwoInputNonBroadcastStreamProcessFunction; +import org.apache.flink.datastream.api.function.TwoOutputStreamProcessFunction; +import org.apache.flink.datastream.api.stream.BroadcastStream; +import org.apache.flink.datastream.api.stream.GlobalStream; +import org.apache.flink.datastream.api.stream.KeyedPartitionStream; +import org.apache.flink.datastream.api.stream.NonKeyedPartitionStream; +import org.apache.flink.datastream.impl.ExecutionEnvironmentImpl; +import org.apache.flink.datastream.impl.operators.ProcessOperator; +import org.apache.flink.datastream.impl.operators.TwoInputBroadcastProcessOperator; +import org.apache.flink.datastream.impl.operators.TwoInputNonBroadcastProcessOperator; +import org.apache.flink.datastream.impl.operators.TwoOutputProcessOperator; +import org.apache.flink.datastream.impl.utils.StreamUtils; +import org.apache.flink.streaming.api.transformations.OneInputTransformation; +import org.apache.flink.streaming.api.transformations.PartitionTransformation; +import org.apache.flink.streaming.runtime.partitioner.GlobalPartitioner; +import org.apache.flink.streaming.runtime.partitioner.ShufflePartitioner; +import org.apache.flink.util.OutputTag; + +/** The implementation of {@link NonKeyedPartitionStream}. */ +public class NonKeyedPartitionStreamImpl<T> extends AbstractDataStream<T> + implements NonKeyedPartitionStream<T> { + public NonKeyedPartitionStreamImpl( + ExecutionEnvironmentImpl environment, Transformation<T> transformation) { + super(environment, transformation); + } + + @Override + public <OUT> NonKeyedPartitionStream<OUT> process( + OneInputStreamProcessFunction<T, OUT> processFunction) { + TypeInformation<OUT> outType = + StreamUtils.getOutputTypeForOneInputProcessFunction(processFunction, getType()); + ProcessOperator<T, OUT> operator = new ProcessOperator<>(processFunction); + OneInputTransformation<T, OUT> outputTransform = + StreamUtils.getOneInputTransformation("Process", this, outType, operator); + environment.addOperator(outputTransform); + return new NonKeyedPartitionStreamImpl<>(environment, outputTransform); + } + + @Override + public <OUT1, OUT2> TwoNonKeyedPartitionStreams<OUT1, OUT2> process( + TwoOutputStreamProcessFunction<T, OUT1, OUT2> processFunction) { + Tuple2<TypeInformation<OUT1>, TypeInformation<OUT2>> twoOutputType = + StreamUtils.getOutputTypesForTwoOutputProcessFunction(processFunction, getType()); + TypeInformation<OUT1> firstOutputType = twoOutputType.f0; + TypeInformation<OUT2> secondOutputType = twoOutputType.f1; + OutputTag<OUT2> secondOutputTag = new OutputTag<>("Second-Output", secondOutputType); + + TwoOutputProcessOperator<T, OUT1, OUT2> operator = + new TwoOutputProcessOperator<>(processFunction, secondOutputTag); + OneInputTransformation<T, OUT1> outTransformation = + StreamUtils.getOneInputTransformation( + "Two-Output-Operator", this, firstOutputType, operator); + NonKeyedPartitionStreamImpl<OUT1> firstStream = + new NonKeyedPartitionStreamImpl<>(environment, outTransformation); + NonKeyedPartitionStreamImpl<OUT2> secondStream = + new NonKeyedPartitionStreamImpl<>( + environment, firstStream.getSideOutputTransform(secondOutputTag)); + environment.addOperator(outTransformation); + return TwoNonKeyedPartitionStreamsImpl.of(firstStream, secondStream); + } + + @Override + public <T_OTHER, OUT> NonKeyedPartitionStream<OUT> connectAndProcess( + NonKeyedPartitionStream<T_OTHER> other, + TwoInputNonBroadcastStreamProcessFunction<T, T_OTHER, OUT> processFunction) { + TypeInformation<OUT> outTypeInfo = + StreamUtils.getOutputTypeForTwoInputNonBroadcastProcessFunction( + processFunction, + getType(), + ((NonKeyedPartitionStreamImpl<T_OTHER>) other).getType()); + + TwoInputNonBroadcastProcessOperator<T, T_OTHER, OUT> processOperator = + new TwoInputNonBroadcastProcessOperator<>(processFunction); + Transformation<OUT> outTransformation = + StreamUtils.getTwoInputTransformation( + "TwoInput-Process", + this, + (NonKeyedPartitionStreamImpl<T_OTHER>) other, + outTypeInfo, + processOperator); + environment.addOperator(outTransformation); + return new NonKeyedPartitionStreamImpl<>(environment, outTransformation); + } + + @Override + public <T_OTHER, OUT> NonKeyedPartitionStream<OUT> connectAndProcess( + BroadcastStream<T_OTHER> other, + TwoInputBroadcastStreamProcessFunction<T, T_OTHER, OUT> processFunction) { + TypeInformation<OUT> outTypeInfo = + StreamUtils.getOutputTypeForTwoInputBroadcastProcessFunction( + processFunction, + getType(), + ((BroadcastStreamImpl<T_OTHER>) other).getType()); + + TwoInputBroadcastProcessOperator<T, T_OTHER, OUT> processOperator = + new TwoInputBroadcastProcessOperator<>(processFunction); + Transformation<OUT> outTransformation = + StreamUtils.getTwoInputTransformation( + "Broadcast-TwoInput-Process", + this, + // we should always take the broadcast input as second input. + (BroadcastStreamImpl<T_OTHER>) other, + outTypeInfo, + processOperator); + environment.addOperator(outTransformation); + return new NonKeyedPartitionStreamImpl<>(environment, outTransformation); + } + + // TODO add toSink method. + + // --------------------- + // Partitioning + // --------------------- + + @Override + public GlobalStream<T> global() { + return new GlobalStreamImpl<>( + environment, + new PartitionTransformation<>(transformation, new GlobalPartitioner<>())); + } + + @Override + public <K> KeyedPartitionStream<K, T> keyBy(KeySelector<T, K> keySelector) { + return new KeyedPartitionStreamImpl<>(this, keySelector); + } + + @Override + public NonKeyedPartitionStream<T> shuffle() { + return new NonKeyedPartitionStreamImpl<>( + environment, + new PartitionTransformation<>(getTransformation(), new ShufflePartitioner<>())); + } + + @Override + public BroadcastStream<T> broadcast() { + return new BroadcastStreamImpl<>(environment, getTransformation()); + } + + static class TwoNonKeyedPartitionStreamsImpl<OUT1, OUT2> + implements TwoNonKeyedPartitionStreams<OUT1, OUT2> { + + private final NonKeyedPartitionStream<OUT1> firstStream; + + private final NonKeyedPartitionStream<OUT2> secondStream; + + public static <OUT1, OUT2> TwoNonKeyedPartitionStreamsImpl<OUT1, OUT2> of( + NonKeyedPartitionStreamImpl<OUT1> firstStream, + NonKeyedPartitionStreamImpl<OUT2> secondStream) { + return new TwoNonKeyedPartitionStreamsImpl<>(firstStream, secondStream); + } + + private TwoNonKeyedPartitionStreamsImpl( + NonKeyedPartitionStreamImpl<OUT1> firstStream, + NonKeyedPartitionStreamImpl<OUT2> secondStream) { + this.firstStream = firstStream; + this.secondStream = secondStream; + } + + @Override + public NonKeyedPartitionStream<OUT1> getFirst() { + return firstStream; + } + + @Override + public NonKeyedPartitionStream<OUT2> getSecond() { + return secondStream; + } + } +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/utils/StreamUtils.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/utils/StreamUtils.java new file mode 100644 index 00000000000..0471c0f9a0a --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/utils/StreamUtils.java @@ -0,0 +1,231 @@ +/* + * 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.flink.datastream.impl.utils; + +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.java.Utils; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.api.java.tuple.Tuple2; +import org.apache.flink.api.java.typeutils.TypeExtractor; +import org.apache.flink.datastream.api.function.OneInputStreamProcessFunction; +import org.apache.flink.datastream.api.function.TwoInputBroadcastStreamProcessFunction; +import org.apache.flink.datastream.api.function.TwoInputNonBroadcastStreamProcessFunction; +import org.apache.flink.datastream.api.function.TwoOutputStreamProcessFunction; +import org.apache.flink.datastream.impl.stream.AbstractDataStream; +import org.apache.flink.datastream.impl.stream.KeyedPartitionStreamImpl; +import org.apache.flink.streaming.api.operators.OneInputStreamOperator; +import org.apache.flink.streaming.api.operators.SimpleOperatorFactory; +import org.apache.flink.streaming.api.operators.SimpleUdfStreamOperatorFactory; +import org.apache.flink.streaming.api.operators.TwoInputStreamOperator; +import org.apache.flink.streaming.api.transformations.OneInputTransformation; +import org.apache.flink.streaming.api.transformations.TwoInputTransformation; + +/** + * This class encapsulates the common logic for all type of streams. It can be used to handle things + * like extract type information, create a new transformation and so on for AbstractDataStream. + */ +public final class StreamUtils { + /** + * Get the output type information for {@link OneInputStreamProcessFunction} from input type + * information. + */ + public static <IN, OUT> TypeInformation<OUT> getOutputTypeForOneInputProcessFunction( + OneInputStreamProcessFunction<IN, OUT> processFunction, + TypeInformation<IN> inTypeInformation) { + return TypeExtractor.getUnaryOperatorReturnType( + processFunction, + OneInputStreamProcessFunction.class, + 0, + 1, + new int[] {1, 0}, + inTypeInformation, + Utils.getCallLocationName(), + true); + } + + /** + * Get the output type information for {@link TwoInputNonBroadcastStreamProcessFunction} from + * two input type information . + */ + public static <IN1, IN2, OUT> + TypeInformation<OUT> getOutputTypeForTwoInputNonBroadcastProcessFunction( + TwoInputNonBroadcastStreamProcessFunction<IN1, IN2, OUT> processFunction, + TypeInformation<IN1> in1TypeInformation, + TypeInformation<IN2> in2TypeInformation) { + return TypeExtractor.getBinaryOperatorReturnType( + processFunction, + TwoInputNonBroadcastStreamProcessFunction.class, + 0, + 1, + 2, + TypeExtractor.NO_INDEX, + in1TypeInformation, + in2TypeInformation, + Utils.getCallLocationName(), + true); + } + + /** + * Get the output type information for {@link TwoInputBroadcastStreamProcessFunction} from two + * input type information . + */ + public static <IN1, IN2, OUT> + TypeInformation<OUT> getOutputTypeForTwoInputBroadcastProcessFunction( + TwoInputBroadcastStreamProcessFunction<IN1, IN2, OUT> processFunction, + TypeInformation<IN1> in1TypeInformation, + TypeInformation<IN2> in2TypeInformation) { + return TypeExtractor.getBinaryOperatorReturnType( + processFunction, + TwoInputBroadcastStreamProcessFunction.class, + 0, + 1, + 2, + TypeExtractor.NO_INDEX, + in1TypeInformation, + in2TypeInformation, + Utils.getCallLocationName(), + true); + } + + /** + * Get output types information for {@link TwoOutputStreamProcessFunction} from the input type + * information. + */ + public static <IN, OUT1, OUT2> + Tuple2<TypeInformation<OUT1>, TypeInformation<OUT2>> + getOutputTypesForTwoOutputProcessFunction( + TwoOutputStreamProcessFunction<IN, OUT1, OUT2> + twoOutputStreamProcessFunction, + TypeInformation<IN> inTypeInformation) { + TypeInformation<OUT1> firstOutputType = + TypeExtractor.getUnaryOperatorReturnType( + twoOutputStreamProcessFunction, + TwoOutputStreamProcessFunction.class, + 0, + 1, + new int[] {1, 0}, + inTypeInformation, + Utils.getCallLocationName(), + true); + TypeInformation<OUT2> secondOutputType = + TypeExtractor.getUnaryOperatorReturnType( + twoOutputStreamProcessFunction, + TwoOutputStreamProcessFunction.class, + 0, + 2, + new int[] {2, 0}, + inTypeInformation, + Utils.getCallLocationName(), + true); + return Tuple2.of(firstOutputType, secondOutputType); + } + + /** Construct and return a {@link OneInputTransformation} from non-keyed input streams. */ + public static <T, R> OneInputTransformation<T, R> getOneInputTransformation( + String operatorName, + AbstractDataStream<T> inputStream, + TypeInformation<R> outTypeInformation, + OneInputStreamOperator<T, R> operator) { + // read the output type of the input Transform to coax out errors about MissingTypeInfo + inputStream.getTransformation().getOutputType(); + + OneInputTransformation<T, R> resultTransform = + new OneInputTransformation<>( + inputStream.getTransformation(), + operatorName, + SimpleUdfStreamOperatorFactory.of(operator), + outTypeInformation, + inputStream.getEnvironment().getParallelism(), + false); + + return resultTransform; + } + + /** Construct and return a {@link OneInputTransformation} from keyed input streams. */ + public static <T, R, K> OneInputTransformation<T, R> getOneInputKeyedTransformation( + String operatorName, + AbstractDataStream<T> inputStream, + TypeInformation<R> outTypeInformation, + OneInputStreamOperator<T, R> operator, + KeySelector<T, K> keySelector, + TypeInformation<K> keyType) { + OneInputTransformation<T, R> resultTransform = + getOneInputTransformation(operatorName, inputStream, outTypeInformation, operator); + + // inject the key selector and key type + resultTransform.setStateKeySelector(keySelector); + resultTransform.setStateKeyType(keyType); + + return resultTransform; + } + + /** Construct and return a {@link TwoInputTransformation} from two input streams. */ + public static <IN1, IN2, OUT> TwoInputTransformation<IN1, IN2, OUT> getTwoInputTransformation( + String operatorName, + AbstractDataStream<IN1> inputStream1, + AbstractDataStream<IN2> inputStream2, + TypeInformation<OUT> outTypeInformation, + TwoInputStreamOperator<IN1, IN2, OUT> operator) { + TwoInputTransformation<IN1, IN2, OUT> transform = + new TwoInputTransformation<>( + inputStream1.getTransformation(), + inputStream2.getTransformation(), + operatorName, + SimpleOperatorFactory.of(operator), + outTypeInformation, + // inputStream1 & 2 share the same env. + inputStream1.getEnvironment().getParallelism(), + false); + + TypeInformation<?> keyType = null; + if (inputStream1 instanceof KeyedPartitionStreamImpl) { + KeyedPartitionStreamImpl<?, IN1> keyedInput1 = + (KeyedPartitionStreamImpl<?, IN1>) inputStream1; + + keyType = keyedInput1.getKeyType(); + + transform.setStateKeySelectors(keyedInput1.getKeySelector(), null); + transform.setStateKeyType(keyType); + } + if (inputStream2 instanceof KeyedPartitionStreamImpl) { + KeyedPartitionStreamImpl<?, IN2> keyedInput2 = + (KeyedPartitionStreamImpl<?, IN2>) inputStream2; + + TypeInformation<?> keyType2 = keyedInput2.getKeyType(); + + if (keyType != null && !(keyType.canEqual(keyType2) && keyType.equals(keyType2))) { + throw new UnsupportedOperationException( + "Key types if input KeyedStreams " + + "don't match: " + + keyType + + " and " + + keyType2 + + "."); + } + + transform.setStateKeySelectors( + transform.getStateKeySelector1(), keyedInput2.getKeySelector()); + + // we might be overwriting the one that's already set, but it's the same + transform.setStateKeyType(keyType2); + } + + return transform; + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/BroadcastStreamImplTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/BroadcastStreamImplTest.java new file mode 100644 index 00000000000..0bd867276bc --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/BroadcastStreamImplTest.java @@ -0,0 +1,86 @@ +/* + * 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.flink.datastream.impl.stream; + +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.dag.Transformation; +import org.apache.flink.datastream.api.stream.KeyedPartitionStream; +import org.apache.flink.datastream.impl.ExecutionEnvironmentImpl; +import org.apache.flink.datastream.impl.TestingTransformation; +import org.apache.flink.streaming.api.transformations.TwoInputTransformation; + +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static org.apache.flink.datastream.impl.stream.StreamTestUtils.assertProcessType; +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link BroadcastStreamImpl}. */ +class BroadcastStreamImplTest { + @Test + void testConnectNonKeyedStream() throws Exception { + ExecutionEnvironmentImpl env = StreamTestUtils.getEnv(); + BroadcastStreamImpl<Integer> stream = + new BroadcastStreamImpl<>(env, new TestingTransformation<>("t1", Types.INT, 1)); + NonKeyedPartitionStreamImpl<Long> nonKeyedStream = + new NonKeyedPartitionStreamImpl<>( + env, new TestingTransformation<>("t2", Types.LONG, 2)); + stream.connectAndProcess( + nonKeyedStream, new StreamTestUtils.NoOpTwoInputBroadcastStreamProcessFunction()); + List<Transformation<?>> transformations = env.getTransformations(); + assertThat(transformations).hasSize(1); + assertProcessType(transformations.get(0), TwoInputTransformation.class, Types.LONG); + } + + @Test + void testConnectKeyedStream() throws Exception { + ExecutionEnvironmentImpl env = StreamTestUtils.getEnv(); + BroadcastStreamImpl<Integer> stream = + new BroadcastStreamImpl<>(env, new TestingTransformation<>("t1", Types.INT, 1)); + NonKeyedPartitionStreamImpl<Long> nonKeyedStream = + new NonKeyedPartitionStreamImpl<>( + env, new TestingTransformation<>("t2", Types.LONG, 2)); + stream.connectAndProcess( + nonKeyedStream.keyBy(x -> x), + new StreamTestUtils.NoOpTwoInputBroadcastStreamProcessFunction()); + List<Transformation<?>> transformations = env.getTransformations(); + assertThat(transformations).hasSize(1); + assertProcessType(transformations.get(0), TwoInputTransformation.class, Types.LONG); + } + + @Test + void testConnectKeyedStreamWithOutputKeySelector() throws Exception { + ExecutionEnvironmentImpl env = StreamTestUtils.getEnv(); + BroadcastStreamImpl<Integer> stream = + new BroadcastStreamImpl<>(env, new TestingTransformation<>("t1", Types.INT, 1)); + NonKeyedPartitionStreamImpl<Long> nonKeyedStream = + new NonKeyedPartitionStreamImpl<>( + env, new TestingTransformation<>("t2", Types.LONG, 2)); + KeyedPartitionStream<Long, Long> resultStream = + stream.connectAndProcess( + nonKeyedStream.keyBy(x -> x), + new StreamTestUtils.NoOpTwoInputBroadcastStreamProcessFunction(), + x -> x); + assertThat(resultStream).isInstanceOf(KeyedPartitionStream.class); + List<Transformation<?>> transformations = env.getTransformations(); + assertThat(transformations).hasSize(1); + assertProcessType(transformations.get(0), TwoInputTransformation.class, Types.LONG); + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/GlobalStreamImplTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/GlobalStreamImplTest.java new file mode 100644 index 00000000000..c63edd74872 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/GlobalStreamImplTest.java @@ -0,0 +1,75 @@ +/* + * 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.flink.datastream.impl.stream; + +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.dag.Transformation; +import org.apache.flink.datastream.impl.ExecutionEnvironmentImpl; +import org.apache.flink.datastream.impl.TestingTransformation; +import org.apache.flink.datastream.impl.stream.StreamTestUtils.NoOpOneInputStreamProcessFunction; +import org.apache.flink.datastream.impl.stream.StreamTestUtils.NoOpTwoInputBroadcastStreamProcessFunction; +import org.apache.flink.datastream.impl.stream.StreamTestUtils.NoOpTwoOutputStreamProcessFunction; +import org.apache.flink.streaming.api.transformations.PartitionTransformation; + +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link GlobalStreamImpl}. */ +class GlobalStreamImplTest { + @Test + void testParallelism() throws Exception { + ExecutionEnvironmentImpl env = StreamTestUtils.getEnv(); + GlobalStreamImpl<Integer> stream = + new GlobalStreamImpl<>(env, new TestingTransformation<>("t1", Types.INT, 1)); + stream.process(new NoOpOneInputStreamProcessFunction()); + stream.process(new NoOpTwoOutputStreamProcessFunction()); + stream.connectAndProcess( + new GlobalStreamImpl<>(env, new TestingTransformation<>("t2", Types.LONG, 1)), + new StreamTestUtils.NoOpTwoInputNonBroadcastStreamProcessFunction()); + List<Transformation<?>> transformations = env.getTransformations(); + assertThat(transformations).hasSize(3); + assertThat(transformations.get(0).getParallelism()).isOne(); + assertThat(transformations.get(1).getParallelism()).isOne(); + assertThat(transformations.get(2).getParallelism()).isOne(); + } + + @Test + void testPartitioning() throws Exception { + ExecutionEnvironmentImpl env = StreamTestUtils.getEnv(); + GlobalStreamImpl<Integer> stream = + new GlobalStreamImpl<>(env, new TestingTransformation<>("t1", Types.INT, 1)); + stream.keyBy((data) -> 1).process(new NoOpOneInputStreamProcessFunction()); + stream.shuffle().process(new NoOpOneInputStreamProcessFunction()); + stream.broadcast() + .connectAndProcess( + new NonKeyedPartitionStreamImpl<>( + env, new TestingTransformation<>("t2", Types.LONG, 1)), + new NoOpTwoInputBroadcastStreamProcessFunction()); + + List<Transformation<?>> transformations = env.getTransformations(); + Transformation<?> keyedTransform = transformations.get(0).getInputs().get(0); + assertThat(keyedTransform).isInstanceOf(PartitionTransformation.class); + Transformation<?> shuffleTransform = transformations.get(1).getInputs().get(0); + assertThat(shuffleTransform).isInstanceOf(PartitionTransformation.class); + assertThat(transformations.get(2).getParallelism()).isOne(); + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/KeyedPartitionStreamImplTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/KeyedPartitionStreamImplTest.java new file mode 100644 index 00000000000..12a7361af46 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/KeyedPartitionStreamImplTest.java @@ -0,0 +1,152 @@ +/* + * 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.flink.datastream.impl.stream; + +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.dag.Transformation; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.datastream.api.stream.KeyedPartitionStream; +import org.apache.flink.datastream.api.stream.NonKeyedPartitionStream; +import org.apache.flink.datastream.impl.ExecutionEnvironmentImpl; +import org.apache.flink.datastream.impl.TestingTransformation; +import org.apache.flink.datastream.impl.stream.StreamTestUtils.NoOpOneInputStreamProcessFunction; +import org.apache.flink.datastream.impl.stream.StreamTestUtils.NoOpTwoOutputStreamProcessFunction; +import org.apache.flink.streaming.api.transformations.OneInputTransformation; +import org.apache.flink.streaming.api.transformations.PartitionTransformation; +import org.apache.flink.streaming.api.transformations.TwoInputTransformation; + +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static org.apache.flink.datastream.impl.stream.StreamTestUtils.assertProcessType; +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link KeyedPartitionStreamImpl}. */ +public class KeyedPartitionStreamImplTest { + @Test + void testPartitioning() throws Exception { + ExecutionEnvironmentImpl env = StreamTestUtils.getEnv(); + KeyedPartitionStream<Integer, Integer> stream = createKeyedStream(env); + stream.keyBy((data) -> 1).process(new NoOpOneInputStreamProcessFunction()); + stream.shuffle().process(new NoOpOneInputStreamProcessFunction()); + stream.broadcast() + .connectAndProcess( + new NonKeyedPartitionStreamImpl<>( + env, new TestingTransformation<>("t2", Types.LONG, 1)), + new StreamTestUtils.NoOpTwoInputBroadcastStreamProcessFunction()); + + List<Transformation<?>> transformations = env.getTransformations(); + Transformation<?> keyedTransform = transformations.get(0).getInputs().get(0); + assertThat(keyedTransform).isInstanceOf(PartitionTransformation.class); + Transformation<?> shuffleTransform = transformations.get(1).getInputs().get(0); + assertThat(shuffleTransform).isInstanceOf(PartitionTransformation.class); + assertThat(transformations.get(2).getParallelism()).isOne(); + } + + @Test + void testProcess() throws Exception { + ExecutionEnvironmentImpl env = StreamTestUtils.getEnv(); + KeyedPartitionStream<Integer, Integer> stream = createKeyedStream(env); + stream.process(new NoOpOneInputStreamProcessFunction()); + KeyedPartitionStream<Integer, Long> resultStream = + stream.process(new NoOpOneInputStreamProcessFunction(), Math::toIntExact); + assertThat(resultStream).isInstanceOf(KeyedPartitionStream.class); + List<Transformation<?>> transformations = env.getTransformations(); + assertThat(transformations).hasSize(2); + assertProcessType(transformations.get(0), OneInputTransformation.class, Types.LONG); + assertProcessType(transformations.get(1), OneInputTransformation.class, Types.LONG); + } + + @Test + void testProcessTwoOutput() throws Exception { + ExecutionEnvironmentImpl env = StreamTestUtils.getEnv(); + KeyedPartitionStream<Integer, Integer> stream = createKeyedStream(env); + NonKeyedPartitionStream.TwoNonKeyedPartitionStreams<Integer, Long> resultStream1 = + stream.process(new NoOpTwoOutputStreamProcessFunction()); + assertThat(resultStream1.getFirst()).isInstanceOf(NonKeyedPartitionStream.class); + assertThat(resultStream1.getSecond()).isInstanceOf(NonKeyedPartitionStream.class); + KeyedPartitionStream.TwoKeyedPartitionStreams<Integer, Integer, Long> resultStream2 = + stream.process(new NoOpTwoOutputStreamProcessFunction(), x -> x, Math::toIntExact); + assertThat(resultStream2.getFirst()).isInstanceOf(KeyedPartitionStream.class); + assertThat(resultStream2.getSecond()).isInstanceOf(KeyedPartitionStream.class); + List<Transformation<?>> transformations = env.getTransformations(); + assertThat(transformations).hasSize(2); + assertProcessType(transformations.get(0), OneInputTransformation.class, Types.INT); + assertProcessType(transformations.get(1), OneInputTransformation.class, Types.INT); + } + + @Test + void testConnectKeyedStream() throws Exception { + ExecutionEnvironmentImpl env = StreamTestUtils.getEnv(); + KeyedPartitionStream<Integer, Integer> stream = createKeyedStream(env); + stream.connectAndProcess( + createKeyedStream( + env, + new TestingTransformation<>("t2", Types.LONG, 1), + (KeySelector<Long, Integer>) Math::toIntExact), + new StreamTestUtils.NoOpTwoInputNonBroadcastStreamProcessFunction()); + stream.connectAndProcess( + createKeyedStream( + env, + new TestingTransformation<>("t3", Types.LONG, 1), + (KeySelector<Long, Integer>) Math::toIntExact), + new StreamTestUtils.NoOpTwoInputNonBroadcastStreamProcessFunction(), + (KeySelector<Long, Integer>) Math::toIntExact); + + List<Transformation<?>> transformations = env.getTransformations(); + assertThat(transformations).hasSize(2); + assertProcessType(transformations.get(0), TwoInputTransformation.class, Types.LONG); + assertProcessType(transformations.get(1), TwoInputTransformation.class, Types.LONG); + } + + @Test + void testConnectBroadcastStream() throws Exception { + ExecutionEnvironmentImpl env = StreamTestUtils.getEnv(); + KeyedPartitionStream<Long, Long> stream = + createKeyedStream(env, new TestingTransformation<>("t1", Types.LONG, 1), x -> x); + BroadcastStreamImpl<Integer> s = + new BroadcastStreamImpl<>(env, new TestingTransformation<>("t2", Types.INT, 1)); + stream.connectAndProcess( + s, new StreamTestUtils.NoOpTwoInputBroadcastStreamProcessFunction()); + stream.connectAndProcess( + s, new StreamTestUtils.NoOpTwoInputBroadcastStreamProcessFunction(), x -> x); + List<Transformation<?>> transformations = env.getTransformations(); + assertThat(transformations).hasSize(2); + assertProcessType(transformations.get(0), TwoInputTransformation.class, Types.LONG); + assertProcessType(transformations.get(1), TwoInputTransformation.class, Types.LONG); + } + + private static KeyedPartitionStream<Integer, Integer> createKeyedStream( + ExecutionEnvironmentImpl env) { + NonKeyedPartitionStreamImpl<Integer> stream = + new NonKeyedPartitionStreamImpl<>( + env, new TestingTransformation<>("t1", Types.INT, 1)); + return stream.keyBy(x -> x); + } + + private static <OUT, KEY> KeyedPartitionStream<KEY, OUT> createKeyedStream( + ExecutionEnvironmentImpl env, + Transformation<OUT> transformation, + KeySelector<OUT, KEY> keySelector) { + NonKeyedPartitionStreamImpl<OUT> stream = + new NonKeyedPartitionStreamImpl<>(env, transformation); + return stream.keyBy(keySelector); + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/NonKeyedPartitionStreamImplTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/NonKeyedPartitionStreamImplTest.java new file mode 100644 index 00000000000..04a9ae8d3fb --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/NonKeyedPartitionStreamImplTest.java @@ -0,0 +1,118 @@ +/* + * 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.flink.datastream.impl.stream; + +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.dag.Transformation; +import org.apache.flink.datastream.api.stream.NonKeyedPartitionStream; +import org.apache.flink.datastream.impl.ExecutionEnvironmentImpl; +import org.apache.flink.datastream.impl.TestingTransformation; +import org.apache.flink.streaming.api.transformations.OneInputTransformation; +import org.apache.flink.streaming.api.transformations.PartitionTransformation; +import org.apache.flink.streaming.api.transformations.TwoInputTransformation; + +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static org.apache.flink.datastream.impl.stream.StreamTestUtils.assertProcessType; +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link NonKeyedPartitionStreamImpl}. */ +class NonKeyedPartitionStreamImplTest { + @Test + void testPartitioning() throws Exception { + ExecutionEnvironmentImpl env = StreamTestUtils.getEnv(); + NonKeyedPartitionStreamImpl<Integer> stream = + new NonKeyedPartitionStreamImpl<>( + env, new TestingTransformation<>("t1", Types.INT, 1)); + stream.keyBy((data) -> 1).process(new StreamTestUtils.NoOpOneInputStreamProcessFunction()); + stream.shuffle().process(new StreamTestUtils.NoOpOneInputStreamProcessFunction()); + stream.broadcast() + .connectAndProcess( + new NonKeyedPartitionStreamImpl<>( + env, new TestingTransformation<>("t2", Types.LONG, 1)), + new StreamTestUtils.NoOpTwoInputBroadcastStreamProcessFunction()); + + List<Transformation<?>> transformations = env.getTransformations(); + Transformation<?> keyedTransform = transformations.get(0).getInputs().get(0); + assertThat(keyedTransform).isInstanceOf(PartitionTransformation.class); + Transformation<?> shuffleTransform = transformations.get(1).getInputs().get(0); + assertThat(shuffleTransform).isInstanceOf(PartitionTransformation.class); + assertThat(transformations.get(2).getParallelism()).isOne(); + } + + @Test + void testProcess() throws Exception { + ExecutionEnvironmentImpl env = StreamTestUtils.getEnv(); + NonKeyedPartitionStreamImpl<Integer> stream = + new NonKeyedPartitionStreamImpl<>( + env, new TestingTransformation<>("t1", Types.INT, 1)); + stream.process(new StreamTestUtils.NoOpOneInputStreamProcessFunction()); + List<Transformation<?>> transformations = env.getTransformations(); + assertThat(transformations).hasSize(1); + assertProcessType(transformations.get(0), OneInputTransformation.class, Types.LONG); + } + + @Test + void testProcessTwoOutput() throws Exception { + ExecutionEnvironmentImpl env = StreamTestUtils.getEnv(); + NonKeyedPartitionStreamImpl<Integer> stream = + new NonKeyedPartitionStreamImpl<>( + env, new TestingTransformation<>("t1", Types.INT, 1)); + NonKeyedPartitionStream.TwoNonKeyedPartitionStreams<Integer, Long> resultStream = + stream.process(new StreamTestUtils.NoOpTwoOutputStreamProcessFunction()); + assertThat(resultStream.getFirst()).isInstanceOf(NonKeyedPartitionStream.class); + assertThat(resultStream.getSecond()).isInstanceOf(NonKeyedPartitionStream.class); + List<Transformation<?>> transformations = env.getTransformations(); + assertThat(transformations).hasSize(1); + assertProcessType(transformations.get(0), OneInputTransformation.class, Types.INT); + } + + @Test + void testConnectNonKeyedStream() throws Exception { + ExecutionEnvironmentImpl env = StreamTestUtils.getEnv(); + NonKeyedPartitionStreamImpl<Integer> stream = + new NonKeyedPartitionStreamImpl<>( + env, new TestingTransformation<>("t1", Types.INT, 1)); + stream.connectAndProcess( + new NonKeyedPartitionStreamImpl<>( + env, new TestingTransformation<>("t2", Types.LONG, 1)), + new StreamTestUtils.NoOpTwoInputNonBroadcastStreamProcessFunction()); + + List<Transformation<?>> transformations = env.getTransformations(); + assertThat(transformations).hasSize(1); + assertProcessType(transformations.get(0), TwoInputTransformation.class, Types.LONG); + } + + @Test + void testConnectBroadcastStream() throws Exception { + ExecutionEnvironmentImpl env = StreamTestUtils.getEnv(); + NonKeyedPartitionStreamImpl<Long> stream = + new NonKeyedPartitionStreamImpl<>( + env, new TestingTransformation<>("t1", Types.LONG, 1)); + stream.connectAndProcess( + new BroadcastStreamImpl<>(env, new TestingTransformation<>("t2", Types.INT, 1)), + new StreamTestUtils.NoOpTwoInputBroadcastStreamProcessFunction()); + + List<Transformation<?>> transformations = env.getTransformations(); + assertThat(transformations).hasSize(1); + assertProcessType(transformations.get(0), TwoInputTransformation.class, Types.LONG); + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/StreamTestUtils.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/StreamTestUtils.java new file mode 100644 index 00000000000..f4fd0738757 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/stream/StreamTestUtils.java @@ -0,0 +1,119 @@ +/* + * 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.flink.datastream.impl.stream; + +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.dag.Transformation; +import org.apache.flink.datastream.api.ExecutionEnvironment; +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.datastream.api.context.NonPartitionedContext; +import org.apache.flink.datastream.api.context.RuntimeContext; +import org.apache.flink.datastream.api.function.OneInputStreamProcessFunction; +import org.apache.flink.datastream.api.function.TwoInputBroadcastStreamProcessFunction; +import org.apache.flink.datastream.api.function.TwoInputNonBroadcastStreamProcessFunction; +import org.apache.flink.datastream.api.function.TwoOutputStreamProcessFunction; +import org.apache.flink.datastream.impl.ExecutionEnvironmentImpl; +import org.apache.flink.streaming.api.graph.StreamGraph; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Test utils for steam. */ +public final class StreamTestUtils { + public static StreamGraph getStreamGraph(ExecutionEnvironment env) { + assertThat(env).isInstanceOf(ExecutionEnvironmentImpl.class); + ExecutionEnvironmentImpl envImpl = (ExecutionEnvironmentImpl) env; + return envImpl.getStreamGraph(); + } + + /** Assert a transformation has a specific class and type information. */ + public static <OUT> void assertProcessType( + Transformation<?> transformation, + Class<?> transformationClass, + TypeInformation<OUT> typeInformation) { + assertThat(transformation) + .isInstanceOf(transformationClass) + .extracting(Transformation::getOutputType) + .isEqualTo(typeInformation); + } + + public static ExecutionEnvironmentImpl getEnv() throws Exception { + return (ExecutionEnvironmentImpl) ExecutionEnvironment.getInstance(); + } + + /** An implementation of the {@link OneInputStreamProcessFunction} that does nothing. */ + public static class NoOpOneInputStreamProcessFunction + implements OneInputStreamProcessFunction<Integer, Long> { + + @Override + public void processRecord(Integer record, Collector<Long> output, RuntimeContext ctx) { + // do nothing. + } + } + + /** An implementation of the {@link TwoOutputStreamProcessFunction} that does nothing. */ + public static class NoOpTwoOutputStreamProcessFunction + implements TwoOutputStreamProcessFunction<Integer, Integer, Long> { + + @Override + public void processRecord( + Integer record, + Collector<Integer> output1, + Collector<Long> output2, + RuntimeContext ctx) { + // do nothing. + } + } + + /** + * An implementation of the {@link TwoInputNonBroadcastStreamProcessFunction} that does nothing. + */ + public static class NoOpTwoInputNonBroadcastStreamProcessFunction + implements TwoInputNonBroadcastStreamProcessFunction<Integer, Long, Long> { + + @Override + public void processRecordFromFirstInput( + Integer record, Collector<Long> output, RuntimeContext ctx) { + // do nothing. + } + + @Override + public void processRecordFromSecondInput( + Long record, Collector<Long> output, RuntimeContext ctx) throws Exception { + // do nothing. + } + } + + /** + * An implementation of the {@link TwoInputBroadcastStreamProcessFunction} that does nothing. + */ + public static class NoOpTwoInputBroadcastStreamProcessFunction + implements TwoInputBroadcastStreamProcessFunction<Long, Integer, Long> { + @Override + public void processRecordFromNonBroadcastInput( + Long record, Collector<Long> output, RuntimeContext ctx) { + // do nothing. + } + + @Override + public void processRecordFromBroadcastInput( + Integer record, NonPartitionedContext<Long> ctx) { + // do nothing. + } + } +} diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/utils/StreamUtilsTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/utils/StreamUtilsTest.java new file mode 100644 index 00000000000..03484237256 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/utils/StreamUtilsTest.java @@ -0,0 +1,197 @@ +/* + * 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.flink.datastream.impl.utils; + +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.api.java.tuple.Tuple2; +import org.apache.flink.api.java.typeutils.MissingTypeInfo; +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.datastream.api.context.NonPartitionedContext; +import org.apache.flink.datastream.api.context.RuntimeContext; +import org.apache.flink.datastream.api.function.OneInputStreamProcessFunction; +import org.apache.flink.datastream.api.function.TwoInputBroadcastStreamProcessFunction; +import org.apache.flink.datastream.api.function.TwoInputNonBroadcastStreamProcessFunction; +import org.apache.flink.datastream.api.function.TwoOutputStreamProcessFunction; +import org.apache.flink.datastream.impl.ExecutionEnvironmentImpl; +import org.apache.flink.datastream.impl.TestingTransformation; +import org.apache.flink.datastream.impl.operators.ProcessOperator; +import org.apache.flink.datastream.impl.operators.TwoInputNonBroadcastProcessOperator; +import org.apache.flink.datastream.impl.stream.NonKeyedPartitionStreamImpl; +import org.apache.flink.datastream.impl.stream.StreamTestUtils; +import org.apache.flink.streaming.api.transformations.OneInputTransformation; +import org.apache.flink.streaming.api.transformations.TwoInputTransformation; + +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link StreamUtils}. */ +class StreamUtilsTest { + @Test + void testGetOneInputOutputType() { + TypeInformation<Long> outputType = + StreamUtils.getOutputTypeForOneInputProcessFunction( + new OneInputStreamProcessFunction<Integer, Long>() { + @Override + public void processRecord( + Integer record, Collector<Long> output, RuntimeContext ctx) + throws Exception { + // ignore + } + }, + Types.INT); + assertThat(outputType).isEqualTo(Types.LONG); + } + + @Test + void testLambdaMissingType() { + TypeInformation<Long> outputType = + StreamUtils.getOutputTypeForOneInputProcessFunction( + (OneInputStreamProcessFunction<Integer, Long>) + (record, output, ctx) -> { + // ignore + }, + Types.INT); + assertThat(outputType).isInstanceOf(MissingTypeInfo.class); + } + + @Test + void testGetTwoInputOutputType() { + TypeInformation<String> outputType = + StreamUtils.getOutputTypeForTwoInputNonBroadcastProcessFunction( + new TwoInputNonBroadcastStreamProcessFunction<Integer, Long, String>() { + @Override + public void processRecordFromFirstInput( + Integer record, Collector<String> output, RuntimeContext ctx) + throws Exception { + // ignore + } + + @Override + public void processRecordFromSecondInput( + Long record, Collector<String> output, RuntimeContext ctx) + throws Exception { + // ignore + } + }, + Types.INT, + Types.LONG); + assertThat(outputType).isEqualTo(Types.STRING); + + outputType = + StreamUtils.getOutputTypeForTwoInputBroadcastProcessFunction( + new TwoInputBroadcastStreamProcessFunction<Integer, Long, String>() { + @Override + public void processRecordFromNonBroadcastInput( + Integer record, Collector<String> output, RuntimeContext ctx) + throws Exception { + // ignore + } + + @Override + public void processRecordFromBroadcastInput( + Long record, NonPartitionedContext<String> ctx) + throws Exception { + // ignore + } + }, + Types.INT, + Types.LONG); + assertThat(outputType).isEqualTo(Types.STRING); + } + + @Test + void testTwoOutputType() { + Tuple2<TypeInformation<Long>, TypeInformation<String>> outputType = + StreamUtils.getOutputTypesForTwoOutputProcessFunction( + new TwoOutputStreamProcessFunction<Integer, Long, String>() { + @Override + public void processRecord( + Integer record, + Collector<Long> output1, + Collector<String> output2, + RuntimeContext ctx) + throws Exception { + // ignore + } + }, + Types.INT); + assertThat(outputType.f0).isEqualTo(Types.LONG); + assertThat(outputType.f1).isEqualTo(Types.STRING); + } + + @Test + void testGetOneInputTransformation() throws Exception { + ExecutionEnvironmentImpl env = StreamTestUtils.getEnv(); + ProcessOperator<Integer, Long> operator = + new ProcessOperator<>(new StreamTestUtils.NoOpOneInputStreamProcessFunction()); + OneInputTransformation<Integer, Long> transformation = + StreamUtils.getOneInputTransformation( + "op", + new NonKeyedPartitionStreamImpl<>( + env, new TestingTransformation<>("t", Types.INT, 1)), + Types.LONG, + operator); + assertThat(transformation.getOperator()).isEqualTo(operator); + assertThat(transformation.getOutputType()).isEqualTo(Types.LONG); + assertThat(transformation.getStateKeySelector()).isNull(); + } + + @Test + void testGetOneInputKeyedTransformation() throws Exception { + ExecutionEnvironmentImpl env = StreamTestUtils.getEnv(); + ProcessOperator<Integer, Long> operator = + new ProcessOperator<>(new StreamTestUtils.NoOpOneInputStreamProcessFunction()); + OneInputTransformation<Integer, Long> transformation = + StreamUtils.getOneInputKeyedTransformation( + "op", + new NonKeyedPartitionStreamImpl<>( + env, new TestingTransformation<>("t", Types.INT, 1)), + Types.LONG, + operator, + (KeySelector<Integer, Integer>) value -> value, + Types.INT); + assertThat(transformation.getOperator()).isEqualTo(operator); + assertThat(transformation.getOutputType()).isEqualTo(Types.LONG); + assertThat(transformation.getStateKeySelector()).isNotNull(); + } + + @Test + void testGetTwoInputTransformation() throws Exception { + ExecutionEnvironmentImpl env = StreamTestUtils.getEnv(); + TwoInputNonBroadcastProcessOperator<Integer, Long, Long> operator = + new TwoInputNonBroadcastProcessOperator<>( + new StreamTestUtils.NoOpTwoInputNonBroadcastStreamProcessFunction()); + TwoInputTransformation<Integer, Long, Long> transformation = + StreamUtils.getTwoInputTransformation( + "op", + new NonKeyedPartitionStreamImpl<>( + env, new TestingTransformation<>("t1", Types.INT, 1)), + new NonKeyedPartitionStreamImpl<>( + env, new TestingTransformation<>("t2", Types.LONG, 1)), + Types.LONG, + operator); + assertThat(transformation.getOperator()).isEqualTo(operator); + assertThat(transformation.getOutputType()).isEqualTo(Types.LONG); + assertThat(transformation.getStateKeySelector1()).isNull(); + assertThat(transformation.getStateKeySelector2()).isNull(); + } +}
