Search in sources :

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() {
    });
}
Also used : DoubleCoder(org.apache.beam.sdk.coders.DoubleCoder) IntervalWindowCoder(org.apache.beam.sdk.transforms.windowing.IntervalWindow.IntervalWindowCoder) ByteCoder(org.apache.beam.sdk.coders.ByteCoder) KvCoder(org.apache.beam.sdk.coders.KvCoder) VarLongCoder(org.apache.beam.sdk.coders.VarLongCoder) BooleanCoder(org.apache.beam.sdk.coders.BooleanCoder) TimestampPrefixingWindowCoder(org.apache.beam.sdk.coders.TimestampPrefixingWindowCoder) Coder(org.apache.beam.sdk.coders.Coder) RowCoder(org.apache.beam.sdk.coders.RowCoder) StringUtf8Coder(org.apache.beam.sdk.coders.StringUtf8Coder) IterableCoder(org.apache.beam.sdk.coders.IterableCoder) IterableLikeCoder(org.apache.beam.sdk.coders.IterableLikeCoder) ModelCoderRegistrar(org.apache.beam.runners.core.construction.ModelCoderRegistrar) BeamFnStateClient(org.apache.beam.fn.harness.state.BeamFnStateClient) ArrayList(java.util.ArrayList) StateResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.StateResponse) TranslationContext(org.apache.beam.runners.core.construction.CoderTranslation.TranslationContext)

Aggregations

BeamFnStateClient (org.apache.beam.fn.harness.state.BeamFnStateClient)6 BeamFnApi (org.apache.beam.model.fnexecution.v1.BeamFnApi)5 ProcessBundleDescriptor (org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor)5 ShortIdMap (org.apache.beam.runners.core.metrics.ShortIdMap)5 ByteString (org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString)5 ArrayList (java.util.ArrayList)4 PTransformRunnerFactory (org.apache.beam.fn.harness.PTransformRunnerFactory)4 BundleProcessorCache (org.apache.beam.fn.harness.control.ProcessBundleHandler.BundleProcessorCache)4 StateResponse (org.apache.beam.model.fnexecution.v1.BeamFnApi.StateResponse)4 Test (org.junit.Test)4 HashSet (java.util.HashSet)3 Map (java.util.Map)3 CompletableFuture (java.util.concurrent.CompletableFuture)3 ProgressRequestCallback (org.apache.beam.fn.harness.PTransformRunnerFactory.ProgressRequestCallback)3 CallbackRegistration (org.apache.beam.fn.harness.control.FinalizeBundleHandler.CallbackRegistration)3 PCollectionConsumerRegistry (org.apache.beam.fn.harness.data.PCollectionConsumerRegistry)3 PTransformFunctionRegistry (org.apache.beam.fn.harness.data.PTransformFunctionRegistry)3 BeamFnStateGrpcClientCache (org.apache.beam.fn.harness.state.BeamFnStateGrpcClientCache)3 StateRequest (org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest)3 IOException (java.io.IOException)2