Java tutorial
/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.beam.sdk.io; import static org.hamcrest.Matchers.is; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertThat; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import com.google.common.collect.Lists; import java.io.BufferedReader; import java.io.File; import java.io.FileInputStream; import java.io.FileNotFoundException; import java.io.FileOutputStream; import java.io.FileReader; import java.io.IOException; import java.io.InputStreamReader; import java.io.PrintWriter; import java.nio.ByteBuffer; import java.nio.channels.Channels; import java.nio.channels.WritableByteChannel; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.zip.GZIPInputStream; import org.apache.beam.sdk.io.FileBasedSink.CompressionType; import org.apache.beam.sdk.io.FileBasedSink.FileResult; import org.apache.beam.sdk.io.FileBasedSink.FilenamePolicy; import org.apache.beam.sdk.io.FileBasedSink.FilenamePolicy.Context; import org.apache.beam.sdk.io.FileBasedSink.WritableByteChannelFactory; import org.apache.beam.sdk.io.FileBasedSink.WriteOperation; import org.apache.beam.sdk.io.FileBasedSink.Writer; import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions; import org.apache.beam.sdk.io.fs.ResourceId; import org.apache.commons.compress.compressors.bzip2.BZip2CompressorInputStream; import org.apache.commons.compress.compressors.deflate.DeflateCompressorInputStream; import org.junit.Rule; import org.junit.Test; import org.junit.rules.TemporaryFolder; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; /** * Tests for {@link FileBasedSink}. */ @RunWith(JUnit4.class) public class FileBasedSinkTest { @Rule public TemporaryFolder tmpFolder = new TemporaryFolder(); private final String tempDirectoryName = "temp"; private ResourceId getTemporaryFolder() { return LocalResources.fromFile(tmpFolder.getRoot(), /* isDirectory */ true); } private ResourceId getBaseOutputDirectory() { String baseOutputDirname = "output"; return getTemporaryFolder().resolve(baseOutputDirname, StandardResolveOptions.RESOLVE_DIRECTORY); } private ResourceId getBaseTempDirectory() { return getTemporaryFolder().resolve(tempDirectoryName, StandardResolveOptions.RESOLVE_DIRECTORY); } /** * Writer opens the correct file, writes the header, footer, and elements in the correct * order, and returns the correct filename. */ @Test public void testWriter() throws Exception { String testUid = "testId"; ResourceId expectedTempFile = getBaseTempDirectory().resolve(testUid, StandardResolveOptions.RESOLVE_FILE); List<String> values = Arrays.asList("sympathetic vulture", "boresome hummingbird"); List<String> expected = new ArrayList<>(); expected.add(SimpleSink.SimpleWriter.HEADER); expected.addAll(values); expected.add(SimpleSink.SimpleWriter.FOOTER); SimpleSink.SimpleWriter writer = buildWriteOperationWithTempDir(getBaseTempDirectory()).createWriter(); writer.openUnwindowed(testUid, -1); for (String value : values) { writer.write(value); } FileResult result = writer.close(); FileBasedSink sink = writer.getWriteOperation().getSink(); assertEquals(expectedTempFile, result.getTempFilename()); assertFileContains(expected, expectedTempFile); } /** * Assert that a file contains the lines provided, in the same order as expected. */ private void assertFileContains(List<String> expected, ResourceId file) throws Exception { try (BufferedReader reader = new BufferedReader(new FileReader(file.toString()))) { List<String> actual = new ArrayList<>(); for (;;) { String line = reader.readLine(); if (line == null) { break; } actual.add(line); } assertEquals("contents for " + file, expected, actual); } } /** Write lines to a file. */ private void writeFile(List<String> lines, File file) throws Exception { try (PrintWriter writer = new PrintWriter(new FileOutputStream(file))) { for (String line : lines) { writer.println(line); } } } /** * Removes temporary files when temporary and output directories differ. */ @Test public void testRemoveWithTempFilename() throws Exception { testRemoveTemporaryFiles(3, getBaseTempDirectory()); } /** Finalize copies temporary files to output files and removes any temporary files. */ @Test public void testFinalize() throws Exception { List<File> files = generateTemporaryFilesForFinalize(3); runFinalize(buildWriteOperation(), files); } /** Finalize can be called repeatedly. */ @Test public void testFinalizeMultipleCalls() throws Exception { List<File> files = generateTemporaryFilesForFinalize(3); SimpleSink.SimpleWriteOperation writeOp = buildWriteOperation(); runFinalize(writeOp, files); runFinalize(writeOp, files); } /** Finalize can be called when some temporary files do not exist and output files exist. */ @Test public void testFinalizeWithIntermediateState() throws Exception { SimpleSink.SimpleWriteOperation writeOp = buildWriteOperation(); List<File> files = generateTemporaryFilesForFinalize(3); runFinalize(writeOp, files); // create a temporary file and then rerun finalize tmpFolder.newFolder(tempDirectoryName); tmpFolder.newFile(tempDirectoryName + "/1"); runFinalize(writeOp, files); } /** Generate n temporary files using the temporary file pattern of Writer. */ private List<File> generateTemporaryFilesForFinalize(int numFiles) throws Exception { List<File> temporaryFiles = new ArrayList<>(); for (int i = 0; i < numFiles; i++) { ResourceId temporaryFile = WriteOperation.buildTemporaryFilename(getBaseTempDirectory(), "" + i); File tmpFile = new File(tmpFolder.getRoot(), temporaryFile.toString()); tmpFile.getParentFile().mkdirs(); assertTrue(tmpFile.createNewFile()); temporaryFiles.add(tmpFile); } return temporaryFiles; } /** Finalize and verify that files are copied and temporary files are optionally removed. */ private void runFinalize(SimpleSink.SimpleWriteOperation writeOp, List<File> temporaryFiles) throws Exception { int numFiles = temporaryFiles.size(); List<FileResult> fileResults = new ArrayList<>(); // Create temporary output bundles and output File objects. for (int i = 0; i < numFiles; i++) { fileResults.add(new FileResult(LocalResources.fromFile(temporaryFiles.get(i), false), WriteFiles.UNKNOWN_SHARDNUM, null, null)); } writeOp.finalize(fileResults); ResourceId outputDirectory = writeOp.getSink().getBaseOutputDirectoryProvider().get(); for (int i = 0; i < numFiles; i++) { ResourceId outputFilename = writeOp.getSink().getFilenamePolicy().unwindowedFilename(outputDirectory, new Context(i, numFiles), ""); assertTrue(new File(outputFilename.toString()).exists()); assertFalse(temporaryFiles.get(i).exists()); } assertFalse(new File(writeOp.tempDirectory.get().toString()).exists()); // Test that repeated requests of the temp directory return a stable result. assertEquals(writeOp.tempDirectory.get(), writeOp.tempDirectory.get()); } /** * Create n temporary and output files and verify that removeTemporaryFiles only removes temporary * files. */ private void testRemoveTemporaryFiles(int numFiles, ResourceId tempDirectory) throws Exception { String prefix = "file"; SimpleSink sink = new SimpleSink(getBaseOutputDirectory(), prefix, "", ""); WriteOperation<String> writeOp = new SimpleSink.SimpleWriteOperation(sink, tempDirectory); List<File> temporaryFiles = new ArrayList<>(); List<File> outputFiles = new ArrayList<>(); for (int i = 0; i < numFiles; i++) { ResourceId tempResource = WriteOperation.buildTemporaryFilename(tempDirectory, prefix + i); File tmpFile = new File(tempResource.toString()); tmpFile.getParentFile().mkdirs(); assertTrue("not able to create new temp file", tmpFile.createNewFile()); temporaryFiles.add(tmpFile); ResourceId outputFileId = getBaseOutputDirectory().resolve(prefix + i, StandardResolveOptions.RESOLVE_FILE); File outputFile = new File(outputFileId.toString()); outputFile.getParentFile().mkdirs(); assertTrue("not able to create new output file", outputFile.createNewFile()); outputFiles.add(outputFile); } writeOp.removeTemporaryFiles(Collections.<ResourceId>emptySet(), true); for (int i = 0; i < numFiles; i++) { File temporaryFile = temporaryFiles.get(i); assertThat(String.format("temp file %s exists", temporaryFile), temporaryFile.exists(), is(false)); File outputFile = outputFiles.get(i); assertThat(String.format("output file %s exists", outputFile), outputFile.exists(), is(true)); } } /** Output files are copied to the destination location with the correct names and contents. */ @Test public void testCopyToOutputFiles() throws Exception { SimpleSink.SimpleWriteOperation writeOp = buildWriteOperation(); ResourceId outputDirectory = writeOp.getSink().getBaseOutputDirectoryProvider().get(); List<String> inputFilenames = Arrays.asList("input-1", "input-2", "input-3"); List<String> inputContents = Arrays.asList("1", "2", "3"); List<String> expectedOutputFilenames = Arrays.asList("file-00-of-03.test", "file-01-of-03.test", "file-02-of-03.test"); Map<ResourceId, ResourceId> inputFilePaths = new HashMap<>(); List<ResourceId> expectedOutputPaths = new ArrayList<>(); for (int i = 0; i < inputFilenames.size(); i++) { // Generate output paths. expectedOutputPaths.add(getBaseOutputDirectory().resolve(expectedOutputFilenames.get(i), StandardResolveOptions.RESOLVE_FILE)); // Generate and write to input paths. File inputTmpFile = tmpFolder.newFile(inputFilenames.get(i)); List<String> lines = Collections.singletonList(inputContents.get(i)); writeFile(lines, inputTmpFile); inputFilePaths.put(LocalResources.fromFile(inputTmpFile, false), writeOp.getSink().getFilenamePolicy() .unwindowedFilename(outputDirectory, new Context(i, inputFilenames.size()), "")); } // Copy input files to output files. writeOp.copyToOutputFiles(inputFilePaths); // Assert that the contents were copied. for (int i = 0; i < expectedOutputPaths.size(); i++) { assertFileContains(Collections.singletonList(inputContents.get(i)), expectedOutputPaths.get(i)); } } public List<ResourceId> generateDestinationFilenames(ResourceId outputDirectory, FilenamePolicy policy, int numFiles) { List<ResourceId> filenames = new ArrayList<>(); for (int i = 0; i < numFiles; i++) { filenames.add(policy.unwindowedFilename(outputDirectory, new Context(i, numFiles), "")); } return filenames; } /** * Output filenames are generated correctly when an extension is supplied. */ @Test public void testGenerateOutputFilenames() { List<ResourceId> expected; List<ResourceId> actual; ResourceId root = getBaseOutputDirectory(); SimpleSink sink = new SimpleSink(root, "file", ".SSSSS.of.NNNNN", ".test"); FilenamePolicy policy = sink.getFilenamePolicy(); expected = Arrays.asList(root.resolve("file.00000.of.00003.test", StandardResolveOptions.RESOLVE_FILE), root.resolve("file.00001.of.00003.test", StandardResolveOptions.RESOLVE_FILE), root.resolve("file.00002.of.00003.test", StandardResolveOptions.RESOLVE_FILE)); actual = generateDestinationFilenames(root, policy, 3); assertEquals(expected, actual); expected = Collections .singletonList(root.resolve("file.00000.of.00001.test", StandardResolveOptions.RESOLVE_FILE)); actual = generateDestinationFilenames(root, policy, 1); assertEquals(expected, actual); expected = new ArrayList<>(); actual = generateDestinationFilenames(root, policy, 0); assertEquals(expected, actual); } /** Reject non-distinct output filenames. */ @Test public void testCollidingOutputFilenames() throws IOException { ResourceId root = getBaseOutputDirectory(); SimpleSink sink = new SimpleSink(root, "file", "-NN", "test"); SimpleSink.SimpleWriteOperation writeOp = new SimpleSink.SimpleWriteOperation(sink); ResourceId temp1 = root.resolve("temp1", StandardResolveOptions.RESOLVE_FILE); ResourceId temp2 = root.resolve("temp2", StandardResolveOptions.RESOLVE_FILE); ResourceId temp3 = root.resolve("temp3", StandardResolveOptions.RESOLVE_FILE); ResourceId output = root.resolve("file-03.test", StandardResolveOptions.RESOLVE_FILE); // More than one shard does. try { Iterable<FileResult> results = Lists.newArrayList(new FileResult(temp1, 1, null, null), new FileResult(temp2, 1, null, null), new FileResult(temp3, 1, null, null)); writeOp.buildOutputFilenames(results); fail("Should have failed."); } catch (IllegalStateException exn) { assertEquals("Only generated 1 distinct file names for 3 files.", exn.getMessage()); } } /** Output filenames are generated correctly when an extension is not supplied. */ @Test public void testGenerateOutputFilenamesWithoutExtension() { List<ResourceId> expected; List<ResourceId> actual; ResourceId root = getBaseOutputDirectory(); SimpleSink sink = new SimpleSink(root, "file", "-SSSSS-of-NNNNN", ""); FilenamePolicy policy = sink.getFilenamePolicy(); expected = Arrays.asList(root.resolve("file-00000-of-00003", StandardResolveOptions.RESOLVE_FILE), root.resolve("file-00001-of-00003", StandardResolveOptions.RESOLVE_FILE), root.resolve("file-00002-of-00003", StandardResolveOptions.RESOLVE_FILE)); actual = generateDestinationFilenames(root, policy, 3); assertEquals(expected, actual); expected = Collections .singletonList(root.resolve("file-00000-of-00001", StandardResolveOptions.RESOLVE_FILE)); actual = generateDestinationFilenames(root, policy, 1); assertEquals(expected, actual); expected = new ArrayList<>(); actual = generateDestinationFilenames(root, policy, 0); assertEquals(expected, actual); } /** {@link CompressionType#BZIP2} correctly writes BZip2 data. */ @Test public void testCompressionTypeBZIP2() throws FileNotFoundException, IOException { final File file = writeValuesWithWritableByteChannelFactory(CompressionType.BZIP2, "abc", "123"); // Read Bzip2ed data back in using Apache commons API (de facto standard). assertReadValues( new BufferedReader(new InputStreamReader(new BZip2CompressorInputStream(new FileInputStream(file)), StandardCharsets.UTF_8.name())), "abc", "123"); } /** {@link CompressionType#GZIP} correctly writes Gzipped data. */ @Test public void testCompressionTypeGZIP() throws FileNotFoundException, IOException { final File file = writeValuesWithWritableByteChannelFactory(CompressionType.GZIP, "abc", "123"); // Read Gzipped data back in using standard API. assertReadValues(new BufferedReader(new InputStreamReader(new GZIPInputStream(new FileInputStream(file)), StandardCharsets.UTF_8.name())), "abc", "123"); } /** {@link CompressionType#DEFLATE} correctly writes deflate data. */ @Test public void testCompressionTypeDEFLATE() throws FileNotFoundException, IOException { final File file = writeValuesWithWritableByteChannelFactory(CompressionType.DEFLATE, "abc", "123"); // Read Gzipped data back in using standard API. assertReadValues(new BufferedReader(new InputStreamReader( new DeflateCompressorInputStream(new FileInputStream(file)), StandardCharsets.UTF_8.name())), "abc", "123"); } /** {@link CompressionType#UNCOMPRESSED} correctly writes uncompressed data. */ @Test public void testCompressionTypeUNCOMPRESSED() throws FileNotFoundException, IOException { final File file = writeValuesWithWritableByteChannelFactory(CompressionType.UNCOMPRESSED, "abc", "123"); // Read uncompressed data back in using standard API. assertReadValues( new BufferedReader(new InputStreamReader(new FileInputStream(file), StandardCharsets.UTF_8.name())), "abc", "123"); } private void assertReadValues(final BufferedReader br, String... values) throws IOException { try (final BufferedReader _br = br) { for (String value : values) { assertEquals(String.format("Line should read '%s'", value), value, _br.readLine()); } } } private File writeValuesWithWritableByteChannelFactory(final WritableByteChannelFactory factory, String... values) throws IOException { final File file = tmpFolder.newFile("test.gz"); final WritableByteChannel channel = factory.create(Channels.newChannel(new FileOutputStream(file))); for (String value : values) { channel.write(ByteBuffer.wrap((value + "\n").getBytes(StandardCharsets.UTF_8))); } channel.close(); return file; } /** * {@link Writer} writes to the {@link WritableByteChannel} provided by {@link * DrunkWritableByteChannelFactory}. */ @Test public void testFileBasedWriterWithWritableByteChannelFactory() throws Exception { final String testUid = "testId"; ResourceId root = getBaseOutputDirectory(); WriteOperation<String> writeOp = new SimpleSink(root, "file", "-SS-of-NN", "txt", new DrunkWritableByteChannelFactory()).createWriteOperation(); final Writer<String> writer = writeOp.createWriter(); final ResourceId expectedFile = writeOp.tempDirectory.get().resolve(testUid, StandardResolveOptions.RESOLVE_FILE); final List<String> expected = new ArrayList<>(); expected.add("header"); expected.add("header"); expected.add("a"); expected.add("a"); expected.add("b"); expected.add("b"); expected.add("footer"); expected.add("footer"); writer.openUnwindowed(testUid, -1); writer.write("a"); writer.write("b"); final FileResult result = writer.close(); assertEquals(expectedFile, result.getTempFilename()); assertFileContains(expected, expectedFile); } /** Build a SimpleSink with default options. */ private SimpleSink buildSink() { return new SimpleSink(getBaseOutputDirectory(), "file", "-SS-of-NN", ".test"); } /** * Build a SimpleWriteOperation with default options and the given temporary directory. */ private SimpleSink.SimpleWriteOperation buildWriteOperationWithTempDir(ResourceId tempDirectory) { SimpleSink sink = buildSink(); return new SimpleSink.SimpleWriteOperation(sink, tempDirectory); } /** Build a write operation with the default options for it and its parent sink. */ private SimpleSink.SimpleWriteOperation buildWriteOperation() { return buildSink().createWriteOperation(); } }