beam-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From dhalp...@apache.org
Subject [41/50] [abbrv] incubator-beam git commit: Add header/footer support to TextIO.Write
Date Tue, 13 Sep 2016 00:41:12 GMT
Add header/footer support to TextIO.Write


Project: http://git-wip-us.apache.org/repos/asf/incubator-beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-beam/commit/b2355957
Tree: http://git-wip-us.apache.org/repos/asf/incubator-beam/tree/b2355957
Diff: http://git-wip-us.apache.org/repos/asf/incubator-beam/diff/b2355957

Branch: refs/heads/gearpump-runner
Commit: b23559578f2a7acda477b4ebccccb6a6f7c9b03e
Parents: 50c1c88
Author: Stas Levin <staslevin@gmail.com>
Authored: Mon Sep 5 20:26:12 2016 +0300
Committer: Dan Halperin <dhalperi@google.com>
Committed: Mon Sep 12 17:40:13 2016 -0700

----------------------------------------------------------------------
 .../java/org/apache/beam/sdk/io/TextIO.java     | 114 ++++++++++++++++---
 .../java/org/apache/beam/sdk/io/TextIOTest.java |  60 +++++++++-
 2 files changed, 152 insertions(+), 22 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/b2355957/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java
index 242470b..c754a0b 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java
@@ -448,7 +448,15 @@ public class TextIO {
       return new Bound<>(DEFAULT_TEXT_CODER).withoutValidation();
     }
 
-    // TODO: appendingNewlines, header, footer, etc.
+    public static Bound<String> withHeader(String header) {
+      return new Bound<>(DEFAULT_TEXT_CODER).withHeader(header);
+    }
+
+    public static Bound<String> withFooter(String footer) {
+      return new Bound<>(DEFAULT_TEXT_CODER).withFooter(footer);
+    }
+
+    // TODO: appendingNewlines, etc.
 
     /**
      * A PTransform that writes a bounded PCollection to a text file (or
@@ -465,6 +473,12 @@ public class TextIO {
       /** The suffix of each file written, combined with prefix and shardTemplate. */
       private final String filenameSuffix;
 
+      /** An optional header to add to each file. */
+      private final String header;
+
+      /** An optional footer to add to each file. */
+      private final String footer;
+
       /** The Coder to use to decode each line. */
       private final Coder<T> coder;
 
@@ -478,12 +492,15 @@ public class TextIO {
       private final boolean validate;
 
       Bound(Coder<T> coder) {
-        this(null, null, "", coder, 0, DEFAULT_SHARD_TEMPLATE, true);
+        this(null, null, "", null, null, coder, 0, DEFAULT_SHARD_TEMPLATE, true);
       }
 
-      private Bound(String name, String filenamePrefix, String filenameSuffix, Coder<T>
coder,
-          int numShards, String shardTemplate, boolean validate) {
+      private Bound(String name, String filenamePrefix, String filenameSuffix, String header,
+                    String footer, Coder<T> coder, int numShards, String shardTemplate,
+                    boolean validate) {
         super(name);
+        this.header = header;
+        this.footer = footer;
         this.coder = coder;
         this.filenamePrefix = filenamePrefix;
         this.filenameSuffix = filenameSuffix;
@@ -502,7 +519,7 @@ public class TextIO {
        */
       public Bound<T> to(String filenamePrefix) {
         validateOutputComponent(filenamePrefix);
-        return new Bound<>(name, filenamePrefix, filenameSuffix, coder, numShards,
+        return new Bound<>(name, filenamePrefix, filenameSuffix, header, footer, coder,
numShards,
             shardTemplate, validate);
       }
 
@@ -516,7 +533,7 @@ public class TextIO {
        */
       public Bound<T> withSuffix(String nameExtension) {
         validateOutputComponent(nameExtension);
-        return new Bound<>(name, filenamePrefix, nameExtension, coder, numShards,
+        return new Bound<>(name, filenamePrefix, nameExtension, header, footer, coder,
numShards,
             shardTemplate, validate);
       }
 
@@ -536,7 +553,7 @@ public class TextIO {
        */
       public Bound<T> withNumShards(int numShards) {
         checkArgument(numShards >= 0);
-        return new Bound<>(name, filenamePrefix, filenameSuffix, coder, numShards,
+        return new Bound<>(name, filenamePrefix, filenameSuffix, header, footer, coder,
numShards,
             shardTemplate, validate);
       }
 
@@ -549,7 +566,7 @@ public class TextIO {
        * @see ShardNameTemplate
        */
       public Bound<T> withShardNameTemplate(String shardTemplate) {
-        return new Bound<>(name, filenamePrefix, filenameSuffix, coder, numShards,
+        return new Bound<>(name, filenamePrefix, filenameSuffix, header, footer, coder,
numShards,
             shardTemplate, validate);
       }
 
@@ -567,7 +584,8 @@ public class TextIO {
        * <p>Does not modify this object.
        */
       public Bound<T> withoutSharding() {
-        return new Bound<>(name, filenamePrefix, filenameSuffix, coder, 1, "", validate);
+        return new Bound<>(name, filenamePrefix, filenameSuffix, header, footer, coder,
1, "",
+            validate);
       }
 
       /**
@@ -579,7 +597,7 @@ public class TextIO {
        * @param <X> the type of the elements of the input {@link PCollection}
        */
       public <X> Bound<X> withCoder(Coder<X> coder) {
-        return new Bound<>(name, filenamePrefix, filenameSuffix, coder, numShards,
+        return new Bound<>(name, filenamePrefix, filenameSuffix, header, footer, coder,
numShards,
             shardTemplate, validate);
       }
 
@@ -594,7 +612,17 @@ public class TextIO {
        * <p>Does not modify this object.
        */
       public Bound<T> withoutValidation() {
-        return new Bound<>(name, filenamePrefix, filenameSuffix, coder, numShards,
+        return new Bound<>(name, filenamePrefix, filenameSuffix, header, footer, coder,
numShards,
+            shardTemplate, false);
+      }
+
+      public Bound<T> withHeader(String header) {
+        return new Bound<>(name, filenamePrefix, filenameSuffix, header, footer, coder,
numShards,
+            shardTemplate, false);
+      }
+
+      public Bound<T> withFooter(String footer) {
+        return new Bound<>(name, filenamePrefix, filenameSuffix, header, footer, coder,
numShards,
             shardTemplate, false);
       }
 
@@ -607,7 +635,8 @@ public class TextIO {
 
         org.apache.beam.sdk.io.Write.Bound<T> write =
             org.apache.beam.sdk.io.Write.to(
-                new TextSink<>(filenamePrefix, filenameSuffix, shardTemplate, coder));
+                new TextSink<>(filenamePrefix, filenameSuffix, header, footer, shardTemplate,
+                    coder));
         if (getNumShards() > 0) {
           write = write.withNumShards(getNumShards());
         }
@@ -629,7 +658,11 @@ public class TextIO {
             .addIfNotDefault(DisplayData.item("validation", validate)
               .withLabel("Validation Enabled"), true)
             .addIfNotDefault(DisplayData.item("numShards", numShards)
-              .withLabel("Maximum Output Shards"), 0);
+              .withLabel("Maximum Output Shards"), 0)
+            .addIfNotNull(DisplayData.item("fileHeader", header)
+              .withLabel("Output file header"))
+            .addIfNotNull(DisplayData.item("fileFooter", footer)
+                .withLabel("Output file footer"));
       }
 
       /**
@@ -664,6 +697,14 @@ public class TextIO {
         return coder;
       }
 
+      public String getHeader() {
+        return header;
+      }
+
+      public String getFooter() {
+        return footer;
+      }
+
       public boolean needsValidation() {
         return validate;
       }
@@ -946,17 +987,22 @@ public class TextIO {
   @VisibleForTesting
   static class TextSink<T> extends FileBasedSink<T> {
     private final Coder<T> coder;
+    private final String header;
+    private final String footer;
 
     @VisibleForTesting
     TextSink(
-        String baseOutputFilename, String extension, String fileNameTemplate, Coder<T>
coder) {
+        String baseOutputFilename, String extension, String header, String footer,
+        String fileNameTemplate, Coder<T> coder) {
       super(baseOutputFilename, extension, fileNameTemplate);
       this.coder = coder;
+      this.header = header;
+      this.footer = footer;
     }
 
     @Override
     public FileBasedSink.FileBasedWriteOperation<T> createWriteOperation(PipelineOptions
options) {
-      return new TextWriteOperation<>(this, coder);
+      return new TextWriteOperation<>(this, coder, header, footer);
     }
 
     /**
@@ -965,15 +1011,19 @@ public class TextIO {
      */
     private static class TextWriteOperation<T> extends FileBasedWriteOperation<T>
{
       private final Coder<T> coder;
+      private final String header;
+      private final String footer;
 
-      private TextWriteOperation(TextSink<T> sink, Coder<T> coder) {
+      private TextWriteOperation(TextSink<T> sink, Coder<T> coder, String header,
String footer) {
         super(sink);
         this.coder = coder;
+        this.header = header;
+        this.footer = footer;
       }
 
       @Override
       public FileBasedWriter<T> createWriter(PipelineOptions options) throws Exception
{
-        return new TextWriter<>(this, coder);
+        return new TextWriter<>(this, coder, header, footer);
       }
     }
 
@@ -984,20 +1034,50 @@ public class TextIO {
     private static class TextWriter<T> extends FileBasedWriter<T> {
       private static final byte[] NEWLINE = "\n".getBytes(StandardCharsets.UTF_8);
       private final Coder<T> coder;
+      private final String header;
+      private final String footer;
       private OutputStream out;
 
       public TextWriter(FileBasedWriteOperation<T> writeOperation, Coder<T> coder)
{
+        this(writeOperation, coder, null, null);
+      }
+
+      public TextWriter(FileBasedWriteOperation<T> writeOperation, Coder<T> coder,
String header) {
+        this(writeOperation, coder, header, null);
+      }
+
+      public TextWriter(FileBasedWriteOperation<T> writeOperation, Coder<T> coder,
String header,
+                        String footer) {
         super(writeOperation);
+        this.header = header;
+        this.footer = footer;
         this.mimeType = MimeTypes.TEXT;
         this.coder = coder;
       }
 
+      private void writeLine(String line) throws IOException {
+        if (line != null) {
+          out.write(line.getBytes(StandardCharsets.UTF_8));
+          out.write(NEWLINE);
+        }
+      }
+
       @Override
       protected void prepareWrite(WritableByteChannel channel) throws Exception {
         out = Channels.newOutputStream(channel);
       }
 
       @Override
+      protected void writeHeader() throws Exception {
+        writeLine(header);
+      }
+
+      @Override
+      protected void writeFooter() throws Exception {
+        writeLine(footer);
+      }
+
+      @Override
       public void write(T value) throws Exception {
         coder.encode(value, out, Context.OUTER);
         out.write(NEWLINE);

http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/b2355957/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TextIOTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TextIOTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TextIOTest.java
index 8f94766..2ab2683 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TextIOTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TextIOTest.java
@@ -34,6 +34,8 @@ import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Lists;
+
 import java.io.BufferedReader;
 import java.io.File;
 import java.io.FileOutputStream;
@@ -48,6 +50,7 @@ import java.nio.file.Files;
 import java.nio.file.StandardOpenOption;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Set;
 import java.util.zip.GZIPOutputStream;
@@ -101,6 +104,9 @@ import org.mockito.stubbing.Answer;
 @SuppressWarnings("unchecked")
 public class TextIOTest {
 
+  private static final String MY_HEADER = "myHeader";
+  private static final String MY_FOOTER = "myFooter";
+
   @Rule
   public TemporaryFolder tmpFolder = new TemporaryFolder();
   @Rule
@@ -216,10 +222,19 @@ public class TextIOTest {
   }
 
   <T> void runTestWrite(T[] elems, Coder<T> coder) throws Exception {
-    runTestWrite(elems, coder, 1);
+    runTestWrite(elems, null, null, coder, 1);
   }
 
   <T> void runTestWrite(T[] elems, Coder<T> coder, int numShards) throws Exception
{
+    runTestWrite(elems, null, null, coder, numShards);
+  }
+
+  <T> void runTestWrite(T[] elems, Coder<T> coder, String header, String footer)
throws Exception {
+    runTestWrite(elems, header, footer, coder, 1);
+  }
+
+  <T> void runTestWrite(T[] elems, String header, String footer, Coder<T> coder,
int numShards)
+      throws Exception {
     String outputName = "file.txt";
     String baseFilename = tmpFolder.newFile(outputName).getPath();
 
@@ -235,6 +250,8 @@ public class TextIOTest {
     } else {
       write = TextIO.Write.to(baseFilename).withCoder(coder);
     }
+    write = write.withHeader(header).withFooter(footer);
+
     if (numShards == 1) {
       write = write.withoutSharding();
     } else if (numShards > 0) {
@@ -244,11 +261,14 @@ public class TextIOTest {
 
     p.run();
 
-    assertOutputFiles(elems, coder, numShards, tmpFolder, outputName, write.getShardNameTemplate());
+    assertOutputFiles(elems, header, footer, coder, numShards, tmpFolder, outputName,
+        write.getShardNameTemplate());
   }
 
   public static <T> void assertOutputFiles(
       T[] elems,
+      String header,
+      String footer,
       Coder<T> coder,
       int numShards,
       TemporaryFolder rootLocation,
@@ -284,15 +304,23 @@ public class TextIOTest {
       }
     }
 
-    String[] expected = new String[elems.length];
+    LinkedList<String> expected = Lists.newLinkedList();
+
     for (int i = 0; i < elems.length; i++) {
       T elem = elems[i];
       byte[] encodedElem = CoderUtils.encodeToByteArray(coder, elem);
       String line = new String(encodedElem);
-      expected[i] = line;
+      expected.add(line);
+    }
+
+    if (header != null) {
+      expected.addFirst(header);
+    }
+    if (footer != null) {
+      expected.addLast(footer);
     }
 
-    assertThat(actual, containsInAnyOrder(expected));
+    assertThat(actual, containsInAnyOrder(expected.toArray()));
   }
 
   @Test
@@ -332,18 +360,40 @@ public class TextIOTest {
   }
 
   @Test
+  @Category(NeedsRunner.class)
+  public void testWriteWithHeader() throws Exception {
+    runTestWrite(LINES_ARRAY, StringUtf8Coder.of(), MY_HEADER, null);
+  }
+
+  @Test
+  @Category(NeedsRunner.class)
+  public void testWriteWithFooter() throws Exception {
+    runTestWrite(LINES_ARRAY, StringUtf8Coder.of(), null, MY_FOOTER);
+  }
+
+  @Test
+  @Category(NeedsRunner.class)
+  public void testWriteWithHeaderAndFooter() throws Exception {
+    runTestWrite(LINES_ARRAY, StringUtf8Coder.of(), MY_HEADER, MY_FOOTER);
+  }
+
+  @Test
   public void testWriteDisplayData() {
     TextIO.Write.Bound<?> write = TextIO.Write
         .to("foo")
         .withSuffix("bar")
         .withShardNameTemplate("-SS-of-NN-")
         .withNumShards(100)
+        .withFooter("myFooter")
+        .withHeader("myHeader")
         .withoutValidation();
 
     DisplayData displayData = DisplayData.from(write);
 
     assertThat(displayData, hasDisplayItem("filePrefix", "foo"));
     assertThat(displayData, hasDisplayItem("fileSuffix", "bar"));
+    assertThat(displayData, hasDisplayItem("fileHeader", "myHeader"));
+    assertThat(displayData, hasDisplayItem("fileFooter", "myFooter"));
     assertThat(displayData, hasDisplayItem("shardNameTemplate", "-SS-of-NN-"));
     assertThat(displayData, hasDisplayItem("numShards", 100));
     assertThat(displayData, hasDisplayItem("validation", false));


Mime
View raw message