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

Example 1 with PInput

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

the class TransformHierarchyTest method replaceSucceeds.

@Test
public void replaceSucceeds() {
    PTransform<?, ?> enclosingPT = new PTransform<PInput, POutput>() {

        @Override
        public POutput expand(PInput input) {
            return PDone.in(input.getPipeline());
        }
    };
    TransformHierarchy.Node enclosing = hierarchy.pushNode("Enclosing", PBegin.in(pipeline), enclosingPT);
    Create.Values<Long> originalTransform = Create.of(1L);
    TransformHierarchy.Node original = hierarchy.pushNode("Create", PBegin.in(pipeline), originalTransform);
    assertThat(hierarchy.getCurrent(), equalTo(original));
    PCollection<Long> originalOutput = pipeline.apply(originalTransform);
    hierarchy.setOutput(originalOutput);
    hierarchy.popNode();
    assertThat(original.finishedSpecifying, is(true));
    hierarchy.setOutput(PDone.in(pipeline));
    hierarchy.popNode();
    assertThat(hierarchy.getCurrent(), not(equalTo(enclosing)));
    Read.Bounded<Long> replacementTransform = Read.from(CountingSource.upTo(1L));
    PCollection<Long> replacementOutput = pipeline.apply(replacementTransform);
    Node replacement = hierarchy.replaceNode(original, PBegin.in(pipeline), replacementTransform);
    assertThat(hierarchy.getCurrent(), equalTo(replacement));
    hierarchy.setOutput(replacementOutput);
    TaggedPValue taggedReplacement = TaggedPValue.ofExpandedValue(replacementOutput);
    Map<PValue, ReplacementOutput> replacementOutputs = Collections.<PValue, ReplacementOutput>singletonMap(replacementOutput, ReplacementOutput.of(TaggedPValue.ofExpandedValue(originalOutput), taggedReplacement));
    hierarchy.replaceOutputs(replacementOutputs);
    assertThat(replacement.getInputs(), equalTo(original.getInputs()));
    assertThat(replacement.getEnclosingNode(), equalTo(original.getEnclosingNode()));
    assertThat(replacement.getEnclosingNode(), equalTo(enclosing));
    assertThat(replacement.getTransform(), Matchers.<PTransform<?, ?>>equalTo(replacementTransform));
    // THe tags of the replacement transform are matched to the appropriate PValues of the original
    assertThat(replacement.getOutputs().keySet(), Matchers.<TupleTag<?>>contains(taggedReplacement.getTag()));
    assertThat(replacement.getOutputs().values(), Matchers.<PValue>contains(originalOutput));
    hierarchy.popNode();
}

Example 2 with PInput

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

the class TransformHierarchy method finishSpecifying.

/**
   * Finish specifying any remaining nodes within the {@link TransformHierarchy}. These are {@link
   * PValue PValues} that are produced as output of some {@link PTransform} but are never consumed
   * as input. These values must still be finished specifying.
   */
private void finishSpecifying() {
    for (Entry<PValue, PInput> producerInputEntry : producerInput.entrySet()) {
        PValue value = producerInputEntry.getKey();
        value.finishSpecifying(producerInputEntry.getValue(), getProducer(value).getTransform());
    }
    producerInput.clear();
}

Example 3 with PInput

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

the class TransformHierarchy method finishSpecifyingInput.

/**
   * Finish specifying all of the input {@link PValue PValues} of the current {@link
   * Node}. Ensures that all of the inputs to the current node have been fully
   * specified, and have been produced by a node in this graph.
   */
public void finishSpecifyingInput() {
    // Inputs must be completely specified before they are consumed by a transform.
    for (PValue inputValue : current.getInputs().values()) {
        Node producerNode = getProducer(inputValue);
        PInput input = producerInput.remove(inputValue);
        inputValue.finishSpecifying(input, producerNode.getTransform());
        checkState(producers.get(inputValue) != null, "Producer unknown for input %s", inputValue);
        checkState(producers.get(inputValue) != null, "Producer unknown for input %s", inputValue);
    }
}

Example 4 with PInput

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

the class TranslationContext method getViewInput.

public <InputT extends PInput> InputT getViewInput(PCollectionView<?> view) {
    PInput input = this.viewInputs.get(view);
    checkArgument(input != null, "unknown view " + view.getName());
    return (InputT) input;
}

Example 5 with PInput

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

the class DirectGraphVisitorTest method getStepNamesContainsAllTransforms.

@Test
public void getStepNamesContainsAllTransforms() {
    PCollection<String> created = p.apply(Create.of("1", "2", "3"));
    PCollection<String> transformed = created.apply(ParDo.of(new DoFn<String, String>() {

        @ProcessElement
        public void processElement(DoFn<String, String>.ProcessContext<String, String> c) throws Exception {
            c.output(Integer.toString(c.element().length()));
        }
    }));
    PDone finished = transformed.apply(new PTransform<PInput, PDone>() {

        @Override
        public PDone expand(PInput input) {
            return PDone.in(input.getPipeline());
        }
    });
    p.traverseTopologically(visitor);
    DirectGraph graph = visitor.getGraph();
    assertThat(graph.getStepName(graph.getProducer(created)), equalTo("s0"));
    assertThat(graph.getStepName(graph.getProducer(transformed)), equalTo("s1"));
// finished doesn't have a producer, because it's not a PValue.
// TODO: Demonstrate that PCollectionList/Tuple and other composite PValues are either safe to
// use, or make them so.
}