org.apache.flink.runtime.checkpoint.StateAssignmentOperation Java Examples

The following examples show how to use org.apache.flink.runtime.checkpoint.StateAssignmentOperation. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example #1
Source File: RestoreStreamTaskTest.java    From Flink-CEPplus with Apache License 2.0 5 votes vote down vote up
@Test
public void testRestoreAfterScaleUp() throws Exception {
	OperatorID headOperatorID = new OperatorID(42L, 42L);
	OperatorID tailOperatorID = new OperatorID(44L, 44L);

	JobManagerTaskRestore restore = createRunAndCheckpointOperatorChain(
		headOperatorID,
		new CounterOperator(),
		tailOperatorID,
		new CounterOperator(),
		Optional.empty());

	TaskStateSnapshot stateHandles = restore.getTaskStateSnapshot();

	assertEquals(2, stateHandles.getSubtaskStateMappings().size());

	// test empty state in case of scale up
	OperatorSubtaskState emptyHeadOperatorState = StateAssignmentOperation.operatorSubtaskStateFrom(
		new OperatorInstanceID(0, headOperatorID),
		Collections.emptyMap(),
		Collections.emptyMap(),
		Collections.emptyMap(),
		Collections.emptyMap());

	stateHandles.putSubtaskStateByOperatorID(headOperatorID, emptyHeadOperatorState);

	createRunAndCheckpointOperatorChain(
		headOperatorID,
		new CounterOperator(),
		tailOperatorID,
		new CounterOperator(),
		Optional.of(restore));

	assertEquals(new HashSet<>(Arrays.asList(headOperatorID, tailOperatorID)), RESTORED_OPERATORS);
}
 
Example #2
Source File: OperatorStateInputFormat.java    From flink with Apache License 2.0 5 votes vote down vote up
private OperatorStateInputSplit[] getOperatorStateInputSplits(int minNumSplits) {
	final Map<OperatorInstanceID, List<OperatorStateHandle>> newManagedOperatorStates = new HashMap<>();

	StateAssignmentOperation.reDistributePartitionableStates(
		Collections.singletonList(operatorState),
		minNumSplits,
		Collections.singletonList(operatorState.getOperatorID()),
		newManagedOperatorStates,
		new HashMap<>());

	return CollectionUtil.mapWithIndex(
		newManagedOperatorStates.values(),
		(handles, index) -> new OperatorStateInputSplit(new StateObjectCollection<>(handles), index)
	).toArray(OperatorStateInputSplit[]::new);
}
 
Example #3
Source File: KeyedStateInputFormat.java    From flink with Apache License 2.0 5 votes vote down vote up
private static KeyGroupRangeInputSplit createKeyGroupRangeInputSplit(
	OperatorState operatorState,
	int maxParallelism,
	KeyGroupRange keyGroupRange,
	Integer index) {

	final List<KeyedStateHandle> managedKeyedState = StateAssignmentOperation.getManagedKeyedStateHandles(operatorState, keyGroupRange);
	final List<KeyedStateHandle> rawKeyedState = StateAssignmentOperation.getRawKeyedStateHandles(operatorState, keyGroupRange);

	return new KeyGroupRangeInputSplit(managedKeyedState, rawKeyedState, maxParallelism, index);
}
 
Example #4
Source File: KeyedStateInputFormat.java    From flink with Apache License 2.0 5 votes vote down vote up
@Nonnull
private static List<KeyGroupRange> sortedKeyGroupRanges(int minNumSplits, int maxParallelism) {
	List<KeyGroupRange> keyGroups = StateAssignmentOperation.createKeyGroupPartitions(
		maxParallelism,
		Math.min(minNumSplits, maxParallelism));

	keyGroups.sort(Comparator.comparing(KeyGroupRange::getStartKeyGroup));
	return keyGroups;
}
 
Example #5
Source File: RestoreStreamTaskTest.java    From flink with Apache License 2.0 5 votes vote down vote up
@Test
public void testRestoreAfterScaleUp() throws Exception {
	OperatorID headOperatorID = new OperatorID(42L, 42L);
	OperatorID tailOperatorID = new OperatorID(44L, 44L);

	JobManagerTaskRestore restore = createRunAndCheckpointOperatorChain(
		headOperatorID,
		new CounterOperator(),
		tailOperatorID,
		new CounterOperator(),
		Optional.empty());

	TaskStateSnapshot stateHandles = restore.getTaskStateSnapshot();

	assertEquals(2, stateHandles.getSubtaskStateMappings().size());

	// test empty state in case of scale up
	OperatorSubtaskState emptyHeadOperatorState = StateAssignmentOperation.operatorSubtaskStateFrom(
		new OperatorInstanceID(0, headOperatorID),
		Collections.emptyMap(),
		Collections.emptyMap(),
		Collections.emptyMap(),
		Collections.emptyMap());

	stateHandles.putSubtaskStateByOperatorID(headOperatorID, emptyHeadOperatorState);

	createRunAndCheckpointOperatorChain(
		headOperatorID,
		new CounterOperator(),
		tailOperatorID,
		new CounterOperator(),
		Optional.of(restore));

	assertEquals(new HashSet<>(Arrays.asList(headOperatorID, tailOperatorID)), RESTORED_OPERATORS);
}
 
Example #6
Source File: KeyedStateInputFormat.java    From flink with Apache License 2.0 5 votes vote down vote up
private static KeyGroupRangeInputSplit createKeyGroupRangeInputSplit(
	OperatorState operatorState,
	int maxParallelism,
	KeyGroupRange keyGroupRange,
	Integer index) {

	final List<KeyedStateHandle> managedKeyedState = StateAssignmentOperation.getManagedKeyedStateHandles(operatorState, keyGroupRange);
	final List<KeyedStateHandle> rawKeyedState = StateAssignmentOperation.getRawKeyedStateHandles(operatorState, keyGroupRange);

	return new KeyGroupRangeInputSplit(managedKeyedState, rawKeyedState, maxParallelism, index);
}
 
Example #7
Source File: KeyedStateInputFormat.java    From flink with Apache License 2.0 5 votes vote down vote up
@Nonnull
private static List<KeyGroupRange> sortedKeyGroupRanges(int minNumSplits, int maxParallelism) {
	List<KeyGroupRange> keyGroups = StateAssignmentOperation.createKeyGroupPartitions(
		maxParallelism,
		Math.min(minNumSplits, maxParallelism));

	keyGroups.sort(Comparator.comparing(KeyGroupRange::getStartKeyGroup));
	return keyGroups;
}
 
Example #8
Source File: RestoreStreamTaskTest.java    From flink with Apache License 2.0 5 votes vote down vote up
@Test
public void testRestoreAfterScaleUp() throws Exception {
	OperatorID headOperatorID = new OperatorID(42L, 42L);
	OperatorID tailOperatorID = new OperatorID(44L, 44L);

	JobManagerTaskRestore restore = createRunAndCheckpointOperatorChain(
		headOperatorID,
		new CounterOperator(),
		tailOperatorID,
		new CounterOperator(),
		Optional.empty());

	TaskStateSnapshot stateHandles = restore.getTaskStateSnapshot();

	assertEquals(2, stateHandles.getSubtaskStateMappings().size());

	// test empty state in case of scale up
	OperatorSubtaskState emptyHeadOperatorState = StateAssignmentOperation.operatorSubtaskStateFrom(
		new OperatorInstanceID(0, headOperatorID),
		Collections.emptyMap(),
		Collections.emptyMap(),
		Collections.emptyMap(),
		Collections.emptyMap(),
		Collections.emptyMap(),
		Collections.emptyMap());

	stateHandles.putSubtaskStateByOperatorID(headOperatorID, emptyHeadOperatorState);

	createRunAndCheckpointOperatorChain(
		headOperatorID,
		new CounterOperator(),
		tailOperatorID,
		new CounterOperator(),
		Optional.of(restore));

	assertEquals(new HashSet<>(Arrays.asList(headOperatorID, tailOperatorID)), RESTORED_OPERATORS);
}
 
Example #9
Source File: StateBackendTestBase.java    From Flink-CEPplus with Apache License 2.0 4 votes vote down vote up
/**
 * This test verifies that state is correctly assigned to key groups and that restore
 * restores the relevant key groups in the backend.
 *
 * <p>We have ten key groups. Initially, one backend is responsible for all ten key groups.
 * Then we snapshot, split up the state and restore in to backends where each is responsible
 * for five key groups. Then we make sure that the state is only available in the correct
 * backend.
 * @throws Exception
 */
@Test
public void testKeyGroupSnapshotRestore() throws Exception {
	final int MAX_PARALLELISM = 10;

	CheckpointStreamFactory streamFactory = createStreamFactory();
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	final AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(
			IntSerializer.INSTANCE,
			MAX_PARALLELISM,
			new KeyGroupRange(0, MAX_PARALLELISM - 1),
			new DummyEnvironment());

	ValueStateDescriptor<String> kvId = new ValueStateDescriptor<>("id", String.class);

	ValueState<String> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);

	// keys that fall into the first half/second half of the key groups, respectively
	int keyInFirstHalf = 17;
	int keyInSecondHalf = 42;
	Random rand = new Random(0);

	// for each key, determine into which half of the key-group space they fall
	int firstKeyHalf = KeyGroupRangeAssignment.assignKeyToParallelOperator(keyInFirstHalf, MAX_PARALLELISM, 2);
	int secondKeyHalf = KeyGroupRangeAssignment.assignKeyToParallelOperator(keyInFirstHalf, MAX_PARALLELISM, 2);

	while (firstKeyHalf == secondKeyHalf) {
		keyInSecondHalf = rand.nextInt();
		secondKeyHalf = KeyGroupRangeAssignment.assignKeyToParallelOperator(keyInSecondHalf, MAX_PARALLELISM, 2);
	}

	backend.setCurrentKey(keyInFirstHalf);
	state.update("ShouldBeInFirstHalf");

	backend.setCurrentKey(keyInSecondHalf);
	state.update("ShouldBeInSecondHalf");


	KeyedStateHandle snapshot = runSnapshot(
		backend.snapshot(0, 0, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation()),
		sharedStateRegistry);

	List<KeyedStateHandle> firstHalfKeyGroupStates = StateAssignmentOperation.getKeyedStateHandles(
			Collections.singletonList(snapshot),
			KeyGroupRangeAssignment.computeKeyGroupRangeForOperatorIndex(MAX_PARALLELISM, 2, 0));

	List<KeyedStateHandle> secondHalfKeyGroupStates = StateAssignmentOperation.getKeyedStateHandles(
			Collections.singletonList(snapshot),
			KeyGroupRangeAssignment.computeKeyGroupRangeForOperatorIndex(MAX_PARALLELISM, 2, 1));

	backend.dispose();

	// backend for the first half of the key group range
	final AbstractKeyedStateBackend<Integer> firstHalfBackend = restoreKeyedBackend(
			IntSerializer.INSTANCE,
			MAX_PARALLELISM,
			new KeyGroupRange(0, 4),
			firstHalfKeyGroupStates,
			new DummyEnvironment());

	// backend for the second half of the key group range
	final AbstractKeyedStateBackend<Integer> secondHalfBackend = restoreKeyedBackend(
			IntSerializer.INSTANCE,
			MAX_PARALLELISM,
			new KeyGroupRange(5, 9),
			secondHalfKeyGroupStates,
			new DummyEnvironment());


	ValueState<String> firstHalfState = firstHalfBackend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);

	firstHalfBackend.setCurrentKey(keyInFirstHalf);
	assertTrue(firstHalfState.value().equals("ShouldBeInFirstHalf"));

	firstHalfBackend.setCurrentKey(keyInSecondHalf);
	assertTrue(firstHalfState.value() == null);

	ValueState<String> secondHalfState = secondHalfBackend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);

	secondHalfBackend.setCurrentKey(keyInFirstHalf);
	assertTrue(secondHalfState.value() == null);

	secondHalfBackend.setCurrentKey(keyInSecondHalf);
	assertTrue(secondHalfState.value().equals("ShouldBeInSecondHalf"));

	firstHalfBackend.dispose();
	secondHalfBackend.dispose();
}
 
Example #10
Source File: AbstractStreamOperatorTestHarness.java    From Flink-CEPplus with Apache License 2.0 4 votes vote down vote up
/**
 * Returns the reshaped 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 static OperatorSubtaskState repartitionOperatorState(
	final OperatorSubtaskState operatorStateHandles,
	final int numKeyGroups,
	final int oldParallelism,
	final int newParallelism,
	final int subtaskIndex) {

	Preconditions.checkNotNull(operatorStateHandles, "the previous operatorStateHandles should not be null.");

	// create a new OperatorStateHandles that only contains the state for our key-groups

	List<KeyGroupRange> keyGroupPartitions = StateAssignmentOperation.createKeyGroupPartitions(
		numKeyGroups,
		newParallelism);

	KeyGroupRange localKeyGroupRange = keyGroupPartitions.get(subtaskIndex);

	List<KeyedStateHandle> localManagedKeyGroupState = StateAssignmentOperation.getKeyedStateHandles(
		operatorStateHandles.getManagedKeyedState(),
		localKeyGroupRange);

	List<KeyedStateHandle> localRawKeyGroupState = StateAssignmentOperation.getKeyedStateHandles(
		operatorStateHandles.getRawKeyedState(),
		localKeyGroupRange);

	StateObjectCollection<OperatorStateHandle> managedOperatorStates = operatorStateHandles.getManagedOperatorState();
	Collection<OperatorStateHandle> localManagedOperatorState;

	if (!managedOperatorStates.isEmpty()) {
		List<List<OperatorStateHandle>> managedOperatorState =
			managedOperatorStates.stream().map(Collections::singletonList).collect(Collectors.toList());

		localManagedOperatorState = operatorStateRepartitioner.repartitionState(
			managedOperatorState,
			oldParallelism,
			newParallelism).get(subtaskIndex);
	} else {
		localManagedOperatorState = Collections.emptyList();
	}

	StateObjectCollection<OperatorStateHandle> rawOperatorStates = operatorStateHandles.getRawOperatorState();
	Collection<OperatorStateHandle> localRawOperatorState;

	if (!rawOperatorStates.isEmpty()) {
		List<List<OperatorStateHandle>> rawOperatorState =
			rawOperatorStates.stream().map(Collections::singletonList).collect(Collectors.toList());

		localRawOperatorState = operatorStateRepartitioner.repartitionState(
			rawOperatorState,
			oldParallelism,
			newParallelism).get(subtaskIndex);
	} else {
		localRawOperatorState = Collections.emptyList();
	}

	return new OperatorSubtaskState(
		new StateObjectCollection<>(nullToEmptyCollection(localManagedOperatorState)),
		new StateObjectCollection<>(nullToEmptyCollection(localRawOperatorState)),
		new StateObjectCollection<>(nullToEmptyCollection(localManagedKeyGroupState)),
		new StateObjectCollection<>(nullToEmptyCollection(localRawKeyGroupState)));
}
 
Example #11
Source File: StateBackendTestBase.java    From flink with Apache License 2.0 4 votes vote down vote up
/**
 * This test verifies that state is correctly assigned to key groups and that restore
 * restores the relevant key groups in the backend.
 *
 * <p>We have ten key groups. Initially, one backend is responsible for all ten key groups.
 * Then we snapshot, split up the state and restore in to backends where each is responsible
 * for five key groups. Then we make sure that the state is only available in the correct
 * backend.
 * @throws Exception
 */
@Test
public void testKeyGroupSnapshotRestore() throws Exception {
	final int MAX_PARALLELISM = 10;

	CheckpointStreamFactory streamFactory = createStreamFactory();
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	final AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(
			IntSerializer.INSTANCE,
			MAX_PARALLELISM,
			new KeyGroupRange(0, MAX_PARALLELISM - 1),
			new DummyEnvironment());

	ValueStateDescriptor<String> kvId = new ValueStateDescriptor<>("id", String.class);

	ValueState<String> state = backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);

	// keys that fall into the first half/second half of the key groups, respectively
	int keyInFirstHalf = 17;
	int keyInSecondHalf = 42;
	Random rand = new Random(0);

	// for each key, determine into which half of the key-group space they fall
	int firstKeyHalf = KeyGroupRangeAssignment.assignKeyToParallelOperator(keyInFirstHalf, MAX_PARALLELISM, 2);
	int secondKeyHalf = KeyGroupRangeAssignment.assignKeyToParallelOperator(keyInFirstHalf, MAX_PARALLELISM, 2);

	while (firstKeyHalf == secondKeyHalf) {
		keyInSecondHalf = rand.nextInt();
		secondKeyHalf = KeyGroupRangeAssignment.assignKeyToParallelOperator(keyInSecondHalf, MAX_PARALLELISM, 2);
	}

	backend.setCurrentKey(keyInFirstHalf);
	state.update("ShouldBeInFirstHalf");

	backend.setCurrentKey(keyInSecondHalf);
	state.update("ShouldBeInSecondHalf");


	KeyedStateHandle snapshot = runSnapshot(
		backend.snapshot(0, 0, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation()),
		sharedStateRegistry);

	List<KeyedStateHandle> firstHalfKeyGroupStates = StateAssignmentOperation.getKeyedStateHandles(
			Collections.singletonList(snapshot),
			KeyGroupRangeAssignment.computeKeyGroupRangeForOperatorIndex(MAX_PARALLELISM, 2, 0));

	List<KeyedStateHandle> secondHalfKeyGroupStates = StateAssignmentOperation.getKeyedStateHandles(
			Collections.singletonList(snapshot),
			KeyGroupRangeAssignment.computeKeyGroupRangeForOperatorIndex(MAX_PARALLELISM, 2, 1));

	backend.dispose();

	// backend for the first half of the key group range
	final AbstractKeyedStateBackend<Integer> firstHalfBackend = restoreKeyedBackend(
			IntSerializer.INSTANCE,
			MAX_PARALLELISM,
			new KeyGroupRange(0, 4),
			firstHalfKeyGroupStates,
			new DummyEnvironment());

	// backend for the second half of the key group range
	final AbstractKeyedStateBackend<Integer> secondHalfBackend = restoreKeyedBackend(
			IntSerializer.INSTANCE,
			MAX_PARALLELISM,
			new KeyGroupRange(5, 9),
			secondHalfKeyGroupStates,
			new DummyEnvironment());


	ValueState<String> firstHalfState = firstHalfBackend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);

	firstHalfBackend.setCurrentKey(keyInFirstHalf);
	assertTrue(firstHalfState.value().equals("ShouldBeInFirstHalf"));

	firstHalfBackend.setCurrentKey(keyInSecondHalf);
	assertTrue(firstHalfState.value() == null);

	ValueState<String> secondHalfState = secondHalfBackend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);

	secondHalfBackend.setCurrentKey(keyInFirstHalf);
	assertTrue(secondHalfState.value() == null);

	secondHalfBackend.setCurrentKey(keyInSecondHalf);
	assertTrue(secondHalfState.value().equals("ShouldBeInSecondHalf"));

	firstHalfBackend.dispose();
	secondHalfBackend.dispose();
}
 
Example #12
Source File: AbstractStreamOperatorTestHarness.java    From flink with Apache License 2.0 4 votes vote down vote up
/**
 * Returns the reshaped 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 static OperatorSubtaskState repartitionOperatorState(
	final OperatorSubtaskState operatorStateHandles,
	final int numKeyGroups,
	final int oldParallelism,
	final int newParallelism,
	final int subtaskIndex) {

	Preconditions.checkNotNull(operatorStateHandles, "the previous operatorStateHandles should not be null.");

	// create a new OperatorStateHandles that only contains the state for our key-groups

	List<KeyGroupRange> keyGroupPartitions = StateAssignmentOperation.createKeyGroupPartitions(
		numKeyGroups,
		newParallelism);

	KeyGroupRange localKeyGroupRange = keyGroupPartitions.get(subtaskIndex);

	List<KeyedStateHandle> localManagedKeyGroupState = StateAssignmentOperation.getKeyedStateHandles(
		operatorStateHandles.getManagedKeyedState(),
		localKeyGroupRange);

	List<KeyedStateHandle> localRawKeyGroupState = StateAssignmentOperation.getKeyedStateHandles(
		operatorStateHandles.getRawKeyedState(),
		localKeyGroupRange);

	StateObjectCollection<OperatorStateHandle> managedOperatorStates = operatorStateHandles.getManagedOperatorState();
	Collection<OperatorStateHandle> localManagedOperatorState;

	if (!managedOperatorStates.isEmpty()) {
		List<List<OperatorStateHandle>> managedOperatorState =
			managedOperatorStates.stream().map(Collections::singletonList).collect(Collectors.toList());

		localManagedOperatorState = operatorStateRepartitioner.repartitionState(
			managedOperatorState,
			oldParallelism,
			newParallelism).get(subtaskIndex);
	} else {
		localManagedOperatorState = Collections.emptyList();
	}

	StateObjectCollection<OperatorStateHandle> rawOperatorStates = operatorStateHandles.getRawOperatorState();
	Collection<OperatorStateHandle> localRawOperatorState;

	if (!rawOperatorStates.isEmpty()) {
		List<List<OperatorStateHandle>> rawOperatorState =
			rawOperatorStates.stream().map(Collections::singletonList).collect(Collectors.toList());

		localRawOperatorState = operatorStateRepartitioner.repartitionState(
			rawOperatorState,
			oldParallelism,
			newParallelism).get(subtaskIndex);
	} else {
		localRawOperatorState = Collections.emptyList();
	}

	return new OperatorSubtaskState(
		new StateObjectCollection<>(nullToEmptyCollection(localManagedOperatorState)),
		new StateObjectCollection<>(nullToEmptyCollection(localRawOperatorState)),
		new StateObjectCollection<>(nullToEmptyCollection(localManagedKeyGroupState)),
		new StateObjectCollection<>(nullToEmptyCollection(localRawKeyGroupState)));
}
 
Example #13
Source File: AbstractStreamOperatorTestHarness.java    From flink with Apache License 2.0 4 votes vote down vote up
/**
 * Returns the reshaped 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 static OperatorSubtaskState repartitionOperatorState(
	final OperatorSubtaskState operatorStateHandles,
	final int numKeyGroups,
	final int oldParallelism,
	final int newParallelism,
	final int subtaskIndex) {

	Preconditions.checkNotNull(operatorStateHandles, "the previous operatorStateHandles should not be null.");

	// create a new OperatorStateHandles that only contains the state for our key-groups

	List<KeyGroupRange> keyGroupPartitions = StateAssignmentOperation.createKeyGroupPartitions(
		numKeyGroups,
		newParallelism);

	KeyGroupRange localKeyGroupRange = keyGroupPartitions.get(subtaskIndex);

	List<KeyedStateHandle> localManagedKeyGroupState = new ArrayList<>();
	StateAssignmentOperation.extractIntersectingState(
		operatorStateHandles.getManagedKeyedState(),
		localKeyGroupRange,
		localManagedKeyGroupState);

	List<KeyedStateHandle> localRawKeyGroupState = new ArrayList<>();
	StateAssignmentOperation.extractIntersectingState(
		operatorStateHandles.getRawKeyedState(),
		localKeyGroupRange,
		localRawKeyGroupState);

	StateObjectCollection<OperatorStateHandle> managedOperatorStates = operatorStateHandles.getManagedOperatorState();
	Collection<OperatorStateHandle> localManagedOperatorState;

	if (!managedOperatorStates.isEmpty()) {
		List<List<OperatorStateHandle>> managedOperatorState =
			managedOperatorStates.stream().map(Collections::singletonList).collect(Collectors.toList());

		localManagedOperatorState = operatorStateRepartitioner.repartitionState(
			managedOperatorState,
			oldParallelism,
			newParallelism).get(subtaskIndex);
	} else {
		localManagedOperatorState = Collections.emptyList();
	}

	StateObjectCollection<OperatorStateHandle> rawOperatorStates = operatorStateHandles.getRawOperatorState();
	Collection<OperatorStateHandle> localRawOperatorState;

	if (!rawOperatorStates.isEmpty()) {
		List<List<OperatorStateHandle>> rawOperatorState =
			rawOperatorStates.stream().map(Collections::singletonList).collect(Collectors.toList());

		localRawOperatorState = operatorStateRepartitioner.repartitionState(
			rawOperatorState,
			oldParallelism,
			newParallelism).get(subtaskIndex);
	} else {
		localRawOperatorState = Collections.emptyList();
	}

	return new OperatorSubtaskState(
		new StateObjectCollection<>(nullToEmptyCollection(localManagedOperatorState)),
		new StateObjectCollection<>(nullToEmptyCollection(localRawOperatorState)),
		new StateObjectCollection<>(nullToEmptyCollection(localManagedKeyGroupState)),
		new StateObjectCollection<>(nullToEmptyCollection(localRawKeyGroupState)));
}