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

Example 1 with StateKey

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

the class FakeBeamFnStateClient method handle.

@Override
public CompletableFuture<StateResponse> handle(StateRequest.Builder requestBuilder) {
    // The id should never be filled out
    assertEquals("", requestBuilder.getId());
    requestBuilder.setId(generateId());
    StateRequest request = requestBuilder.build();
    StateKey key = request.getStateKey();
    StateResponse.Builder response;
    assertNotEquals(RequestCase.REQUEST_NOT_SET, request.getRequestCase());
    assertNotEquals(TypeCase.TYPE_NOT_SET, key.getTypeCase());
    // multimap side input and runner based state keys only support get requests
    if (key.getTypeCase() == TypeCase.MULTIMAP_SIDE_INPUT || key.getTypeCase() == TypeCase.RUNNER) {
        assertEquals(RequestCase.GET, request.getRequestCase());
    }
    switch(request.getRequestCase()) {
        case GET:
            List<ByteString> byteStrings = data.getOrDefault(request.getStateKey(), Collections.singletonList(ByteString.EMPTY));
            int block = 0;
            if (request.getGet().getContinuationToken().size() > 0) {
                block = Integer.parseInt(request.getGet().getContinuationToken().toStringUtf8());
            }
            ByteString returnBlock = byteStrings.get(block);
            ByteString continuationToken = ByteString.EMPTY;
            if (byteStrings.size() > block + 1) {
                continuationToken = ByteString.copyFromUtf8(Integer.toString(block + 1));
            }
            response = StateResponse.newBuilder().setGet(StateGetResponse.newBuilder().setData(returnBlock).setContinuationToken(continuationToken));
            break;
        case CLEAR:
            data.remove(request.getStateKey());
            response = StateResponse.newBuilder().setClear(StateClearResponse.getDefaultInstance());
            break;
        case APPEND:
            List<ByteString> previousValue = data.computeIfAbsent(request.getStateKey(), (unused) -> new ArrayList<>());
            previousValue.add(request.getAppend().getData());
            response = StateResponse.newBuilder().setAppend(StateAppendResponse.getDefaultInstance());
            break;
        default:
            throw new IllegalStateException(String.format("Unknown request type %s", request.getRequestCase()));
    }
    return CompletableFuture.completedFuture(response.setId(requestBuilder.getId()).build());
}

Example 2 with StateKey

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

the class RegisterAndProcessBundleOperationTest method testProcessingBundleHandlesMultimapSideInputRequests.

@Test
public void testProcessingBundleHandlesMultimapSideInputRequests() throws Exception {
    IdGenerator idGenerator = makeIdGeneratorStartingFrom(777L);
    DataflowStepContext mockStepContext = mock(DataflowStepContext.class);
    DataflowStepContext mockUserStepContext = mock(DataflowStepContext.class);
    when(mockStepContext.namespacedToUser()).thenReturn(mockUserStepContext);
    CountDownLatch waitForStateHandler = new CountDownLatch(1);
    // Issues state calls to the Runner after a process bundle request is sent.
    InstructionRequestHandler fakeClient = new TestInstructionRequestHandler() {

        @Override
        public CompletionStage<InstructionResponse> handle(InstructionRequest request) {
            switch(request.getRequestCase()) {
                case REGISTER:
                    return CompletableFuture.completedFuture(responseFor(request).build());
                case PROCESS_BUNDLE:
                    return MoreFutures.supplyAsync(() -> {
                        StateKey getKey = StateKey.newBuilder().setMultimapSideInput(StateKey.MultimapSideInput.newBuilder().setTransformId("testPTransformId").setSideInputId("testSideInputId").setWindow(ByteString.copyFrom(CoderUtils.encodeToByteArray(GlobalWindow.Coder.INSTANCE, GlobalWindow.INSTANCE))).setKey(ByteString.copyFrom(CoderUtils.encodeToByteArray(ByteArrayCoder.of(), "ABC".getBytes(StandardCharsets.UTF_8), Coder.Context.NESTED)))).build();
                        StateRequest getRequest = StateRequest.newBuilder().setStateKey(getKey).setGet(StateGetRequest.getDefaultInstance()).build();
                        waitForStateHandler.await();
                        StateRequestHandler stateHandler = stateHandlerCaptor.getValue();
                        StateResponse.Builder getResponse = MoreFutures.get(stateHandler.handle(getRequest));
                        assertEquals(encodeAndConcat(Arrays.asList("X", "Y", "Z"), StringUtf8Coder.of()), getResponse.getGet().getData());
                        return responseFor(request).build();
                    });
                default:
                    // block forever on other request types
                    return new CompletableFuture<>();
            }
        }
    };
    SideInputReader fakeSideInputReader = new SideInputReader() {

        @Override
        @Nullable
        public <T> T get(PCollectionView<T> view, BoundedWindow window) {
            assertEquals(GlobalWindow.INSTANCE, window);
            assertEquals("testSideInputId", view.getTagInternal().getId());
            return (T) InMemoryMultimapSideInputView.fromIterable(ByteArrayCoder.of(), ImmutableList.of(KV.of("ABC".getBytes(StandardCharsets.UTF_8), "X"), KV.of("ABC".getBytes(StandardCharsets.UTF_8), "Y"), KV.of("ABC".getBytes(StandardCharsets.UTF_8), "Z")));
        }

        @Override
        public <T> boolean contains(PCollectionView<T> view) {
            return "testSideInputId".equals(view.getTagInternal().getId());
        }

        @Override
        public boolean isEmpty() {
            return false;
        }
    };
    RegisterAndProcessBundleOperation operation = new RegisterAndProcessBundleOperation(idGenerator, fakeClient, mockBeamFnStateDelegator, REGISTER_REQUEST, ImmutableMap.of(), ImmutableMap.of("testPTransformId", mockStepContext), ImmutableMap.of("testPTransformId", fakeSideInputReader), ImmutableTable.of("testPTransformId", "testSideInputId", DataflowPortabilityPCollectionView.with(new TupleTag<>("testSideInputId"), FullWindowedValueCoder.of(KvCoder.of(ByteArrayCoder.of(), StringUtf8Coder.of()), GlobalWindow.Coder.INSTANCE))), ImmutableMap.of(), mockContext);
    operation.start();
    verify(mockBeamFnStateDelegator).registerForProcessBundleInstructionId(eq("778"), stateHandlerCaptor.capture());
    waitForStateHandler.countDown();
    // This method blocks till the requests are completed
    operation.finish();
    // Ensure that the number of reigstrations matches the number of deregistrations
    assertEquals(stateServiceRegisterCounter.get(), stateServiceDeregisterCounter.get());
    assertEquals(0, stateServiceAbortCounter.get());
}

Example 3 with StateKey

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey 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);
}