@Test
public void testCompressionTypeIsSet() throws Exception {
TextIO.Read.Bound<String> read = TextIO.Read.from("gs://bucket/test");
assertEquals(CompressionType.AUTO, read.getCompressionType());
read = TextIO.Read.from("gs://bucket/test").withCompressionType(CompressionType.GZIP);
assertEquals(CompressionType.GZIP, read.getCompressionType());
}
<T> void runTestRead(T[] expected, Coder<T> coder) throws Exception {
File tmpFile = tmpFolder.newFile("file.txt");
String filename = tmpFile.getPath();
try (PrintStream writer = new PrintStream(new FileOutputStream(tmpFile))) {
for (T elem : expected) {
byte[] encodedElem = CoderUtils.encodeToByteArray(coder, elem);
String line = new String(encodedElem);
writer.println(line);
}
}
Pipeline p = TestPipeline.create();
TextIO.Read.Bound<T> read;
if (coder.equals(StringUtf8Coder.of())) {
TextIO.Read.Bound<String> readStrings = TextIO.Read.from(filename);
// T==String
read = (TextIO.Read.Bound<T>) readStrings;
} else {
read = TextIO.Read.from(filename).withCoder(coder);
}
PCollection<T> output = p.apply(read);
DataflowAssert.that(output).containsInAnyOrder(expected);
p.run();
}
@Test
public void testCompressedRead() throws Exception {
String[] lines = {"Irritable eagle", "Optimistic jay", "Fanciful hawk"};
File tmpFile = tmpFolder.newFile("test");
String filename = tmpFile.getPath();
List<String> expected = new ArrayList<>();
try (PrintStream writer =
new PrintStream(new GZIPOutputStream(new FileOutputStream(tmpFile)))) {
for (String line : lines) {
writer.println(line);
expected.add(line);
}
}
Pipeline p = TestPipeline.create();
TextIO.Read.Bound<String> read =
TextIO.Read.from(filename).withCompressionType(CompressionType.GZIP);
PCollection<String> output = p.apply(read);
DataflowAssert.that(output).containsInAnyOrder(expected);
p.run();
tmpFile.delete();
}
@Test
public void testReadNamed() {
Pipeline p = TestPipeline.create();
{
PCollection<String> output1 = p.apply(TextIO.Read.from("/tmp/file.txt"));
assertEquals("TextIO.Read.out", output1.getName());
}
{
PCollection<String> output2 = p.apply(TextIO.Read.named("MyRead").from("/tmp/file.txt"));
assertEquals("MyRead.out", output2.getName());
}
{
PCollection<String> output3 = p.apply(TextIO.Read.from("/tmp/file.txt").named("HerRead"));
assertEquals("HerRead.out", output3.getName());
}
}
/** Recursive wildcards are not supported. This tests "**". */
@Test
public void testBadWildcardRecursive() throws Exception {
Pipeline pipeline = TestPipeline.create();
pipeline.apply(TextIO.Read.from("gs://bucket/foo**/baz"));
// Check that running does fail.
expectedException.expect(IllegalArgumentException.class);
expectedException.expectMessage("wildcard");
pipeline.run();
}
@Test
public void testReadWithoutValidationFlag() throws Exception {
TextIO.Read.Bound<String> read = TextIO.Read.from("gs://bucket/foo*/baz");
assertTrue(read.needsValidation());
assertFalse(read.withoutValidation().needsValidation());
}
private void applyRead(Pipeline pipeline, String path) {
pipeline.apply("Read(" + path + ")", TextIO.Read.from(path));
}
