/** * Constructs a {@link SingletonAssert} for the value of the provided {@code PCollection * PCollection<T>} with the specified reason. The provided PCollection must be a singleton. */ public static <T> SingletonAssert<T> thatSingleton(String reason, PCollection<T> actual) { return new PCollectionViewAssert<>( actual, View.asSingleton(), actual.getCoder(), PAssertionSite.capture(reason)); }
@Test public void testViewNonmergingAsSingletonDirect() { testViewNonmerging(pipeline, View.asSingleton()); }
@Test public void testViewUnboundedAsSingletonDirect() { testViewUnbounded(pipeline, View.asSingleton()); }
@Test public void testViewUnboundedAsSingletonStreaming() { testViewUnbounded(createTestStreamingRunner(), View.asSingleton()); }
@Test public void testViewNonmergingAsSingletonStreaming() { testViewNonmerging(createTestStreamingRunner(), View.asSingleton()); }
@Test public void testViewNonmergingAsSingletonBatch() { testViewNonmerging(createTestBatchRunner(), View.asSingleton()); }
@Test public void testViewUnboundedAsSingletonBatch() { testViewUnbounded(createTestBatchRunner(), View.asSingleton()); }
@Override public PCollectionView<Transaction> expand(PBegin input) { getSpannerConfig().validate(); return input .apply(Create.of(1)) .apply("Create transaction", ParDo.of(new CreateTransactionFn(this))) .apply("As PCollectionView", View.asSingleton()); }
@Override public PCollectionView<Integer> expand(PCollection<T> records) { return records .apply(Window.into(new GlobalWindows())) .apply("CountRecords", Count.globally()) .apply("GenerateShardCount", ParDo.of(new CalculateShardsFn())) .apply(View.asSingleton()); } }
@Override public PCollectionView<Configuration> expand(PCollection<? extends KV<KeyT, ValueT>> input) { Configuration conf = createWriteConf( SequenceFileOutputFormat.class, keyClass, valueClass, outputDirPath, REDUCERS_COUNT, String.valueOf(windowNum++)); return input .getPipeline() .apply(Create.<Configuration>of(conf)) .apply(View.<Configuration>asSingleton().withDefaultValue(conf)); } }
@Test @Category(NeedsRunner.class) public void testParDoReadingFromUnknownSideInput() { List<Integer> inputs = Arrays.asList(3, -42, 666); PCollectionView<Integer> sideView = pipeline.apply("Create3", Create.of(3)).apply(View.asSingleton()); pipeline .apply("CreateMain", Create.of(inputs)) .apply(ParDo.of(new TestDoFn(Arrays.asList(sideView), Arrays.asList()))); thrown.expect(RuntimeException.class); thrown.expectMessage("calling sideInput() with unknown view"); pipeline.run(); }
@Override public PCollectionView<Integer> expand(PCollection<String> input) { return input .apply( ParDo.of( new DoFn<String, Integer>() { @ProcessElement public void toInteger(ProcessContext ctxt) { ctxt.output(Integer.valueOf(ctxt.element())); } })) .apply(Top.largest(1)) .apply(Flatten.iterables()) .apply(View.asSingleton()); } }
private void testSideInput(IsBounded bounded) { PCollectionView<String> sideInput = p.apply("side input", Create.of("foo")).apply(View.asSingleton()); PCollection<String> res = p.apply("input", Create.of(0, 1, 2)) .apply(ParDo.of(sdfWithSideInput(bounded, sideInput)).withSideInputs(sideInput)); PAssert.that(res).containsInAnyOrder(Arrays.asList("foo:0", "foo:1", "foo:2")); p.run(); }
@Test public void fnWithSideInputDefault() throws Exception { PCollection<Integer> pCollection = p.apply(Create.empty(VarIntCoder.of())); final PCollectionView<Integer> value = pCollection.apply(View.<Integer>asSingleton().withDefaultValue(0)); try (DoFnTester<Integer, Integer> tester = DoFnTester.of(new SideInputDoFn(value))) { tester.processElement(1); tester.processElement(2); tester.processElement(4); tester.processElement(8); assertThat(tester.peekOutputElements(), containsInAnyOrder(0, 0, 0, 0)); } }
@Test public void fnWithSideInputExplicit() throws Exception { PCollection<Integer> pCollection = p.apply(Create.of(-2)); final PCollectionView<Integer> value = pCollection.apply(View.<Integer>asSingleton().withDefaultValue(0)); try (DoFnTester<Integer, Integer> tester = DoFnTester.of(new SideInputDoFn(value))) { tester.setSideInput(value, GlobalWindow.INSTANCE, -2); tester.processElement(16); tester.processElement(32); tester.processElement(64); tester.processElement(128); tester.finishBundle(); assertThat(tester.peekOutputElements(), containsInAnyOrder(-2, -2, -2, -2)); } }
@Test public void testViewGetName() { assertEquals("View.AsSingleton", View.<Integer>asSingleton().getName()); assertEquals("View.AsIterable", View.<Integer>asIterable().getName()); assertEquals("View.AsMap", View.<String, Integer>asMap().getName()); assertEquals("View.AsMultimap", View.<String, Integer>asMultimap().getName()); }
@Test public void testPassThroughThenCleanup() throws Exception { PCollection<Integer> output = p.apply(Create.of(1, 2, 3)) .apply( new PassThroughThenCleanup<>( new PassThroughThenCleanup.CleanupOperation() { @Override void cleanup(PassThroughThenCleanup.ContextContainer c) throws Exception { // no-op } }, p.apply("Create1", Create.of("")).apply(View.asSingleton()))); PAssert.that(output).containsInAnyOrder(1, 2, 3); p.run(); }
@Test public void testPassThroughThenCleanupExecuted() throws Exception { p.apply(Create.empty(VarIntCoder.of())) .apply( new PassThroughThenCleanup<>( new PassThroughThenCleanup.CleanupOperation() { @Override void cleanup(PassThroughThenCleanup.ContextContainer c) throws Exception { throw new RuntimeException("cleanup executed"); } }, p.apply("Create1", Create.of("")).apply(View.asSingleton()))); thrown.expect(RuntimeException.class); thrown.expectMessage("cleanup executed"); p.run(); }
/** Creates a simple pipeline with a {@link Combine.PerKey} with side inputs. */ private static TestPipeline createCombinePerKeyWithSideInputsPipeline() { TestPipeline pipeline = TestPipeline.create().enableAbandonedNodeEnforcement(false); PCollection<KV<String, Integer>> input = pipeline .apply(Create.of(KV.of("key", 1))) .setCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())); PCollection<String> sideInput = pipeline.apply(Create.of("side input")); PCollectionView<String> sideInputView = sideInput.apply(View.asSingleton()); input.apply( Combine.<String, Integer, Integer>perKey(new SumCombineFnWithContext()) .withSideInputs(sideInputView)); return pipeline; }
/** Basic test of {@link MapElements} with a {@link Fn} and a side input. */ @Test @Category(NeedsRunner.class) public void testMapBasicWithSideInput() throws Exception { final PCollectionView<Integer> view = pipeline.apply("Create base", Create.of(40)).apply(View.asSingleton()); PCollection<Integer> output = pipeline .apply(Create.of(0, 1, 2)) .apply( MapElements.into(integers()) .via( fn((element, c) -> element + c.sideInput(view), requiresSideInputs(view)))); PAssert.that(output).containsInAnyOrder(40, 41, 42); pipeline.run(); }