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

Example 1 with OperatorStateHandle

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

the class SavepointV1Test method createTaskStates.

static Collection<TaskState> createTaskStates(int numTaskStates, int numSubtasksPerTask) throws IOException {
    Random random = new Random(numTaskStates * 31 + numSubtasksPerTask);
    List<TaskState> taskStates = new ArrayList<>(numTaskStates);
    for (int stateIdx = 0; stateIdx < numTaskStates; ++stateIdx) {
        int chainLength = 1 + random.nextInt(8);
        TaskState taskState = new TaskState(new JobVertexID(), numSubtasksPerTask, 128, chainLength);
        int noNonPartitionableStateAtIndex = random.nextInt(chainLength);
        int noOperatorStateBackendAtIndex = random.nextInt(chainLength);
        int noOperatorStateStreamAtIndex = random.nextInt(chainLength);
        boolean hasKeyedBackend = random.nextInt(4) != 0;
        boolean hasKeyedStream = random.nextInt(4) != 0;
        for (int subtaskIdx = 0; subtaskIdx < numSubtasksPerTask; subtaskIdx++) {
            List<StreamStateHandle> nonPartitionableStates = new ArrayList<>(chainLength);
            List<OperatorStateHandle> operatorStatesBackend = new ArrayList<>(chainLength);
            List<OperatorStateHandle> operatorStatesStream = new ArrayList<>(chainLength);
            for (int chainIdx = 0; chainIdx < chainLength; ++chainIdx) {
                StreamStateHandle nonPartitionableState = new TestByteStreamStateHandleDeepCompare("a-" + chainIdx, ("Hi-" + chainIdx).getBytes(ConfigConstants.DEFAULT_CHARSET));
                StreamStateHandle operatorStateBackend = new TestByteStreamStateHandleDeepCompare("b-" + chainIdx, ("Beautiful-" + chainIdx).getBytes(ConfigConstants.DEFAULT_CHARSET));
                StreamStateHandle operatorStateStream = new TestByteStreamStateHandleDeepCompare("b-" + chainIdx, ("Beautiful-" + chainIdx).getBytes(ConfigConstants.DEFAULT_CHARSET));
                Map<String, OperatorStateHandle.StateMetaInfo> offsetsMap = new HashMap<>();
                offsetsMap.put("A", new OperatorStateHandle.StateMetaInfo(new long[] { 0, 10, 20 }, OperatorStateHandle.Mode.SPLIT_DISTRIBUTE));
                offsetsMap.put("B", new OperatorStateHandle.StateMetaInfo(new long[] { 30, 40, 50 }, OperatorStateHandle.Mode.SPLIT_DISTRIBUTE));
                offsetsMap.put("C", new OperatorStateHandle.StateMetaInfo(new long[] { 60, 70, 80 }, OperatorStateHandle.Mode.BROADCAST));
                if (chainIdx != noNonPartitionableStateAtIndex) {
                    nonPartitionableStates.add(nonPartitionableState);
                }
                if (chainIdx != noOperatorStateBackendAtIndex) {
                    OperatorStateHandle operatorStateHandleBackend = new OperatorStateHandle(offsetsMap, operatorStateBackend);
                    operatorStatesBackend.add(operatorStateHandleBackend);
                }
                if (chainIdx != noOperatorStateStreamAtIndex) {
                    OperatorStateHandle operatorStateHandleStream = new OperatorStateHandle(offsetsMap, operatorStateStream);
                    operatorStatesStream.add(operatorStateHandleStream);
                }
            }
            KeyGroupsStateHandle keyedStateBackend = null;
            KeyGroupsStateHandle keyedStateStream = null;
            if (hasKeyedBackend) {
                keyedStateBackend = new KeyGroupsStateHandle(new KeyGroupRangeOffsets(1, 1, new long[] { 42 }), new TestByteStreamStateHandleDeepCompare("c", "Hello".getBytes(ConfigConstants.DEFAULT_CHARSET)));
            }
            if (hasKeyedStream) {
                keyedStateStream = new KeyGroupsStateHandle(new KeyGroupRangeOffsets(1, 1, new long[] { 23 }), new TestByteStreamStateHandleDeepCompare("d", "World".getBytes(ConfigConstants.DEFAULT_CHARSET)));
            }
            taskState.putState(subtaskIdx, new SubtaskState(new ChainedStateHandle<>(nonPartitionableStates), new ChainedStateHandle<>(operatorStatesBackend), new ChainedStateHandle<>(operatorStatesStream), keyedStateStream, keyedStateBackend));
        }
        taskStates.add(taskState);
    }
    return taskStates;
}

Example 2 with OperatorStateHandle

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

the class AbstractStreamOperatorTestHarness method initializeStateFromLegacyCheckpoint.

public void initializeStateFromLegacyCheckpoint(String checkpointFilename) throws Exception {
    FileInputStream fin = new FileInputStream(checkpointFilename);
    StreamTaskState state = MigrationInstantiationUtil.deserializeObject(fin, ClassLoader.getSystemClassLoader());
    fin.close();
    if (!setupCalled) {
        setup();
    }
    StreamStateHandle stateHandle = SavepointV0Serializer.convertOperatorAndFunctionState(state);
    List<KeyGroupsStateHandle> keyGroupStatesList = new ArrayList<>();
    if (state.getKvStates() != null) {
        KeyGroupsStateHandle keyedStateHandle = SavepointV0Serializer.convertKeyedBackendState(state.getKvStates(), environment.getTaskInfo().getIndexOfThisSubtask(), 0);
        keyGroupStatesList.add(keyedStateHandle);
    }
    // finally calling the initializeState() with the legacy operatorStateHandles
    initializeState(new OperatorStateHandles(0, stateHandle, keyGroupStatesList, Collections.<KeyGroupsStateHandle>emptyList(), Collections.<OperatorStateHandle>emptyList(), Collections.<OperatorStateHandle>emptyList()));
}

Example 3 with OperatorStateHandle

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

the class AbstractStreamOperatorTestHarness method initializeState.

/**
	 * Calls {@link org.apache.flink.streaming.api.operators.StreamOperator#initializeState(OperatorStateHandles)}.
	 * Calls {@link org.apache.flink.streaming.api.operators.StreamOperator#setup(StreamTask, StreamConfig, Output)}
	 * if it was not called before.
	 *
	 * <p>This will 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.
	 */
public void initializeState(OperatorStateHandles operatorStateHandles) throws Exception {
    if (!setupCalled) {
        setup();
    }
    if (operatorStateHandles != null) {
        int numKeyGroups = getEnvironment().getTaskInfo().getMaxNumberOfParallelSubtasks();
        int numSubtasks = getEnvironment().getTaskInfo().getNumberOfParallelSubtasks();
        int subtaskIndex = getEnvironment().getTaskInfo().getIndexOfThisSubtask();
        // create a new OperatorStateHandles that only contains the state for our key-groups
        List<KeyGroupRange> keyGroupPartitions = StateAssignmentOperation.createKeyGroupPartitions(numKeyGroups, numSubtasks);
        KeyGroupRange localKeyGroupRange = keyGroupPartitions.get(subtaskIndex);
        List<KeyGroupsStateHandle> localManagedKeyGroupState = null;
        if (operatorStateHandles.getManagedKeyedState() != null) {
            localManagedKeyGroupState = StateAssignmentOperation.getKeyGroupsStateHandles(operatorStateHandles.getManagedKeyedState(), localKeyGroupRange);
        }
        List<KeyGroupsStateHandle> localRawKeyGroupState = null;
        if (operatorStateHandles.getRawKeyedState() != null) {
            localRawKeyGroupState = StateAssignmentOperation.getKeyGroupsStateHandles(operatorStateHandles.getRawKeyedState(), localKeyGroupRange);
        }
        List<OperatorStateHandle> managedOperatorState = new ArrayList<>();
        if (operatorStateHandles.getManagedOperatorState() != null) {
            managedOperatorState.addAll(operatorStateHandles.getManagedOperatorState());
        }
        Collection<OperatorStateHandle> localManagedOperatorState = operatorStateRepartitioner.repartitionState(managedOperatorState, numSubtasks).get(subtaskIndex);
        List<OperatorStateHandle> rawOperatorState = new ArrayList<>();
        if (operatorStateHandles.getRawOperatorState() != null) {
            rawOperatorState.addAll(operatorStateHandles.getRawOperatorState());
        }
        Collection<OperatorStateHandle> localRawOperatorState = operatorStateRepartitioner.repartitionState(rawOperatorState, numSubtasks).get(subtaskIndex);
        OperatorStateHandles massagedOperatorStateHandles = new OperatorStateHandles(0, operatorStateHandles.getLegacyOperatorState(), localManagedKeyGroupState, localRawKeyGroupState, localManagedOperatorState, localRawOperatorState);
        operator.initializeState(massagedOperatorStateHandles);
    } else {
        operator.initializeState(null);
    }
    initializeCalled = true;
}

Example 4 with OperatorStateHandle

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

the class AbstractStreamOperatorTestHarness method snapshot.

/**
	 * Calls {@link StreamOperator#snapshotState(long, long, CheckpointOptions)}.
	 */
public OperatorStateHandles snapshot(long checkpointId, long timestamp) throws Exception {
    CheckpointStreamFactory streamFactory = stateBackend.createStreamFactory(new JobID(), "test_op");
    OperatorSnapshotResult operatorStateResult = operator.snapshotState(checkpointId, timestamp, CheckpointOptions.forFullCheckpoint());
    KeyGroupsStateHandle keyedManaged = FutureUtil.runIfNotDoneAndGet(operatorStateResult.getKeyedStateManagedFuture());
    KeyGroupsStateHandle keyedRaw = FutureUtil.runIfNotDoneAndGet(operatorStateResult.getKeyedStateRawFuture());
    OperatorStateHandle opManaged = FutureUtil.runIfNotDoneAndGet(operatorStateResult.getOperatorStateManagedFuture());
    OperatorStateHandle opRaw = FutureUtil.runIfNotDoneAndGet(operatorStateResult.getOperatorStateRawFuture());
    return new OperatorStateHandles(0, null, keyedManaged != null ? Collections.singletonList(keyedManaged) : null, keyedRaw != null ? Collections.singletonList(keyedRaw) : null, opManaged != null ? Collections.singletonList(opManaged) : null, opRaw != null ? Collections.singletonList(opRaw) : null);
}

Example 5 with OperatorStateHandle

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

the class SavepointV0Serializer method convertSubtaskState.

private org.apache.flink.runtime.checkpoint.SubtaskState convertSubtaskState(SubtaskState subtaskState, int parallelInstanceIdx, ClassLoader userClassLoader, long checkpointID) throws Exception {
    SerializedValue<StateHandle<?>> serializedValue = subtaskState.getState();
    StreamTaskStateList stateList = (StreamTaskStateList) serializedValue.deserializeValue(userClassLoader);
    StreamTaskState[] streamTaskStates = stateList.getState(userClassLoader);
    List<StreamStateHandle> newChainStateList = Arrays.asList(new StreamStateHandle[streamTaskStates.length]);
    KeyGroupsStateHandle newKeyedState = null;
    for (int chainIdx = 0; chainIdx < streamTaskStates.length; ++chainIdx) {
        StreamTaskState streamTaskState = streamTaskStates[chainIdx];
        if (streamTaskState == null) {
            continue;
        }
        newChainStateList.set(chainIdx, convertOperatorAndFunctionState(streamTaskState));
        HashMap<String, KvStateSnapshot<?, ?, ?, ?>> oldKeyedState = streamTaskState.getKvStates();
        if (null != oldKeyedState) {
            Preconditions.checkState(null == newKeyedState, "Found more than one keyed state in chain");
            newKeyedState = convertKeyedBackendState(oldKeyedState, parallelInstanceIdx, checkpointID);
        }
    }
    ChainedStateHandle<StreamStateHandle> newChainedState = new ChainedStateHandle<>(newChainStateList);
    ChainedStateHandle<OperatorStateHandle> nopChain = new ChainedStateHandle<>(Arrays.asList(new OperatorStateHandle[newChainedState.getLength()]));
    return new org.apache.flink.runtime.checkpoint.SubtaskState(newChainedState, nopChain, nopChain, newKeyedState, null);
}