Examples with StateDescriptor org.apache.flink.api.common.state.StateDescriptor used on opensource projects

Example 1 with StateDescriptor

use of org.apache.flink.api.common.state.StateDescriptor in project flink by apache.

the class RocksDBKeyedStateBackend method restoreOldSavepointKeyedState.

/**
	 * For backwards compatibility, remove again later!
	 */
@Deprecated
private void restoreOldSavepointKeyedState(Collection<KeyGroupsStateHandle> restoreState) throws Exception {
    if (restoreState.isEmpty()) {
        return;
    }
    Preconditions.checkState(1 == restoreState.size(), "Only one element expected here.");
    HashMap<String, RocksDBStateBackend.FinalFullyAsyncSnapshot> namedStates;
    try (FSDataInputStream inputStream = restoreState.iterator().next().openInputStream()) {
        namedStates = InstantiationUtil.deserializeObject(inputStream, userCodeClassLoader);
    }
    Preconditions.checkState(1 == namedStates.size(), "Only one element expected here.");
    DataInputView inputView = namedStates.values().iterator().next().stateHandle.getState(userCodeClassLoader);
    // clear k/v state information before filling it
    kvStateInformation.clear();
    // first get the column family mapping
    int numColumns = inputView.readInt();
    Map<Byte, StateDescriptor<?, ?>> columnFamilyMapping = new HashMap<>(numColumns);
    for (int i = 0; i < numColumns; i++) {
        byte mappingByte = inputView.readByte();
        ObjectInputStream ooIn = new InstantiationUtil.ClassLoaderObjectInputStream(new DataInputViewStream(inputView), userCodeClassLoader);
        StateDescriptor stateDescriptor = (StateDescriptor) ooIn.readObject();
        columnFamilyMapping.put(mappingByte, stateDescriptor);
        // this will fill in the k/v state information
        getColumnFamily(stateDescriptor, MigrationNamespaceSerializerProxy.INSTANCE);
    }
    // try and read until EOF
    try {
        // the EOFException will get us out of this...
        while (true) {
            byte mappingByte = inputView.readByte();
            ColumnFamilyHandle handle = getColumnFamily(columnFamilyMapping.get(mappingByte), MigrationNamespaceSerializerProxy.INSTANCE);
            byte[] keyAndNamespace = BytePrimitiveArraySerializer.INSTANCE.deserialize(inputView);
            ByteArrayInputStreamWithPos bis = new ByteArrayInputStreamWithPos(keyAndNamespace);
            K reconstructedKey = keySerializer.deserialize(new DataInputViewStreamWrapper(bis));
            int len = bis.getPosition();
            int keyGroup = (byte) KeyGroupRangeAssignment.assignToKeyGroup(reconstructedKey, numberOfKeyGroups);
            if (keyGroupPrefixBytes == 1) {
                // copy and override one byte (42) between key and namespace
                System.arraycopy(keyAndNamespace, 0, keyAndNamespace, 1, len);
                keyAndNamespace[0] = (byte) keyGroup;
            } else {
                byte[] largerKey = new byte[1 + keyAndNamespace.length];
                // write key-group
                largerKey[0] = (byte) ((keyGroup >> 8) & 0xFF);
                largerKey[1] = (byte) (keyGroup & 0xFF);
                // write key
                System.arraycopy(keyAndNamespace, 0, largerKey, 2, len);
                //skip one byte (42), write namespace
                System.arraycopy(keyAndNamespace, 1 + len, largerKey, 2 + len, keyAndNamespace.length - len - 1);
                keyAndNamespace = largerKey;
            }
            byte[] value = BytePrimitiveArraySerializer.INSTANCE.deserialize(inputView);
            db.put(handle, keyAndNamespace, value);
        }
    } catch (EOFException e) {
    // expected
    }
}

Example 2 with StateDescriptor

use of org.apache.flink.api.common.state.StateDescriptor in project flink by apache.

the class SimpleStateRequestHandler method getMapState.

private MapState<ByteArrayWrapper, byte[]> getMapState(BeamFnApi.StateRequest request) throws Exception {
    BeamFnApi.StateKey.MultimapSideInput mapUserState = request.getStateKey().getMultimapSideInput();
    byte[] data = Base64.getDecoder().decode(mapUserState.getSideInputId());
    FlinkFnApi.StateDescriptor stateDescriptor = FlinkFnApi.StateDescriptor.parseFrom(data);
    String stateName = PYTHON_STATE_PREFIX + stateDescriptor.getStateName();
    StateDescriptor cachedStateDescriptor = stateDescriptorCache.get(stateName);
    MapStateDescriptor<ByteArrayWrapper, byte[]> mapStateDescriptor;
    if (cachedStateDescriptor instanceof MapStateDescriptor) {
        mapStateDescriptor = (MapStateDescriptor<ByteArrayWrapper, byte[]>) cachedStateDescriptor;
    } else if (cachedStateDescriptor == null) {
        mapStateDescriptor = new MapStateDescriptor<>(stateName, ByteArrayWrapperSerializer.INSTANCE, valueSerializer);
        if (stateDescriptor.hasStateTtlConfig()) {
            FlinkFnApi.StateDescriptor.StateTTLConfig stateTtlConfigProto = stateDescriptor.getStateTtlConfig();
            StateTtlConfig stateTtlConfig = ProtoUtils.parseStateTtlConfigFromProto(stateTtlConfigProto);
            mapStateDescriptor.enableTimeToLive(stateTtlConfig);
        }
        stateDescriptorCache.put(stateName, mapStateDescriptor);
    } else {
        throw new RuntimeException(String.format("State name corrupt detected: " + "'%s' is used both as MAP state and '%s' state at the same time.", stateName, cachedStateDescriptor.getType()));
    }
    byte[] windowBytes = mapUserState.getWindow().toByteArray();
    if (windowBytes.length != 0) {
        bais.setBuffer(windowBytes, 0, windowBytes.length);
        Object namespace = namespaceSerializer.deserialize(baisWrapper);
        return (MapState<ByteArrayWrapper, byte[]>) keyedStateBackend.getPartitionedState(namespace, namespaceSerializer, mapStateDescriptor);
    } else {
        return (MapState<ByteArrayWrapper, byte[]>) keyedStateBackend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, mapStateDescriptor);
    }
}

Example 3 with StateDescriptor

use of org.apache.flink.api.common.state.StateDescriptor in project flink by apache.

the class SimpleStateRequestHandler method getListState.

private ListState<byte[]> getListState(BeamFnApi.StateRequest request) throws Exception {
    BeamFnApi.StateKey.BagUserState bagUserState = request.getStateKey().getBagUserState();
    byte[] data = Base64.getDecoder().decode(bagUserState.getUserStateId());
    FlinkFnApi.StateDescriptor stateDescriptor = FlinkFnApi.StateDescriptor.parseFrom(data);
    String stateName = PYTHON_STATE_PREFIX + stateDescriptor.getStateName();
    ListStateDescriptor<byte[]> listStateDescriptor;
    StateDescriptor cachedStateDescriptor = stateDescriptorCache.get(stateName);
    if (cachedStateDescriptor instanceof ListStateDescriptor) {
        listStateDescriptor = (ListStateDescriptor<byte[]>) cachedStateDescriptor;
    } else if (cachedStateDescriptor == null) {
        listStateDescriptor = new ListStateDescriptor<>(stateName, valueSerializer);
        if (stateDescriptor.hasStateTtlConfig()) {
            FlinkFnApi.StateDescriptor.StateTTLConfig stateTtlConfigProto = stateDescriptor.getStateTtlConfig();
            StateTtlConfig stateTtlConfig = ProtoUtils.parseStateTtlConfigFromProto(stateTtlConfigProto);
            listStateDescriptor.enableTimeToLive(stateTtlConfig);
        }
        stateDescriptorCache.put(stateName, listStateDescriptor);
    } else {
        throw new RuntimeException(String.format("State name corrupt detected: " + "'%s' is used both as LIST state and '%s' state at the same time.", stateName, cachedStateDescriptor.getType()));
    }
    byte[] windowBytes = bagUserState.getWindow().toByteArray();
    if (windowBytes.length != 0) {
        bais.setBuffer(windowBytes, 0, windowBytes.length);
        Object namespace = namespaceSerializer.deserialize(baisWrapper);
        return (ListState<byte[]>) keyedStateBackend.getPartitionedState(namespace, namespaceSerializer, listStateDescriptor);
    } else {
        return (ListState<byte[]>) keyedStateBackend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, listStateDescriptor);
    }
}

Example 4 with StateDescriptor

use of org.apache.flink.api.common.state.StateDescriptor in project flink by apache.

the class StreamingRuntimeContextTest method testValueStateInstantiation.

@Test
public void testValueStateInstantiation() throws Exception {
    final ExecutionConfig config = new ExecutionConfig();
    config.registerKryoType(Path.class);
    final AtomicReference<Object> descriptorCapture = new AtomicReference<>();
    StreamingRuntimeContext context = createRuntimeContext(descriptorCapture, config);
    ValueStateDescriptor<TaskInfo> descr = new ValueStateDescriptor<>("name", TaskInfo.class);
    context.getState(descr);
    StateDescriptor<?, ?> descrIntercepted = (StateDescriptor<?, ?>) descriptorCapture.get();
    TypeSerializer<?> serializer = descrIntercepted.getSerializer();
    // check that the Path class is really registered, i.e., the execution config was applied
    assertTrue(serializer instanceof KryoSerializer);
    assertTrue(((KryoSerializer<?>) serializer).getKryo().getRegistration(Path.class).getId() > 0);
}

Example 5 with StateDescriptor

use of org.apache.flink.api.common.state.StateDescriptor 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;
}