@Override public void populateDisplayData(DisplayData.Builder builder) { super.populateDisplayData(builder); builder .add(DisplayData.item("filePrefix", getOutputPrefix()).withLabel("Output File Prefix")) .addIfNotNull( DisplayData.item("fileSuffix", getFilenameSuffix()).withLabel("Output File Suffix")) .addIfNotNull( DisplayData.item("shardNameTemplate", getShardTemplate()) .withLabel("Output Shard Name Template")) .addIfNotDefault( DisplayData.item("numShards", getNumShards()).withLabel("Maximum Output Shards"), 0) .add( DisplayData.item("compressionType", getCompression().toString()) .withLabel("Compression Type")); } }
@Override public PDone expand(PCollection<byte[]> input) { checkState( getOutputPrefix() != null, "need to set the output prefix of a TFRecordIO.Write transform"); WriteFiles<byte[], Void, byte[]> write = WriteFiles.to( new TFRecordSink( getOutputPrefix(), getShardTemplate(), getFilenameSuffix(), getCompression())); if (getNumShards() > 0) { write = write.withNumShards(getNumShards()); } input.apply("Write", write); return PDone.in(input.getPipeline()); }