private KeyedStateStore checkPreconditionsAndGetKeyedStateStore(StateDescriptor<?, ?> stateDescriptor) { Preconditions.checkNotNull(stateDescriptor, "The state properties must not be null"); KeyedStateStore keyedStateStore = operator.getKeyedStateStore(); Preconditions.checkNotNull(keyedStateStore, "Keyed state can only be used on a 'keyed stream', i.e., after a 'keyBy()' operation."); return keyedStateStore; }
@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 keyDone() throws Exception { contextOperator.setCurrentKey(currentKey); context.globalState = contextOperator.getKeyedStateStore(); function.process(currentKey, window, context, unionIterator, out); } }
private KeyedStateStore checkPreconditionsAndGetKeyedStateStore(StateDescriptor<?, ?> stateDescriptor) { Preconditions.checkNotNull(stateDescriptor, "The state properties must not be null"); KeyedStateStore keyedStateStore = operator.getKeyedStateStore(); Preconditions.checkNotNull(keyedStateStore, "Keyed state can only be used on a 'keyed stream', i.e., after a 'keyBy()' operation."); return keyedStateStore; }
private KeyedStateStore checkPreconditionsAndGetKeyedStateStore(StateDescriptor<?, ?> stateDescriptor) { Preconditions.checkNotNull(stateDescriptor, "The state properties must not be null"); KeyedStateStore keyedStateStore = operator.getKeyedStateStore(); Preconditions.checkNotNull(keyedStateStore, "Keyed state can only be used on a 'keyed stream', i.e., after a 'keyBy()' operation."); return keyedStateStore; }
private KeyedStateStore checkPreconditionsAndGetKeyedStateStore(StateDescriptor<?, ?> stateDescriptor) { Preconditions.checkNotNull(stateDescriptor, "The state properties must not be null"); KeyedStateStore keyedStateStore = operator.getKeyedStateStore(); Preconditions.checkNotNull(keyedStateStore, "Keyed state can only be used on a 'keyed stream', i.e., after a 'keyBy()' operation."); return keyedStateStore; }
@Override public void evaluateWindow(Collector<Result> out, final TimeWindow window, AbstractStreamOperator<Result> operator) throws Exception { if (previousPanes.isEmpty()) { // optimized path for single pane case (tumbling window) for (KeyMap.Entry<Key, ArrayList<Type>> entry : latestPane) { Key key = entry.getKey(); operator.setCurrentKey(key); context.globalState = operator.getKeyedStateStore(); function.process(entry.getKey(), window, context, entry.getValue(), out); } } else { // general code path for multi-pane case WindowFunctionTraversal<Key, Type, Result> evaluator = new WindowFunctionTraversal<>( function, window, out, operator, context); traverseAllPanes(evaluator, evaluationPass); } evaluationPass++; }