Examples with BeamFnStateClient org.apache.beam.fn.harness.state.BeamFnStateClient used on opensource projects

Example 6 with BeamFnStateClient

use of org.apache.beam.fn.harness.state.BeamFnStateClient in project beam by apache.

the class CommonCoderTest method instantiateCoder.

private static Coder<?> instantiateCoder(CommonCoder coder) {
    List<Coder<?>> components = new ArrayList<>();
    for (CommonCoder innerCoder : coder.getComponents()) {
        components.add(instantiateCoder(innerCoder));
    }
    // translator since we need to interact with a fake state client.
    if (coder.getUrn().equals(getUrn(StandardCoders.Enum.STATE_BACKED_ITERABLE))) {
        BeamFnStateClient stateClient = new BeamFnStateClient() {

            @Override
            public CompletableFuture<StateResponse> handle(StateRequest.Builder requestBuilder) {
                checkState(requestBuilder.hasGet());
                checkState(requestBuilder.hasStateKey());
                checkState(requestBuilder.getStateKey().hasRunner());
                StateResponse.Builder rval = StateResponse.newBuilder();
                rval.setId(requestBuilder.getId());
                rval.setGet(StateGetResponse.newBuilder().setData(coder.getState().getOrDefault(requestBuilder.getStateKey().getRunner().getKey(), ByteString.EMPTY)));
                return CompletableFuture.completedFuture(rval.build());
            }
        };
        return new StateBackedIterable.Coder<>(() -> Caches.noop(), stateClient, () -> "instructionId", (Coder) Iterables.getOnlyElement(components));
    }
    Class<? extends Coder> coderType = ImmutableBiMap.copyOf(new ModelCoderRegistrar().getCoderURNs()).inverse().get(coder.getUrn());
    checkNotNull(coderType, "Unknown coder URN: " + coder.getUrn());
    CoderTranslator<?> translator = new ModelCoderRegistrar().getCoderTranslators().get(coderType);
    checkNotNull(translator, "No translator found for common coder class: " + coderType.getSimpleName());
    return translator.fromComponents(components, coder.getPayload(), new TranslationContext() {
    });
}