org.apache.flink.runtime.state.StateSnapshotKeyGroupReader Java Examples

/**
 * Creates a new StateTableByKeyGroupReader that inserts de-serialized mappings into the given table, using the
 * de-serialization algorithm that matches the given version.
 *
 * @param <K> type of key.
 * @param <N> type of namespace.
 * @param <S> type of state.
 * @param stateTable the {@link StateTable} into which de-serialized mappings are inserted.
 * @param version version for the de-serialization algorithm.
 * @return the appropriate reader.
 */
static <K, N, S> StateSnapshotKeyGroupReader readerForVersion(
	StateTable<K, N, S> stateTable,
	int version) {
	switch (version) {
		case 1:
			return new StateTableByKeyGroupReaderV1<>(stateTable);
		case 2:
		case 3:
		case 4:
		case 5:
		case 6:
			return createV2PlusReader(stateTable);
		default:
			throw new IllegalArgumentException("Unknown version: " + version);
	}
}
 

/**
 * Creates a new StateTableByKeyGroupReader that inserts de-serialized mappings into the given table, using the
 * de-serialization algorithm that matches the given version.
 *
 * @param <K> type of key.
 * @param <N> type of namespace.
 * @param <S> type of state.
 * @param stateTable the {@link StateTable} into which de-serialized mappings are inserted.
 * @param version version for the de-serialization algorithm.
 * @return the appropriate reader.
 */
static <K, N, S> StateSnapshotKeyGroupReader readerForVersion(
	StateTable<K, N, S> stateTable,
	int version) {
	switch (version) {
		case 1:
			return new StateTableByKeyGroupReaderV1<>(stateTable);
		case 2:
		case 3:
		case 4:
		case 5:
		case 6:
			return createV2PlusReader(stateTable);
		default:
			throw new IllegalArgumentException("Unknown version: " + version);
	}
}
 

@Nonnull
@Override
public StateSnapshotKeyGroupReader keyGroupReader(int readVersionHint) {
	final TypeSerializer<T> elementSerializer = metaInfo.getElementSerializer();
	return KeyGroupPartitioner.createKeyGroupPartitionReader(
		elementSerializer::deserialize, //we know that this does not deliver nulls, because we never write nulls
		(element, keyGroupId) -> priorityQueue.add(element));
}
 

/**
 * Creates a new StateTableByKeyGroupReader that inserts de-serialized mappings into the given table, using the
 * de-serialization algorithm that matches the given version.
 *
 * @param stateTable the {@link StateTable} into which de-serialized mappings are inserted.
 * @param version version for the de-serialization algorithm.
 * @param <K> type of key.
 * @param <N> type of namespace.
 * @param <S> type of state.
 * @return the appropriate reader.
 */
static <K, N, S> StateSnapshotKeyGroupReader readerForVersion(StateTable<K, N, S> stateTable, int version) {
	switch (version) {
		case 1:
			return new StateTableByKeyGroupReaderV1<>(stateTable);
		case 2:
		case 3:
		case 4:
		case 5:
		case 6:
			return createV2PlusReader(stateTable);
		default:
			throw new IllegalArgumentException("Unknown version: " + version);
	}
}
 

private static <K, N, S> StateSnapshotKeyGroupReader createV2PlusReader(StateTable<K, N, S> stateTable) {
	final TypeSerializer<K> keySerializer = stateTable.keyContext.getKeySerializer();
	final TypeSerializer<N> namespaceSerializer = stateTable.getNamespaceSerializer();
	final TypeSerializer<S> stateSerializer = stateTable.getStateSerializer();
	final Tuple3<N, K, S> buffer = new Tuple3<>();
	return KeyGroupPartitioner.createKeyGroupPartitionReader((in) -> {
		buffer.f0 = namespaceSerializer.deserialize(in);
		buffer.f1 = keySerializer.deserialize(in);
		buffer.f2 = stateSerializer.deserialize(in);
		return buffer;
	}, (element, keyGroupId1) -> stateTable.put(element.f1, keyGroupId1, element.f0, element.f2));
}
 

private void readKeyGroupStateData(
	InputStream inputStream,
	Map<Integer, StateMetaInfoSnapshot> kvStatesById,
	int keyGroupIndex,
	int numStates,
	int readVersion) throws IOException {

	DataInputViewStreamWrapper inView =
		new DataInputViewStreamWrapper(inputStream);

	for (int i = 0; i < numStates; i++) {

		final int kvStateId = inView.readShort();
		final StateMetaInfoSnapshot stateMetaInfoSnapshot = kvStatesById.get(kvStateId);
		final StateSnapshotRestore registeredState;

		switch (stateMetaInfoSnapshot.getBackendStateType()) {
			case KEY_VALUE:
				registeredState = registeredKVStates.get(stateMetaInfoSnapshot.getName());
				break;
			case PRIORITY_QUEUE:
				registeredState = registeredPQStates.get(stateMetaInfoSnapshot.getName());
				break;
			default:
				throw new IllegalStateException("Unexpected state type: " +
					stateMetaInfoSnapshot.getBackendStateType() + ".");
		}

		StateSnapshotKeyGroupReader keyGroupReader = registeredState.keyGroupReader(readVersion);
		keyGroupReader.readMappingsInKeyGroup(inView, keyGroupIndex);
	}
}
 

@Nonnull
@Override
public StateSnapshotKeyGroupReader keyGroupReader(int readVersionHint) {
	final TypeSerializer<T> elementSerializer = metaInfo.getElementSerializer();
	return KeyGroupPartitioner.createKeyGroupPartitionReader(
		elementSerializer::deserialize, //we know that this does not deliver nulls, because we never write nulls
		(element, keyGroupId) -> priorityQueue.add(element));
}
 

private static <K, N, S> StateSnapshotKeyGroupReader createV2PlusReader(
	StateTable<K, N, S> stateTable) {
	final TypeSerializer<N> namespaceSerializer = stateTable.getNamespaceSerializer();
	final TypeSerializer<S> stateSerializer = stateTable.getStateSerializer();
	final TypeSerializer<K> keySerializer = stateTable.keySerializer;
	final Tuple3<N, K, S> buffer = new Tuple3<>();
	return KeyGroupPartitioner.createKeyGroupPartitionReader((in) -> {
		buffer.f0 = namespaceSerializer.deserialize(in);
		buffer.f1 = keySerializer.deserialize(in);
		buffer.f2 = stateSerializer.deserialize(in);
		return buffer;
	}, (element, keyGroupId1) -> stateTable.put(element.f1, keyGroupId1, element.f0, element.f2));
}
 

private void readKeyGroupStateData(
	InputStream inputStream,
	Map<Integer, StateMetaInfoSnapshot> kvStatesById,
	int keyGroupIndex,
	int numStates,
	int readVersion) throws IOException {

	DataInputViewStreamWrapper inView =
		new DataInputViewStreamWrapper(inputStream);

	for (int i = 0; i < numStates; i++) {

		final int kvStateId = inView.readShort();
		final StateMetaInfoSnapshot stateMetaInfoSnapshot = kvStatesById.get(kvStateId);
		final StateSnapshotRestore registeredState;

		switch (stateMetaInfoSnapshot.getBackendStateType()) {
			case KEY_VALUE:
				registeredState = registeredKVStates.get(stateMetaInfoSnapshot.getName());
				break;
			case PRIORITY_QUEUE:
				registeredState = registeredPQStates.get(stateMetaInfoSnapshot.getName());
				break;
			default:
				throw new IllegalStateException("Unexpected state type: " +
					stateMetaInfoSnapshot.getBackendStateType() + ".");
		}

		StateSnapshotKeyGroupReader keyGroupReader = registeredState.keyGroupReader(readVersion);
		keyGroupReader.readMappingsInKeyGroup(inView, keyGroupIndex);
	}
}
 

@Nonnull
@Override
public StateSnapshotKeyGroupReader keyGroupReader(int readVersionHint) {
	final TypeSerializer<T> elementSerializer = metaInfo.getElementSerializer();
	return KeyGroupPartitioner.createKeyGroupPartitionReader(
		elementSerializer::deserialize, //we know that this does not deliver nulls, because we never write nulls
		(element, keyGroupId) -> priorityQueue.add(element));
}
 

private static <K, N, S> StateSnapshotKeyGroupReader createV2PlusReader(
	StateTable<K, N, S> stateTable) {
	final TypeSerializer<N> namespaceSerializer = stateTable.getNamespaceSerializer();
	final TypeSerializer<S> stateSerializer = stateTable.getStateSerializer();
	final TypeSerializer<K> keySerializer = stateTable.keySerializer;
	final Tuple3<N, K, S> buffer = new Tuple3<>();
	return KeyGroupPartitioner.createKeyGroupPartitionReader((in) -> {
		buffer.f0 = namespaceSerializer.deserialize(in);
		buffer.f1 = keySerializer.deserialize(in);
		buffer.f2 = stateSerializer.deserialize(in);
		return buffer;
	}, (element, keyGroupId1) -> stateTable.put(element.f1, keyGroupId1, element.f0, element.f2));
}
 

private void readKeyGroupStateData(
	InputStream inputStream,
	Map<Integer, StateMetaInfoSnapshot> kvStatesById,
	int keyGroupIndex,
	int numStates,
	int readVersion) throws IOException {

	DataInputViewStreamWrapper inView =
		new DataInputViewStreamWrapper(inputStream);

	for (int i = 0; i < numStates; i++) {

		final int kvStateId = inView.readShort();
		final StateMetaInfoSnapshot stateMetaInfoSnapshot = kvStatesById.get(kvStateId);
		final StateSnapshotRestore registeredState;

		switch (stateMetaInfoSnapshot.getBackendStateType()) {
			case KEY_VALUE:
				registeredState = registeredKVStates.get(stateMetaInfoSnapshot.getName());
				break;
			case PRIORITY_QUEUE:
				registeredState = registeredPQStates.get(stateMetaInfoSnapshot.getName());
				break;
			default:
				throw new IllegalStateException("Unexpected state type: " +
					stateMetaInfoSnapshot.getBackendStateType() + ".");
		}

		StateSnapshotKeyGroupReader keyGroupReader = registeredState.keyGroupReader(readVersion);
		keyGroupReader.readMappingsInKeyGroup(inView, keyGroupIndex);
	}
}
 

@Nonnull
@Override
public StateSnapshotKeyGroupReader keyGroupReader(int readVersion) {
	return StateTableByKeyGroupReaders.readerForVersion(this, readVersion);
}
 

@Nonnull
@Override
public StateSnapshotKeyGroupReader keyGroupReader(int readVersion) {
	return StateTableByKeyGroupReaders.readerForVersion(this, readVersion);
}
 

@Nonnull
@Override
public StateSnapshotKeyGroupReader keyGroupReader(int readVersion) {
	return StateTableByKeyGroupReaders.readerForVersion(this, readVersion);
}