use of org.apache.flink.runtime.state.StateSnapshotTransformer in project flink by apache.
the class CopyOnWriteSkipListStateMapTestUtils method verifySnapshotWithTransform.
static <K, N, S> void verifySnapshotWithTransform(@Nonnull Map<N, Map<K, S>> referenceStates, @Nonnull CopyOnWriteSkipListStateMapSnapshot<K, N, S> snapshot, StateSnapshotTransformer<S> transformer, TypeSerializer<K> keySerializer, TypeSerializer<N> namespaceSerializer, TypeSerializer<S> stateSerializer, SnapshotVerificationMode verificationMode) throws IOException {
ByteArrayOutputStreamWithPos outputStream = new ByteArrayOutputStreamWithPos();
DataOutputView outputView = new DataOutputViewStreamWrapper(outputStream);
Map<N, Map<K, S>> transformedStates = new HashMap<>();
for (Map.Entry<N, Map<K, S>> namespaceEntry : referenceStates.entrySet()) {
for (Map.Entry<K, S> keyEntry : namespaceEntry.getValue().entrySet()) {
S state = transformer.filterOrTransform(keyEntry.getValue());
if (state != null) {
transformedStates.computeIfAbsent(namespaceEntry.getKey(), (none) -> new HashMap<>()).put(keyEntry.getKey(), state);
}
}
}
if (verificationMode == SnapshotVerificationMode.SERIALIZED) {
snapshot.writeState(keySerializer, namespaceSerializer, stateSerializer, outputView, transformer);
Map<N, Map<K, S>> actualStates = readStateFromSnapshot(outputStream.toByteArray(), keySerializer, namespaceSerializer, stateSerializer);
assertEquals(transformedStates, actualStates);
} else {
Iterator<StateEntry<K, N, S>> iterator = snapshot.getIterator(keySerializer, namespaceSerializer, stateSerializer, null);
assertThat(() -> iterator, containsInAnyOrder(toMatchers(referenceStates)));
}
}
use of org.apache.flink.runtime.state.StateSnapshotTransformer in project flink by apache.
the class MockKeyedStateBackend method createInternalState.
@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.getType());
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;
}
Aggregations