Examples with AppliedPTransform org.apache.beam.sdk.runners.AppliedPTransform used on opensource projects

Example 16 with AppliedPTransform

use of org.apache.beam.sdk.runners.AppliedPTransform in project beam by apache.

the class WindowIntoTranslationTest method testToFromProto.

@Test
public void testToFromProto() throws InvalidProtocolBufferException {
    pipeline.apply(GenerateSequence.from(0)).apply(Window.<Long>into((WindowFn) windowFn));
    final AtomicReference<AppliedPTransform<?, ?, Assign<?>>> assign = new AtomicReference<>(null);
    pipeline.traverseTopologically(new PipelineVisitor.Defaults() {

        @Override
        public void visitPrimitiveTransform(Node node) {
            if (node.getTransform() instanceof Window.Assign) {
                checkState(assign.get() == null);
                assign.set((AppliedPTransform<?, ?, Assign<?>>) node.toAppliedPTransform(getPipeline()));
            }
        }
    });
    checkState(assign.get() != null);
    SdkComponents components = SdkComponents.create();
    WindowIntoPayload payload = WindowIntoTranslation.toProto(assign.get().getTransform(), components);
    assertEquals(windowFn, WindowIntoTranslation.getWindowFn(payload));
}

Example 17 with AppliedPTransform

use of org.apache.beam.sdk.runners.AppliedPTransform 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 18 with AppliedPTransform

use of org.apache.beam.sdk.runners.AppliedPTransform in project beam by apache.

the class PTransformMatchersTest method flattenWithDuplicateInputsWithoutDuplicates.

@Test
public void flattenWithDuplicateInputsWithoutDuplicates() {
    AppliedPTransform application = AppliedPTransform.<PCollectionList<Object>, PCollection<Object>, Flatten.PCollections<Object>>of("Flatten", Collections.<TupleTag<?>, PValue>singletonMap(new TupleTag<Object>(), PCollection.createPrimitiveOutputInternal(p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED)), Collections.<TupleTag<?>, PValue>singletonMap(new TupleTag<Object>(), PCollection.createPrimitiveOutputInternal(p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED)), Flatten.pCollections(), p);
    assertThat(PTransformMatchers.flattenWithDuplicateInputs().matches(application), is(false));
}

Example 19 with AppliedPTransform

use of org.apache.beam.sdk.runners.AppliedPTransform in project beam by apache.

the class PTransformMatchersTest method emptyFlattenWithEmptyFlatten.

@Test
public void emptyFlattenWithEmptyFlatten() {
    AppliedPTransform application = AppliedPTransform.<PCollectionList<Object>, PCollection<Object>, Flatten.PCollections<Object>>of("EmptyFlatten", Collections.<TupleTag<?>, PValue>emptyMap(), Collections.<TupleTag<?>, PValue>singletonMap(new TupleTag<Object>(), PCollection.createPrimitiveOutputInternal(p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED)), Flatten.pCollections(), p);
    assertThat(PTransformMatchers.emptyFlatten().matches(application), is(true));
}

Example 20 with AppliedPTransform

use of org.apache.beam.sdk.runners.AppliedPTransform in project beam by apache.

the class PTransformMatchersTest method flattenWithDuplicateInputsWithDuplicates.

@Test
public void flattenWithDuplicateInputsWithDuplicates() {
    PCollection<Object> duplicate = PCollection.createPrimitiveOutputInternal(p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED);
    AppliedPTransform application = AppliedPTransform.<PCollectionList<Object>, PCollection<Object>, Flatten.PCollections<Object>>of("Flatten", ImmutableMap.<TupleTag<?>, PValue>builder().put(new TupleTag<Object>(), duplicate).put(new TupleTag<Object>(), duplicate).build(), Collections.<TupleTag<?>, PValue>singletonMap(new TupleTag<Object>(), PCollection.createPrimitiveOutputInternal(p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED)), Flatten.pCollections(), p);
    assertThat(PTransformMatchers.flattenWithDuplicateInputs().matches(application), is(true));
}