Examples with MapStateDescriptor org.apache.flink.api.common.state.MapStateDescriptor used on opensource projects

Example 16 with MapStateDescriptor

use of org.apache.flink.api.common.state.MapStateDescriptor in project flink by apache.

the class SortingBoundedInputITCase method testBatchExecutionWithTimersTwoInput.

@Test
public void testBatchExecutionWithTimersTwoInput() {
    StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
    // set parallelism to 1 to have consistent order of results
    env.setParallelism(1);
    Configuration config = new Configuration();
    config.set(ExecutionOptions.RUNTIME_MODE, RuntimeExecutionMode.BATCH);
    env.configure(config, this.getClass().getClassLoader());
    WatermarkStrategy<Tuple2<Integer, Integer>> watermarkStrategy = WatermarkStrategy.forGenerator(ctx -> GENERATE_WATERMARK_AFTER_4_14_TIMESTAMP).withTimestampAssigner((r, previousTimestamp) -> r.f1);
    SingleOutputStreamOperator<Integer> elements1 = env.fromElements(Tuple2.of(1, 3), Tuple2.of(1, 1), Tuple2.of(2, 1), Tuple2.of(1, 4), // late element
    Tuple2.of(2, 3), // late element
    Tuple2.of(1, 2), Tuple2.of(1, 13), Tuple2.of(1, 11), Tuple2.of(2, 14), // late element
    Tuple2.of(1, 11)).assignTimestampsAndWatermarks(watermarkStrategy).map(element -> element.f0);
    SingleOutputStreamOperator<Integer> elements2 = env.fromElements(Tuple2.of(1, 3), Tuple2.of(1, 1), Tuple2.of(2, 1), Tuple2.of(1, 4), // late element
    Tuple2.of(2, 3), // late element
    Tuple2.of(1, 2), Tuple2.of(1, 13), Tuple2.of(1, 11), Tuple2.of(2, 14), // late element
    Tuple2.of(1, 11)).assignTimestampsAndWatermarks(watermarkStrategy).map(element -> element.f0);
    OutputTag<Integer> lateElements = new OutputTag<>("late_elements", BasicTypeInfo.INT_TYPE_INFO);
    SingleOutputStreamOperator<Tuple3<Long, Integer, Integer>> sums = elements1.connect(elements2).keyBy(element -> element, element -> element).process(new KeyedCoProcessFunction<Integer, Integer, Integer, Tuple3<Long, Integer, Integer>>() {

        private MapState<Long, Integer> countState;

        private ValueState<Long> previousTimestampState;

        @Override
        public void open(Configuration parameters) {
            countState = getRuntimeContext().getMapState(new MapStateDescriptor<>("sum", BasicTypeInfo.LONG_TYPE_INFO, BasicTypeInfo.INT_TYPE_INFO));
            previousTimestampState = getRuntimeContext().getState(new ValueStateDescriptor<>("previousTimestamp", BasicTypeInfo.LONG_TYPE_INFO));
        }

        @Override
        public void processElement1(Integer value, Context ctx, Collector<Tuple3<Long, Integer, Integer>> out) throws Exception {
            processElement(value, ctx);
        }

        @Override
        public void processElement2(Integer value, Context ctx, Collector<Tuple3<Long, Integer, Integer>> out) throws Exception {
            processElement(value, ctx);
        }

        private void processElement(Integer value, Context ctx) throws Exception {
            Long elementTimestamp = ctx.timestamp();
            long nextTen = ((elementTimestamp + 10) / 10) * 10;
            ctx.timerService().registerEventTimeTimer(nextTen);
            if (elementTimestamp < ctx.timerService().currentWatermark()) {
                ctx.output(lateElements, value);
            } else {
                Long previousTimestamp = Optional.ofNullable(previousTimestampState.value()).orElse(0L);
                assertThat(elementTimestamp, greaterThanOrEqualTo(previousTimestamp));
                previousTimestampState.update(elementTimestamp);
                Integer currentCount = Optional.ofNullable(countState.get(nextTen)).orElse(0);
                countState.put(nextTen, currentCount + 1);
            }
        }

        @Override
        public void onTimer(long timestamp, OnTimerContext ctx, Collector<Tuple3<Long, Integer, Integer>> out) throws Exception {
            out.collect(Tuple3.of(timestamp, ctx.getCurrentKey(), countState.get(timestamp)));
            countState.remove(timestamp);
            // this would go in infinite loop if we did not quiesce the
            // timer service.
            ctx.timerService().registerEventTimeTimer(timestamp + 1);
        }
    });
    DataStream<Integer> lateStream = sums.getSideOutput(lateElements);
    List<Integer> lateRecordsCollected = CollectionUtil.iteratorToList(DataStreamUtils.collect(lateStream));
    List<Tuple3<Long, Integer, Integer>> sumsCollected = CollectionUtil.iteratorToList(DataStreamUtils.collect(sums));
    assertTrue(lateRecordsCollected.isEmpty());
    assertThat(sumsCollected, equalTo(Arrays.asList(Tuple3.of(10L, 1, 8), Tuple3.of(20L, 1, 6), Tuple3.of(10L, 2, 4), Tuple3.of(20L, 2, 2))));
}

Example 17 with MapStateDescriptor

use of org.apache.flink.api.common.state.MapStateDescriptor in project flink by apache.

the class OperatorStateBackendTest method testCorrectClassLoaderUsedOnSnapshot.

@SuppressWarnings("unchecked")
@Test
public void testCorrectClassLoaderUsedOnSnapshot() throws Exception {
    AbstractStateBackend abstractStateBackend = new MemoryStateBackend(4096);
    final Environment env = createMockEnvironment();
    CloseableRegistry cancelStreamRegistry = new CloseableRegistry();
    OperatorStateBackend operatorStateBackend = abstractStateBackend.createOperatorStateBackend(env, "test-op-name", emptyStateHandles, cancelStreamRegistry);
    AtomicInteger copyCounter = new AtomicInteger(0);
    TypeSerializer<Integer> serializer = new VerifyingIntSerializer(env.getUserCodeClassLoader().asClassLoader(), copyCounter);
    // write some state
    ListStateDescriptor<Integer> stateDescriptor = new ListStateDescriptor<>("test", serializer);
    ListState<Integer> listState = operatorStateBackend.getListState(stateDescriptor);
    listState.add(42);
    AtomicInteger keyCopyCounter = new AtomicInteger(0);
    AtomicInteger valueCopyCounter = new AtomicInteger(0);
    TypeSerializer<Integer> keySerializer = new VerifyingIntSerializer(env.getUserCodeClassLoader().asClassLoader(), keyCopyCounter);
    TypeSerializer<Integer> valueSerializer = new VerifyingIntSerializer(env.getUserCodeClassLoader().asClassLoader(), valueCopyCounter);
    MapStateDescriptor<Integer, Integer> broadcastStateDesc = new MapStateDescriptor<>("test-broadcast", keySerializer, valueSerializer);
    BroadcastState<Integer, Integer> broadcastState = operatorStateBackend.getBroadcastState(broadcastStateDesc);
    broadcastState.put(1, 2);
    broadcastState.put(3, 4);
    broadcastState.put(5, 6);
    CheckpointStreamFactory streamFactory = new MemCheckpointStreamFactory(4096);
    RunnableFuture<SnapshotResult<OperatorStateHandle>> runnableFuture = operatorStateBackend.snapshot(1, 1, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());
    FutureUtils.runIfNotDoneAndGet(runnableFuture);
    // make sure that the copy method has been called
    assertTrue(copyCounter.get() > 0);
    assertTrue(keyCopyCounter.get() > 0);
    assertTrue(valueCopyCounter.get() > 0);
}

Example 18 with MapStateDescriptor

use of org.apache.flink.api.common.state.MapStateDescriptor in project flink by apache.

the class OperatorStateBackendTest method testSnapshotRestoreAsync.

@Test
public void testSnapshotRestoreAsync() throws Exception {
    OperatorStateBackend operatorStateBackend = new DefaultOperatorStateBackendBuilder(OperatorStateBackendTest.class.getClassLoader(), new ExecutionConfig(), true, emptyStateHandles, new CloseableRegistry()).build();
    ListStateDescriptor<MutableType> stateDescriptor1 = new ListStateDescriptor<>("test1", new JavaSerializer<MutableType>());
    ListStateDescriptor<MutableType> stateDescriptor2 = new ListStateDescriptor<>("test2", new JavaSerializer<MutableType>());
    ListStateDescriptor<MutableType> stateDescriptor3 = new ListStateDescriptor<>("test3", new JavaSerializer<MutableType>());
    MapStateDescriptor<MutableType, MutableType> broadcastStateDescriptor1 = new MapStateDescriptor<>("test4", new JavaSerializer<MutableType>(), new JavaSerializer<MutableType>());
    MapStateDescriptor<MutableType, MutableType> broadcastStateDescriptor2 = new MapStateDescriptor<>("test5", new JavaSerializer<MutableType>(), new JavaSerializer<MutableType>());
    MapStateDescriptor<MutableType, MutableType> broadcastStateDescriptor3 = new MapStateDescriptor<>("test6", new JavaSerializer<MutableType>(), new JavaSerializer<MutableType>());
    ListState<MutableType> listState1 = operatorStateBackend.getListState(stateDescriptor1);
    ListState<MutableType> listState2 = operatorStateBackend.getListState(stateDescriptor2);
    ListState<MutableType> listState3 = operatorStateBackend.getUnionListState(stateDescriptor3);
    BroadcastState<MutableType, MutableType> broadcastState1 = operatorStateBackend.getBroadcastState(broadcastStateDescriptor1);
    BroadcastState<MutableType, MutableType> broadcastState2 = operatorStateBackend.getBroadcastState(broadcastStateDescriptor2);
    BroadcastState<MutableType, MutableType> broadcastState3 = operatorStateBackend.getBroadcastState(broadcastStateDescriptor3);
    listState1.add(MutableType.of(42));
    listState1.add(MutableType.of(4711));
    listState2.add(MutableType.of(7));
    listState2.add(MutableType.of(13));
    listState2.add(MutableType.of(23));
    listState3.add(MutableType.of(17));
    listState3.add(MutableType.of(18));
    listState3.add(MutableType.of(19));
    listState3.add(MutableType.of(20));
    broadcastState1.put(MutableType.of(1), MutableType.of(2));
    broadcastState1.put(MutableType.of(2), MutableType.of(5));
    broadcastState2.put(MutableType.of(2), MutableType.of(5));
    BlockerCheckpointStreamFactory streamFactory = new BlockerCheckpointStreamFactory(1024 * 1024);
    OneShotLatch waiterLatch = new OneShotLatch();
    OneShotLatch blockerLatch = new OneShotLatch();
    streamFactory.setWaiterLatch(waiterLatch);
    streamFactory.setBlockerLatch(blockerLatch);
    RunnableFuture<SnapshotResult<OperatorStateHandle>> runnableFuture = operatorStateBackend.snapshot(1, 1, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());
    ExecutorService executorService = Executors.newFixedThreadPool(1);
    executorService.submit(runnableFuture);
    // wait until the async checkpoint is in the write code, then continue
    waiterLatch.await();
    // do some mutations to the state, to test if our snapshot will NOT reflect them
    listState1.add(MutableType.of(77));
    broadcastState1.put(MutableType.of(32), MutableType.of(97));
    int n = 0;
    for (MutableType mutableType : listState2.get()) {
        if (++n == 2) {
            // allow the write code to continue, so that we could do changes while state is
            // written in parallel.
            blockerLatch.trigger();
        }
        mutableType.setValue(mutableType.getValue() + 10);
    }
    listState3.clear();
    broadcastState2.clear();
    operatorStateBackend.getListState(new ListStateDescriptor<>("test4", new JavaSerializer<MutableType>()));
    // run the snapshot
    SnapshotResult<OperatorStateHandle> snapshotResult = runnableFuture.get();
    OperatorStateHandle stateHandle = snapshotResult.getJobManagerOwnedSnapshot();
    try {
        operatorStateBackend.close();
        operatorStateBackend.dispose();
        AbstractStateBackend abstractStateBackend = new MemoryStateBackend(4096);
        CloseableRegistry cancelStreamRegistry = new CloseableRegistry();
        operatorStateBackend = abstractStateBackend.createOperatorStateBackend(createMockEnvironment(), "testOperator", StateObjectCollection.singleton(stateHandle), cancelStreamRegistry);
        assertEquals(3, operatorStateBackend.getRegisteredStateNames().size());
        assertEquals(3, operatorStateBackend.getRegisteredBroadcastStateNames().size());
        listState1 = operatorStateBackend.getListState(stateDescriptor1);
        listState2 = operatorStateBackend.getListState(stateDescriptor2);
        listState3 = operatorStateBackend.getUnionListState(stateDescriptor3);
        broadcastState1 = operatorStateBackend.getBroadcastState(broadcastStateDescriptor1);
        broadcastState2 = operatorStateBackend.getBroadcastState(broadcastStateDescriptor2);
        broadcastState3 = operatorStateBackend.getBroadcastState(broadcastStateDescriptor3);
        assertEquals(3, operatorStateBackend.getRegisteredStateNames().size());
        assertEquals(3, operatorStateBackend.getRegisteredBroadcastStateNames().size());
        Iterator<MutableType> it = listState1.get().iterator();
        assertEquals(42, it.next().value);
        assertEquals(4711, it.next().value);
        assertFalse(it.hasNext());
        it = listState2.get().iterator();
        assertEquals(7, it.next().value);
        assertEquals(13, it.next().value);
        assertEquals(23, it.next().value);
        assertFalse(it.hasNext());
        it = listState3.get().iterator();
        assertEquals(17, it.next().value);
        assertEquals(18, it.next().value);
        assertEquals(19, it.next().value);
        assertEquals(20, it.next().value);
        assertFalse(it.hasNext());
        Iterator<Map.Entry<MutableType, MutableType>> bIt = broadcastState1.iterator();
        assertTrue(bIt.hasNext());
        Map.Entry<MutableType, MutableType> entry = bIt.next();
        assertEquals(1, entry.getKey().value);
        assertEquals(2, entry.getValue().value);
        assertTrue(bIt.hasNext());
        entry = bIt.next();
        assertEquals(2, entry.getKey().value);
        assertEquals(5, entry.getValue().value);
        assertFalse(bIt.hasNext());
        bIt = broadcastState2.iterator();
        assertTrue(bIt.hasNext());
        entry = bIt.next();
        assertEquals(2, entry.getKey().value);
        assertEquals(5, entry.getValue().value);
        assertFalse(bIt.hasNext());
        bIt = broadcastState3.iterator();
        assertFalse(bIt.hasNext());
        operatorStateBackend.close();
        operatorStateBackend.dispose();
    } finally {
        stateHandle.discardState();
    }
    executorService.shutdown();
}

Example 19 with MapStateDescriptor

use of org.apache.flink.api.common.state.MapStateDescriptor in project flink by apache.

the class OperatorStateBackendTest method testSnapshotRestoreSync.

@Test
public void testSnapshotRestoreSync() throws Exception {
    AbstractStateBackend abstractStateBackend = new MemoryStateBackend(2 * 4096);
    OperatorStateBackend operatorStateBackend = abstractStateBackend.createOperatorStateBackend(createMockEnvironment(), "test-op-name", emptyStateHandles, new CloseableRegistry());
    ListStateDescriptor<Serializable> stateDescriptor1 = new ListStateDescriptor<>("test1", new JavaSerializer<>());
    ListStateDescriptor<Serializable> stateDescriptor2 = new ListStateDescriptor<>("test2", new JavaSerializer<>());
    ListStateDescriptor<Serializable> stateDescriptor3 = new ListStateDescriptor<>("test3", new JavaSerializer<>());
    MapStateDescriptor<Serializable, Serializable> broadcastStateDescriptor1 = new MapStateDescriptor<>("test4", new JavaSerializer<>(), new JavaSerializer<>());
    MapStateDescriptor<Serializable, Serializable> broadcastStateDescriptor2 = new MapStateDescriptor<>("test5", new JavaSerializer<>(), new JavaSerializer<>());
    MapStateDescriptor<Serializable, Serializable> broadcastStateDescriptor3 = new MapStateDescriptor<>("test6", new JavaSerializer<>(), new JavaSerializer<>());
    ListState<Serializable> listState1 = operatorStateBackend.getListState(stateDescriptor1);
    ListState<Serializable> listState2 = operatorStateBackend.getListState(stateDescriptor2);
    ListState<Serializable> listState3 = operatorStateBackend.getUnionListState(stateDescriptor3);
    BroadcastState<Serializable, Serializable> broadcastState1 = operatorStateBackend.getBroadcastState(broadcastStateDescriptor1);
    BroadcastState<Serializable, Serializable> broadcastState2 = operatorStateBackend.getBroadcastState(broadcastStateDescriptor2);
    BroadcastState<Serializable, Serializable> broadcastState3 = operatorStateBackend.getBroadcastState(broadcastStateDescriptor3);
    listState1.add(42);
    listState1.add(4711);
    listState2.add(7);
    listState2.add(13);
    listState2.add(23);
    listState3.add(17);
    listState3.add(18);
    listState3.add(19);
    listState3.add(20);
    broadcastState1.put(1, 2);
    broadcastState1.put(2, 5);
    broadcastState2.put(2, 5);
    CheckpointStreamFactory streamFactory = new MemCheckpointStreamFactory(2 * 4096);
    RunnableFuture<SnapshotResult<OperatorStateHandle>> snapshot = operatorStateBackend.snapshot(1L, 1L, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());
    SnapshotResult<OperatorStateHandle> snapshotResult = FutureUtils.runIfNotDoneAndGet(snapshot);
    OperatorStateHandle stateHandle = snapshotResult.getJobManagerOwnedSnapshot();
    try {
        operatorStateBackend.close();
        operatorStateBackend.dispose();
        operatorStateBackend = abstractStateBackend.createOperatorStateBackend(createMockEnvironment(), "testOperator", StateObjectCollection.singleton(stateHandle), new CloseableRegistry());
        assertEquals(3, operatorStateBackend.getRegisteredStateNames().size());
        assertEquals(3, operatorStateBackend.getRegisteredBroadcastStateNames().size());
        listState1 = operatorStateBackend.getListState(stateDescriptor1);
        listState2 = operatorStateBackend.getListState(stateDescriptor2);
        listState3 = operatorStateBackend.getUnionListState(stateDescriptor3);
        broadcastState1 = operatorStateBackend.getBroadcastState(broadcastStateDescriptor1);
        broadcastState2 = operatorStateBackend.getBroadcastState(broadcastStateDescriptor2);
        broadcastState3 = operatorStateBackend.getBroadcastState(broadcastStateDescriptor3);
        assertEquals(3, operatorStateBackend.getRegisteredStateNames().size());
        assertEquals(3, operatorStateBackend.getRegisteredBroadcastStateNames().size());
        Iterator<Serializable> it = listState1.get().iterator();
        assertEquals(42, it.next());
        assertEquals(4711, it.next());
        assertFalse(it.hasNext());
        it = listState2.get().iterator();
        assertEquals(7, it.next());
        assertEquals(13, it.next());
        assertEquals(23, it.next());
        assertFalse(it.hasNext());
        it = listState3.get().iterator();
        assertEquals(17, it.next());
        assertEquals(18, it.next());
        assertEquals(19, it.next());
        assertEquals(20, it.next());
        assertFalse(it.hasNext());
        Iterator<Map.Entry<Serializable, Serializable>> bIt = broadcastState1.iterator();
        assertTrue(bIt.hasNext());
        Map.Entry<Serializable, Serializable> entry = bIt.next();
        assertEquals(1, entry.getKey());
        assertEquals(2, entry.getValue());
        assertTrue(bIt.hasNext());
        entry = bIt.next();
        assertEquals(2, entry.getKey());
        assertEquals(5, entry.getValue());
        assertFalse(bIt.hasNext());
        bIt = broadcastState2.iterator();
        assertTrue(bIt.hasNext());
        entry = bIt.next();
        assertEquals(2, entry.getKey());
        assertEquals(5, entry.getValue());
        assertFalse(bIt.hasNext());
        bIt = broadcastState3.iterator();
        assertFalse(bIt.hasNext());
        operatorStateBackend.close();
        operatorStateBackend.dispose();
    } finally {
        stateHandle.discardState();
    }
}

Example 20 with MapStateDescriptor

use of org.apache.flink.api.common.state.MapStateDescriptor in project flink by apache.

the class OperatorStateBackendTest method testSnapshotAsyncClose.

@Test
public void testSnapshotAsyncClose() throws Exception {
    DefaultOperatorStateBackend operatorStateBackend = new DefaultOperatorStateBackendBuilder(OperatorStateBackendTest.class.getClassLoader(), new ExecutionConfig(), true, emptyStateHandles, new CloseableRegistry()).build();
    ListStateDescriptor<MutableType> stateDescriptor1 = new ListStateDescriptor<>("test1", new JavaSerializer<MutableType>());
    ListState<MutableType> listState1 = operatorStateBackend.getListState(stateDescriptor1);
    listState1.add(MutableType.of(42));
    listState1.add(MutableType.of(4711));
    MapStateDescriptor<MutableType, MutableType> broadcastStateDescriptor1 = new MapStateDescriptor<>("test4", new JavaSerializer<MutableType>(), new JavaSerializer<MutableType>());
    BroadcastState<MutableType, MutableType> broadcastState1 = operatorStateBackend.getBroadcastState(broadcastStateDescriptor1);
    broadcastState1.put(MutableType.of(1), MutableType.of(2));
    broadcastState1.put(MutableType.of(2), MutableType.of(5));
    BlockerCheckpointStreamFactory streamFactory = new BlockerCheckpointStreamFactory(1024 * 1024);
    OneShotLatch waiterLatch = new OneShotLatch();
    OneShotLatch blockerLatch = new OneShotLatch();
    streamFactory.setWaiterLatch(waiterLatch);
    streamFactory.setBlockerLatch(blockerLatch);
    RunnableFuture<SnapshotResult<OperatorStateHandle>> runnableFuture = operatorStateBackend.snapshot(1, 1, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());
    ExecutorService executorService = Executors.newFixedThreadPool(1);
    executorService.submit(runnableFuture);
    // wait until the async checkpoint is in the write code, then continue
    waiterLatch.await();
    operatorStateBackend.close();
    blockerLatch.trigger();
    try {
        runnableFuture.get(60, TimeUnit.SECONDS);
        Assert.fail();
    } catch (CancellationException expected) {
    }
}