@Test public void testSourceTransform() { PTransform<? super PBegin, ? extends POutput> myTransform = TextIO.read().from("foo.*"); DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(); Set<DisplayData> displayData = evaluator.displayDataForPrimitiveSourceTransforms(myTransform); assertThat(displayData, hasItem(hasDisplayItem("filePattern", "foo.*"))); } }
/** * 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); }
/** Create a new {@link DisplayDataEvaluator} using the specified {@link PipelineOptions}. */ public static DisplayDataEvaluator create(PipelineOptions pipelineOptions) { return new DisplayDataEvaluator(pipelineOptions); }
/** * Create a new {@link DisplayDataEvaluator} using options returned from {@link * #getDefaultOptions()}. */ public static DisplayDataEvaluator create() { return create(getDefaultOptions()); }
@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"))); }
/** * Traverse the specified source {@link PTransform}, collecting {@link DisplayData} registered on * the inner primitive {@link PTransform PTransforms}. * * @param root The source root {@link PTransform} to traverse * @return the set of {@link DisplayData} for primitive source {@link PTransform PTransforms}. */ public Set<DisplayData> displayDataForPrimitiveSourceTransforms( final PTransform<? super PBegin, ? extends POutput> root) { Pipeline pipeline = Pipeline.create(options); pipeline.apply("SourceTransform", root); return displayDataForPipeline(pipeline, root); }
@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()))); }
/** * Traverse the specified {@link PTransform}, collecting {@link DisplayData} registered on the * inner primitive {@link PTransform PTransforms}. * * @param root The root {@link PTransform} to traverse * @param inputCoder The coder to set for the {@link PTransform} input, or null to infer the * default coder. * @return the set of {@link DisplayData} for primitive {@link PTransform PTransforms}. */ public <InputT> Set<DisplayData> displayDataForPrimitiveTransforms( final PTransform<? super PCollection<InputT>, ? extends POutput> root, Coder<InputT> inputCoder) { Create.Values<InputT> input; if (inputCoder != null) { input = Create.empty(inputCoder); } else { // These types don't actually work, but the pipeline will never be run input = (Create.Values<InputT>) Create.empty(VoidCoder.of()); } Pipeline pipeline = Pipeline.create(options); pipeline.apply("Input", input).apply("Transform", root); return displayDataForPipeline(pipeline, root); }
@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 @Category(ValidatesRunner.class) public void testPrimitiveReadDisplayData() { DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(); AvroIO.Read<GenericRecord> read = AvroIO.readGenericRecords(Schema.create(Schema.Type.STRING)).from("/foo.*"); Set<DisplayData> displayData = evaluator.displayDataForPrimitiveSourceTransforms(read); assertThat( "AvroIO.Read should include the file pattern in its primitive transform", displayData, hasItem(hasDisplayItem("filePattern"))); }
@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 testPrimitiveReadDisplayData() { DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(); TextIO.Read read = TextIO.read().from("foobar"); Set<DisplayData> displayData = evaluator.displayDataForPrimitiveSourceTransforms(read); assertThat( "TextIO.Read should include the file prefix in its primitive display data", displayData, hasItem(hasDisplayItem(hasValue(startsWith("foobar"))))); }
@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 @Category({ValidatesRunner.class, UsesUnboundedPCollections.class}) public void testPrimitiveReadDisplayData() { DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(); Set<DisplayData> displayData; PubsubIO.Read<String> baseRead = PubsubIO.readStrings(); // Reading from a subscription. PubsubIO.Read<String> read = baseRead.fromSubscription("projects/project/subscriptions/subscription"); displayData = evaluator.displayDataForPrimitiveSourceTransforms(read); assertThat( "PubsubIO.Read should include the subscription in its primitive display data", displayData, hasItem(hasDisplayItem("subscription"))); // Reading from a topic. read = baseRead.fromTopic("projects/project/topics/topic"); displayData = evaluator.displayDataForPrimitiveSourceTransforms(read); assertThat( "PubsubIO.Read should include the topic in its primitive display data", displayData, hasItem(hasDisplayItem("topic"))); }
@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 public void testSourcePrimitiveDisplayData() { DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(); int numSplits = 98; PTransform<PBegin, PCollection<Entity>> read = DatastoreIO.v1() .read() .withProjectId(PROJECT_ID) .withQuery(Query.newBuilder().build()) .withNumQuerySplits(numSplits); String assertMessage = "DatastoreIO read should include the '%s' in its primitive display data"; Set<DisplayData> displayData = evaluator.displayDataForPrimitiveSourceTransforms(read); assertThat( String.format(assertMessage, "project id"), displayData, hasItem(hasDisplayItem("projectId", PROJECT_ID))); assertThat( String.format(assertMessage, "number of query splits"), displayData, hasItem(hasDisplayItem("numQuerySplits", numSplits))); }
@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 public void testTableSourcePrimitiveDisplayData() throws IOException, InterruptedException { DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(); BigQueryIO.Read read = BigQueryIO.read() .from("project:dataset.tableId") .withTestServices( new FakeBigQueryServices() .withDatasetService(new FakeDatasetService()) .withJobService(new FakeJobService())) .withoutValidation(); Set<DisplayData> displayData = evaluator.displayDataForPrimitiveSourceTransforms(read); assertThat( "BigQueryIO.Read should include the table spec in its primitive display data", displayData, hasItem(hasDisplayItem("table"))); }
@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 testQuerySourcePrimitiveDisplayData() throws IOException, InterruptedException { DisplayDataEvaluator evaluator = DisplayDataEvaluator.create(); BigQueryIO.Read read = BigQueryIO.read() .fromQuery("foobar") .withTestServices( new FakeBigQueryServices() .withDatasetService(new FakeDatasetService()) .withJobService(new FakeJobService())) .withoutValidation(); Set<DisplayData> displayData = evaluator.displayDataForPrimitiveSourceTransforms(read); assertThat( "BigQueryIO.Read should include the query in its primitive display data", displayData, hasItem(hasDisplayItem("query"))); }