Examples with MessageWithComponents org.apache.beam.model.pipeline.v1.RunnerApi.MessageWithComponents used on opensource projects

Example 1 with MessageWithComponents

use of org.apache.beam.model.pipeline.v1.RunnerApi.MessageWithComponents in project beam by apache.

the class ModelCodersTest method windowedValueCoderComponentsToConstructor.

@Test
public void windowedValueCoderComponentsToConstructor() throws IOException {
    FullWindowedValueCoder<Iterable<KV<String, Integer>>> javaCoder = FullWindowedValueCoder.of(IterableCoder.of(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())), IntervalWindowCoder.of());
    MessageWithComponents coderAndComponents = CoderTranslation.toProto(javaCoder);
    WindowedValueCoderComponents windowedValueCoderComponents = ModelCoders.getWindowedValueCoderComponents(coderAndComponents.getCoder());
    Coder windowedCoder = ModelCoders.windowedValueCoder(windowedValueCoderComponents.elementCoderId(), windowedValueCoderComponents.windowCoderId());
    assertThat(windowedCoder, equalTo(coderAndComponents.getCoder()));
}

Example 2 with MessageWithComponents

use of org.apache.beam.model.pipeline.v1.RunnerApi.MessageWithComponents in project beam by apache.

the class ModelCodersTest method kvCoderComponentsToConstructor.

@Test
public void kvCoderComponentsToConstructor() throws IOException {
    KvCoder<byte[], Iterable<Long>> javaCoder = KvCoder.of(ByteArrayCoder.of(), IterableCoder.of(LengthPrefixCoder.of(VarLongCoder.of())));
    MessageWithComponents coderAndComponents = CoderTranslation.toProto(javaCoder);
    KvCoderComponents kvCoderComponents = ModelCoders.getKvCoderComponents(coderAndComponents.getCoder());
    Coder kvCoder = ModelCoders.kvCoder(kvCoderComponents.keyCoderId(), kvCoderComponents.valueCoderId());
    assertThat(kvCoder, equalTo(coderAndComponents.getCoder()));
}

Example 3 with MessageWithComponents

use of org.apache.beam.model.pipeline.v1.RunnerApi.MessageWithComponents in project beam by apache.

the class ProtoOverrides method updateTransform.

/**
 * Update all composites present in the {@code originalPipeline} with an URN equal to the provided
 * {@code urn} using the provided {@link TransformReplacement}.
 */
public static Pipeline updateTransform(String urn, Pipeline originalPipeline, TransformReplacement compositeBuilder) {
    Components.Builder resultComponents = originalPipeline.getComponents().toBuilder();
    for (Map.Entry<String, PTransform> pt : originalPipeline.getComponents().getTransformsMap().entrySet()) {
        if (pt.getValue().getSpec() != null && urn.equals(pt.getValue().getSpec().getUrn())) {
            MessageWithComponents updated = compositeBuilder.getReplacement(pt.getKey(), originalPipeline.getComponents());
            if (updated == null) {
                continue;
            }
            checkArgument(updated.getPtransform().getOutputsMap().equals(pt.getValue().getOutputsMap()), "A %s must produce all of the outputs of the original %s", TransformReplacement.class.getSimpleName(), PTransform.class.getSimpleName());
            removeSubtransforms(pt.getValue(), resultComponents);
            resultComponents.mergeFrom(updated.getComponents()).putTransforms(pt.getKey(), updated.getPtransform());
        }
    }
    return originalPipeline.toBuilder().setComponents(resultComponents).build();
}