WindowMergingFnRunner { static <T, W extends BoundedWindow> ThrowingFunction<KV<T, Iterable<W>>, KV<T, KV<Iterable<W>, Iterable<KV<W, Iterable<W>>>>>> createMapFunctionForPTransform(String ptransformId, PTransform ptransform) throws IOException { RunnerApi.FunctionSpec payload = RunnerApi.FunctionSpec.parseFrom(ptransform.getSpec().getPayload()); WindowFn<?, W> windowFn = (WindowFn<?, W>) WindowingStrategyTranslation.windowFnFromProto(payload); return WindowMergingFnRunner.<T, W>create(windowFn)::mergeWindows; }    }
@Test public void testWindowMergingWithNonMergingWindowFn() throws Exception { ThrowingFunction< KV<Object, Iterable<BoundedWindow>>, KV< Object, KV<Iterable<BoundedWindow>, Iterable<KV<BoundedWindow, Iterable<BoundedWindow>>>>>> mapFunction = WindowMergingFnRunner.createMapFunctionForPTransform( "ptransformId", createMergeTransformForWindowFn(new GlobalWindows())); KV<Object, Iterable<BoundedWindow>> input = KV.of( "abc", ImmutableList.of(new IntervalWindow(Instant.now(), Duration.standardMinutes(1)))); assertEquals( KV.of(input.getKey(), KV.of(input.getValue(), Collections.emptyList())), mapFunction.apply(input)); }
@Test public void testWindowMergingWithMergingWindowFn() throws Exception { ThrowingFunction< KV<Object, Iterable<BoundedWindow>>, KV< Object, KV<Iterable<BoundedWindow>, Iterable<KV<BoundedWindow, Iterable<BoundedWindow>>>>>> mapFunction = WindowMergingFnRunner.createMapFunctionForPTransform( "ptransformId", createMergeTransformForWindowFn(Sessions.withGapDuration(Duration.millis(5L)))); BoundedWindow[] expectedToBeMerged = new BoundedWindow[] { new IntervalWindow(new Instant(9L), new Instant(11L)), new IntervalWindow(new Instant(10L), new Instant(10L)), new IntervalWindow(new Instant(7L), new Instant(10L)) }; Iterable<BoundedWindow> expectedToBeUnmerged = Sets.newHashSet( new IntervalWindow(new Instant(1L), new Instant(1L)), new IntervalWindow(new Instant(20L), new Instant(20L))); KV<Object, Iterable<BoundedWindow>> input = KV.of( "abc", ImmutableList.<BoundedWindow>builder() .add(expectedToBeMerged) .addAll(expectedToBeUnmerged) .build()); KV<Object, KV<Iterable<BoundedWindow>, Iterable<KV<BoundedWindow, Iterable<BoundedWindow>>>>> output = mapFunction.apply(input); assertEquals(input.getKey(), output.getKey()); assertEquals(expectedToBeUnmerged, output.getValue().getKey()); KV<BoundedWindow, Iterable<BoundedWindow>> mergedOutput = Iterables.getOnlyElement(output.getValue().getValue()); assertEquals(new IntervalWindow(new Instant(7L), new Instant(11L)), mergedOutput.getKey()); assertThat(mergedOutput.getValue(), containsInAnyOrder(expectedToBeMerged)); }