Search in sources :

Example 1 with WindowIntoPayload

use of org.apache.beam.sdk.common.runner.v1.RunnerApi.WindowIntoPayload 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)

Aggregations

AtomicReference (java.util.concurrent.atomic.AtomicReference)1 PipelineVisitor (org.apache.beam.sdk.Pipeline.PipelineVisitor)1 WindowIntoPayload (org.apache.beam.sdk.common.runner.v1.RunnerApi.WindowIntoPayload)1 AppliedPTransform (org.apache.beam.sdk.runners.AppliedPTransform)1 Node (org.apache.beam.sdk.runners.TransformHierarchy.Node)1 BoundedWindow (org.apache.beam.sdk.transforms.windowing.BoundedWindow)1 GlobalWindow (org.apache.beam.sdk.transforms.windowing.GlobalWindow)1 PartitioningWindowFn (org.apache.beam.sdk.transforms.windowing.PartitioningWindowFn)1 Window (org.apache.beam.sdk.transforms.windowing.Window)1 Assign (org.apache.beam.sdk.transforms.windowing.Window.Assign)1 WindowFn (org.apache.beam.sdk.transforms.windowing.WindowFn)1 Test (org.junit.Test)1