Examples with TranslationContext org.apache.beam.runners.core.construction.CoderTranslation.TranslationContext used on opensource projects

Example 1 with TranslationContext

use of org.apache.beam.runners.core.construction.CoderTranslation.TranslationContext in project beam by apache.

the class CoderTranslators method row.

static CoderTranslator<RowCoder> row() {
    return new CoderTranslator<RowCoder>() {

        @Override
        public List<? extends Coder<?>> getComponents(RowCoder from) {
            return ImmutableList.of();
        }

        @Override
        public byte[] getPayload(RowCoder from) {
            return SchemaTranslation.schemaToProto(from.getSchema(), true).toByteArray();
        }

        @Override
        public RowCoder fromComponents(List<Coder<?>> components, byte[] payload, TranslationContext context) {
            checkArgument(components.isEmpty(), "Expected empty component list, but received: " + components);
            Schema schema;
            try {
                schema = SchemaTranslation.schemaFromProto(SchemaApi.Schema.parseFrom(payload));
            } catch (InvalidProtocolBufferException e) {
                throw new RuntimeException("Unable to parse schema for RowCoder: ", e);
            }
            return RowCoder.of(schema);
        }
    };
}

Example 2 with TranslationContext

use of org.apache.beam.runners.core.construction.CoderTranslation.TranslationContext 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() {
    });
}