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

Example 6 with StateRequest

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

the class MultimapSideInput method get.

@Override
public Iterable<V> get(K k) {
    ByteString.Output output = ByteString.newOutput();
    try {
        keyCoder.encode(k, output);
    } catch (IOException e) {
        throw new IllegalStateException(String.format("Failed to encode key %s for side input id %s.", k, keysRequest.getStateKey().getMultimapKeysSideInput().getSideInputId()), e);
    }
    ByteString encodedKey = output.toByteString();
    StateKey stateKey = StateKey.newBuilder().setMultimapSideInput(StateKey.MultimapSideInput.newBuilder().setTransformId(keysRequest.getStateKey().getMultimapKeysSideInput().getTransformId()).setSideInputId(keysRequest.getStateKey().getMultimapKeysSideInput().getSideInputId()).setWindow(keysRequest.getStateKey().getMultimapKeysSideInput().getWindow()).setKey(encodedKey)).build();
    StateRequest request = keysRequest.toBuilder().setStateKey(stateKey).build();
    return StateFetchingIterators.readAllAndDecodeStartingFrom(Caches.subCache(cache, "ValuesForKey", encodedKey), beamFnStateClient, request, valueCoder);
}

Example 7 with StateRequest

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

the class RegisterAndProcessBundleOperation method handleBagUserState.

private CompletionStage<BeamFnApi.StateResponse.Builder> handleBagUserState(StateRequest stateRequest) {
    StateKey.BagUserState bagUserStateKey = stateRequest.getStateKey().getBagUserState();
    DataflowStepContext userStepContext = ptransformIdToUserStepContext.get(bagUserStateKey.getTransformId());
    checkState(userStepContext != null, String.format("Unknown PTransform id '%s'", bagUserStateKey.getTransformId()));
    // TODO: We should not be required to hold onto a pointer to the bag states for the
    // user. InMemoryStateInternals assumes that the Java garbage collector does the clean-up work
    // but instead StateInternals should hold its own references and write out any data and
    // clear references when the MapTask within Dataflow completes like how WindmillStateInternals
    // works.
    BagState<ByteString> state = userStateData.computeIfAbsent(stateRequest.getStateKey(), unused -> userStepContext.stateInternals().state(// window.
    StateNamespaces.window(GlobalWindow.Coder.INSTANCE, GlobalWindow.INSTANCE), StateTags.bag(bagUserStateKey.getUserStateId(), ByteStringCoder.of())));
    switch(stateRequest.getRequestCase()) {
        case GET:
            return CompletableFuture.completedFuture(StateResponse.newBuilder().setGet(StateGetResponse.newBuilder().setData(concat(state.read()))));
        case APPEND:
            state.add(stateRequest.getAppend().getData());
            return CompletableFuture.completedFuture(StateResponse.newBuilder().setAppend(StateAppendResponse.getDefaultInstance()));
        case CLEAR:
            state.clear();
            return CompletableFuture.completedFuture(StateResponse.newBuilder().setClear(StateClearResponse.getDefaultInstance()));
        default:
            throw new IllegalArgumentException(String.format("Unknown request type %s", stateRequest.getRequestCase()));
    }
}

Example 8 with StateRequest

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

the class RegisterAndProcessBundleOperation method handleMultimapSideInput.

private CompletionStage<BeamFnApi.StateResponse.Builder> handleMultimapSideInput(StateRequest stateRequest) {
    checkState(stateRequest.getRequestCase() == RequestCase.GET, String.format("MultimapSideInput state requests only support '%s' requests, received '%s'", RequestCase.GET, stateRequest.getRequestCase()));
    StateKey.MultimapSideInput multimapSideInputStateKey = stateRequest.getStateKey().getMultimapSideInput();
    SideInputReader sideInputReader = ptransformIdToSideInputReader.get(multimapSideInputStateKey.getTransformId());
    checkState(sideInputReader != null, String.format("Unknown PTransform '%s'", multimapSideInputStateKey.getTransformId()));
    PCollectionView<Materializations.MultimapView<Object, Object>> view = (PCollectionView<Materializations.MultimapView<Object, Object>>) ptransformIdToSideInputIdToPCollectionView.get(multimapSideInputStateKey.getTransformId(), multimapSideInputStateKey.getSideInputId());
    checkState(view != null, String.format("Unknown side input '%s' on PTransform '%s'", multimapSideInputStateKey.getSideInputId(), multimapSideInputStateKey.getTransformId()));
    checkState(Materializations.MULTIMAP_MATERIALIZATION_URN.equals(view.getViewFn().getMaterialization().getUrn()), String.format("Unknown materialization for side input '%s' on PTransform '%s' with urn '%s'", multimapSideInputStateKey.getSideInputId(), multimapSideInputStateKey.getTransformId(), view.getViewFn().getMaterialization().getUrn()));
    checkState(view.getCoderInternal() instanceof KvCoder, String.format("Materialization of side input '%s' on PTransform '%s' expects %s but received %s.", multimapSideInputStateKey.getSideInputId(), multimapSideInputStateKey.getTransformId(), KvCoder.class.getSimpleName(), view.getCoderInternal().getClass().getSimpleName()));
    Coder<Object> keyCoder = ((KvCoder) view.getCoderInternal()).getKeyCoder();
    Coder<Object> valueCoder = ((KvCoder) view.getCoderInternal()).getValueCoder();
    BoundedWindow window;
    try {
        // TODO: Use EncodedWindow instead of decoding the window.
        window = view.getWindowingStrategyInternal().getWindowFn().windowCoder().decode(multimapSideInputStateKey.getWindow().newInput());
    } catch (IOException e) {
        throw new IllegalArgumentException(String.format("Unable to decode window for side input '%s' on PTransform '%s'.", multimapSideInputStateKey.getSideInputId(), multimapSideInputStateKey.getTransformId()), e);
    }
    Object userKey;
    try {
        // TODO: Use the encoded representation of the key.
        userKey = keyCoder.decode(multimapSideInputStateKey.getKey().newInput());
    } catch (IOException e) {
        throw new IllegalArgumentException(String.format("Unable to decode user key for side input '%s' on PTransform '%s'.", multimapSideInputStateKey.getSideInputId(), multimapSideInputStateKey.getTransformId()), e);
    }
    Materializations.MultimapView<Object, Object> sideInput = sideInputReader.get(view, window);
    Iterable<Object> values = sideInput.get(userKey);
    try {
        // TODO: Use the raw value so we don't go through a decode/encode cycle for no reason.
        return CompletableFuture.completedFuture(StateResponse.newBuilder().setGet(StateGetResponse.newBuilder().setData(encodeAndConcat(values, valueCoder))));
    } catch (IOException e) {
        throw new IllegalArgumentException(String.format("Unable to encode values for side input '%s' on PTransform '%s'.", multimapSideInputStateKey.getSideInputId(), multimapSideInputStateKey.getTransformId()), e);
    }
}

Example 9 with StateRequest

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

the class StateRequestHandlersTest method testDelegatingStateHandlerDelegates.

@Test
public void testDelegatingStateHandlerDelegates() throws Exception {
    StateRequestHandler mockHandler = Mockito.mock(StateRequestHandler.class);
    StateRequestHandler mockHandler2 = Mockito.mock(StateRequestHandler.class);
    EnumMap<StateKey.TypeCase, StateRequestHandler> handlers = new EnumMap<>(StateKey.TypeCase.class);
    handlers.put(StateKey.TypeCase.TYPE_NOT_SET, mockHandler);
    handlers.put(TypeCase.MULTIMAP_SIDE_INPUT, mockHandler2);
    StateRequest request = StateRequest.getDefaultInstance();
    StateRequest request2 = StateRequest.newBuilder().setStateKey(StateKey.newBuilder().setMultimapSideInput(MultimapSideInput.getDefaultInstance())).build();
    StateRequestHandlers.delegateBasedUponType(handlers).handle(request);
    StateRequestHandlers.delegateBasedUponType(handlers).handle(request2);
    verify(mockHandler).handle(request);
    verify(mockHandler2).handle(request2);
    verifyNoMoreInteractions(mockHandler, mockHandler2);
}