Examples with PTransform org.apache.beam.sdk.transforms.PTransform used on opensource projects

Example 31 with PTransform

use of org.apache.beam.sdk.transforms.PTransform in project beam by apache.

the class TransformHierarchyTest method visitIsTopologicallyOrdered.

@Test
public void visitIsTopologicallyOrdered() {
    PCollection<String> one = PCollection.createPrimitiveOutputInternal(pipeline, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, StringUtf8Coder.of());
    final PCollection<Integer> two = PCollection.createPrimitiveOutputInternal(pipeline, WindowingStrategy.globalDefault(), IsBounded.UNBOUNDED, VarIntCoder.of());
    final PDone done = PDone.in(pipeline);
    final TupleTag<String> oneTag = new TupleTag<String>() {
    };
    final TupleTag<Integer> twoTag = new TupleTag<Integer>() {
    };
    final PCollectionTuple oneAndTwo = PCollectionTuple.of(oneTag, one).and(twoTag, two);
    PTransform<PCollection<String>, PDone> multiConsumer = new PTransform<PCollection<String>, PDone>() {

        @Override
        public PDone expand(PCollection<String> input) {
            return done;
        }

        @Override
        public Map<TupleTag<?>, PValue> getAdditionalInputs() {
            return Collections.singletonMap(twoTag, two);
        }
    };
    hierarchy.pushNode("consumes_both", one, multiConsumer);
    hierarchy.setOutput(done);
    hierarchy.popNode();
    final PTransform<PBegin, PCollectionTuple> producer = new PTransform<PBegin, PCollectionTuple>() {

        @Override
        public PCollectionTuple expand(PBegin input) {
            return oneAndTwo;
        }
    };
    hierarchy.pushNode("encloses_producer", PBegin.in(pipeline), new PTransform<PBegin, PCollectionTuple>() {

        @Override
        public PCollectionTuple expand(PBegin input) {
            return input.apply(producer);
        }
    });
    hierarchy.pushNode("creates_one_and_two", PBegin.in(pipeline), producer);
    hierarchy.setOutput(oneAndTwo);
    hierarchy.popNode();
    hierarchy.setOutput(oneAndTwo);
    hierarchy.popNode();
    hierarchy.pushNode("second_copy_of_consumes_both", one, multiConsumer);
    hierarchy.setOutput(done);
    hierarchy.popNode();
    final Set<Node> visitedNodes = new HashSet<>();
    final Set<Node> exitedNodes = new HashSet<>();
    final Set<PValue> visitedValues = new HashSet<>();
    hierarchy.visit(new PipelineVisitor.Defaults() {

        @Override
        public CompositeBehavior enterCompositeTransform(Node node) {
            for (PValue input : node.getInputs().values()) {
                assertThat(visitedValues, hasItem(input));
            }
            assertThat("Nodes should not be visited more than once", visitedNodes, not(hasItem(node)));
            if (!node.isRootNode()) {
                assertThat("Nodes should always be visited after their enclosing nodes", visitedNodes, hasItem(node.getEnclosingNode()));
            }
            visitedNodes.add(node);
            return CompositeBehavior.ENTER_TRANSFORM;
        }

        @Override
        public void leaveCompositeTransform(Node node) {
            assertThat(visitedNodes, hasItem(node));
            if (!node.isRootNode()) {
                assertThat("Nodes should always be left before their enclosing nodes are left", exitedNodes, not(hasItem(node.getEnclosingNode())));
            }
            assertThat(exitedNodes, not(hasItem(node)));
            exitedNodes.add(node);
        }

        @Override
        public void visitPrimitiveTransform(Node node) {
            assertThat(visitedNodes, hasItem(node.getEnclosingNode()));
            assertThat(exitedNodes, not(hasItem(node.getEnclosingNode())));
            assertThat("Nodes should not be visited more than once", visitedNodes, not(hasItem(node)));
            for (PValue input : node.getInputs().values()) {
                assertThat(visitedValues, hasItem(input));
            }
            visitedNodes.add(node);
        }

        @Override
        public void visitValue(PValue value, Node producer) {
            assertThat(visitedNodes, hasItem(producer));
            assertThat(visitedValues, not(hasItem(value)));
            visitedValues.add(value);
        }
    });
    assertThat("Should have visited all the nodes", visitedNodes.size(), equalTo(5));
    assertThat("Should have left all of the visited composites", exitedNodes.size(), equalTo(2));
}

Example 32 with PTransform

use of org.apache.beam.sdk.transforms.PTransform 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<?>, PCollection<?>> replacementLongs = (Map.Entry) Iterables.getOnlyElement(replacementOutput.expand().entrySet());
    hierarchy.replaceOutputs(Collections.singletonMap(replacementOutput.get(longs), ReplacementOutput.of(TaggedPValue.ofExpandedValue(output), TaggedPValue.of(replacementLongs.getKey(), replacementLongs.getValue()))));
    assertThat(replacementParNode.getOutputs().keySet(), Matchers.contains(replacementLongs.getKey()));
    assertThat(replacementParNode.getOutputs().values(), Matchers.contains(output));
    assertThat(compositeNode.getOutputs().keySet(), equalTo(replacementOutput.get(longs).expand().keySet()));
    assertThat(compositeNode.getOutputs().values(), Matchers.contains(output));
    hierarchy.popNode();
}

Example 33 with PTransform

use of org.apache.beam.sdk.transforms.PTransform in project beam by apache.

the class TransformTreeTest method testCompositeCapture.

@Test
public void testCompositeCapture() throws Exception {
    p.enableAbandonedNodeEnforcement(false);
    File inputFile = tmpFolder.newFile();
    File outputFile = tmpFolder.newFile();
    final PTransform<PCollection<String>, PCollection<Iterable<String>>> sample = Sample.fixedSizeGlobally(10);
    p.apply("ReadMyFile", TextIO.read().from(inputFile.getPath())).apply(sample).apply(Flatten.iterables()).apply("WriteMyFile", TextIO.write().to(outputFile.getPath()));
    final EnumSet<TransformsSeen> visited = EnumSet.noneOf(TransformsSeen.class);
    final EnumSet<TransformsSeen> left = EnumSet.noneOf(TransformsSeen.class);
    p.traverseTopologically(new Pipeline.PipelineVisitor.Defaults() {

        @Override
        public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) {
            if (node.isRootNode()) {
                return CompositeBehavior.ENTER_TRANSFORM;
            }
            PTransform<?, ?> transform = node.getTransform();
            if (sample.getClass().equals(transform.getClass())) {
                assertTrue(visited.add(TransformsSeen.SAMPLE));
                assertNotNull(node.getEnclosingNode());
                assertTrue(node.isCompositeNode());
            } else if (transform instanceof WriteFiles) {
                assertTrue(visited.add(TransformsSeen.WRITE));
                assertNotNull(node.getEnclosingNode());
                assertTrue(node.isCompositeNode());
            }
            assertThat(transform, not(instanceOf(Impulse.class)));
            return CompositeBehavior.ENTER_TRANSFORM;
        }

        @Override
        public void leaveCompositeTransform(TransformHierarchy.Node node) {
            PTransform<?, ?> transform = node.getTransform();
            if (!node.isRootNode() && transform.getClass().equals(sample.getClass())) {
                assertTrue(left.add(TransformsSeen.SAMPLE));
            }
        }

        @Override
        public void visitPrimitiveTransform(TransformHierarchy.Node node) {
            PTransform<?, ?> transform = node.getTransform();
            // Composites should not be visited here.
            assertThat(transform, not(instanceOf(Combine.Globally.class)));
            assertThat(transform, not(instanceOf(WriteFiles.class)));
            assertThat(transform, not(instanceOf(TextIO.Read.class)));
            // There are multiple impulses in the graph so we don't validate that we haven't
            // seen one before.
            visited.add(TransformsSeen.IMPULSE);
        }
    });
    assertEquals(visited, EnumSet.allOf(TransformsSeen.class));
    assertEquals(left, EnumSet.of(TransformsSeen.SAMPLE));
}

Example 34 with PTransform

use of org.apache.beam.sdk.transforms.PTransform in project beam by apache.

the class DeadLetteredTransformTest method testDeadLettersOnlyFailures.

@Test
@Category(NeedsRunner.class)
public void testDeadLettersOnlyFailures() throws Exception {
    resetFailures();
    PCollection<Long> elements = p.apply(Create.of(10L, 20L).withCoder(VarLongCoder.of()));
    PCollection<Long> results = elements.apply(new DeadLetteredTransform<>(SimpleFunction.fromSerializableFunctionWithOutputType(x -> {
        if (x == 10L) {
            throw new RuntimeException(FAILURE_KEY);
        }
        return x;
    }, TypeDescriptor.of(Long.class)), new PTransform<PCollection<Failure>, PDone>() {

        @Override
        public PDone expand(PCollection<Failure> input) {
            input.apply(MapElements.into(TypeDescriptor.of(Void.class)).via(failure -> {
                capture(failure);
                return null;
            }));
            return PDone.in(input.getPipeline());
        }
    }));
    PAssert.that(results).containsInAnyOrder(20L);
    p.run().waitUntilFinish();
    List<Failure> failures = getFailures();
    assertEquals(1, failures.size());
    Failure failure = failures.iterator().next();
    assertEquals(10L, VarLongCoder.of().decode(new ByteArrayInputStream(failure.getPayload())).longValue());
    assertTrue(failure.getError().contains(FAILURE_KEY));
}

Example 35 with PTransform

use of org.apache.beam.sdk.transforms.PTransform in project beam by apache.

the class DisplayDataTest method testTypicalUsage.

@Test
public void testTypicalUsage() {
    final HasDisplayData subComponent1 = new HasDisplayData() {

        @Override
        public void populateDisplayData(DisplayData.Builder builder) {
            builder.add(DisplayData.item("ExpectedAnswer", 42));
        }
    };
    final HasDisplayData subComponent2 = new HasDisplayData() {

        @Override
        public void populateDisplayData(DisplayData.Builder builder) {
            builder.add(DisplayData.item("Location", "Seattle")).add(DisplayData.item("Forecast", "Rain"));
        }
    };
    PTransform<?, ?> transform = new PTransform<PCollection<String>, PCollection<String>>() {

        final Instant defaultStartTime = new Instant(0);

        Instant startTime = defaultStartTime;

        @Override
        public PCollection<String> expand(PCollection<String> begin) {
            throw new IllegalArgumentException("Should never be applied");
        }

        @Override
        public void populateDisplayData(DisplayData.Builder builder) {
            builder.include("p1", subComponent1).include("p2", subComponent2).add(DisplayData.item("minSproggles", 200).withLabel("Minimum Required Sproggles")).add(DisplayData.item("fireLasers", true)).addIfNotDefault(DisplayData.item("startTime", startTime), defaultStartTime).add(DisplayData.item("timeBomb", Instant.now().plus(Duration.standardDays(1)))).add(DisplayData.item("filterLogic", subComponent1.getClass())).add(DisplayData.item("serviceUrl", "google.com/fizzbang").withLinkUrl("http://www.google.com/fizzbang"));
        }
    };
    DisplayData data = DisplayData.from(transform);
    assertThat(data.items(), not(empty()));
    assertThat(data.items(), everyItem(allOf(hasKey(not(isEmptyOrNullString())), hasNamespace(Matchers.<Class<?>>isOneOf(transform.getClass(), subComponent1.getClass(), subComponent2.getClass())), hasType(notNullValue(DisplayData.Type.class)), hasValue(not(isEmptyOrNullString())))));
}