Examples with TypeCase org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey.TypeCase used on opensource projects

Example 1 with TypeCase

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey.TypeCase in project beam by apache.

the class SparkExecutableStageFunction method getStateRequestHandler.

private StateRequestHandler getStateRequestHandler(ExecutableStage executableStage, ProcessBundleDescriptors.ExecutableProcessBundleDescriptor processBundleDescriptor) {
    EnumMap<TypeCase, StateRequestHandler> handlerMap = new EnumMap<>(StateKey.TypeCase.class);
    final StateRequestHandler sideInputHandler;
    StateRequestHandlers.SideInputHandlerFactory sideInputHandlerFactory = BatchSideInputHandlerFactory.forStage(executableStage, new BatchSideInputHandlerFactory.SideInputGetter() {

        @Override
        public <T> List<T> getSideInput(String pCollectionId) {
            Tuple2<Broadcast<List<byte[]>>, WindowedValueCoder<SideInputT>> tuple2 = sideInputs.get(pCollectionId);
            Broadcast<List<byte[]>> broadcast = tuple2._1;
            WindowedValueCoder<SideInputT> coder = tuple2._2;
            return (List<T>) broadcast.value().stream().map(bytes -> CoderHelpers.fromByteArray(bytes, coder)).collect(Collectors.toList());
        }
    });
    try {
        sideInputHandler = StateRequestHandlers.forSideInputHandlerFactory(ProcessBundleDescriptors.getSideInputs(executableStage), sideInputHandlerFactory);
    } catch (IOException e) {
        throw new RuntimeException("Failed to setup state handler", e);
    }
    if (bagUserStateHandlerFactory == null) {
        bagUserStateHandlerFactory = new InMemoryBagUserStateFactory();
    }
    final StateRequestHandler userStateHandler;
    if (executableStage.getUserStates().size() > 0) {
        // Need to discard the old key's state
        bagUserStateHandlerFactory.resetForNewKey();
        userStateHandler = StateRequestHandlers.forBagUserStateHandlerFactory(processBundleDescriptor, bagUserStateHandlerFactory);
    } else {
        userStateHandler = StateRequestHandler.unsupported();
    }
    handlerMap.put(StateKey.TypeCase.ITERABLE_SIDE_INPUT, sideInputHandler);
    handlerMap.put(StateKey.TypeCase.MULTIMAP_SIDE_INPUT, sideInputHandler);
    handlerMap.put(StateKey.TypeCase.MULTIMAP_KEYS_SIDE_INPUT, sideInputHandler);
    handlerMap.put(StateKey.TypeCase.BAG_USER_STATE, userStateHandler);
    return StateRequestHandlers.delegateBasedUponType(handlerMap);
}

Example 2 with TypeCase

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey.TypeCase in project beam by apache.

the class ExecutableStageDoFnOperator method getStateRequestHandler.

private StateRequestHandler getStateRequestHandler(ExecutableStage executableStage) {
    final StateRequestHandler sideInputStateHandler;
    if (executableStage.getSideInputs().size() > 0) {
        checkNotNull(super.sideInputHandler);
        StateRequestHandlers.SideInputHandlerFactory sideInputHandlerFactory = Preconditions.checkNotNull(StreamingSideInputHandlerFactory.forStage(executableStage, sideInputIds, super.sideInputHandler));
        try {
            sideInputStateHandler = StateRequestHandlers.forSideInputHandlerFactory(ProcessBundleDescriptors.getSideInputs(executableStage), sideInputHandlerFactory);
        } catch (IOException e) {
            throw new RuntimeException("Failed to initialize SideInputHandler", e);
        }
    } else {
        sideInputStateHandler = StateRequestHandler.unsupported();
    }
    final StateRequestHandler userStateRequestHandler;
    if (!executableStage.getUserStates().isEmpty()) {
        if (keyedStateInternals == null) {
            throw new IllegalStateException("Input must be keyed when user state is used");
        }
        userStateRequestHandler = StateRequestHandlers.forBagUserStateHandlerFactory(stageBundleFactory.getProcessBundleDescriptor(), new BagUserStateFactory(keyedStateInternals, getKeyedStateBackend(), stateBackendLock, keyCoder));
    } else {
        userStateRequestHandler = StateRequestHandler.unsupported();
    }
    EnumMap<TypeCase, StateRequestHandler> handlerMap = new EnumMap<>(TypeCase.class);
    handlerMap.put(TypeCase.ITERABLE_SIDE_INPUT, sideInputStateHandler);
    handlerMap.put(TypeCase.MULTIMAP_SIDE_INPUT, sideInputStateHandler);
    handlerMap.put(TypeCase.MULTIMAP_KEYS_SIDE_INPUT, sideInputStateHandler);
    handlerMap.put(TypeCase.BAG_USER_STATE, userStateRequestHandler);
    return StateRequestHandlers.delegateBasedUponType(handlerMap);
}