GroupIntoBatches extends PTransform<PCollection<KV<K, InputT>>, PCollection<KV<K, Iterable<InputT>>>> { public static <K, InputT> GroupIntoBatches<K, InputT> ofSize(long batchSize) { return new GroupIntoBatches<>(batchSize, null); } private  GroupIntoBatches(long batchSize, @Nullable Duration maxBufferingDuration); static GroupIntoBatches<K, InputT> ofSize(long batchSize); long getBatchSize(); GroupIntoBatches<K, InputT> withMaxBufferingDuration(Duration duration); @Override PCollection<KV<K, Iterable<InputT>>> expand(PCollection<KV<K, InputT>> input);  }
@Test @Category({NeedsRunner.class, UsesTimersInParDo.class, UsesStatefulParDo.class}) public void testInGlobalWindow() { PCollection<KV<String, Iterable<String>>> collection = pipeline .apply("Input data", Create.of(data)) .apply(GroupIntoBatches.ofSize(BATCH_SIZE)) .setCoder(KvCoder.of(StringUtf8Coder.of(), IterableCoder.of(StringUtf8Coder.of()))); PAssert.that("Incorrect batch size in one or more elements", collection) .satisfies( new SerializableFunction<Iterable<KV<String, Iterable<String>>>, Void>() { private boolean checkBatchSizes(Iterable<KV<String, Iterable<String>>> listToCheck) { for (KV<String, Iterable<String>> element : listToCheck) { if (Iterables.size(element.getValue()) != BATCH_SIZE) { return false; } } return true; } @Override public Void apply(Iterable<KV<String, Iterable<String>>> input) { assertTrue(checkBatchSizes(input)); return null; } }); PAssert.thatSingleton("Incorrect collection size", collection.apply("Count", Count.globally())) .isEqualTo(EVEN_NUM_ELEMENTS / BATCH_SIZE); pipeline.run(); }
@Test @Category({NeedsRunner.class, UsesTimersInParDo.class, UsesStatefulParDo.class}) public void testWithUnevenBatches() { PCollection<KV<String, Iterable<String>>> collection = pipeline .apply("Input data", Create.of(createTestData(ODD_NUM_ELEMENTS))) .apply(GroupIntoBatches.ofSize(BATCH_SIZE)) .setCoder(KvCoder.of(StringUtf8Coder.of(), IterableCoder.of(StringUtf8Coder.of()))); PAssert.that("Incorrect batch size in one or more elements", collection) .satisfies( new SerializableFunction<Iterable<KV<String, Iterable<String>>>, Void>() { private boolean checkBatchSizes(Iterable<KV<String, Iterable<String>>> listToCheck) { for (KV<String, Iterable<String>> element : listToCheck) { if (Iterables.size(element.getValue()) > BATCH_SIZE) { return false; } } return true; } @Override public Void apply(Iterable<KV<String, Iterable<String>>> input) { assertTrue(checkBatchSizes(input)); return null; } }); PAssert.thatSingleton("Incorrect collection size", collection.apply("Count", Count.globally())) .isEqualTo( (ODD_NUM_ELEMENTS + BATCH_SIZE - 1) / BATCH_SIZE); pipeline.run(); }
@Test @Category({ NeedsRunner.class, UsesTimersInParDo.class, UsesTestStream.class, UsesStatefulParDo.class }) public void testInStreamingMode() { int timestampInterval = 1; Instant startInstant = new Instant(0L); TestStream.Builder<KV<String, String>> streamBuilder = TestStream.create(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())) .advanceWatermarkTo(startInstant); long offset = 0L; for (KV<String, String> element : data) { streamBuilder = streamBuilder.addElements( TimestampedValue.of( element, startInstant.plus(Duration.standardSeconds(offset * timestampInterval)))); offset++; } final long windowDuration = 6; TestStream<KV<String, String>> stream = streamBuilder .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(windowDuration - 1))) .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(windowDuration + 1))) .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(EVEN_NUM_ELEMENTS))) .advanceWatermarkToInfinity(); PCollection<KV<String, String>> inputCollection = pipeline .apply(stream) .apply( Window.<KV<String, String>>into( FixedWindows.of(Duration.standardSeconds(windowDuration))) .withAllowedLateness(Duration.millis(ALLOWED_LATENESS))); inputCollection.apply( ParDo.of( new DoFn<KV<String, String>, Void>() { @ProcessElement public void processElement(ProcessContext c, BoundedWindow window) { LOG.debug( "*** ELEMENT: ({},{}) *** with timestamp {} in window {}", c.element().getKey(), c.element().getValue(), c.timestamp(), window); } })); PCollection<KV<String, Iterable<String>>> outputCollection = inputCollection .apply(GroupIntoBatches.ofSize(BATCH_SIZE)) .setCoder(KvCoder.of(StringUtf8Coder.of(), IterableCoder.of(StringUtf8Coder.of()))); PCollection<KV<String, Long>> countOutput = outputCollection.apply( "Count elements in windows after applying GroupIntoBatches", Count.perKey()); PAssert.that("Wrong number of elements in windows after GroupIntoBatches", countOutput) .satisfies( input -> { Iterator<KV<String, Long>> inputIterator = input.iterator(); long count0 = inputIterator.next().getValue(); assertEquals("Wrong number of elements in first window", 2, count0); long count1 = inputIterator.next().getValue(); assertEquals("Wrong number of elements in second window", 1, count1); return null; }); PAssert.that("Incorrect output collection after GroupIntoBatches", outputCollection) .satisfies( input -> { Iterator<KV<String, Iterable<String>>> inputIterator = input.iterator(); int size0 = Iterables.size(inputIterator.next().getValue()); assertEquals("Wrong first element batch Size", 5, size0); int size1 = Iterables.size(inputIterator.next().getValue()); assertEquals("Wrong second element batch Size", 1, size1); int size2 = Iterables.size(inputIterator.next().getValue()); assertEquals("Wrong third element batch Size", 4, size2); return null; }); pipeline.run().waitUntilFinish(); }