| /* |
| * 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.test.iteration; |
| |
| import org.apache.flink.api.common.functions.JoinFunction; |
| import org.apache.flink.api.common.state.BroadcastState; |
| import org.apache.flink.api.common.state.ListState; |
| import org.apache.flink.api.common.state.ListStateDescriptor; |
| import org.apache.flink.api.common.state.MapStateDescriptor; |
| import org.apache.flink.api.common.typeinfo.BasicTypeInfo; |
| import org.apache.flink.api.common.typeinfo.Types; |
| import org.apache.flink.api.java.tuple.Tuple2; |
| import org.apache.flink.iteration.DataStreamList; |
| import org.apache.flink.iteration.IterationBody; |
| import org.apache.flink.iteration.IterationBodyResult; |
| import org.apache.flink.iteration.IterationConfig; |
| import org.apache.flink.iteration.IterationConfig.OperatorLifeCycle; |
| import org.apache.flink.iteration.Iterations; |
| import org.apache.flink.iteration.ReplayableDataStreamList; |
| import org.apache.flink.ml.common.datastream.EndOfStreamWindows; |
| import org.apache.flink.ml.common.iteration.TerminateOnMaxIter; |
| import org.apache.flink.runtime.jobgraph.JobGraph; |
| import org.apache.flink.runtime.minicluster.MiniCluster; |
| import org.apache.flink.runtime.state.StateInitializationContext; |
| import org.apache.flink.streaming.api.datastream.DataStream; |
| import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator; |
| import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; |
| import org.apache.flink.streaming.api.functions.sink.SinkFunction; |
| import org.apache.flink.streaming.api.operators.AbstractStreamOperator; |
| import org.apache.flink.streaming.api.operators.OneInputStreamOperator; |
| import org.apache.flink.streaming.api.watermark.Watermark; |
| import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; |
| import org.apache.flink.test.iteration.operators.CollectSink; |
| import org.apache.flink.test.iteration.operators.EpochRecord; |
| import org.apache.flink.test.iteration.operators.OutputRecord; |
| import org.apache.flink.test.iteration.operators.SequenceSource; |
| import org.apache.flink.test.iteration.operators.StatefulProcessFunction; |
| import org.apache.flink.test.iteration.operators.TwoInputReducePerRoundOperator; |
| import org.apache.flink.testutils.junit.SharedObjects; |
| import org.apache.flink.testutils.junit.SharedReference; |
| import org.apache.flink.util.TestLogger; |
| |
| import org.junit.After; |
| import org.junit.Before; |
| import org.junit.Rule; |
| import org.junit.Test; |
| |
| import java.util.ArrayList; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.concurrent.BlockingQueue; |
| import java.util.concurrent.LinkedBlockingQueue; |
| |
| import static org.apache.flink.test.iteration.UnboundedStreamIterationITCase.computeRoundStat; |
| import static org.apache.flink.test.iteration.UnboundedStreamIterationITCase.createMiniClusterConfiguration; |
| import static org.apache.flink.test.iteration.UnboundedStreamIterationITCase.verifyResult; |
| import static org.junit.Assert.assertEquals; |
| |
| /** Tests the per-round iterations. */ |
| public class BoundedPerRoundStreamIterationITCase extends TestLogger { |
| |
| @Rule public final SharedObjects sharedObjects = SharedObjects.create(); |
| |
| private MiniCluster miniCluster; |
| |
| private SharedReference<BlockingQueue<OutputRecord<Integer>>> collectedOutputRecord; |
| private SharedReference<BlockingQueue<Long>> collectedWatermarks; |
| private SharedReference<BlockingQueue<Long>> collectedOutputs; |
| |
| @Before |
| public void setup() throws Exception { |
| miniCluster = new MiniCluster(createMiniClusterConfiguration(2, 2)); |
| miniCluster.start(); |
| |
| collectedOutputRecord = sharedObjects.add(new LinkedBlockingQueue<>()); |
| collectedWatermarks = sharedObjects.add(new LinkedBlockingQueue<>()); |
| collectedOutputs = sharedObjects.add(new LinkedBlockingQueue<>()); |
| } |
| |
| @After |
| public void teardown() throws Exception { |
| if (miniCluster != null) { |
| miniCluster.close(); |
| } |
| } |
| |
| @Test |
| public void testPerRoundIteration() throws Exception { |
| JobGraph jobGraph = createPerRoundJobGraph(4, 1000, 5, collectedOutputRecord); |
| miniCluster.executeJobBlocking(jobGraph); |
| |
| assertEquals(5, collectedOutputRecord.get().size()); |
| Map<Integer, Tuple2<Integer, Integer>> roundsStat = |
| computeRoundStat(collectedOutputRecord.get(), OutputRecord.Event.TERMINATED, 5); |
| verifyResult(roundsStat, 5, 1, 4 * (0 + 999) * 1000 / 2); |
| } |
| |
| @Test |
| public void testPerRoundIterationWithJoin() throws Exception { |
| StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); |
| env.setParallelism(4); |
| |
| DataStream<Tuple2<Long, Integer>> input1 = env.fromElements(Tuple2.of(1L, 1)); |
| |
| DataStream<Tuple2<Long, Long>> input2 = env.fromElements(Tuple2.of(1L, 2L)); |
| |
| DataStream<Tuple2<Long, Long>> iterationWithJoinResult = |
| Iterations.iterateBoundedStreamsUntilTermination( |
| DataStreamList.of(input1), |
| ReplayableDataStreamList.replay(input2), |
| IterationConfig.newBuilder() |
| .setOperatorLifeCycle( |
| IterationConfig.OperatorLifeCycle.PER_ROUND) |
| .build(), |
| new IterationBodyWithJoin()) |
| .get(0); |
| DataStream<Long> watermarks = |
| iterationWithJoinResult.transform( |
| "CollectingWatermark", Types.LONG, new CollectingWatermark()); |
| |
| watermarks.addSink(new LongSink(collectedWatermarks)); |
| |
| JobGraph graph = env.getStreamGraph().getJobGraph(); |
| miniCluster.executeJobBlocking(graph); |
| |
| assertEquals(env.getParallelism(), collectedWatermarks.get().size()); |
| collectedWatermarks |
| .get() |
| .iterator() |
| .forEachRemaining(x -> assertEquals(Long.MAX_VALUE, (long) x)); |
| } |
| |
| @Test |
| public void testPerRoundIterationWithState() throws Exception { |
| StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); |
| env.setParallelism(1); |
| DataStream<Long> broadcastStream = env.fromElements(1L); |
| DataStream<Long> inputStream = env.fromElements(1L); |
| DataStreamList outputStream = |
| Iterations.iterateBoundedStreamsUntilTermination( |
| DataStreamList.of(inputStream), |
| ReplayableDataStreamList.replay(broadcastStream), |
| IterationConfig.newBuilder() |
| .setOperatorLifeCycle(OperatorLifeCycle.PER_ROUND) |
| .build(), |
| new PerRoundIterationBodyWithState()); |
| |
| outputStream.<Long>get(0).addSink(new LongSink(collectedOutputs)); |
| JobGraph jobGraph = env.getStreamGraph().getJobGraph(); |
| miniCluster.executeJobBlocking(jobGraph); |
| |
| List<Long> result = new ArrayList<>(3); |
| collectedOutputs.get().drainTo(result); |
| assertEquals(3, result.size()); |
| for (long value : result) { |
| assertEquals(1L, value); |
| } |
| } |
| |
| private static JobGraph createPerRoundJobGraph( |
| int numSources, |
| int numRecordsPerSource, |
| int maxRound, |
| SharedReference<BlockingQueue<OutputRecord<Integer>>> result) { |
| StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); |
| env.getConfig().enableObjectReuse(); |
| env.setParallelism(1); |
| env.setMaxParallelism(5); |
| |
| DataStream<Integer> variableSource = env.fromElements(0); |
| DataStream<Integer> constSource = |
| env.addSource(new SequenceSource(numRecordsPerSource, false, 0)) |
| .setParallelism(numSources) |
| .map(EpochRecord::getValue) |
| .setParallelism(numSources) |
| .name("Constant"); |
| |
| DataStreamList outputs = |
| Iterations.iterateBoundedStreamsUntilTermination( |
| DataStreamList.of(variableSource), |
| ReplayableDataStreamList.replay(constSource), |
| IterationConfig.newBuilder() |
| .setOperatorLifeCycle(IterationConfig.OperatorLifeCycle.PER_ROUND) |
| .build(), |
| (variableStreams, dataStreams) -> { |
| SingleOutputStreamOperator<Integer> reducer = |
| variableStreams |
| .<Integer>get(0) |
| .connect(dataStreams.<Integer>get(0)) |
| .transform( |
| "Reducer", |
| BasicTypeInfo.INT_TYPE_INFO, |
| new TwoInputReducePerRoundOperator()) |
| .setParallelism(1); |
| |
| return new IterationBodyResult( |
| DataStreamList.of( |
| reducer.partitionCustom( |
| (k, numPartitions) -> k % numPartitions, |
| x -> x) |
| .map(x -> x) |
| .keyBy(x -> x) |
| .process( |
| new StatefulProcessFunction< |
| Integer>() {}) |
| .setParallelism(4) |
| .filter(x -> x < maxRound) |
| .setParallelism(1)), |
| DataStreamList.of( |
| reducer.getSideOutput( |
| TwoInputReducePerRoundOperator.OUTPUT_TAG)), |
| reducer.filter(x -> x < maxRound) |
| .map(x -> (double) x) |
| .setParallelism(1)); |
| }); |
| outputs.<OutputRecord<Integer>>get(0).addSink(new CollectSink(result)); |
| |
| return env.getStreamGraph().getJobGraph(); |
| } |
| |
| private static class IterationBodyWithJoin implements IterationBody { |
| @Override |
| public IterationBodyResult process( |
| DataStreamList variableStreams, DataStreamList dataStreams) { |
| DataStream<Tuple2<Long, Integer>> input1 = variableStreams.get(0); |
| DataStream<Tuple2<Long, Long>> input2 = dataStreams.get(0); |
| |
| DataStream<Long> terminationCriteria = input1.flatMap(new TerminateOnMaxIter(1)); |
| |
| DataStream<Tuple2<Long, Long>> res = |
| input1.join(input2) |
| .where(x -> x.f0) |
| .equalTo(x -> x.f0) |
| .window(EndOfStreamWindows.get()) |
| .apply( |
| new JoinFunction< |
| Tuple2<Long, Integer>, |
| Tuple2<Long, Long>, |
| Tuple2<Long, Long>>() { |
| @Override |
| public Tuple2<Long, Long> join( |
| Tuple2<Long, Integer> longIntegerTuple2, |
| Tuple2<Long, Long> longLongTuple2) { |
| return longLongTuple2; |
| } |
| }); |
| |
| return new IterationBodyResult( |
| DataStreamList.of(input1), DataStreamList.of(res), terminationCriteria); |
| } |
| } |
| |
| private static class PerRoundIterationBodyWithState implements IterationBody { |
| |
| @Override |
| public IterationBodyResult process( |
| DataStreamList variableStreams, DataStreamList dataStreams) { |
| DataStream<Long> variableStream = variableStreams.get(0); |
| |
| DataStream<Long> feedback = |
| variableStream.transform("mapWithState", Types.LONG, new MapWithState()); |
| |
| DataStream<Integer> terminationCriteria = |
| feedback.<Long>flatMap(new TerminateOnMaxIter(2)).returns(Types.INT); |
| |
| return new IterationBodyResult( |
| DataStreamList.of(feedback), DataStreamList.of(feedback), terminationCriteria); |
| } |
| } |
| |
| private static class MapWithState extends AbstractStreamOperator<Long> |
| implements OneInputStreamOperator<Long, Long> { |
| private ListState<Long> listState; |
| private ListState<Long> unionState; |
| private BroadcastState<Long, Long> broadcastState; |
| |
| @Override |
| public void processElement(StreamRecord<Long> element) throws Exception { |
| long val = element.getValue(); |
| listState.add(val); |
| unionState.add(val); |
| broadcastState.put(val, val); |
| output.collect(element); |
| } |
| |
| @Override |
| public void initializeState(StateInitializationContext context) throws Exception { |
| super.initializeState(context); |
| listState = |
| context.getOperatorStateStore() |
| .getListState(new ListStateDescriptor<>("longState", Types.LONG)); |
| unionState = |
| context.getOperatorStateStore() |
| .getUnionListState(new ListStateDescriptor<>("unionState", Types.LONG)); |
| broadcastState = |
| context.getOperatorStateStore() |
| .getBroadcastState( |
| new MapStateDescriptor<>( |
| "broadcastState", Types.LONG, Types.LONG)); |
| } |
| } |
| |
| private static class LongSink implements SinkFunction<Long> { |
| private final SharedReference<BlockingQueue<Long>> collectedLong; |
| |
| public LongSink(SharedReference<BlockingQueue<Long>> collectedLong) { |
| this.collectedLong = collectedLong; |
| } |
| |
| @Override |
| public void invoke(Long value, Context context) { |
| collectedLong.get().add(value); |
| } |
| } |
| |
| private static class CollectingWatermark extends AbstractStreamOperator<Long> |
| implements OneInputStreamOperator<Tuple2<Long, Long>, Long> { |
| |
| @Override |
| public void processElement(StreamRecord<Tuple2<Long, Long>> streamRecord) {} |
| |
| @Override |
| public void processWatermark(Watermark mark) throws Exception { |
| super.processWatermark(mark); |
| output.collect(new StreamRecord<>(mark.getTimestamp())); |
| } |
| } |
| } |
