AssignWindowsRunner { static <T, W extends BoundedWindow> AssignWindowsRunner<T, W> create( WindowFn<? super T, W> windowFn) { WindowFn<T, W> typedWindowFn = (WindowFn<T, W>) windowFn; return new AssignWindowsRunner<>(typedWindowFn); } private  AssignWindowsRunner(WindowFn<T, W> windowFn);   }
@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(); } }; Collection<WindowedValue<?>> outputs = new ArrayList<>(); MetricsContainerStepMap metricsContainerRegistry = new MetricsContainerStepMap(); PCollectionConsumerRegistry pCollectionConsumerRegistry = new PCollectionConsumerRegistry( metricsContainerRegistry, mock(ExecutionStateTracker.class)); pCollectionConsumerRegistry.register("output", "ptransform", outputs::add); SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); MapFnRunners.forWindowedValueMapFnFactory(new AssignWindowsMapFnFactory<>()) .createRunnerForPTransform( null , null , null , null , "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(), null , null , null , null , pCollectionConsumerRegistry, null , null , null , null , null , null , null ); 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); pCollectionConsumerRegistry.getMultiplexingConsumer("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))); }
@Test public void factoryCreatesFromJavaWindowFn() throws Exception { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); PTransform windowPTransform = PTransform.newBuilder() .putInputs("in", "input") .putOutputs("out", "output") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN) .setPayload( WindowIntoPayload.newBuilder() .setWindowFn( WindowingStrategyTranslation.toProto( new TestWindowFn(), components)) .build() .toByteString()) .build()) .build(); ThrowingFunction<WindowedValue<?>, WindowedValue<?>> fn = (ThrowingFunction) factory.forPTransform("transform", windowPTransform); assertThat( fn.apply( WindowedValue.of( 22L, new Instant(5), new IntervalWindow(new Instant(0L), new Instant(20027L)), PaneInfo.ON_TIME_AND_ONLY_FIRING)), equalTo( WindowedValue.of( 22L, new Instant(5), new TestWindowFn().assignWindow(new Instant(5)), PaneInfo.ON_TIME_AND_ONLY_FIRING))); }
@Test public void factoryCreatesFromKnownWindowFn() throws Exception { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); PTransform windowPTransform = PTransform.newBuilder() .putInputs("in", "input") .putOutputs("out", "output") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN) .setPayload( WindowIntoPayload.newBuilder() .setWindowFn( WindowingStrategyTranslation.toProto( Sessions.withGapDuration(Duration.standardMinutes(12L)), components)) .build() .toByteString()) .build()) .build(); ThrowingFunction<WindowedValue<?>, WindowedValue<?>> fn = (ThrowingFunction) factory.forPTransform("transform", windowPTransform); WindowedValue<?> output = fn.apply( WindowedValue.of( 22L, new Instant(5), new IntervalWindow(new Instant(0L), new Instant(20027L)), PaneInfo.ON_TIME_AND_ONLY_FIRING)); assertThat( output, equalTo( WindowedValue.of( 22L, new Instant(5), new IntervalWindow(new Instant(5L), Duration.standardMinutes(12L)), PaneInfo.ON_TIME_AND_ONLY_FIRING))); }
@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(); } }; Collection<WindowedValue<?>> outputs = new ArrayList<>(); MetricsContainerStepMap metricsContainerRegistry = new MetricsContainerStepMap(); PCollectionConsumerRegistry pCollectionConsumerRegistry = new PCollectionConsumerRegistry( metricsContainerRegistry, mock(ExecutionStateTracker.class)); pCollectionConsumerRegistry.register("output", "ptransform", outputs::add); SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); MapFnRunners.forWindowedValueMapFnFactory(new AssignWindowsMapFnFactory<>()) .createRunnerForPTransform( null , null , null , "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(), null , null , null , null , pCollectionConsumerRegistry, null , null, null, null ); 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); pCollectionConsumerRegistry.getMultiplexingConsumer("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))); }