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

Example 6 with KeyedStateStore

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

the class WindowOperatorContractTest method testStateTypeIsConsistentFromWindowStateAndGlobalState.

@Test
public void testStateTypeIsConsistentFromWindowStateAndGlobalState() throws Exception {
    class NoOpAggregateFunction implements AggregateFunction<String, String, String> {

        @Override
        public String createAccumulator() {
            return null;
        }

        @Override
        public String add(String value, String accumulator) {
            return null;
        }

        @Override
        public String getResult(String accumulator) {
            return null;
        }

        @Override
        public String merge(String a, String b) {
            return null;
        }
    }
    WindowAssigner<Integer, TimeWindow> mockAssigner = mockTimeWindowAssigner();
    Trigger<Integer, TimeWindow> mockTrigger = mockTrigger();
    InternalWindowFunction<Iterable<Integer>, Void, Integer, TimeWindow> mockWindowFunction = mockWindowFunction();
    KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Void> testHarness = createWindowOperator(mockAssigner, mockTrigger, 20L, mockWindowFunction);
    testHarness.open();
    when(mockTrigger.onElement(anyInt(), anyLong(), anyTimeWindow(), anyTriggerContext())).thenReturn(TriggerResult.FIRE);
    when(mockAssigner.assignWindows(anyInt(), anyLong(), anyAssignerContext())).thenReturn(Arrays.asList(new TimeWindow(0, 20)));
    AtomicBoolean processWasInvoked = new AtomicBoolean(false);
    doAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
            InternalWindowFunction.InternalWindowContext context = (InternalWindowFunction.InternalWindowContext) invocationOnMock.getArguments()[2];
            KeyedStateStore windowKeyedStateStore = context.windowState();
            KeyedStateStore globalKeyedStateStore = context.globalState();
            ListStateDescriptor<String> windowListStateDescriptor = new ListStateDescriptor<String>("windowListState", String.class);
            ListStateDescriptor<String> globalListStateDescriptor = new ListStateDescriptor<String>("globalListState", String.class);
            assertEquals(windowKeyedStateStore.getListState(windowListStateDescriptor).getClass(), globalKeyedStateStore.getListState(globalListStateDescriptor).getClass());
            ValueStateDescriptor<String> windowValueStateDescriptor = new ValueStateDescriptor<String>("windowValueState", String.class);
            ValueStateDescriptor<String> globalValueStateDescriptor = new ValueStateDescriptor<String>("globalValueState", String.class);
            assertEquals(windowKeyedStateStore.getState(windowValueStateDescriptor).getClass(), globalKeyedStateStore.getState(globalValueStateDescriptor).getClass());
            AggregatingStateDescriptor<String, String, String> windowAggStateDesc = new AggregatingStateDescriptor<String, String, String>("windowAgg", new NoOpAggregateFunction(), String.class);
            AggregatingStateDescriptor<String, String, String> globalAggStateDesc = new AggregatingStateDescriptor<String, String, String>("globalAgg", new NoOpAggregateFunction(), String.class);
            assertEquals(windowKeyedStateStore.getAggregatingState(windowAggStateDesc).getClass(), globalKeyedStateStore.getAggregatingState(globalAggStateDesc).getClass());
            ReducingStateDescriptor<String> windowReducingStateDesc = new ReducingStateDescriptor<String>("windowReducing", (a, b) -> a, String.class);
            ReducingStateDescriptor<String> globalReducingStateDesc = new ReducingStateDescriptor<String>("globalReducing", (a, b) -> a, String.class);
            assertEquals(windowKeyedStateStore.getReducingState(windowReducingStateDesc).getClass(), globalKeyedStateStore.getReducingState(globalReducingStateDesc).getClass());
            MapStateDescriptor<String, String> windowMapStateDescriptor = new MapStateDescriptor<String, String>("windowMapState", String.class, String.class);
            MapStateDescriptor<String, String> globalMapStateDescriptor = new MapStateDescriptor<String, String>("globalMapState", String.class, String.class);
            assertEquals(windowKeyedStateStore.getMapState(windowMapStateDescriptor).getClass(), globalKeyedStateStore.getMapState(globalMapStateDescriptor).getClass());
            processWasInvoked.set(true);
            return null;
        }
    }).when(mockWindowFunction).process(anyInt(), anyTimeWindow(), anyInternalWindowContext(), anyIntIterable(), WindowOperatorContractTest.<Void>anyCollector());
    testHarness.processElement(new StreamRecord<>(0, 0L));
    assertTrue(processWasInvoked.get());
}

Example 7 with KeyedStateStore

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

the class StreamingRuntimeContext method getListState.

@Override
public <T> ListState<T> getListState(ListStateDescriptor<T> stateProperties) {
    KeyedStateStore keyedStateStore = checkPreconditionsAndGetKeyedStateStore(stateProperties);
    stateProperties.initializeSerializerUnlessSet(getExecutionConfig());
    return keyedStateStore.getListState(stateProperties);
}

Example 8 with KeyedStateStore

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

the class StreamingRuntimeContext method getMapState.

@Override
public <UK, UV> MapState<UK, UV> getMapState(MapStateDescriptor<UK, UV> stateProperties) {
    KeyedStateStore keyedStateStore = checkPreconditionsAndGetKeyedStateStore(stateProperties);
    stateProperties.initializeSerializerUnlessSet(getExecutionConfig());
    return keyedStateStore.getMapState(stateProperties);
}

Example 9 with KeyedStateStore

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

the class StreamingRuntimeContext method getReducingState.

@Override
public <T> ReducingState<T> getReducingState(ReducingStateDescriptor<T> stateProperties) {
    KeyedStateStore keyedStateStore = checkPreconditionsAndGetKeyedStateStore(stateProperties);
    stateProperties.initializeSerializerUnlessSet(getExecutionConfig());
    return keyedStateStore.getReducingState(stateProperties);
}