Examples with DataInputViewStreamWrapper org.apache.flink.core.memory.DataInputViewStreamWrapper used on opensource projects

Example 31 with DataInputViewStreamWrapper

use of org.apache.flink.core.memory.DataInputViewStreamWrapper in project flink by apache.

the class DefaultOperatorStateBackend method restore.

@Override
public void restore(Collection<OperatorStateHandle> restoreSnapshots) throws Exception {
    if (null == restoreSnapshots) {
        return;
    }
    for (OperatorStateHandle stateHandle : restoreSnapshots) {
        if (stateHandle == null) {
            continue;
        }
        FSDataInputStream in = stateHandle.openInputStream();
        closeStreamOnCancelRegistry.registerClosable(in);
        ClassLoader restoreClassLoader = Thread.currentThread().getContextClassLoader();
        try {
            Thread.currentThread().setContextClassLoader(userClassloader);
            OperatorBackendSerializationProxy backendSerializationProxy = new OperatorBackendSerializationProxy(userClassloader);
            backendSerializationProxy.read(new DataInputViewStreamWrapper(in));
            List<OperatorBackendSerializationProxy.StateMetaInfo<?>> metaInfoList = backendSerializationProxy.getNamedStateSerializationProxies();
            // Recreate all PartitionableListStates from the meta info
            for (OperatorBackendSerializationProxy.StateMetaInfo<?> stateMetaInfo : metaInfoList) {
                PartitionableListState<?> listState = registeredStates.get(stateMetaInfo.getName());
                if (null == listState) {
                    listState = new PartitionableListState<>(stateMetaInfo.getName(), stateMetaInfo.getStateSerializer(), stateMetaInfo.getMode());
                    registeredStates.put(listState.getName(), listState);
                } else {
                    Preconditions.checkState(listState.getPartitionStateSerializer().isCompatibleWith(stateMetaInfo.getStateSerializer()), "Incompatible state serializers found: " + listState.getPartitionStateSerializer() + " is not compatible with " + stateMetaInfo.getStateSerializer());
                }
            }
            // Restore all the state in PartitionableListStates
            for (Map.Entry<String, OperatorStateHandle.StateMetaInfo> nameToOffsets : stateHandle.getStateNameToPartitionOffsets().entrySet()) {
                PartitionableListState<?> stateListForName = registeredStates.get(nameToOffsets.getKey());
                Preconditions.checkState(null != stateListForName, "Found state without " + "corresponding meta info: " + nameToOffsets.getKey());
                deserializeStateValues(stateListForName, in, nameToOffsets.getValue());
            }
        } finally {
            Thread.currentThread().setContextClassLoader(restoreClassLoader);
            closeStreamOnCancelRegistry.unregisterClosable(in);
            IOUtils.closeQuietly(in);
        }
    }
}

Example 32 with DataInputViewStreamWrapper

use of org.apache.flink.core.memory.DataInputViewStreamWrapper in project flink by apache.

the class NFA method readComputationState.

@SuppressWarnings("unchecked")
private ComputationState<T> readComputationState(ObjectInputStream ois) throws IOException, ClassNotFoundException {
    final State<T> state = (State<T>) ois.readObject();
    final long timestamp = ois.readLong();
    final DeweyNumber version = (DeweyNumber) ois.readObject();
    final long startTimestamp = ois.readLong();
    final boolean hasEvent = ois.readBoolean();
    final T event;
    if (hasEvent) {
        DataInputViewStreamWrapper input = new DataInputViewStreamWrapper(ois);
        event = nonDuplicatingTypeSerializer.deserialize(input);
    } else {
        event = null;
    }
    return new ComputationState<>(state, event, timestamp, version, startTimestamp);
}

Example 33 with DataInputViewStreamWrapper

use of org.apache.flink.core.memory.DataInputViewStreamWrapper in project flink by apache.

the class SharedBuffer method readObject.

private void readObject(ObjectInputStream ois) throws IOException, ClassNotFoundException {
    DataInputViewStreamWrapper source = new DataInputViewStreamWrapper(ois);
    ArrayList<SharedBufferEntry<K, V>> entryList = new ArrayList<>();
    ois.defaultReadObject();
    this.pages = new HashMap<>();
    int numberPages = ois.readInt();
    for (int i = 0; i < numberPages; i++) {
        // key of the page
        @SuppressWarnings("unchecked") K key = (K) ois.readObject();
        SharedBufferPage<K, V> page = new SharedBufferPage<>(key);
        pages.put(key, page);
        int numberEntries = ois.readInt();
        for (int j = 0; j < numberEntries; j++) {
            // restore the SharedBufferEntries for the given page
            V value = valueSerializer.deserialize(source);
            long timestamp = ois.readLong();
            ValueTimeWrapper<V> valueTimeWrapper = new ValueTimeWrapper<>(value, timestamp);
            SharedBufferEntry<K, V> sharedBufferEntry = new SharedBufferEntry<K, V>(valueTimeWrapper, page);
            sharedBufferEntry.referenceCounter = ois.readInt();
            page.entries.put(valueTimeWrapper, sharedBufferEntry);
            entryList.add(sharedBufferEntry);
        }
    }
    // read the edges of the shared buffer entries
    int numberEdges = ois.readInt();
    for (int j = 0; j < numberEdges; j++) {
        int sourceIndex = ois.readInt();
        int targetIndex = ois.readInt();
        if (sourceIndex >= entryList.size() || sourceIndex < 0) {
            throw new RuntimeException("Could not find source entry with index " + sourceIndex + ". This indicates a corrupted state.");
        } else {
            // We've already deserialized the shared buffer entry. Simply read its ID and
            // retrieve the buffer entry from the list of entries
            SharedBufferEntry<K, V> sourceEntry = entryList.get(sourceIndex);
            final DeweyNumber version = (DeweyNumber) ois.readObject();
            final SharedBufferEntry<K, V> target;
            if (targetIndex >= 0) {
                if (targetIndex >= entryList.size()) {
                    throw new RuntimeException("Could not find target entry with index " + targetIndex + ". This indicates a corrupted state.");
                } else {
                    target = entryList.get(targetIndex);
                }
            } else {
                target = null;
            }
            sourceEntry.edges.add(new SharedBufferEdge<K, V>(target, version));
        }
    }
}

Example 34 with DataInputViewStreamWrapper

use of org.apache.flink.core.memory.DataInputViewStreamWrapper in project beam by apache.

the class DoFnOperator method initializeState.

@Override
public void initializeState(StateInitializationContext context) throws Exception {
    if (getKeyedStateBackend() != null) {
        int totalKeyGroups = getKeyedStateBackend().getNumberOfKeyGroups();
        KeyGroupsList localKeyGroupRange = getKeyedStateBackend().getKeyGroupRange();
        for (KeyGroupStatePartitionStreamProvider streamProvider : context.getRawKeyedStateInputs()) {
            DataInputViewStreamWrapper div = new DataInputViewStreamWrapper(streamProvider.getStream());
            int keyGroupIdx = streamProvider.getKeyGroupId();
            checkArgument(localKeyGroupRange.contains(keyGroupIdx), "Key Group " + keyGroupIdx + " does not belong to the local range.");
            // if (this instanceof KeyGroupRestoringOperator)
            restoreKeyGroupState(keyGroupIdx, div);
            // We just initialize our timerService
            if (keyCoder != null) {
                if (timerService == null) {
                    timerService = new HeapInternalTimerService<>(totalKeyGroups, localKeyGroupRange, this, getRuntimeContext().getProcessingTimeService());
                }
                timerService.restoreTimersForKeyGroup(div, keyGroupIdx, getUserCodeClassloader());
            }
        }
    }
}

Example 35 with DataInputViewStreamWrapper

use of org.apache.flink.core.memory.DataInputViewStreamWrapper in project flink by apache.

the class AbstractKeyedCEPPatternOperator method restoreState.

//////////////////////			Backwards Compatibility			//////////////////////
@Override
public void restoreState(FSDataInputStream in) throws Exception {
    // this is the flag indicating if we have udf
    // state to restore (not needed here)
    in.read();
    DataInputViewStreamWrapper inputView = new DataInputViewStreamWrapper(in);
    InternalWatermarkCallbackService<KEY> watermarkCallbackService = getInternalWatermarkCallbackService();
    if (migratingFromOldKeyedOperator) {
        int numberEntries = inputView.readInt();
        for (int i = 0; i < numberEntries; i++) {
            watermarkCallbackService.registerKeyForWatermarkCallback(keySerializer.deserialize(inputView));
        }
    } else {
        final ObjectInputStream ois = new ObjectInputStream(in);
        // retrieve the NFA
        @SuppressWarnings("unchecked") NFA<IN> nfa = (NFA<IN>) ois.readObject();
        // retrieve the elements that were pending in the priority queue
        MultiplexingStreamRecordSerializer<IN> recordSerializer = new MultiplexingStreamRecordSerializer<>(inputSerializer);
        PriorityQueue<StreamRecord<IN>> priorityQueue = priorityQueueFactory.createPriorityQueue();
        int entries = ois.readInt();
        for (int i = 0; i < entries; i++) {
            StreamElement streamElement = recordSerializer.deserialize(inputView);
            priorityQueue.offer(streamElement.<IN>asRecord());
        }
        // finally register the retrieved state with the new keyed state.
        setCurrentKey((byte) 0);
        nfaOperatorState.update(nfa);
        priorityQueueOperatorState.update(priorityQueue);
        if (!isProcessingTime) {
            // this is relevant only for event/ingestion time
            // need to work around type restrictions
            InternalWatermarkCallbackService rawWatermarkCallbackService = (InternalWatermarkCallbackService) watermarkCallbackService;
            rawWatermarkCallbackService.registerKeyForWatermarkCallback((byte) 0);
        }
        ois.close();
    }
}