Examples with CombinePayload org.apache.beam.sdk.common.runner.v1.RunnerApi.CombinePayload used on opensource projects

Example 1 with CombinePayload

use of org.apache.beam.sdk.common.runner.v1.RunnerApi.CombinePayload in project beam by apache.

the class PipelineValidator method validateCombine.

private static void validateCombine(String id, PTransform transform, Components components, Set<String> requirements) throws Exception {
    CombinePayload payload = CombinePayload.parseFrom(transform.getSpec().getPayload());
    checkArgument(components.containsCoders(payload.getAccumulatorCoderId()), "Transform %s uses unknown accumulator coder id %s", id, payload.getAccumulatorCoderId());
}

Example 2 with CombinePayload

use of org.apache.beam.sdk.common.runner.v1.RunnerApi.CombinePayload in project beam by apache.

the class CombineTranslationTest method testToFromProto.

@Test
public void testToFromProto() throws Exception {
    PCollection<Integer> input = pipeline.apply(Create.of(1, 2, 3));
    input.apply(Combine.globally(combineFn));
    final AtomicReference<AppliedPTransform<?, ?, Combine.PerKey<?, ?, ?>>> combine = new AtomicReference<>();
    pipeline.traverseTopologically(new PipelineVisitor.Defaults() {

        @Override
        public void leaveCompositeTransform(Node node) {
            if (node.getTransform() instanceof Combine.PerKey) {
                checkState(combine.get() == null);
                combine.set((AppliedPTransform) node.toAppliedPTransform(getPipeline()));
            }
        }
    });
    checkState(combine.get() != null);
    SdkComponents sdkComponents = SdkComponents.create();
    CombinePayload combineProto = CombineTranslation.toProto(combine.get(), sdkComponents);
    RunnerApi.Components componentsProto = sdkComponents.toComponents();
    assertEquals(combineFn.getAccumulatorCoder(pipeline.getCoderRegistry(), input.getCoder()), CombineTranslation.getAccumulatorCoder(combineProto, componentsProto));
    assertEquals(combineFn, CombineTranslation.getCombineFn(combineProto));
}

Example 3 with CombinePayload

use of org.apache.beam.sdk.common.runner.v1.RunnerApi.CombinePayload in project beam by apache.

the class CombineRunners method createCombineGroupedValuesMapFunction.

static <KeyT, InputT, AccumT, OutputT> ThrowingFunction<KV<KeyT, Iterable<InputT>>, KV<KeyT, OutputT>> createCombineGroupedValuesMapFunction(String pTransformId, PTransform pTransform) throws IOException {
    CombinePayload combinePayload = CombinePayload.parseFrom(pTransform.getSpec().getPayload());
    CombineFn<InputT, AccumT, OutputT> combineFn = (CombineFn) SerializableUtils.deserializeFromByteArray(combinePayload.getCombineFn().getPayload().toByteArray(), "CombineFn");
    return (KV<KeyT, Iterable<InputT>> input) -> {
        return KV.of(input.getKey(), combineFn.apply(input.getValue()));
    };
}

Example 4 with CombinePayload

use of org.apache.beam.sdk.common.runner.v1.RunnerApi.CombinePayload in project beam by apache.

the class CombineRunners method createMergeAccumulatorsMapFunction.

static <KeyT, AccumT> ThrowingFunction<KV<KeyT, Iterable<AccumT>>, KV<KeyT, AccumT>> createMergeAccumulatorsMapFunction(String pTransformId, PTransform pTransform) throws IOException {
    CombinePayload combinePayload = CombinePayload.parseFrom(pTransform.getSpec().getPayload());
    CombineFn<?, AccumT, ?> combineFn = (CombineFn) SerializableUtils.deserializeFromByteArray(combinePayload.getCombineFn().getPayload().toByteArray(), "CombineFn");
    return (KV<KeyT, Iterable<AccumT>> input) -> KV.of(input.getKey(), combineFn.mergeAccumulators(input.getValue()));
}

Example 5 with CombinePayload

use of org.apache.beam.sdk.common.runner.v1.RunnerApi.CombinePayload in project beam by apache.

the class CombineRunners method createExtractOutputsMapFunction.

static <KeyT, AccumT, OutputT> ThrowingFunction<KV<KeyT, AccumT>, KV<KeyT, OutputT>> createExtractOutputsMapFunction(String pTransformId, PTransform pTransform) throws IOException {
    CombinePayload combinePayload = CombinePayload.parseFrom(pTransform.getSpec().getPayload());
    CombineFn<?, AccumT, OutputT> combineFn = (CombineFn) SerializableUtils.deserializeFromByteArray(combinePayload.getCombineFn().getPayload().toByteArray(), "CombineFn");
    return (KV<KeyT, AccumT> input) -> KV.of(input.getKey(), combineFn.extractOutput(input.getValue()));
}