Examples with OperatorStateHandle org.apache.flink.runtime.state.OperatorStateHandle used on opensource projects

Example 46 with OperatorStateHandle

use of org.apache.flink.runtime.state.OperatorStateHandle in project flink by apache.

the class AbstractStreamOperatorTestHarness method repackageState.

/**
 * Takes the different {@link OperatorSubtaskState} created by calling {@link #snapshot(long,
 * long)} on different instances of {@link AbstractStreamOperatorTestHarness} (each one
 * representing one subtask) and repacks them into a single {@link OperatorSubtaskState} so that
 * the parallelism of the test can change arbitrarily (i.e. be able to scale both up and down).
 *
 * <p>After repacking the partial states, remember to use {@link
 * #repartitionOperatorState(OperatorSubtaskState, int, int, int, int)} to reshape the state
 * handles to include only those key-group states in the local key-group range and the operator
 * states that would be assigned to the local subtask. Bear in mind that for parallelism greater
 * than one, you have to use the constructor {@link
 * #AbstractStreamOperatorTestHarness(StreamOperator, int, int, int)}.
 *
 * <p><b>NOTE: </b> each of the {@code handles} in the argument list is assumed to be from a
 * single task of a single operator (i.e. chain length of one).
 *
 * <p>For an example of how to use it, have a look at {@link
 * AbstractStreamOperatorTest#testStateAndTimerStateShufflingScalingDown()}.
 *
 * @param handles the different states to be merged.
 * @return the resulting state, or {@code null} if no partial states are specified.
 */
public static OperatorSubtaskState repackageState(OperatorSubtaskState... handles) throws Exception {
    if (handles.length < 1) {
        return null;
    } else if (handles.length == 1) {
        return handles[0];
    }
    List<OperatorStateHandle> mergedManagedOperatorState = new ArrayList<>(handles.length);
    List<OperatorStateHandle> mergedRawOperatorState = new ArrayList<>(handles.length);
    List<KeyedStateHandle> mergedManagedKeyedState = new ArrayList<>(handles.length);
    List<KeyedStateHandle> mergedRawKeyedState = new ArrayList<>(handles.length);
    for (OperatorSubtaskState handle : handles) {
        Collection<OperatorStateHandle> managedOperatorState = handle.getManagedOperatorState();
        Collection<OperatorStateHandle> rawOperatorState = handle.getRawOperatorState();
        Collection<KeyedStateHandle> managedKeyedState = handle.getManagedKeyedState();
        Collection<KeyedStateHandle> rawKeyedState = handle.getRawKeyedState();
        mergedManagedOperatorState.addAll(managedOperatorState);
        mergedRawOperatorState.addAll(rawOperatorState);
        mergedManagedKeyedState.addAll(managedKeyedState);
        mergedRawKeyedState.addAll(rawKeyedState);
    }
    return OperatorSubtaskState.builder().setManagedOperatorState(new StateObjectCollection<>(mergedManagedOperatorState)).setRawOperatorState(new StateObjectCollection<>(mergedRawOperatorState)).setManagedKeyedState(new StateObjectCollection<>(mergedManagedKeyedState)).setRawKeyedState(new StateObjectCollection<>(mergedRawKeyedState)).build();
}

Example 47 with OperatorStateHandle

use of org.apache.flink.runtime.state.OperatorStateHandle in project flink by apache.

the class RoundRobinOperatorStateRepartitioner method groupByStateName.

/**
	 * Group by the different named states.
	 */
@SuppressWarnings("unchecked, rawtype")
private GroupByStateNameResults groupByStateName(List<OperatorStateHandle> previousParallelSubtaskStates) {
    //Reorganize: group by (State Name -> StreamStateHandle + StateMetaInfo)
    EnumMap<OperatorStateHandle.Mode, Map<String, List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>>>> nameToStateByMode = new EnumMap<>(OperatorStateHandle.Mode.class);
    for (OperatorStateHandle.Mode mode : OperatorStateHandle.Mode.values()) {
        Map<String, List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>>> map = new HashMap<>();
        nameToStateByMode.put(mode, new HashMap<String, List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>>>());
    }
    for (OperatorStateHandle psh : previousParallelSubtaskStates) {
        for (Map.Entry<String, OperatorStateHandle.StateMetaInfo> e : psh.getStateNameToPartitionOffsets().entrySet()) {
            OperatorStateHandle.StateMetaInfo metaInfo = e.getValue();
            Map<String, List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>>> nameToState = nameToStateByMode.get(metaInfo.getDistributionMode());
            List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>> stateLocations = nameToState.get(e.getKey());
            if (stateLocations == null) {
                stateLocations = new ArrayList<>();
                nameToState.put(e.getKey(), stateLocations);
            }
            stateLocations.add(new Tuple2<>(psh.getDelegateStateHandle(), e.getValue()));
        }
    }
    return new GroupByStateNameResults(nameToStateByMode);
}

Example 48 with OperatorStateHandle

use of org.apache.flink.runtime.state.OperatorStateHandle in project flink by apache.

the class CheckpointCoordinatorTest method testReplicateModeStateHandle.

@Test
public void testReplicateModeStateHandle() {
    Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap = new HashMap<>(1);
    metaInfoMap.put("t-1", new OperatorStateHandle.StateMetaInfo(new long[] { 0, 23 }, OperatorStateHandle.Mode.BROADCAST));
    metaInfoMap.put("t-2", new OperatorStateHandle.StateMetaInfo(new long[] { 42, 64 }, OperatorStateHandle.Mode.BROADCAST));
    metaInfoMap.put("t-3", new OperatorStateHandle.StateMetaInfo(new long[] { 72, 83 }, OperatorStateHandle.Mode.SPLIT_DISTRIBUTE));
    OperatorStateHandle osh = new OperatorStateHandle(metaInfoMap, new ByteStreamStateHandle("test", new byte[100]));
    OperatorStateRepartitioner repartitioner = RoundRobinOperatorStateRepartitioner.INSTANCE;
    List<Collection<OperatorStateHandle>> repartitionedStates = repartitioner.repartitionState(Collections.singletonList(osh), 3);
    Map<String, Integer> checkCounts = new HashMap<>(3);
    for (Collection<OperatorStateHandle> operatorStateHandles : repartitionedStates) {
        for (OperatorStateHandle operatorStateHandle : operatorStateHandles) {
            for (Map.Entry<String, OperatorStateHandle.StateMetaInfo> stateNameToMetaInfo : operatorStateHandle.getStateNameToPartitionOffsets().entrySet()) {
                String stateName = stateNameToMetaInfo.getKey();
                Integer count = checkCounts.get(stateName);
                if (null == count) {
                    checkCounts.put(stateName, 1);
                } else {
                    checkCounts.put(stateName, 1 + count);
                }
                OperatorStateHandle.StateMetaInfo stateMetaInfo = stateNameToMetaInfo.getValue();
                if (OperatorStateHandle.Mode.SPLIT_DISTRIBUTE.equals(stateMetaInfo.getDistributionMode())) {
                    Assert.assertEquals(1, stateNameToMetaInfo.getValue().getOffsets().length);
                } else {
                    Assert.assertEquals(2, stateNameToMetaInfo.getValue().getOffsets().length);
                }
            }
        }
    }
    Assert.assertEquals(3, checkCounts.size());
    Assert.assertEquals(3, checkCounts.get("t-1").intValue());
    Assert.assertEquals(3, checkCounts.get("t-2").intValue());
    Assert.assertEquals(2, checkCounts.get("t-3").intValue());
}

Example 49 with OperatorStateHandle

use of org.apache.flink.runtime.state.OperatorStateHandle in project flink by apache.

the class AbstractStreamOperator method initializeState.

@Override
public final void initializeState(OperatorStateHandles stateHandles) throws Exception {
    Collection<KeyGroupsStateHandle> keyedStateHandlesRaw = null;
    Collection<OperatorStateHandle> operatorStateHandlesRaw = null;
    Collection<OperatorStateHandle> operatorStateHandlesBackend = null;
    boolean restoring = null != stateHandles;
    //TODO we should move the actual initialization of this from StreamTask to this class
    initKeyedState();
    if (getKeyedStateBackend() != null && timeServiceManager == null) {
        timeServiceManager = new InternalTimeServiceManager<>(getKeyedStateBackend().getNumberOfKeyGroups(), getKeyedStateBackend().getKeyGroupRange(), this, getRuntimeContext().getProcessingTimeService());
    }
    if (restoring) {
        //pass directly
        operatorStateHandlesBackend = stateHandles.getManagedOperatorState();
        operatorStateHandlesRaw = stateHandles.getRawOperatorState();
        if (null != getKeyedStateBackend()) {
            //only use the keyed state if it is meant for us (aka head operator)
            keyedStateHandlesRaw = stateHandles.getRawKeyedState();
        }
    }
    checkpointStreamFactory = container.createCheckpointStreamFactory(this);
    initOperatorState(operatorStateHandlesBackend);
    StateInitializationContext initializationContext = new StateInitializationContextImpl(// information whether we restore or start for the first time
    restoring, // access to operator state backend
    operatorStateBackend, // access to keyed state backend
    keyedStateStore, // access to keyed state stream
    keyedStateHandlesRaw, // access to operator state stream
    operatorStateHandlesRaw, // access to register streams for canceling
    getContainingTask().getCancelables());
    initializeState(initializationContext);
    if (restoring) {
        // finally restore the legacy state in case we are
        // migrating from a previous Flink version.
        restoreStreamCheckpointed(stateHandles);
    }
}

Example 50 with OperatorStateHandle

use of org.apache.flink.runtime.state.OperatorStateHandle in project flink by apache.

the class OperatorSnapshotResultTest method testCancelAndCleanup.

/**
	 * Tests that all runnable futures in an OperatorSnapshotResult are properly cancelled and if
	 * the StreamStateHandle result is retrievable that the state handle are discarded.
	 */
@Test
public void testCancelAndCleanup() throws Exception {
    OperatorSnapshotResult operatorSnapshotResult = new OperatorSnapshotResult();
    operatorSnapshotResult.cancel();
    KeyGroupsStateHandle keyedManagedStateHandle = mock(KeyGroupsStateHandle.class);
    RunnableFuture<KeyGroupsStateHandle> keyedStateManagedFuture = mock(RunnableFuture.class);
    when(keyedStateManagedFuture.get()).thenReturn(keyedManagedStateHandle);
    KeyGroupsStateHandle keyedRawStateHandle = mock(KeyGroupsStateHandle.class);
    RunnableFuture<KeyGroupsStateHandle> keyedStateRawFuture = mock(RunnableFuture.class);
    when(keyedStateRawFuture.get()).thenReturn(keyedRawStateHandle);
    OperatorStateHandle operatorManagedStateHandle = mock(OperatorStateHandle.class);
    RunnableFuture<OperatorStateHandle> operatorStateManagedFuture = mock(RunnableFuture.class);
    when(operatorStateManagedFuture.get()).thenReturn(operatorManagedStateHandle);
    OperatorStateHandle operatorRawStateHandle = mock(OperatorStateHandle.class);
    RunnableFuture<OperatorStateHandle> operatorStateRawFuture = mock(RunnableFuture.class);
    when(operatorStateRawFuture.get()).thenReturn(operatorRawStateHandle);
    operatorSnapshotResult = new OperatorSnapshotResult(keyedStateManagedFuture, keyedStateRawFuture, operatorStateManagedFuture, operatorStateRawFuture);
    operatorSnapshotResult.cancel();
    verify(keyedStateManagedFuture).cancel(true);
    verify(keyedStateRawFuture).cancel(true);
    verify(operatorStateManagedFuture).cancel(true);
    verify(operatorStateRawFuture).cancel(true);
    verify(keyedManagedStateHandle).discardState();
    verify(keyedRawStateHandle).discardState();
    verify(operatorManagedStateHandle).discardState();
    verify(operatorRawStateHandle).discardState();
}