/** * Tests map serialization utils. */ @Test public void testMapSerialization() throws Exception { final long key = 0L; final KeyGroupRange keyGroupRange = new KeyGroupRange(0, 0); // objects for heap state list serialisation final HeapKeyedStateBackend<Long> longHeapKeyedStateBackend = new HeapKeyedStateBackend<>( mock(TaskKvStateRegistry.class), LongSerializer.INSTANCE, ClassLoader.getSystemClassLoader(), keyGroupRange.getNumberOfKeyGroups(), keyGroupRange, async, new ExecutionConfig(), TestLocalRecoveryConfig.disabled(), new HeapPriorityQueueSetFactory(keyGroupRange, keyGroupRange.getNumberOfKeyGroups(), 128), TtlTimeProvider.DEFAULT ); longHeapKeyedStateBackend.setCurrentKey(key); final InternalMapState<Long, VoidNamespace, Long, String> mapState = (InternalMapState<Long, VoidNamespace, Long, String>) longHeapKeyedStateBackend.getPartitionedState( VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, new MapStateDescriptor<>("test", LongSerializer.INSTANCE, StringSerializer.INSTANCE)); testMapSerialization(key, mapState); }
/** * Tests list serialization utils. */ @Test public void testListSerialization() throws Exception { final long key = 0L; final KeyGroupRange keyGroupRange = new KeyGroupRange(0, 0); // objects for heap state list serialisation final HeapKeyedStateBackend<Long> longHeapKeyedStateBackend = new HeapKeyedStateBackend<>( mock(TaskKvStateRegistry.class), LongSerializer.INSTANCE, ClassLoader.getSystemClassLoader(), keyGroupRange.getNumberOfKeyGroups(), keyGroupRange, async, new ExecutionConfig(), TestLocalRecoveryConfig.disabled(), new HeapPriorityQueueSetFactory(keyGroupRange, keyGroupRange.getNumberOfKeyGroups(), 128), TtlTimeProvider.DEFAULT ); longHeapKeyedStateBackend.setCurrentKey(key); final InternalListState<Long, VoidNamespace, Long> listState = longHeapKeyedStateBackend.createInternalState( VoidNamespaceSerializer.INSTANCE, new ListStateDescriptor<>("test", LongSerializer.INSTANCE)); testListSerialization(key, listState); }
@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); } } }
@Override public <N, T, ACC, R> InternalAggregatingState<N, T, R> createAggregatingState( TypeSerializer<N> namespaceSerializer, AggregatingStateDescriptor<T, ACC, R> stateDesc) throws Exception { StateTable<K, N, ACC> stateTable = tryRegisterStateTable(namespaceSerializer, stateDesc); return new HeapAggregatingState<>(stateDesc, stateTable, keySerializer, namespaceSerializer); }
@Override public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend( Environment env, JobID jobID, String operatorIdentifier, TypeSerializer<K> keySerializer, int numberOfKeyGroups, KeyGroupRange keyGroupRange, TaskKvStateRegistry kvStateRegistry) throws IOException { return new HeapKeyedStateBackend<>( kvStateRegistry, keySerializer, env.getUserClassLoader(), numberOfKeyGroups, keyGroupRange, asynchronousSnapshots, env.getExecutionConfig()); }
@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, getStateSnapshotTransformer(stateDesc, snapshotTransformFactory)); return stateFactory.createState(stateDesc, stateTable, keySerializer); }
@SuppressWarnings("deprecation") public void restore(Collection<KeyedStateHandle> restoredState) throws Exception { if (restoredState == null || restoredState.isEmpty()) { return; } LOG.info("Initializing heap keyed state backend from snapshot."); if (LOG.isDebugEnabled()) { LOG.debug("Restoring snapshot from state handles: {}.", restoredState); } restorePartitionedState(restoredState); }
public TriggerTestHarness( Trigger<T, W> trigger, TypeSerializer<W> windowSerializer) throws Exception { this.trigger = trigger; this.windowSerializer = windowSerializer; // we only ever use one key, other tests make sure that windows work across different // keys DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0); MemoryStateBackend backend = new MemoryStateBackend(); @SuppressWarnings("unchecked") HeapKeyedStateBackend<Integer> stateBackend = (HeapKeyedStateBackend<Integer>) backend.createKeyedStateBackend(dummyEnv, new JobID(), "test_op", IntSerializer.INSTANCE, 1, new KeyGroupRange(0, 0), new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID())); this.stateBackend = stateBackend; this.stateBackend.setCurrentKey(KEY); this.internalTimerService = new TestInternalTimerService<>(new KeyContext() { @Override public void setCurrentKey(Object key) { // ignore } @Override public Object getCurrentKey() { return KEY; } }); }
serializationProxy.getStateMetaInfoSnapshots(); createOrCheckStateForMetaInfo(restoredMetaInfos, kvStatesById); readStateHandleStateData( fsDataInputStream, inView,
@SuppressWarnings("deprecation") @Override public void restore(Collection<KeyedStateHandle> restoredState) throws Exception { if (restoredState == null || restoredState.isEmpty()) { return; } LOG.info("Initializing heap keyed state backend from snapshot."); if (LOG.isDebugEnabled()) { LOG.debug("Restoring snapshot from state handles: {}.", restoredState); } if (MigrationUtil.isOldSavepointKeyedState(restoredState)) { restoreOldSavepointKeyedState(restoredState); } else { restorePartitionedState(restoredState); } }
stateTable = newStateTable(newMetaInfo); stateTables.put(stateName, stateTable); } else {
registeredState = registeredPQStates.get(metaInfoSnapshot.getName()); if (registeredState == null) { createInternal(new RegisteredPriorityQueueStateBackendMetaInfo<>(metaInfoSnapshot));
private void readStateHandleStateData( FSDataInputStream fsDataInputStream, DataInputViewStreamWrapper inView, KeyGroupRangeOffsets keyGroupOffsets, Map<Integer, StateMetaInfoSnapshot> kvStatesById, int numStates, int readVersion, boolean isCompressed) throws IOException { final StreamCompressionDecorator streamCompressionDecorator = isCompressed ? SnappyStreamCompressionDecorator.INSTANCE : UncompressedStreamCompressionDecorator.INSTANCE; for (Tuple2<Integer, Long> groupOffset : keyGroupOffsets) { int keyGroupIndex = groupOffset.f0; long offset = groupOffset.f1; // Check that restored key groups all belong to the backend. Preconditions.checkState(keyGroupRange.contains(keyGroupIndex), "The key group must belong to the backend."); fsDataInputStream.seek(offset); int writtenKeyGroupIndex = inView.readInt(); Preconditions.checkState(writtenKeyGroupIndex == keyGroupIndex, "Unexpected key-group in restore."); try (InputStream kgCompressionInStream = streamCompressionDecorator.decorateWithCompression(fsDataInputStream)) { readKeyGroupStateData( kgCompressionInStream, kvStatesById, keyGroupIndex, numStates, readVersion); } } }
@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); } } }
@Override public <N, T> InternalReducingState<N, T> createReducingState( TypeSerializer<N> namespaceSerializer, ReducingStateDescriptor<T> stateDesc) throws Exception { StateTable<K, N, T> stateTable = tryRegisterStateTable(namespaceSerializer, stateDesc); return new HeapReducingState<>(stateDesc, stateTable, keySerializer, namespaceSerializer); }
@Override public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend( Environment env, JobID jobID, String operatorIdentifier, TypeSerializer<K> keySerializer, int numberOfKeyGroups, KeyGroupRange keyGroupRange, TaskKvStateRegistry kvStateRegistry) { return new HeapKeyedStateBackend<>( kvStateRegistry, keySerializer, env.getUserClassLoader(), numberOfKeyGroups, keyGroupRange, asynchronousSnapshots, env.getExecutionConfig()); } }
@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, getStateSnapshotTransformer(stateDesc, snapshotTransformFactory)); return stateFactory.createState(stateDesc, stateTable, keySerializer); }
@SuppressWarnings("deprecation") public void restore(Collection<KeyedStateHandle> restoredState) throws Exception { if (restoredState == null || restoredState.isEmpty()) { return; } LOG.info("Initializing heap keyed state backend from snapshot."); if (LOG.isDebugEnabled()) { LOG.debug("Restoring snapshot from state handles: {}.", restoredState); } restorePartitionedState(restoredState); }
serializationProxy.getStateMetaInfoSnapshots(); createOrCheckStateForMetaInfo(restoredMetaInfos, kvStatesById); readStateHandleStateData( fsDataInputStream, inView,
registeredState = registeredPQStates.get(metaInfoSnapshot.getName()); if (registeredState == null) { createInternal(new RegisteredPriorityQueueStateBackendMetaInfo<>(metaInfoSnapshot));