Flatten { public static <T> PCollections<T> pCollections() { return new PCollections<>(); }  static PCollections<T> pCollections(); static Iterables<T> iterables();  }
@Test @Category(ValidatesRunner.class) public void testFlattenPCollections() { List<List<String>> inputs = Arrays.asList(LINES, NO_LINES, LINES2, NO_LINES, LINES, NO_LINES); PCollection<String> output = makePCollectionListOfStrings(p, inputs).apply(Flatten.pCollections()); PAssert.that(output).containsInAnyOrder(flattenLists(inputs)); p.run(); }
@Test @Category(ValidatesRunner.class) public void testFlattenPCollectionsSingletonList() { PCollection<String> input = p.apply(Create.of(LINES)); PCollection<String> output = PCollectionList.of(input).apply(Flatten.pCollections()); assertThat(output, not(equalTo(input))); PAssert.that(output).containsInAnyOrder(LINES); p.run(); }
@Test @Category(ValidatesRunner.class) public void testFlattenPCollectionsThenParDo() { List<List<String>> inputs = Arrays.asList(LINES, NO_LINES, LINES2, NO_LINES, LINES, NO_LINES); PCollection<String> output = makePCollectionListOfStrings(p, inputs) .apply(Flatten.pCollections()) .apply(ParDo.of(new IdentityFn<>())); PAssert.that(output).containsInAnyOrder(flattenLists(inputs)); p.run(); }
@Test @Category(ValidatesRunner.class) public void testFlattenPCollectionsEmpty() { PCollection<String> output = PCollectionList.<String>empty(p) .apply(Flatten.pCollections()) .setCoder(StringUtf8Coder.of()); PAssert.that(output).empty(); p.run(); }
@Test @Category(ValidatesRunner.class) public void testFlattenInputMultipleCopies() { int count = 5; PCollection<Long> longs = p.apply("mkLines", GenerateSequence.from(0).to(count)); PCollection<Long> biggerLongs = p.apply("mkOtherLines", GenerateSequence.from(0).to(count)) .apply( MapElements.via( new SimpleFunction<Long, Long>() { @Override public Long apply(Long input) { return input + 10L; } })); PCollection<Long> flattened = PCollectionList.of(longs).and(longs).and(biggerLongs).apply(Flatten.pCollections()); List<Long> expectedLongs = new ArrayList<>(); for (int i = 0; i < count; i++) { expectedLongs.add((long) i); expectedLongs.add((long) i); expectedLongs.add(i + 10L); } PAssert.that(flattened).containsInAnyOrder(expectedLongs); p.run(); }
@Test @Category({ValidatesRunner.class, FlattenWithHeterogeneousCoders.class}) public void testFlattenMultipleCoders() throws CannotProvideCoderException { PCollection<Long> bigEndianLongs = p.apply( "BigEndianLongs", Create.of(0L, 1L, 2L, 3L, null, 4L, 5L, null, 6L, 7L, 8L, null, 9L) .withCoder(NullableCoder.of(BigEndianLongCoder.of()))); PCollection<Long> varLongs = p.apply("VarLengthLongs", GenerateSequence.from(0).to(5)).setCoder(VarLongCoder.of()); PCollection<Long> flattened = PCollectionList.of(bigEndianLongs) .and(varLongs) .apply(Flatten.pCollections()) .setCoder(NullableCoder.of(VarLongCoder.of())); PAssert.that(flattened) .containsInAnyOrder( 0L, 0L, 1L, 1L, 2L, 3L, 2L, 4L, 5L, 3L, 6L, 7L, 4L, 8L, 9L, null, null, null); p.run(); }
@Test @Category({ValidatesRunner.class, UsesSideInputs.class}) public void testEmptyFlattenAsSideInput() { final PCollectionView<Iterable<String>> view = PCollectionList.<String>empty(p) .apply(Flatten.pCollections()) .setCoder(StringUtf8Coder.of()) .apply(View.asIterable()); PCollection<String> output = p.apply(Create.of((Void) null).withCoder(VoidCoder.of())) .apply( ParDo.of( new DoFn<Void, String>() { @ProcessElement public void processElement(ProcessContext c) { for (String side : c.sideInput(view)) { c.output(side); } } }) .withSideInputs(view)); PAssert.that(output).empty(); p.run(); }
@Test @Category(ValidatesRunner.class) public void testFlattenPCollectionsEmptyThenParDo() { PCollection<String> output = PCollectionList.<String>empty(p) .apply(Flatten.pCollections()) .setCoder(StringUtf8Coder.of()) .apply(ParDo.of(new IdentityFn<>())); PAssert.that(output).empty(); p.run(); }
@Test @Category(NeedsRunner.class) public void testFlattenNoListsNoCoder() { thrown.expect(IllegalStateException.class); thrown.expectMessage("Unable to return a default Coder"); PCollectionList.<ClassWithoutCoder>empty(p).apply(Flatten.pCollections()); p.run(); }
@Test @Category(ValidatesRunner.class) public void testFlattenMultiplePCollectionsHavingMultipleConsumers() { PCollection<String> input = p.apply(Create.of("AA", "BBB", "CC")); final TupleTag<String> outputEvenLengthTag = new TupleTag<String>() {}; final TupleTag<String> outputOddLengthTag = new TupleTag<String>() {}; PCollectionTuple tuple = input.apply( ParDo.of( new DoFn<String, String>() { @ProcessElement public void processElement(ProcessContext c) { if (c.element().length() % 2 == 0) { c.output(c.element()); } else { c.output(outputOddLengthTag, c.element()); } } }) .withOutputTags(outputEvenLengthTag, TupleTagList.of(outputOddLengthTag))); PCollection<String> outputEvenLength = tuple.get(outputEvenLengthTag); PCollection<String> outputOddLength = tuple.get(outputOddLengthTag); PCollection<String> outputMerged = PCollectionList.of(outputEvenLength).and(outputOddLength).apply(Flatten.pCollections()); PAssert.that(outputMerged).containsInAnyOrder("AA", "BBB", "CC"); PAssert.that(outputEvenLength).containsInAnyOrder("AA", "CC"); PAssert.that(outputOddLength).containsInAnyOrder("BBB"); p.run(); }
@Test @Category(ValidatesRunner.class) public void testFlattenWithDifferentInputAndOutputCoders() { PCollection<KV<String, String>> flattenInput = p.apply(Create.of(LINES)) .apply(WithKeys.of("a")) .setCoder(SerializableCoder.of(new TypeDescriptor<KV<String, String>>() {})); PCollection<String> output = PCollectionList.of(flattenInput) .apply(Flatten.pCollections()) .setCoder(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())) .apply(GroupByKey.create()) .apply(Values.create()) .apply( FlatMapElements.into(TypeDescriptors.strings()) .via((Iterable<String> values) -> values)); PAssert.that(output).containsInAnyOrder(LINES); p.run(); }
@Test @Category(ValidatesRunner.class) public void testFlattenWithDifferentInputAndOutputCoders2() { PCollection<KV<String, Iterable<String>>> flattenInput = p.apply(Create.of(LINES)) .apply(WithKeys.of("a")) .setCoder(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())) .apply(GroupByKey.create()); PCollection<String> output = PCollectionList.of(flattenInput) .apply(Flatten.pCollections()) .setCoder(SerializableCoder.of(new TypeDescriptor<KV<String, Iterable<String>>>() {})) .apply(Values.create()) .setCoder(IterableCoder.of(StringUtf8Coder.of())) .apply( FlatMapElements.into(TypeDescriptors.strings()) .via((Iterable<String> values) -> values)); PAssert.that(output).containsInAnyOrder(LINES); p.run(); }
@Test @Category(NeedsRunner.class) public void testEqualWindowFnPropagation() { PCollection<String> input1 = p.apply("CreateInput1", Create.of("Input1")) .apply("Window1", Window.into(FixedWindows.of(Duration.standardMinutes(1)))); PCollection<String> input2 = p.apply("CreateInput2", Create.of("Input2")) .apply("Window2", Window.into(FixedWindows.of(Duration.standardMinutes(1)))); PCollection<String> output = PCollectionList.of(input1).and(input2).apply(Flatten.pCollections()); p.run(); Assert.assertTrue( output .getWindowingStrategy() .getWindowFn() .isCompatible(FixedWindows.of(Duration.standardMinutes(1)))); }
@Test @Category(NeedsRunner.class) public void testCompatibleWindowFnPropagation() { PCollection<String> input1 = p.apply("CreateInput1", Create.of("Input1")) .apply("Window1", Window.into(Sessions.withGapDuration(Duration.standardMinutes(1)))); PCollection<String> input2 = p.apply("CreateInput2", Create.of("Input2")) .apply("Window2", Window.into(Sessions.withGapDuration(Duration.standardMinutes(2)))); PCollection<String> output = PCollectionList.of(input1).and(input2).apply(Flatten.pCollections()); p.run(); Assert.assertTrue( output .getWindowingStrategy() .getWindowFn() .isCompatible(Sessions.withGapDuration(Duration.standardMinutes(2)))); }
@Test public void testIncompatibleWindowFnPropagationFailure() { p.enableAbandonedNodeEnforcement(false); PCollection<String> input1 = p.apply("CreateInput1", Create.of("Input1")) .apply("Window1", Window.into(FixedWindows.of(Duration.standardMinutes(1)))); PCollection<String> input2 = p.apply("CreateInput2", Create.of("Input2")) .apply("Window2", Window.into(FixedWindows.of(Duration.standardMinutes(2)))); try { PCollectionList.of(input1).and(input2).apply(Flatten.pCollections()); Assert.fail("Exception should have been thrown"); } catch (IllegalStateException e) { Assert.assertTrue( e.getMessage().startsWith("Inputs to Flatten had incompatible window windowFns")); } }