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

Example 36 with TupleTag

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

the class PipelineTest method testTupleInjectionTransform.

/**
 * Tests that Pipeline supports putting an element into a tuple as a transform.
 */
@Test
@Category(ValidatesRunner.class)
public void testTupleInjectionTransform() throws Exception {
    PCollection<Integer> input = pipeline.apply(Create.of(1, 2, 3, 4));
    TupleTag<Integer> tag = new TupleTag<>();
    PCollectionTuple output = input.apply("ProjectTag", new TupleInjectionTransform<>(tag));
    PAssert.that(output.get(tag)).containsInAnyOrder(1, 2, 3, 4);
    pipeline.run();
}

Example 37 with TupleTag

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

the class TransformHierarchyTest method visitAfterReplace.

/**
 * Tests that visiting the {@link TransformHierarchy} after replacing nodes does not visit any of
 * the original nodes or inaccessible values but does visit all of the replacement nodes, new
 * inaccessible replacement values, and the original output values.
 */
@Test
public void visitAfterReplace() {
    Node root = hierarchy.getCurrent();
    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);
    Node upstreamNode = hierarchy.pushNode("Upstream", pipeline.begin(), genUpstream);
    hierarchy.finishSpecifyingInput();
    hierarchy.setOutput(upstream);
    hierarchy.popNode();
    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));
    Map<TupleTag<?>, PCollection<?>> expandedReplacementOutput = (Map) replacementOutput.expand();
    Entry<TupleTag<?>, PCollection<?>> replacementLongs = Iterables.getOnlyElement(expandedReplacementOutput.entrySet());
    hierarchy.replaceOutputs(Collections.singletonMap(replacementOutput.get(longs), ReplacementOutput.of(TaggedPValue.ofExpandedValue(output), TaggedPValue.of(replacementLongs.getKey(), replacementLongs.getValue()))));
    hierarchy.popNode();
    final Set<Node> visitedCompositeNodes = new HashSet<>();
    final Set<Node> visitedPrimitiveNodes = new HashSet<>();
    Set<PValue> visitedValues = hierarchy.visit(new Defaults() {

        @Override
        public CompositeBehavior enterCompositeTransform(Node node) {
            visitedCompositeNodes.add(node);
            return CompositeBehavior.ENTER_TRANSFORM;
        }

        @Override
        public void visitPrimitiveTransform(Node node) {
            visitedPrimitiveNodes.add(node);
        }
    });
    /*
    Final Graph:
    Upstream -> Upstream.out -> Composite -> (ReplacementParDo -> OriginalParDo.out)
    */
    assertThat(visitedCompositeNodes, containsInAnyOrder(root, compositeNode));
    assertThat(visitedPrimitiveNodes, containsInAnyOrder(upstreamNode, replacementParNode));
    assertThat(visitedValues, containsInAnyOrder(upstream, output));
}

Example 38 with TupleTag

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

the class SqlTransform method toTableMap.

@SuppressWarnings("unchecked")
private Map<String, BeamSqlTable> toTableMap(PInput inputs) {
    /**
     * A single PCollection is transformed to a table named PCOLLECTION, other input types are
     * expanded and converted to tables using the tags as names.
     */
    if (inputs instanceof PCollection) {
        PCollection<?> pCollection = (PCollection<?>) inputs;
        return ImmutableMap.of(PCOLLECTION_NAME, new BeamPCollectionTable(pCollection));
    }
    ImmutableMap.Builder<String, BeamSqlTable> tables = ImmutableMap.builder();
    for (Map.Entry<TupleTag<?>, PValue> input : inputs.expand().entrySet()) {
        PCollection<?> pCollection = (PCollection<?>) input.getValue();
        tables.put(input.getKey().getId(), new BeamPCollectionTable(pCollection));
    }
    return tables.build();
}

Example 39 with TupleTag

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

the class BeamSqlDslFilterTest method runCompositeFilter.

private void runCompositeFilter(PCollection<Row> input) throws Exception {
    String sql = "SELECT * FROM TABLE_A" + " WHERE f_int > 1 AND (f_long < 3000 OR f_string = 'string_row3')";
    PCollection<Row> result = PCollectionTuple.of(new TupleTag<>("TABLE_A"), input).apply("testCompositeFilter", SqlTransform.query(sql));
    PAssert.that(result).containsInAnyOrder(rowsInTableA.get(1), rowsInTableA.get(2));
    pipeline.run().waitUntilFinish();
}

Example 40 with TupleTag

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

the class BeamSqlDslFilterTest method runNoReturnFilter.

private void runNoReturnFilter(PCollection<Row> input) throws Exception {
    String sql = "SELECT * FROM TABLE_A WHERE f_int < 1";
    PCollection<Row> result = PCollectionTuple.of(new TupleTag<>("TABLE_A"), input).apply("testNoReturnFilter", SqlTransform.query(sql));
    PAssert.that(result).empty();
    pipeline.run().waitUntilFinish();
}