Examples with AssignWindowsMapFnFactory org.apache.beam.fn.harness.AssignWindowsRunner.AssignWindowsMapFnFactory used on opensource projects

Example 1 with AssignWindowsMapFnFactory

use of org.apache.beam.fn.harness.AssignWindowsRunner.AssignWindowsMapFnFactory in project beam by apache.

the class AssignWindowsRunnerTest method multipleInputWindowsAsMapFnSucceeds.

@Test
public void multipleInputWindowsAsMapFnSucceeds() throws Exception {
    WindowFn<Object, BoundedWindow> windowFn = new WindowFn<Object, BoundedWindow>() {

        @Override
        public Collection<BoundedWindow> assignWindows(AssignContext c) {
            c.window();
            return ImmutableSet.of(GlobalWindow.INSTANCE, new IntervalWindow(new Instant(-500), Duration.standardMinutes(3)));
        }

        @Override
        public void mergeWindows(MergeContext c) {
            throw new UnsupportedOperationException();
        }

        @Override
        public WindowMappingFn<BoundedWindow> getDefaultWindowMappingFn() {
            throw new UnsupportedOperationException();
        }

        @Override
        public boolean isCompatible(WindowFn<?, ?> other) {
            throw new UnsupportedOperationException();
        }

        @Override
        public Coder<BoundedWindow> windowCoder() {
            throw new UnsupportedOperationException();
        }
    };
    SdkComponents components = SdkComponents.create();
    components.registerEnvironment(Environments.createDockerEnvironment("java"));
    RunnerApi.PCollection pCollection = RunnerApi.PCollection.newBuilder().setUniqueName("input").setCoderId("coder-id").build();
    RunnerApi.Coder coder = CoderTranslation.toProto(VarIntCoder.of()).getCoder();
    PTransformRunnerFactoryTestContext context = PTransformRunnerFactoryTestContext.builder("ptransform", PTransform.newBuilder().putInputs("in", "input").putOutputs("out", "output").setSpec(FunctionSpec.newBuilder().setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN).setPayload(WindowIntoPayload.newBuilder().setWindowFn(WindowingStrategyTranslation.toProto(windowFn, components)).build().toByteString())).build()).pCollections(Collections.singletonMap("input", pCollection)).coders(Collections.singletonMap("coder-id", coder)).build();
    Collection<WindowedValue<?>> outputs = new ArrayList<>();
    context.addPCollectionConsumer("output", outputs::add, VarIntCoder.of());
    MapFnRunners.forWindowedValueMapFnFactory(new AssignWindowsMapFnFactory<>()).createRunnerForPTransform(context);
    WindowedValue<Integer> value = WindowedValue.of(2, new Instant(-10L), ImmutableList.of(new IntervalWindow(new Instant(-22L), Duration.standardMinutes(5L)), new IntervalWindow(new Instant(-120000L), Duration.standardMinutes(3L))), PaneInfo.ON_TIME_AND_ONLY_FIRING);
    context.getPCollectionConsumer("input").accept(value);
    assertThat(outputs, containsInAnyOrder(WindowedValue.of(2, new Instant(-10L), ImmutableSet.of(GlobalWindow.INSTANCE, new IntervalWindow(new Instant(-500), Duration.standardMinutes(3))), PaneInfo.ON_TIME_AND_ONLY_FIRING), WindowedValue.of(2, new Instant(-10L), ImmutableSet.of(GlobalWindow.INSTANCE, new IntervalWindow(new Instant(-500), Duration.standardMinutes(3))), PaneInfo.ON_TIME_AND_ONLY_FIRING)));
}