Reshuffle extends PTransform<PCollection<KV<K, V>>, PCollection<KV<K, V>>> { public static <K, V> Reshuffle<K, V> of() { return new Reshuffle<>(); } private  Reshuffle(); static Reshuffle<K, V> of(); @Experimental static ViaRandomKey<T> viaRandomKey(); @Override PCollection<KV<K, V>> expand(PCollection<KV<K, V>> input);  }
@Test @Category(ValidatesRunner.class) public void testJustReshuffle() { PCollection<KV<String, Integer>> input = pipeline.apply( Create.of(ARBITRARY_KVS).withCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))); PCollection<KV<String, Integer>> output = input.apply(Reshuffle.of()); PAssert.that(output).containsInAnyOrder(ARBITRARY_KVS); assertEquals(input.getWindowingStrategy(), output.getWindowingStrategy()); pipeline.run(); }
@Test @Category(ValidatesRunner.class) public void testReshufflePreservesTimestamps() { PCollection<KV<String, TimestampedValue<String>>> input = pipeline .apply( Create.timestamped( TimestampedValue.of("foo", BoundedWindow.TIMESTAMP_MIN_VALUE), TimestampedValue.of("foo", new Instant(0)), TimestampedValue.of("bar", new Instant(33)), TimestampedValue.of("bar", GlobalWindow.INSTANCE.maxTimestamp())) .withCoder(StringUtf8Coder.of())) .apply( WithKeys.<String, String>of(input12 -> input12) .withKeyType(TypeDescriptors.strings())) .apply("ReifyOriginalTimestamps", Reify.timestampsInValue()); PCollection<TimestampedValue<TimestampedValue<String>>> output = input .apply(Reshuffle.of()) .apply("ReifyReshuffledTimestamps", Reify.timestampsInValue()) .apply(Values.create()); PAssert.that(output) .satisfies( input1 -> { for (TimestampedValue<TimestampedValue<String>> elem : input1) { Instant originalTimestamp = elem.getValue().getTimestamp(); Instant afterReshuffleTimestamp = elem.getTimestamp(); assertThat( "Reshuffle must preserve element timestamps", afterReshuffleTimestamp, equalTo(originalTimestamp)); } return null; }); pipeline.run(); }
@Test @Category(ValidatesRunner.class) public void testReshuffleAfterSessionsAndGroupByKey() { PCollection<KV<String, Iterable<Integer>>> input = pipeline .apply( Create.of(GBK_TESTABLE_KVS) .withCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))) .apply(Window.into(Sessions.withGapDuration(Duration.standardMinutes(10)))) .apply(GroupByKey.create()); PCollection<KV<String, Iterable<Integer>>> output = input.apply(Reshuffle.of()); PAssert.that(output).satisfies(new AssertThatHasExpectedContents()); assertEquals(input.getWindowingStrategy(), output.getWindowingStrategy()); pipeline.run(); }
@Test @Category(ValidatesRunner.class) public void testReshuffleAfterFixedWindowsAndGroupByKey() { PCollection<KV<String, Iterable<Integer>>> input = pipeline .apply( Create.of(GBK_TESTABLE_KVS) .withCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))) .apply(Window.into(FixedWindows.of(Duration.standardMinutes(10L)))) .apply(GroupByKey.create()); PCollection<KV<String, Iterable<Integer>>> output = input.apply(Reshuffle.of()); PAssert.that(output).satisfies(new AssertThatHasExpectedContents()); assertEquals(input.getWindowingStrategy(), output.getWindowingStrategy()); pipeline.run(); }
@Test @Category(ValidatesRunner.class) public void testReshuffleAfterSlidingWindowsAndGroupByKey() { PCollection<KV<String, Iterable<Integer>>> input = pipeline .apply( Create.of(GBK_TESTABLE_KVS) .withCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))) .apply(Window.into(FixedWindows.of(Duration.standardMinutes(10L)))) .apply(GroupByKey.create()); PCollection<KV<String, Iterable<Integer>>> output = input.apply(Reshuffle.of()); PAssert.that(output).satisfies(new AssertThatHasExpectedContents()); assertEquals(input.getWindowingStrategy(), output.getWindowingStrategy()); pipeline.run(); }
@Test @Category(ValidatesRunner.class) public void testReshuffleAfterFixedWindows() { PCollection<KV<String, Integer>> input = pipeline .apply( Create.of(ARBITRARY_KVS) .withCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))) .apply(Window.into(FixedWindows.of(Duration.standardMinutes(10L)))); PCollection<KV<String, Integer>> output = input.apply(Reshuffle.of()); PAssert.that(output).containsInAnyOrder(ARBITRARY_KVS); assertEquals(input.getWindowingStrategy(), output.getWindowingStrategy()); pipeline.run(); }
@Test @Category(ValidatesRunner.class) public void testReshuffleAfterSlidingWindows() { PCollection<KV<String, Integer>> input = pipeline .apply( Create.of(ARBITRARY_KVS) .withCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()))) .apply(Window.into(FixedWindows.of(Duration.standardMinutes(10L)))); PCollection<KV<String, Integer>> output = input.apply(Reshuffle.of()); PAssert.that(output).containsInAnyOrder(ARBITRARY_KVS); assertEquals(input.getWindowingStrategy(), output.getWindowingStrategy()); pipeline.run(); }
@Test @Category({ValidatesRunner.class, UsesTestStream.class}) public void testReshuffleWithTimestampsStreaming() { TestStream<Long> stream = TestStream.create(VarLongCoder.of()) .advanceWatermarkTo(new Instant(0L).plus(Duration.standardDays(48L))) .addElements( TimestampedValue.of(0L, new Instant(0L)), TimestampedValue.of(1L, new Instant(0L).plus(Duration.standardDays(48L))), TimestampedValue.of( 2L, BoundedWindow.TIMESTAMP_MAX_VALUE.minus(Duration.standardDays(48L)))) .advanceWatermarkToInfinity(); PCollection<KV<String, Long>> input = pipeline .apply(stream) .apply(WithKeys.of("")) .apply(Window.into(FixedWindows.of(Duration.standardMinutes(10L)))); PCollection<KV<String, Long>> reshuffled = input.apply(Reshuffle.of()); PAssert.that(reshuffled.apply(Values.create())).containsInAnyOrder(0L, 1L, 2L); pipeline.run(); }