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

Example 31 with OperatorStateHandle

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

the class TaskStateSnapshotTest method hasState.

@Test
public void hasState() {
    Random random = new Random(0x42);
    TaskStateSnapshot taskStateSnapshot = new TaskStateSnapshot();
    Assert.assertFalse(taskStateSnapshot.hasState());
    OperatorSubtaskState emptyOperatorSubtaskState = OperatorSubtaskState.builder().build();
    Assert.assertFalse(emptyOperatorSubtaskState.hasState());
    taskStateSnapshot.putSubtaskStateByOperatorID(new OperatorID(), emptyOperatorSubtaskState);
    Assert.assertFalse(taskStateSnapshot.hasState());
    OperatorStateHandle stateHandle = StateHandleDummyUtil.createNewOperatorStateHandle(2, random);
    OperatorSubtaskState nonEmptyOperatorSubtaskState = OperatorSubtaskState.builder().setManagedOperatorState(stateHandle).build();
    Assert.assertTrue(nonEmptyOperatorSubtaskState.hasState());
    taskStateSnapshot.putSubtaskStateByOperatorID(new OperatorID(), nonEmptyOperatorSubtaskState);
    Assert.assertTrue(taskStateSnapshot.hasState());
}

Example 32 with OperatorStateHandle

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

the class TaskStateSnapshotTest method getStateSize.

@Test
public void getStateSize() {
    Random random = new Random(0x42);
    TaskStateSnapshot taskStateSnapshot = new TaskStateSnapshot();
    Assert.assertEquals(0, taskStateSnapshot.getStateSize());
    OperatorSubtaskState emptyOperatorSubtaskState = OperatorSubtaskState.builder().build();
    Assert.assertFalse(emptyOperatorSubtaskState.hasState());
    taskStateSnapshot.putSubtaskStateByOperatorID(new OperatorID(), emptyOperatorSubtaskState);
    Assert.assertEquals(0, taskStateSnapshot.getStateSize());
    OperatorStateHandle stateHandle_1 = StateHandleDummyUtil.createNewOperatorStateHandle(2, random);
    OperatorSubtaskState nonEmptyOperatorSubtaskState_1 = OperatorSubtaskState.builder().setManagedOperatorState(stateHandle_1).build();
    OperatorStateHandle stateHandle_2 = StateHandleDummyUtil.createNewOperatorStateHandle(2, random);
    OperatorSubtaskState nonEmptyOperatorSubtaskState_2 = OperatorSubtaskState.builder().setRawOperatorState(stateHandle_2).build();
    taskStateSnapshot.putSubtaskStateByOperatorID(new OperatorID(), nonEmptyOperatorSubtaskState_1);
    taskStateSnapshot.putSubtaskStateByOperatorID(new OperatorID(), nonEmptyOperatorSubtaskState_2);
    long totalSize = stateHandle_1.getStateSize() + stateHandle_2.getStateSize();
    Assert.assertEquals(totalSize, taskStateSnapshot.getStateSize());
}

Example 33 with OperatorStateHandle

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

the class RoundRobinOperatorStateRepartitioner method groupByStateMode.

/**
 * Group by the different named states.
 */
@SuppressWarnings("unchecked, rawtype")
private GroupByStateNameResults groupByStateMode(List<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()) {
        nameToStateByMode.put(mode, new HashMap<>());
    }
    for (List<OperatorStateHandle> previousParallelSubtaskState : previousParallelSubtaskStates) {
        for (OperatorStateHandle operatorStateHandle : previousParallelSubtaskState) {
            if (operatorStateHandle == null) {
                continue;
            }
            final Set<Map.Entry<String, OperatorStateHandle.StateMetaInfo>> partitionOffsetEntries = operatorStateHandle.getStateNameToPartitionOffsets().entrySet();
            for (Map.Entry<String, OperatorStateHandle.StateMetaInfo> e : partitionOffsetEntries) {
                OperatorStateHandle.StateMetaInfo metaInfo = e.getValue();
                Map<String, List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>>> nameToState = nameToStateByMode.get(metaInfo.getDistributionMode());
                List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>> stateLocations = nameToState.computeIfAbsent(e.getKey(), k -> new ArrayList<>(previousParallelSubtaskStates.size() * partitionOffsetEntries.size()));
                stateLocations.add(Tuple2.of(operatorStateHandle.getDelegateStateHandle(), e.getValue()));
            }
        }
    }
    return new GroupByStateNameResults(nameToStateByMode);
}

Example 34 with OperatorStateHandle

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

the class RoundRobinOperatorStateRepartitioner method collectStates.

/**
 * Collect the states from given parallelSubtaskStates with the specific {@code mode}.
 */
private Map<String, StateEntry> collectStates(List<List<OperatorStateHandle>> parallelSubtaskStates, OperatorStateHandle.Mode mode) {
    Map<String, StateEntry> states = new HashMap<>(parallelSubtaskStates.size());
    for (int i = 0; i < parallelSubtaskStates.size(); ++i) {
        final int subtaskIndex = i;
        List<OperatorStateHandle> subTaskState = parallelSubtaskStates.get(i);
        for (OperatorStateHandle operatorStateHandle : subTaskState) {
            if (operatorStateHandle == null) {
                continue;
            }
            final Set<Map.Entry<String, OperatorStateHandle.StateMetaInfo>> partitionOffsetEntries = operatorStateHandle.getStateNameToPartitionOffsets().entrySet();
            partitionOffsetEntries.stream().filter(entry -> entry.getValue().getDistributionMode().equals(mode)).forEach(entry -> {
                StateEntry stateEntry = states.computeIfAbsent(entry.getKey(), k -> new StateEntry(parallelSubtaskStates.size() * partitionOffsetEntries.size(), parallelSubtaskStates.size()));
                stateEntry.addEntry(subtaskIndex, Tuple2.of(operatorStateHandle.getDelegateStateHandle(), entry.getValue()));
            });
        }
    }
    return states;
}

Example 35 with OperatorStateHandle

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

the class RoundRobinOperatorStateRepartitioner method repartitionBroadcastState.

/**
 * Repartition BROADCAST state.
 */
private void repartitionBroadcastState(Map<String, List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>>> broadcastState, List<Map<StreamStateHandle, OperatorStateHandle>> mergeMapList) {
    int newParallelism = mergeMapList.size();
    for (int i = 0; i < newParallelism; ++i) {
        final Map<StreamStateHandle, OperatorStateHandle> mergeMap = mergeMapList.get(i);
        // for each name, pick the i-th entry
        for (Map.Entry<String, List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>>> e : broadcastState.entrySet()) {
            int previousParallelism = e.getValue().size();
            Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo> handleWithMetaInfo = e.getValue().get(i % previousParallelism);
            OperatorStateHandle operatorStateHandle = mergeMap.get(handleWithMetaInfo.f0);
            if (operatorStateHandle == null) {
                operatorStateHandle = new OperatorStreamStateHandle(new HashMap<>(broadcastState.size()), handleWithMetaInfo.f0);
                mergeMap.put(handleWithMetaInfo.f0, operatorStateHandle);
            }
            operatorStateHandle.getStateNameToPartitionOffsets().put(e.getKey(), handleWithMetaInfo.f1);
        }
    }
}