@Override public <S extends State> S getPartitionedState(StateDescriptor<S, ?> stateDescriptor) { try { return stateBackend.getPartitionedState(window, windowSerializer, stateDescriptor); } catch (Exception e) { throw new RuntimeException("Error getting state", e); } }
@Override public void open() throws Exception { super.open(); this.state = getKeyedStateBackend().getPartitionedState( VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, stateDescriptor); }
@SuppressWarnings("unchecked") private static AbstractStreamOperator<?> createDescriptorCapturingMockOp( final AtomicReference<Object> ref, final ExecutionConfig config) throws Exception { AbstractStreamOperator<?> operatorMock = mock(AbstractStreamOperator.class); KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class); DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config); when(operatorMock.getExecutionConfig()).thenReturn(config); doAnswer(new Answer<Object>() { @Override public Object answer(InvocationOnMock invocationOnMock) throws Throwable { ref.set(invocationOnMock.getArguments()[2]); return null; } }).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(StateDescriptor.class)); when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore); when(operatorMock.getOperatorID()).thenReturn(new OperatorID()); return operatorMock; }
@SuppressWarnings("unchecked") private static AbstractStreamOperator<?> createMapPlainMockOp() throws Exception { AbstractStreamOperator<?> operatorMock = mock(AbstractStreamOperator.class); ExecutionConfig config = new ExecutionConfig(); KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class); DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config); when(operatorMock.getExecutionConfig()).thenReturn(config); doAnswer(new Answer<MapState<Integer, String>>() { @Override public MapState<Integer, String> answer(InvocationOnMock invocationOnMock) throws Throwable { MapStateDescriptor<Integer, String> descr = (MapStateDescriptor<Integer, String>) invocationOnMock.getArguments()[2]; AbstractKeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend( new DummyEnvironment("test_task", 1, 0), new JobID(), "test_op", IntSerializer.INSTANCE, 1, new KeyGroupRange(0, 0), new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID())); backend.setCurrentKey(0); return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr); } }).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(MapStateDescriptor.class)); when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore); when(operatorMock.getOperatorID()).thenReturn(new OperatorID()); return operatorMock; }
@SuppressWarnings("unchecked") private static AbstractStreamOperator<?> createListPlainMockOp() throws Exception { AbstractStreamOperator<?> operatorMock = mock(AbstractStreamOperator.class); ExecutionConfig config = new ExecutionConfig(); KeyedStateBackend keyedStateBackend = mock(KeyedStateBackend.class); DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, config); when(operatorMock.getExecutionConfig()).thenReturn(config); doAnswer(new Answer<ListState<String>>() { @Override public ListState<String> answer(InvocationOnMock invocationOnMock) throws Throwable { ListStateDescriptor<String> descr = (ListStateDescriptor<String>) invocationOnMock.getArguments()[2]; AbstractKeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend( new DummyEnvironment("test_task", 1, 0), new JobID(), "test_op", IntSerializer.INSTANCE, 1, new KeyGroupRange(0, 0), new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID())); backend.setCurrentKey(0); return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr); } }).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(ListStateDescriptor.class)); when(operatorMock.getKeyedStateStore()).thenReturn(keyedStateStore); when(operatorMock.getOperatorID()).thenReturn(new OperatorID()); return operatorMock; }
@Override public void onProcessingTime(InternalTimer<Long, Long> timer) throws Exception { ValueState<Long> state = getKeyedStateBackend().getPartitionedState( timer.getNamespace(), LongSerializer.INSTANCE, stateDescriptor); assertEquals(state.value(), timer.getNamespace()); getRuntimeContext().getAccumulator(SUCCESSFUL_PROCESSING_TIME_CHECK_ACCUMULATOR).add(1); } }
@Override public void onProcessingTime(InternalTimer<Long, Long> timer) throws Exception { ValueState<Long> state = getKeyedStateBackend().getPartitionedState( timer.getNamespace(), LongSerializer.INSTANCE, stateDescriptor); assertEquals(state.value(), timer.getNamespace()); getRuntimeContext().getAccumulator(SUCCESSFUL_PROCESSING_TIME_CHECK_ACCUMULATOR).add(1); } }
@Override public void onEventTime(InternalTimer<Long, Long> timer) throws Exception { ValueState<Long> state = getKeyedStateBackend().getPartitionedState( timer.getNamespace(), LongSerializer.INSTANCE, stateDescriptor); assertEquals(state.value(), timer.getNamespace()); getRuntimeContext().getAccumulator(SUCCESSFUL_EVENT_TIME_CHECK_ACCUMULATOR).add(1); }
@Override public void onEventTime(InternalTimer<Long, Long> timer) throws Exception { ValueState<Long> state = getKeyedStateBackend().getPartitionedState( timer.getNamespace(), LongSerializer.INSTANCE, stateDescriptor); assertEquals(state.value(), timer.getNamespace()); getRuntimeContext().getAccumulator(SUCCESSFUL_EVENT_TIME_CHECK_ACCUMULATOR).add(1); }
@Override public void processElement(StreamRecord<Tuple2<Long, Long>> element) throws Exception { ValueState<Long> state = getKeyedStateBackend().getPartitionedState( element.getValue().f0, LongSerializer.INSTANCE, stateDescriptor); assertEquals(state.value(), element.getValue().f1); getRuntimeContext().getAccumulator(SUCCESSFUL_PROCESS_CHECK_ACCUMULATOR).add(1); output.collect(element); }
@Override public void processElement(StreamRecord<Tuple2<Long, Long>> element) throws Exception { ValueState<Long> state = getKeyedStateBackend().getPartitionedState( element.getValue().f0, LongSerializer.INSTANCE, stateDescriptor); assertEquals(state.value(), element.getValue().f1); getRuntimeContext().getAccumulator(SUCCESSFUL_PROCESS_CHECK_ACCUMULATOR).add(1); output.collect(element); }
desc.setQueryable("vanilla"); backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc);
desc.setQueryable("vanilla"); backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, desc);
protected <S extends State> S getPartitionedState(StateDescriptor<S, ?> stateDescriptor) throws Exception { return keyedStateBackend.getPartitionedState( VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, stateDescriptor); } }
@Override public AccumT getAccum() { try { return flinkStateBackend .getPartitionedState( namespace.stringKey(), StringSerializer.INSTANCE, flinkStateDescriptor) .value(); } catch (Exception e) { throw new RuntimeException("Error reading state.", e); } }
@Override public Boolean read() { try { return flinkStateBackend .getPartitionedState( namespace.stringKey(), StringSerializer.INSTANCE, flinkStateDescriptor) .value() == null; } catch (Exception e) { throw new RuntimeException("Error reading state.", e); } }
@Override public void put(KeyT key, ValueT value) { try { flinkStateBackend .getPartitionedState( namespace.stringKey(), StringSerializer.INSTANCE, flinkStateDescriptor) .put(key, value); } catch (Exception e) { throw new RuntimeException("Error put kv to state.", e); } }
@Override public void clear() { try { flinkStateBackend .getPartitionedState( namespace.stringKey(), StringSerializer.INSTANCE, flinkStateDescriptor) .clear(); } catch (Exception e) { throw new RuntimeException("Error clearing state.", e); } }
@Override public ReadableState<ValueT> get(final KeyT input) { try { return ReadableStates.immediate( flinkStateBackend .getPartitionedState( namespace.stringKey(), StringSerializer.INSTANCE, flinkStateDescriptor) .get(input)); } catch (Exception e) { throw new RuntimeException("Error get from state.", e); } }
@Override public void pushBack(T element) throws Exception { ListState<T> state = backend.getPartitionedState( VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, stateDescriptor); backend.setCurrentKey(keySelector.getKey(element)); state.add(element); }