org.apache.flink.api.common.state.State Java Examples

@Override
@SuppressWarnings("unchecked")
@Nonnull
public <N, SV, SEV, S extends State, IS extends S> IS createInternalState(
	@Nonnull TypeSerializer<N> namespaceSerializer,
	@Nonnull StateDescriptor<S, SV> stateDesc,
	@Nonnull StateSnapshotTransformFactory<SEV> snapshotTransformFactory) throws Exception {
	StateFactory stateFactory = STATE_FACTORIES.get(stateDesc.getClass());
	if (stateFactory == null) {
		String message = String.format("State %s is not supported by %s",
			stateDesc.getClass(), TtlStateFactory.class);
		throw new FlinkRuntimeException(message);
	}
	IS state = stateFactory.createInternalState(namespaceSerializer, stateDesc);
	stateSnapshotFilters.put(stateDesc.getName(),
		(StateSnapshotTransformer<Object>) getStateSnapshotTransformer(stateDesc, snapshotTransformFactory));
	((MockInternalKvState<K, N, SV>) state).values = () -> stateValues
		.computeIfAbsent(stateDesc.getName(), n -> new HashMap<>())
		.computeIfAbsent(getCurrentKey(), k -> new HashMap<>());
	return state;
}
 

@Override
@SuppressWarnings("unchecked")
@Nonnull
public <N, SV, SEV, S extends State, IS extends S> IS createInternalState(
	@Nonnull TypeSerializer<N> namespaceSerializer,
	@Nonnull StateDescriptor<S, SV> stateDesc,
	@Nonnull StateSnapshotTransformFactory<SEV> snapshotTransformFactory) throws Exception {
	StateFactory stateFactory = STATE_FACTORIES.get(stateDesc.getClass());
	if (stateFactory == null) {
		String message = String.format("State %s is not supported by %s",
			stateDesc.getClass(), TtlStateFactory.class);
		throw new FlinkRuntimeException(message);
	}
	IS state = stateFactory.createInternalState(namespaceSerializer, stateDesc);
	stateSnapshotFilters.put(stateDesc.getName(),
		(StateSnapshotTransformer<Object>) getStateSnapshotTransformer(stateDesc, snapshotTransformFactory));
	((MockInternalKvState<K, N, SV>) state).values = () -> stateValues
		.computeIfAbsent(stateDesc.getName(), n -> new HashMap<>())
		.computeIfAbsent(getCurrentKey(), k -> new HashMap<>());
	return state;
}
 

@Override
@Nonnull
public <N, SV, SEV, S extends State, IS extends S> IS createInternalState(
	@Nonnull TypeSerializer<N> namespaceSerializer,
	@Nonnull StateDescriptor<S, SV> stateDesc,
	@Nonnull StateSnapshotTransformFactory<SEV> snapshotTransformFactory) throws Exception {
	StateFactory stateFactory = STATE_FACTORIES.get(stateDesc.getClass());
	if (stateFactory == null) {
		String message = String.format("State %s is not supported by %s",
			stateDesc.getClass(), this.getClass());
		throw new FlinkRuntimeException(message);
	}
	Tuple2<ColumnFamilyHandle, RegisteredKeyValueStateBackendMetaInfo<N, SV>> registerResult = tryRegisterKvStateInformation(
		stateDesc, namespaceSerializer, snapshotTransformFactory);
	return stateFactory.createState(stateDesc, registerResult, RocksDBKeyedStateBackend.this);
}
 

private <N, S extends State, SV> RegisteredKeyValueStateBackendMetaInfo<N, SV> updateRestoredStateMetaInfo(
	Tuple2<ColumnFamilyHandle, RegisteredKeyValueStateBackendMetaInfo<N, SV>> oldStateInfo,
	StateDescriptor<S, SV> stateDesc,
	TypeSerializer<N> namespaceSerializer,
	TypeSerializer<SV> stateSerializer) throws Exception {

	@SuppressWarnings("unchecked")
	RegisteredKeyValueStateBackendMetaInfo<N, SV> restoredKvStateMetaInfo = oldStateInfo.f1;

	TypeSerializerSchemaCompatibility<N> s = restoredKvStateMetaInfo.updateNamespaceSerializer(namespaceSerializer);
	if (s.isCompatibleAfterMigration() || s.isIncompatible()) {
		throw new StateMigrationException("The new namespace serializer must be compatible.");
	}

	restoredKvStateMetaInfo.checkStateMetaInfo(stateDesc);

	TypeSerializerSchemaCompatibility<SV> newStateSerializerCompatibility =
		restoredKvStateMetaInfo.updateStateSerializer(stateSerializer);
	if (newStateSerializerCompatibility.isCompatibleAfterMigration()) {
		migrateStateValues(stateDesc, oldStateInfo);
	} else if (newStateSerializerCompatibility.isIncompatible()) {
		throw new StateMigrationException("The new state serializer cannot be incompatible.");
	}

	return restoredKvStateMetaInfo;
}
 

/**
 * @see KeyedStateBackend
 */
@Override
@SuppressWarnings("unchecked")
public <N, S extends State, V> S getOrCreateKeyedState(
		final TypeSerializer<N> namespaceSerializer,
		StateDescriptor<S, V> stateDescriptor) throws Exception {
	checkNotNull(namespaceSerializer, "Namespace serializer");
	checkNotNull(keySerializerProvider, "State key serializer has not been configured in the config. " +
			"This operation cannot use partitioned state.");

	InternalKvState<K, ?, ?> kvState = keyValueStatesByName.get(stateDescriptor.getName());
	if (kvState == null) {
		if (!stateDescriptor.isSerializerInitialized()) {
			stateDescriptor.initializeSerializerUnlessSet(executionConfig);
		}
		kvState = TtlStateFactory.createStateAndWrapWithTtlIfEnabled(
			namespaceSerializer, stateDescriptor, this, ttlTimeProvider);
		keyValueStatesByName.put(stateDescriptor.getName(), kvState);
		publishQueryableStateIfEnabled(stateDescriptor, kvState);
	}
	return (S) kvState;
}
 

/**
 * @see KeyedStateBackend
 */
@Override
@SuppressWarnings("unchecked")
public <N, S extends State, V> S getOrCreateKeyedState(
		final TypeSerializer<N> namespaceSerializer,
		StateDescriptor<S, V> stateDescriptor) throws Exception {
	checkNotNull(namespaceSerializer, "Namespace serializer");
	checkNotNull(keySerializer, "State key serializer has not been configured in the config. " +
			"This operation cannot use partitioned state.");

	InternalKvState<K, ?, ?> kvState = keyValueStatesByName.get(stateDescriptor.getName());
	if (kvState == null) {
		if (!stateDescriptor.isSerializerInitialized()) {
			stateDescriptor.initializeSerializerUnlessSet(executionConfig);
		}
		kvState = TtlStateFactory.createStateAndWrapWithTtlIfEnabled(
			namespaceSerializer, stateDescriptor, this, ttlTimeProvider);
		keyValueStatesByName.put(stateDescriptor.getName(), kvState);
		publishQueryableStateIfEnabled(stateDescriptor, kvState);
	}
	return (S) kvState;
}
 

@Override
@Nonnull
public <N, SV, SEV, S extends State, IS extends S> IS createInternalState(
	@Nonnull TypeSerializer<N> namespaceSerializer,
	@Nonnull StateDescriptor<S, SV> stateDesc,
	@Nonnull StateSnapshotTransformFactory<SEV> snapshotTransformFactory) throws Exception {
	StateFactory stateFactory = STATE_FACTORIES.get(stateDesc.getClass());
	if (stateFactory == null) {
		String message = String.format("State %s is not supported by %s",
			stateDesc.getClass(), this.getClass());
		throw new FlinkRuntimeException(message);
	}
	StateTable<K, N, SV> stateTable = tryRegisterStateTable(
		namespaceSerializer, stateDesc, getStateSnapshotTransformFactory(stateDesc, snapshotTransformFactory));
	return stateFactory.createState(stateDesc, stateTable, getKeySerializer());
}
 

@Override
public <N, S extends State, T> void applyToAllKeys(
	final N namespace,
	final TypeSerializer<N> namespaceSerializer,
	final StateDescriptor<S, T> stateDescriptor,
	final KeyedStateFunction<K, S> function) throws Exception {

	try (Stream<K> keyStream = getKeys(stateDescriptor.getName(), namespace)) {

		// we copy the keys into list to avoid the concurrency problem
		// when state.clear() is invoked in function.process().
		final List<K> keys = keyStream.collect(Collectors.toList());

		final S state = getPartitionedState(
			namespace,
			namespaceSerializer,
			stateDescriptor);

		for (K key : keys) {
			setCurrentKey(key);
			function.process(key, state);
		}
	}
}
 

@SuppressWarnings("unchecked")
public static <V, S extends State> S createState(
	StateDescriptor<S, V> stateDescriptor,
	byte[] serializedState) throws IOException {
	final V state = KvStateSerializer.deserializeValue(
		serializedState,
		stateDescriptor.getSerializer());
	return (S) new ImmutableReducingState<>(state);
}
 

@SuppressWarnings("unchecked")
public static <V, T, S extends State> S createState(
	StateDescriptor<S, T> stateDescriptor,
	byte[] serializedState) throws IOException {
	final List<V> state = KvStateSerializer.deserializeList(
		serializedState,
		((ListStateDescriptor<V>) stateDescriptor).getElementSerializer());
	return (S) new ImmutableListState<>(state);
}
 

/**
 * Returns a future holding the request result.
 * @param jobId                     JobID of the job the queryable state belongs to.
 * @param queryableStateName        Name under which the state is queryable.
 * @param key			            The key we are interested in.
 * @param keyTypeInfo				The {@link TypeInformation} of the key.
 * @param stateDescriptor			The {@link StateDescriptor} of the state we want to query.
 * @return Future holding the immutable {@link State} object containing the result.
 */
@PublicEvolving
public <K, S extends State, V> CompletableFuture<S> getKvState(
		final JobID jobId,
		final String queryableStateName,
		final K key,
		final TypeInformation<K> keyTypeInfo,
		final StateDescriptor<S, V> stateDescriptor) {

	return getKvState(jobId, queryableStateName, key, VoidNamespace.INSTANCE,
			keyTypeInfo, VoidNamespaceTypeInfo.INSTANCE, stateDescriptor);
}
 

@SuppressWarnings({"unchecked", "unused"})
static <T, N, ACC, S extends State, IS extends S> IS createState(
	TypeSerializer<N> namespaceSerializer,
	StateDescriptor<S, ACC> stateDesc) {
	FoldingStateDescriptor<T, ACC> foldingStateDesc = (FoldingStateDescriptor<T, ACC>) stateDesc;
	return (IS) new MockInternalFoldingState<>(foldingStateDesc.getFoldFunction());
}
 

@SuppressWarnings("unchecked")
static <K, N, SV, S extends State, IS extends S> IS create(
	StateDescriptor<S, SV> stateDesc,
	Tuple2<ColumnFamilyHandle, RegisteredKeyValueStateBackendMetaInfo<N, SV>> registerResult,
	RocksDBKeyedStateBackend<K> backend) {
	return (IS) new RocksDBValueState<>(
		registerResult.f0,
		registerResult.f1.getNamespaceSerializer(),
		registerResult.f1.getStateSerializer(),
		stateDesc.getDefaultValue(),
		backend);
}
 

@Override
public <S extends State> S getPartitionedState(StateDescriptor<S, ?> stateDescriptor) {
	try {
		return stateBackend.getPartitionedState(window, windowSerializer, stateDescriptor);
	} catch (Exception e) {
		throw new RuntimeException("Error getting state", e);
	}
}
 

@SuppressWarnings({"unchecked", "unused"})
static <IN, OUT, N, ACC, S extends State, IS extends S> IS createState(
	TypeSerializer<N> namespaceSerializer,
	StateDescriptor<S, ACC> stateDesc) {
	AggregatingStateDescriptor<IN, ACC, OUT> aggregatingStateDesc =
		(AggregatingStateDescriptor<IN, ACC, OUT>) stateDesc;
	return (IS) new MockInternalAggregatingState<>(aggregatingStateDesc.getAggregateFunction());
}
 

@SuppressWarnings("unchecked")
static <E, K, N, SV, S extends State, IS extends S> IS create(
	StateDescriptor<S, SV> stateDesc,
	Tuple2<ColumnFamilyHandle, RegisteredKeyValueStateBackendMetaInfo<N, SV>> registerResult,
	RocksDBKeyedStateBackend<K> backend) {
	return (IS) new RocksDBListState<>(
		registerResult.f0,
		registerResult.f1.getNamespaceSerializer(),
		(TypeSerializer<List<E>>) registerResult.f1.getStateSerializer(),
		(List<E>) stateDesc.getDefaultValue(),
		backend);
}
 

public <N, S extends State, T> S getOrCreateKeyedState(
		TypeSerializer<N> namespaceSerializer,
		StateDescriptor<S, T> stateDescriptor) throws Exception {

	if (keyedStateStore != null) {
		return keyedStateBackend.getOrCreateKeyedState(namespaceSerializer, stateDescriptor);
	}
	else {
		throw new IllegalStateException("Cannot create partitioned state. " +
				"The keyed state backend has not been set." +
				"This indicates that the operator is not partitioned/keyed.");
	}
}
 

@Override
public <S extends State> S getPartitionedState(StateDescriptor<S, ?> stateDescriptor) {
	try {
		return stateBackend.getPartitionedState(window, windowSerializer, stateDescriptor);
	} catch (Exception e) {
		throw new RuntimeException("Error getting state", e);
	}
}
 

@Override
public <VS, S extends State> void applyToKeyedState(
		final StateDescriptor<S, VS> stateDescriptor,
		final KeyedStateFunction<KS, S> function) throws Exception {

	keyedStateBackend.applyToAllKeys(
			VoidNamespace.INSTANCE,
			VoidNamespaceSerializer.INSTANCE,
			Preconditions.checkNotNull(stateDescriptor),
			Preconditions.checkNotNull(function));
}
 

@SuppressWarnings("unchecked")
static <UK, UV, K, N, SV, S extends State, IS extends S> IS create(
	StateDescriptor<S, SV> stateDesc,
	StateTable<K, N, SV> stateTable,
	TypeSerializer<K> keySerializer) {
	return (IS) new HeapMapState<>(
		(StateTable<K, N, Map<UK, UV>>) stateTable,
		keySerializer,
		(TypeSerializer<Map<UK, UV>>) stateTable.getStateSerializer(),
		stateTable.getNamespaceSerializer(),
		(Map<UK, UV>) stateDesc.getDefaultValue());
}
 

@SuppressWarnings("unchecked")
public static <K, V, T, S extends State> S createState(
	StateDescriptor<S, T> stateDescriptor,
	byte[] serializedState) throws IOException {
	MapStateDescriptor<K, V> mapStateDescriptor = (MapStateDescriptor<K, V>) stateDescriptor;
	final Map<K, V> state = KvStateSerializer.deserializeMap(
		serializedState,
		mapStateDescriptor.getKeySerializer(),
		mapStateDescriptor.getValueSerializer());
	return (S) new ImmutableMapState<>(state);
}
 

/**
 * @see KeyedStateBackend
 */
@Override
public <N, S extends State, T> void applyToAllKeys(
		final N namespace,
		final TypeSerializer<N> namespaceSerializer,
		final StateDescriptor<S, T> stateDescriptor,
		final KeyedStateFunction<K, S> function) throws Exception {

	try (Stream<K> keyStream = getKeys(stateDescriptor.getName(), namespace)) {

		final S state = getPartitionedState(
			namespace,
			namespaceSerializer,
			stateDescriptor);

		keyStream.forEach((K key) -> {
			setCurrentKey(key);
			try {
				function.process(key, state);
			} catch (Throwable e) {
				// we wrap the checked exception in an unchecked
				// one and catch it (and re-throw it) later.
				throw new RuntimeException(e);
			}
		});
	}
}
 

protected <N, S extends State, T> S getOrCreateKeyedState(
		TypeSerializer<N> namespaceSerializer,
		StateDescriptor<S, T> stateDescriptor) throws Exception {

	if (keyedStateStore != null) {
		return keyedStateBackend.getOrCreateKeyedState(namespaceSerializer, stateDescriptor);
	}
	else {
		throw new IllegalStateException("Cannot create partitioned state. " +
				"The keyed state backend has not been set." +
				"This indicates that the operator is not partitioned/keyed.");
	}
}
 

@SuppressWarnings("unchecked")
@Override
@Nonnull
public State createState(@Nonnull FunctionInitializationContext context, @Nonnull StateTtlConfig ttlConfig) {
	stateDesc.enableTimeToLive(ttlConfig);
	return createState(context);
}
 

@SuppressWarnings({"unchecked", "unused"})
static <N, T, S extends State, IS extends S> IS createState(
	TypeSerializer<N> namespaceSerializer,
	StateDescriptor<S, T> stateDesc) {
	ReducingStateDescriptor<T> reducingStateDesc = (ReducingStateDescriptor<T>) stateDesc;
	return (IS) new MockInternalReducingState<>(reducingStateDesc.getReduceFunction());
}
 

public static <K, N, SV, TTLSV, S extends State, IS extends S> IS createStateAndWrapWithTtlIfEnabled(
	TypeSerializer<N> namespaceSerializer,
	StateDescriptor<S, SV> stateDesc,
	KeyedStateBackend<K> stateBackend,
	TtlTimeProvider timeProvider) throws Exception {
	Preconditions.checkNotNull(namespaceSerializer);
	Preconditions.checkNotNull(stateDesc);
	Preconditions.checkNotNull(stateBackend);
	Preconditions.checkNotNull(timeProvider);
	return  stateDesc.getTtlConfig().isEnabled() ?
		new TtlStateFactory<K, N, SV, TTLSV, S, IS>(
			namespaceSerializer, stateDesc, stateBackend, timeProvider)
			.createState() :
		stateBackend.createInternalState(namespaceSerializer, stateDesc);
}
 

@SuppressWarnings({"unchecked", "unused"})
static <T, N, ACC, S extends State, IS extends S> IS createState(
	TypeSerializer<N> namespaceSerializer,
	StateDescriptor<S, ACC> stateDesc) {
	FoldingStateDescriptor<T, ACC> foldingStateDesc = (FoldingStateDescriptor<T, ACC>) stateDesc;
	return (IS) new MockInternalFoldingState<>(foldingStateDesc.getFoldFunction());
}
 

@SuppressWarnings("unchecked")
public <S extends State> S getPartitionedState(StateDescriptor<S, ?> stateDescriptor) {
	try {
		return WindowOperator.this.getPartitionedState(window, windowSerializer, stateDescriptor);
	} catch (Exception e) {
		throw new RuntimeException("Could not retrieve state", e);
	}
}
 

@SuppressWarnings("unchecked")
public static <ACC, S extends State> S createState(
	StateDescriptor<S, ACC> stateDescriptor,
	byte[] serializedState) throws IOException {
	final ACC state = KvStateSerializer.deserializeValue(
		serializedState,
		stateDescriptor.getSerializer());
	return (S) new ImmutableFoldingState<>(state);
}
 

@SuppressWarnings("unchecked")
<N, S extends State, V> S createState(
	StateDescriptor<S, V> stateDescriptor,
	@SuppressWarnings("SameParameterValue") N defaultNamespace) throws Exception {
	S state = keyedStateBackend.getOrCreateKeyedState(StringSerializer.INSTANCE, stateDescriptor);
	((InternalKvState<?, N, ?>) state).setCurrentNamespace(defaultNamespace);
	return state;
}