Search in sources :

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));
}
Also used : Window(org.apache.beam.sdk.transforms.windowing.Window) GlobalWindow(org.apache.beam.sdk.transforms.windowing.GlobalWindow) BoundedWindow(org.apache.beam.sdk.transforms.windowing.BoundedWindow) WindowIntoPayload(org.apache.beam.sdk.common.runner.v1.RunnerApi.WindowIntoPayload) AppliedPTransform(org.apache.beam.sdk.runners.AppliedPTransform) WindowFn(org.apache.beam.sdk.transforms.windowing.WindowFn) PartitioningWindowFn(org.apache.beam.sdk.transforms.windowing.PartitioningWindowFn) Node(org.apache.beam.sdk.runners.TransformHierarchy.Node) Assign(org.apache.beam.sdk.transforms.windowing.Window.Assign) PipelineVisitor(org.apache.beam.sdk.Pipeline.PipelineVisitor) AtomicReference(java.util.concurrent.atomic.AtomicReference) Test(org.junit.Test)

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));
}
Also used : Combine(org.apache.beam.sdk.transforms.Combine) Node(org.apache.beam.sdk.runners.TransformHierarchy.Node) AtomicReference(java.util.concurrent.atomic.AtomicReference) RunnerApi(org.apache.beam.sdk.common.runner.v1.RunnerApi) CombinePayload(org.apache.beam.sdk.common.runner.v1.RunnerApi.CombinePayload) AppliedPTransform(org.apache.beam.sdk.runners.AppliedPTransform) PipelineVisitor(org.apache.beam.sdk.Pipeline.PipelineVisitor) Test(org.junit.Test)

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));
}
Also used : PCollectionList(org.apache.beam.sdk.values.PCollectionList) PCollection(org.apache.beam.sdk.values.PCollection) AppliedPTransform(org.apache.beam.sdk.runners.AppliedPTransform) TupleTag(org.apache.beam.sdk.values.TupleTag) Test(org.junit.Test)

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));
}
Also used : PCollectionList(org.apache.beam.sdk.values.PCollectionList) PCollection(org.apache.beam.sdk.values.PCollection) AppliedPTransform(org.apache.beam.sdk.runners.AppliedPTransform) TupleTag(org.apache.beam.sdk.values.TupleTag) Test(org.junit.Test)

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));
}
Also used : PCollectionList(org.apache.beam.sdk.values.PCollectionList) PCollection(org.apache.beam.sdk.values.PCollection) AppliedPTransform(org.apache.beam.sdk.runners.AppliedPTransform) TupleTag(org.apache.beam.sdk.values.TupleTag) PValue(org.apache.beam.sdk.values.PValue) Test(org.junit.Test)

Aggregations

AppliedPTransform (org.apache.beam.sdk.runners.AppliedPTransform)21 Test (org.junit.Test)13 PCollection (org.apache.beam.sdk.values.PCollection)12 TupleTag (org.apache.beam.sdk.values.TupleTag)9 PTransformMatcher (org.apache.beam.sdk.runners.PTransformMatcher)6 DoFnSignature (org.apache.beam.sdk.transforms.reflect.DoFnSignature)4 PCollectionList (org.apache.beam.sdk.values.PCollectionList)4 PipelineVisitor (org.apache.beam.sdk.Pipeline.PipelineVisitor)3 Node (org.apache.beam.sdk.runners.TransformHierarchy.Node)3 PTransform (org.apache.beam.sdk.transforms.PTransform)3 KV (org.apache.beam.sdk.values.KV)3 PValue (org.apache.beam.sdk.values.PValue)3 Instant (org.joda.time.Instant)3 ArrayList (java.util.ArrayList)2 AtomicReference (java.util.concurrent.atomic.AtomicReference)2 KeyedWorkItem (org.apache.beam.runners.core.KeyedWorkItem)2 StatefulParDo (org.apache.beam.runners.direct.ParDoMultiOverrideFactory.StatefulParDo)2 ValueState (org.apache.beam.sdk.state.ValueState)2 BoundedWindow (org.apache.beam.sdk.transforms.windowing.BoundedWindow)2 IntervalWindow (org.apache.beam.sdk.transforms.windowing.IntervalWindow)2