Examples with PValue org.apache.beam.sdk.values.PValue used on opensource projects

Example 21 with PValue

use of org.apache.beam.sdk.values.PValue in project beam by apache.

the class TransformHierarchyTest method replaceWithCompositeSucceeds.

@Test
public void replaceWithCompositeSucceeds() {
    final SingleOutput<Long, Long> originalParDo = ParDo.of(new DoFn<Long, Long>() {

        @ProcessElement
        public void processElement(ProcessContext ctxt) {
            ctxt.output(ctxt.element() + 1L);
        }
    });
    GenerateSequence genUpstream = GenerateSequence.from(0);
    PCollection<Long> upstream = pipeline.apply(genUpstream);
    PCollection<Long> output = upstream.apply("Original", originalParDo);
    hierarchy.pushNode("Upstream", pipeline.begin(), genUpstream);
    hierarchy.finishSpecifyingInput();
    hierarchy.setOutput(upstream);
    hierarchy.popNode();
    TransformHierarchy.Node original = hierarchy.pushNode("Original", upstream, originalParDo);
    hierarchy.finishSpecifyingInput();
    hierarchy.setOutput(output);
    hierarchy.popNode();
    final TupleTag<Long> longs = new TupleTag<>();
    final MultiOutput<Long, Long> replacementParDo = ParDo.of(new DoFn<Long, Long>() {

        @ProcessElement
        public void processElement(ProcessContext ctxt) {
            ctxt.output(ctxt.element() + 1L);
        }
    }).withOutputTags(longs, TupleTagList.empty());
    PTransform<PCollection<Long>, PCollection<Long>> replacementComposite = new PTransform<PCollection<Long>, PCollection<Long>>() {

        @Override
        public PCollection<Long> expand(PCollection<Long> input) {
            return input.apply("Contained", replacementParDo).get(longs);
        }
    };
    PCollectionTuple replacementOutput = upstream.apply("Contained", replacementParDo);
    Node compositeNode = hierarchy.replaceNode(original, upstream, replacementComposite);
    Node replacementParNode = hierarchy.pushNode("Original/Contained", upstream, replacementParDo);
    hierarchy.finishSpecifyingInput();
    hierarchy.setOutput(replacementOutput);
    hierarchy.popNode();
    hierarchy.setOutput(replacementOutput.get(longs));
    Entry<TupleTag<?>, PValue> replacementLongs = Iterables.getOnlyElement(replacementOutput.expand().entrySet());
    hierarchy.replaceOutputs(Collections.<PValue, ReplacementOutput>singletonMap(replacementOutput.get(longs), ReplacementOutput.of(TaggedPValue.ofExpandedValue(output), TaggedPValue.of(replacementLongs.getKey(), replacementLongs.getValue()))));
    assertThat(replacementParNode.getOutputs().keySet(), Matchers.<TupleTag<?>>contains(replacementLongs.getKey()));
    assertThat(replacementParNode.getOutputs().values(), Matchers.<PValue>contains(output));
    assertThat(compositeNode.getOutputs().keySet(), equalTo(replacementOutput.get(longs).expand().keySet()));
    assertThat(compositeNode.getOutputs().values(), Matchers.<PValue>contains(output));
    hierarchy.popNode();
}

Example 22 with PValue

use of org.apache.beam.sdk.values.PValue in project beam by apache.

the class DirectGraphVisitor method visitPrimitiveTransform.

@Override
public void visitPrimitiveTransform(TransformHierarchy.Node node) {
    AppliedPTransform<?, ?, ?> appliedTransform = getAppliedTransform(node);
    stepNames.put(appliedTransform, genStepName());
    if (node.getInputs().isEmpty()) {
        rootTransforms.add(appliedTransform);
    } else {
        for (PValue value : node.getInputs().values()) {
            primitiveConsumers.put(value, appliedTransform);
        }
        if (node.getTransform() instanceof ViewOverrideFactory.WriteView) {
            viewWriters.put(((ViewOverrideFactory.WriteView<?, ?>) node.getTransform()).getView(), node.toAppliedPTransform(getPipeline()));
        }
    }
}

Example 23 with PValue

use of org.apache.beam.sdk.values.PValue in project beam by apache.

the class WatermarkManager method getInputProcessingWatermarks.

private Collection<Watermark> getInputProcessingWatermarks(AppliedPTransform<?, ?, ?> transform) {
    ImmutableList.Builder<Watermark> inputWmsBuilder = ImmutableList.builder();
    Map<TupleTag<?>, PValue> inputs = transform.getInputs();
    if (inputs.isEmpty()) {
        inputWmsBuilder.add(THE_END_OF_TIME);
    }
    for (PValue pvalue : inputs.values()) {
        Watermark producerOutputWatermark = getValueWatermark(pvalue).synchronizedProcessingOutputWatermark;
        inputWmsBuilder.add(producerOutputWatermark);
    }
    return inputWmsBuilder.build();
}

Example 24 with PValue

use of org.apache.beam.sdk.values.PValue in project beam by apache.

the class WatermarkManager method getInputWatermarks.

private List<Watermark> getInputWatermarks(AppliedPTransform<?, ?, ?> transform) {
    ImmutableList.Builder<Watermark> inputWatermarksBuilder = ImmutableList.builder();
    Map<TupleTag<?>, PValue> inputs = transform.getInputs();
    if (inputs.isEmpty()) {
        inputWatermarksBuilder.add(THE_END_OF_TIME);
    }
    for (PValue pvalue : inputs.values()) {
        Watermark producerOutputWatermark = getValueWatermark(pvalue).outputWatermark;
        inputWatermarksBuilder.add(producerOutputWatermark);
    }
    List<Watermark> inputCollectionWatermarks = inputWatermarksBuilder.build();
    return inputCollectionWatermarks;
}

Example 25 with PValue

use of org.apache.beam.sdk.values.PValue in project beam by apache.

the class DeduplicatedFlattenFactory method getInput.

/**
   * {@inheritDoc}.
   *
   * <p>The input {@link PCollectionList} that is constructed will have the same values in the same
   */
private PCollectionList<T> getInput(Map<TupleTag<?>, PValue> inputs, Pipeline p) {
    PCollectionList<T> pCollections = PCollectionList.empty(p);
    for (PValue input : inputs.values()) {
        PCollection<T> pcollection = (PCollection<T>) input;
        pCollections = pCollections.and(pcollection);
    }
    return pCollections;
}