Java KeyGroupRange类代码示例

import org.apache.flink.runtime.state.KeyGroupRange; //导入依赖的package包/类
@Override
protected StateInternals createStateInternals() {
  MemoryStateBackend backend = new MemoryStateBackend();
  try {
    AbstractKeyedStateBackend<ByteBuffer> keyedStateBackend = backend.createKeyedStateBackend(
        new DummyEnvironment("test", 1, 0),
        new JobID(),
        "test_op",
        new GenericTypeInfo<>(ByteBuffer.class).createSerializer(new ExecutionConfig()),
        1,
        new KeyGroupRange(0, 0),
        new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()));

    keyedStateBackend.setCurrentKey(
        ByteBuffer.wrap(CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "Hello")));

    return new FlinkStateInternals<>(keyedStateBackend, StringUtf8Coder.of());
  } catch (Exception e) {
    throw new RuntimeException(e);
  }
}
 

import org.apache.flink.runtime.state.KeyGroupRange; //导入依赖的package包/类
private static KeyedStateBackend<ByteBuffer> getKeyedStateBackend(int numberOfKeyGroups,
                                                 KeyGroupRange keyGroupRange) {
  MemoryStateBackend backend = new MemoryStateBackend();
  try {
    AbstractKeyedStateBackend<ByteBuffer> keyedStateBackend = backend.createKeyedStateBackend(
        new DummyEnvironment("test", 1, 0),
        new JobID(),
        "test_op",
        new GenericTypeInfo<>(ByteBuffer.class).createSerializer(new ExecutionConfig()),
        numberOfKeyGroups,
        keyGroupRange,
        new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID()));
    keyedStateBackend.setCurrentKey(ByteBuffer.wrap(
        CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "1")));
    return keyedStateBackend;
  } catch (Exception e) {
    throw new RuntimeException(e);
  }
}
 

import org.apache.flink.runtime.state.KeyGroupRange; //导入依赖的package包/类
RocksDBIncrementalKeyedStateHandle(
		JobID jobId,
		String operatorIdentifier,
		KeyGroupRange keyGroupRange,
		long checkpointId,
		Map<String, StreamStateHandle> newSstFiles,
		Map<String, StreamStateHandle> oldSstFiles,
		Map<String, StreamStateHandle> miscFiles,
		StreamStateHandle metaStateHandle) {

	this.jobId = Preconditions.checkNotNull(jobId);
	this.operatorIdentifier = Preconditions.checkNotNull(operatorIdentifier);
	this.keyGroupRange = Preconditions.checkNotNull(keyGroupRange);
	this.checkpointId = checkpointId;
	this.newSstFiles = Preconditions.checkNotNull(newSstFiles);
	this.oldSstFiles = Preconditions.checkNotNull(oldSstFiles);
	this.miscFiles = Preconditions.checkNotNull(miscFiles);
	this.metaStateHandle = Preconditions.checkNotNull(metaStateHandle);
	this.registered = false;
}
 

import org.apache.flink.runtime.state.KeyGroupRange; //导入依赖的package包/类
RocksDBKeyedStateBackend2(
		final String operatorIdentifier,
		final ClassLoader userCodeClassLoader,
		final File instanceBasePath,
		final DBOptions dbOptions,
		final ColumnFamilyOptions columnFamilyOptions,
		final TaskKvStateRegistry kvStateRegistry,
		final TypeSerializer<K> keySerializer,
		final int numberOfKeyGroups,
		final KeyGroupRange keyGroupRange,
		final ExecutionConfig executionConfig) throws Exception {

	super(operatorIdentifier, userCodeClassLoader,
		instanceBasePath,
		dbOptions, columnFamilyOptions, kvStateRegistry, keySerializer,
		numberOfKeyGroups, keyGroupRange, executionConfig, false);
}
 

import org.apache.flink.runtime.state.KeyGroupRange; //导入依赖的package包/类
@VisibleForTesting
public static KeyedStateHandle deserializeKeyedStateHandle(DataInputStream dis) throws IOException {
	final int type = dis.readByte();
	if (NULL_HANDLE == type) {
		return null;
	} else if (KEY_GROUPS_HANDLE == type) {
		int startKeyGroup = dis.readInt();
		int numKeyGroups = dis.readInt();
		KeyGroupRange keyGroupRange = KeyGroupRange.of(startKeyGroup, startKeyGroup + numKeyGroups - 1);
		long[] offsets = new long[numKeyGroups];
		for (int i = 0; i < numKeyGroups; ++i) {
			offsets[i] = dis.readLong();
		}
		KeyGroupRangeOffsets keyGroupRangeOffsets = new KeyGroupRangeOffsets(
			keyGroupRange, offsets);
		StreamStateHandle stateHandle = deserializeStreamStateHandle(dis);
		return new KeyGroupsStateHandle(keyGroupRangeOffsets, stateHandle);
	} else {
		throw new IllegalStateException("Reading invalid KeyedStateHandle, type: " + type);
	}
}
 

import org.apache.flink.runtime.state.KeyGroupRange; //导入依赖的package包/类
/**
 * Collect {@link KeyGroupsStateHandle  rawKeyedStateHandles} which have intersection with given
 * {@link KeyGroupRange} from {@link TaskState operatorState}
 *
 * @param operatorState        all state handles of a operator
 * @param subtaskKeyGroupRange the KeyGroupRange of a subtask
 * @return all rawKeyedStateHandles which have intersection with given KeyGroupRange
 */
public static List<KeyedStateHandle> getRawKeyedStateHandles(
	OperatorState operatorState,
	KeyGroupRange subtaskKeyGroupRange) {

	List<KeyedStateHandle> extractedKeyedStateHandles = new ArrayList<>();

	for (int i = 0; i < operatorState.getParallelism(); i++) {
		if (operatorState.getState(i) != null) {
			Collection<KeyedStateHandle> rawKeyedState = operatorState.getState(i).getRawKeyedState();
			extractIntersectingState(
				rawKeyedState,
				subtaskKeyGroupRange,
				extractedKeyedStateHandles);
		}
	}

	return extractedKeyedStateHandles;
}
 

import org.apache.flink.runtime.state.KeyGroupRange; //导入依赖的package包/类
/**
 * Extracts certain key group ranges from the given state handles and adds them to the collector.
 */
private static void extractIntersectingState(
	Collection<KeyedStateHandle> originalSubtaskStateHandles,
	KeyGroupRange rangeToExtract,
	List<KeyedStateHandle> extractedStateCollector) {

	for (KeyedStateHandle keyedStateHandle : originalSubtaskStateHandles) {

		if (keyedStateHandle != null) {

			KeyedStateHandle intersectedKeyedStateHandle = keyedStateHandle.getIntersection(rangeToExtract);

			if (intersectedKeyedStateHandle != null) {
				extractedStateCollector.add(intersectedKeyedStateHandle);
			}
		}
	}
}
 

import org.apache.flink.runtime.state.KeyGroupRange; //导入依赖的package包/类
/**
 * Determine the subset of {@link KeyGroupsStateHandle KeyGroupsStateHandles} with correct
 * key group index for the given subtask {@link KeyGroupRange}.
 * <p>
 * <p>This is publicly visible to be used in tests.
 */
public static List<KeyedStateHandle> getKeyedStateHandles(
	Collection<? extends KeyedStateHandle> keyedStateHandles,
	KeyGroupRange subtaskKeyGroupRange) {

	List<KeyedStateHandle> subtaskKeyedStateHandles = new ArrayList<>();

	for (KeyedStateHandle keyedStateHandle : keyedStateHandles) {
		KeyedStateHandle intersectedKeyedStateHandle = keyedStateHandle.getIntersection(subtaskKeyGroupRange);

		if (intersectedKeyedStateHandle != null) {
			subtaskKeyedStateHandles.add(intersectedKeyedStateHandle);
		}
	}

	return subtaskKeyedStateHandles;
}
 

import org.apache.flink.runtime.state.KeyGroupRange; //导入依赖的package包/类
@Override
public void notifyKvStateRegistered(
	JobID jobId,
	JobVertexID jobVertexId,
	KeyGroupRange keyGroupRange,
	String registrationName,
	KvStateID kvStateId) {

	Object msg = new KvStateMessage.NotifyKvStateRegistered(
		jobId,
		jobVertexId,
		keyGroupRange,
		registrationName,
		kvStateId,
		kvStateServerAddress);

	jobManager.tell(msg);
}
 

import org.apache.flink.runtime.state.KeyGroupRange; //导入依赖的package包/类
/**
 * Registers a KvState instance for the given key group index.
 *
 * @param keyGroupRange  Key group range to register
 * @param kvStateId      ID of the KvState instance at the key group index.
 * @param kvStateAddress Server address of the KvState instance at the key group index.
 * @throws IndexOutOfBoundsException If key group range start < 0 or key group range end >= Number of key groups
 */
public void registerKvState(KeyGroupRange keyGroupRange, KvStateID kvStateId, InetSocketAddress kvStateAddress) {

	if (keyGroupRange.getStartKeyGroup() < 0 || keyGroupRange.getEndKeyGroup() >= numKeyGroups) {
		throw new IndexOutOfBoundsException("Key group index");
	}

	for (int kgIdx = keyGroupRange.getStartKeyGroup(); kgIdx <= keyGroupRange.getEndKeyGroup(); ++kgIdx) {

		if (kvStateIds[kgIdx] == null && kvStateAddresses[kgIdx] == null) {
			numRegisteredKeyGroups++;
		}

		kvStateIds[kgIdx] = kvStateId;
		kvStateAddresses[kgIdx] = kvStateAddress;
	}
}
 

import org.apache.flink.runtime.state.KeyGroupRange; //导入依赖的package包/类
/**
 * Registers a KvState instance for the given key group index.
 *
 * @param keyGroupRange Key group range to unregister.
 * @throws IndexOutOfBoundsException If key group range start < 0 or key group range end >= Number of key groups
 * @throws IllegalArgumentException  If no location information registered for a key group index in the range.
 */
void unregisterKvState(KeyGroupRange keyGroupRange) {
	if (keyGroupRange.getStartKeyGroup() < 0 || keyGroupRange.getEndKeyGroup() >= numKeyGroups) {
		throw new IndexOutOfBoundsException("Key group index");
	}

	for (int kgIdx = keyGroupRange.getStartKeyGroup(); kgIdx <= keyGroupRange.getEndKeyGroup(); ++kgIdx) {
		if (kvStateIds[kgIdx] == null || kvStateAddresses[kgIdx] == null) {
			throw new IllegalArgumentException("Not registered. Probably registration/unregistration race.");
		}

		numRegisteredKeyGroups--;

		kvStateIds[kgIdx] = null;
		kvStateAddresses[kgIdx] = null;
	}
}
 

import org.apache.flink.runtime.state.KeyGroupRange; //导入依赖的package包/类
/**
 * Notifies the JobManager about a registered {@link InternalKvState} instance.
 *
 * @param jobId                JobID the KvState instance belongs to
 * @param jobVertexId          JobVertexID the KvState instance belongs to
 * @param keyGroupRange        Key group range the KvState instance belongs to
 * @param registrationName     Name under which the KvState has been registered
 * @param kvStateId            ID of the registered KvState instance
 * @param kvStateServerAddress Server address where to find the KvState instance
 */
public NotifyKvStateRegistered(
		JobID jobId,
		JobVertexID jobVertexId,
		KeyGroupRange keyGroupRange,
		String registrationName,
		KvStateID kvStateId,
		InetSocketAddress kvStateServerAddress) {

	this.jobId = Preconditions.checkNotNull(jobId, "JobID");
	this.jobVertexId = Preconditions.checkNotNull(jobVertexId, "JobVertexID");
	Preconditions.checkArgument(keyGroupRange != KeyGroupRange.EMPTY_KEY_GROUP_RANGE);
	this.keyGroupRange = Preconditions.checkNotNull(keyGroupRange);
	this.registrationName = Preconditions.checkNotNull(registrationName, "Registration name");
	this.kvStateId = Preconditions.checkNotNull(kvStateId, "KvStateID");
	this.kvStateServerAddress = Preconditions.checkNotNull(kvStateServerAddress, "ServerAddress");
}
 

import org.apache.flink.runtime.state.KeyGroupRange; //导入依赖的package包/类
/**
 * Unregisters the KvState instance identified by the given KvStateID.
 *
 * @param jobId     JobId the KvState instance belongs to
 * @param kvStateId KvStateID to identify the KvState instance
 * @param keyGroupRange    Key group range the KvState instance belongs to
 */
public void unregisterKvState(
		JobID jobId,
		JobVertexID jobVertexId,
		KeyGroupRange keyGroupRange,
		String registrationName,
		KvStateID kvStateId) {

	if (registeredKvStates.remove(kvStateId) != null) {
		final KvStateRegistryListener listener = listenerRef.get();
		if (listener != null) {
			listener.notifyKvStateUnregistered(
					jobId,
					jobVertexId,
					keyGroupRange,
					registrationName);
		}
	}
}
 

import org.apache.flink.runtime.state.KeyGroupRange; //导入依赖的package包/类
/**
 * Notifies the registry about an unregistered KvState instance.
 *
 * @param jobVertexId JobVertexID the KvState instance belongs to
 * @param keyGroupRange Key group index the KvState instance belongs to
 * @param registrationName Name under which the KvState has been registered
 * @throws IllegalArgumentException If another operator registered the state instance
 * @throws IllegalArgumentException If the registration name is not known
 */
public void notifyKvStateUnregistered(
		JobVertexID jobVertexId,
		KeyGroupRange keyGroupRange,
		String registrationName) {

	KvStateLocation location = lookupTable.get(registrationName);

	if (location != null) {
		// Duplicate name if vertex IDs don't match
		if (!location.getJobVertexId().equals(jobVertexId)) {
			throw new IllegalArgumentException("Another operator (" +
					location.getJobVertexId() + ") registered the KvState " +
					"under '" + registrationName + "'.");
		}

		location.unregisterKvState(keyGroupRange);

		if (location.getNumRegisteredKeyGroups() == 0) {
			lookupTable.remove(registrationName);
		}
	} else {
		throw new IllegalArgumentException("Unknown registration name '" +
				registrationName + "'. " + "Probably registration/unregistration race.");
	}
}
 

import org.apache.flink.runtime.state.KeyGroupRange; //导入依赖的package包/类
@Override
public void notifyKvStateRegistered(
		final JobVertexID jobVertexId,
		final KeyGroupRange keyGroupRange,
		final String registrationName,
		final KvStateID kvStateId,
		final InetSocketAddress kvStateServerAddress) {
	if (log.isDebugEnabled()) {
		log.debug("Key value state registered for job {} under name {}.",
				jobGraph.getJobID(), registrationName);
	}

	try {
		executionGraph.getKvStateLocationRegistry().notifyKvStateRegistered(
				jobVertexId, keyGroupRange, registrationName, kvStateId, kvStateServerAddress);
	} catch (Exception e) {
		log.error("Failed to notify KvStateRegistry about registration {}.", registrationName);
	}
}