/** * Traverse the specified {@link PTransform}, collecting {@link DisplayData} registered on the * inner primitive {@link PTransform PTransforms}. * * @return the set of {@link DisplayData} for primitive {@link PTransform PTransforms}. */ public <InputT> Set<DisplayData> displayDataForPrimitiveTransforms( final PTransform<? super PCollection<InputT>, ? extends POutput> root) { return displayDataForPrimitiveTransforms(root, null); }
@Test public void testPrimitiveTransform() { PTransform<? super PCollection<Integer>, ? super PCollection<Integer>> myTransform = ParDo.of( new DoFn<Integer, Integer>() { @ProcessElement public void processElement(ProcessContext c) throws Exception {} @Override public void populateDisplayData(DisplayData.Builder builder) { builder.add(DisplayData.item("foo", "bar")); } }); DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(); Set<DisplayData> displayData = evaluator.displayDataForPrimitiveTransforms(myTransform); assertThat(displayData, hasItem(hasDisplayItem("foo"))); }
@Test @Category(ValidatesRunner.class) public void testPrimitiveDisplayData() { SimpleFunction<Integer, ?> mapFn = new SimpleFunction<Integer, Integer>() { @Override public Integer apply(Integer input) { return input; } }; MapElements<Integer, ?> map = MapElements.via(mapFn); DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(); Set<DisplayData> displayData = evaluator.displayDataForPrimitiveTransforms(map); assertThat( "MapElements should include the mapFn in its primitive display data", displayData, hasItem(hasDisplayItem("class", mapFn.getClass()))); }
@Test public void testCompositeTransform() { PTransform<? super PCollection<String>, ? super POutput> myTransform = new PTransform<PCollection<String>, POutput>() { @Override public PCollection<String> expand(PCollection<String> input) { return input.apply( ParDo.of( new DoFn<String, String>() { @ProcessElement public void processElement(ProcessContext c) throws Exception { c.output(c.element()); } @Override public void populateDisplayData(DisplayData.Builder builder) { builder.add(DisplayData.item("primitiveKey", "primitiveValue")); } })); } @Override public void populateDisplayData(DisplayData.Builder builder) { builder.add(DisplayData.item("compositeKey", "compositeValue")); } }; DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(); Set<DisplayData> displayData = evaluator.displayDataForPrimitiveTransforms(myTransform); assertThat(displayData, not(hasItem(hasDisplayItem("compositeKey", "compositeValue")))); assertThat(displayData, hasItem(hasDisplayItem("primitiveKey", "primitiveValue"))); }
@Test @Category(ValidatesRunner.class) public void testPrimitiveWriteDisplayData() { DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(); PubsubIO.Write<?> write = PubsubIO.writeStrings().to("projects/project/topics/topic"); Set<DisplayData> displayData = evaluator.displayDataForPrimitiveTransforms(write); assertThat( "PubsubIO.Write should include the topic in its primitive display data", displayData, hasItem(hasDisplayItem("topic"))); } }
@Test public void testDeleteKeyPrimitiveDisplayData() { DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(); PTransform<PCollection<Key>, ?> write = DatastoreIO.v1().deleteKey().withProjectId("myProject"); Set<DisplayData> displayData = evaluator.displayDataForPrimitiveTransforms(write); assertThat( "DatastoreIO write should include the project in its primitive display data", displayData, hasItem(hasDisplayItem("projectId"))); assertThat( "DatastoreIO write should include the deleteKeyFn in its primitive display data", displayData, hasItem(hasDisplayItem("deleteKeyFn"))); }
@Test public void testWritePrimitiveDisplayData() { DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(); PTransform<PCollection<Entity>, ?> write = DatastoreIO.v1().write().withProjectId("myProject"); Set<DisplayData> displayData = evaluator.displayDataForPrimitiveTransforms(write); assertThat( "DatastoreIO write should include the project in its primitive display data", displayData, hasItem(hasDisplayItem("projectId"))); assertThat( "DatastoreIO write should include the upsertFn in its primitive display data", displayData, hasItem(hasDisplayItem("upsertFn"))); }
@Test @Category(ValidatesRunner.class) public void testCombinePerKeyPrimitiveDisplayData() { DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(); UniqueInts combineFn = new UniqueInts(); PTransform<PCollection<KV<Integer, Integer>>, ? extends POutput> combine = Combine.perKey(combineFn); Set<DisplayData> displayData = evaluator.displayDataForPrimitiveTransforms( combine, KvCoder.of(VarIntCoder.of(), VarIntCoder.of())); assertThat( "Combine.perKey should include the combineFn in its primitive transform", displayData, hasItem(hasDisplayItem("combineFn", combineFn.getClass()))); }
@Test public void testDeleteEntityPrimitiveDisplayData() { DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(); PTransform<PCollection<Entity>, ?> write = DatastoreIO.v1().deleteEntity().withProjectId("myProject"); Set<DisplayData> displayData = evaluator.displayDataForPrimitiveTransforms(write); assertThat( "DatastoreIO write should include the project in its primitive display data", displayData, hasItem(hasDisplayItem("projectId"))); assertThat( "DatastoreIO write should include the deleteEntityFn in its primitive display data", displayData, hasItem(hasDisplayItem("deleteEntityFn"))); }
@Test @Category(ValidatesRunner.class) public void testCombinePerKeyWithHotKeyFanoutPrimitiveDisplayData() { int hotKeyFanout = 2; DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(); UniqueInts combineFn = new UniqueInts(); PTransform<PCollection<KV<Integer, Integer>>, PCollection<KV<Integer, Set<Integer>>>> combine = Combine.<Integer, Integer, Set<Integer>>perKey(combineFn) .withHotKeyFanout(hotKeyFanout); Set<DisplayData> displayData = evaluator.displayDataForPrimitiveTransforms( combine, KvCoder.of(VarIntCoder.of(), VarIntCoder.of())); assertThat( "Combine.perKey.withHotKeyFanout should include the combineFn in its primitive " + "transform", displayData, hasItem(hasDisplayItem("combineFn", combineFn.getClass()))); assertThat( "Combine.perKey.withHotKeyFanout(int) should include the fanout in its primitive " + "transform", displayData, hasItem(hasDisplayItem("fanout", hotKeyFanout))); }
/** * A test that demonstrates that the replacement transform has the Display Data of the {@link * ParDo.SingleOutput} it replaces. */ @Test public void getReplacementTransformPopulateDisplayData() { ParDo.SingleOutput<Integer, Long> originalTransform = ParDo.of(new ToLongFn()); DisplayData originalDisplayData = DisplayData.from(originalTransform); PCollection<? extends Integer> input = pipeline.apply(Create.of(1, 2, 3)); AppliedPTransform< PCollection<? extends Integer>, PCollection<Long>, ParDo.SingleOutput<Integer, Long>> application = AppliedPTransform.of( "original", input.expand(), input.apply(originalTransform).expand(), originalTransform, pipeline); PTransformReplacement<PCollection<? extends Integer>, PCollection<Long>> replacement = factory.getReplacementTransform(application); DisplayData replacementDisplayData = DisplayData.from(replacement.getTransform()); assertThat(replacementDisplayData, equalTo(originalDisplayData)); DisplayData primitiveDisplayData = Iterables.getOnlyElement( DisplayDataEvaluator.create() .displayDataForPrimitiveTransforms(replacement.getTransform(), VarIntCoder.of())); assertThat(primitiveDisplayData, equalTo(replacementDisplayData)); }
@Test @Category(ValidatesRunner.class) public void testPrimitiveDisplayData() { FixedWindows windowFn = FixedWindows.of(Duration.standardHours(5)); AfterWatermark.FromEndOfWindow triggerBuilder = AfterWatermark.pastEndOfWindow(); Duration allowedLateness = Duration.standardMinutes(10); Window.ClosingBehavior closingBehavior = Window.ClosingBehavior.FIRE_IF_NON_EMPTY; TimestampCombiner timestampCombiner = TimestampCombiner.END_OF_WINDOW; Window<?> window = Window.into(windowFn) .triggering(triggerBuilder) .accumulatingFiredPanes() .withAllowedLateness(allowedLateness, closingBehavior) .withTimestampCombiner(timestampCombiner); DisplayData primitiveDisplayData = Iterables.getOnlyElement( DisplayDataEvaluator.create().displayDataForPrimitiveTransforms(window)); assertThat(primitiveDisplayData, hasDisplayItem("windowFn", windowFn.getClass())); assertThat(primitiveDisplayData, includesDisplayDataFor("windowFn", windowFn)); assertThat(primitiveDisplayData, hasDisplayItem("trigger", triggerBuilder.toString())); assertThat( primitiveDisplayData, hasDisplayItem("accumulationMode", AccumulationMode.ACCUMULATING_FIRED_PANES.toString())); assertThat(primitiveDisplayData, hasDisplayItem("allowedLateness", allowedLateness)); assertThat(primitiveDisplayData, hasDisplayItem("closingBehavior", closingBehavior.toString())); assertThat( primitiveDisplayData, hasDisplayItem("timestampCombiner", timestampCombiner.toString())); }