Backport Beam PR-828

Set Gcs upload buffer size to 1M in DataflowRunner streaming mode
feczo · Aug 17, 2016 · b9b2c74 · b9b2c74
1 parent 39e57d6
commit b9b2c74
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 0 deletions.
diff --git a/sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/DataflowPipelineRunner.java b/sdk/src/main/java/com/google/cloud/dataflow/sdk/runners/DataflowPipelineRunner.java
@@ -133,6 +133,7 @@
 import com.google.cloud.dataflow.sdk.values.PValue;
 import com.google.cloud.dataflow.sdk.values.TupleTag;
 import com.google.cloud.dataflow.sdk.values.TupleTagList;
+import com.google.cloud.hadoop.util.AbstractGoogleAsyncWriteChannel;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Function;
 import com.google.common.base.Joiner;
@@ -229,6 +230,9 @@ public class DataflowPipelineRunner extends PipelineRunner<DataflowPipelineJob>
   // The limit of CreateJob request size.
   private static final int CREATE_JOB_REQUEST_LIMIT_BYTES = 10 * 1024 * 1024;
 
+  @VisibleForTesting
+  static final int GCS_UPLOAD_BUFFER_SIZE_BYTES_DEFAULT = 1 * 1024 * 1024;
+
   private final Set<PCollection<?>> pcollectionsRequiringIndexedFormat;
 
   /**
@@ -334,6 +338,9 @@ public static DataflowPipelineRunner fromOptions(PipelineOptions options) {
           + "' invalid. Please make sure the value is non-negative.");
     }
 
+    if (dataflowOptions.isStreaming() && dataflowOptions.getGcsUploadBufferSizeBytes() == null) {
+      dataflowOptions.setGcsUploadBufferSizeBytes(GCS_UPLOAD_BUFFER_SIZE_BYTES_DEFAULT);
+    }
     return new DataflowPipelineRunner(dataflowOptions);
   }
 

diff --git a/sdk/src/test/java/com/google/cloud/dataflow/sdk/runners/DataflowPipelineRunnerTest.java b/sdk/src/test/java/com/google/cloud/dataflow/sdk/runners/DataflowPipelineRunnerTest.java
@@ -735,6 +735,43 @@ public void testValidJobName() throws IOException {
     }
   }
 
+  @Test
+  public void testGcsUploadBufferSizeIsUnsetForBatchWhenDefault() throws IOException {
+    DataflowPipelineOptions batchOptions = buildPipelineOptions();
+    batchOptions.setRunner(DataflowPipelineRunner.class);
+    Pipeline.create(batchOptions);
+    assertNull(batchOptions.getGcsUploadBufferSizeBytes());
+  }
+
+  @Test
+  public void testGcsUploadBufferSizeIsSetForStreamingWhenDefault() throws IOException {
+    DataflowPipelineOptions streamingOptions = buildPipelineOptions();
+    streamingOptions.setStreaming(true);
+    streamingOptions.setRunner(DataflowPipelineRunner.class);
+    Pipeline.create(streamingOptions);
+    assertEquals(
+        DataflowPipelineRunner.GCS_UPLOAD_BUFFER_SIZE_BYTES_DEFAULT,
+        streamingOptions.getGcsUploadBufferSizeBytes().intValue());
+  }
+
+  @Test
+  public void testGcsUploadBufferSizeUnchangedWhenNotDefault() throws IOException {
+    int gcsUploadBufferSizeBytes = 12345678;
+    DataflowPipelineOptions batchOptions = buildPipelineOptions();
+    batchOptions.setGcsUploadBufferSizeBytes(gcsUploadBufferSizeBytes);
+    batchOptions.setRunner(DataflowPipelineRunner.class);
+    Pipeline.create(batchOptions);
+    assertEquals(gcsUploadBufferSizeBytes, batchOptions.getGcsUploadBufferSizeBytes().intValue());
+
+    DataflowPipelineOptions streamingOptions = buildPipelineOptions();
+    streamingOptions.setStreaming(true);
+    streamingOptions.setGcsUploadBufferSizeBytes(gcsUploadBufferSizeBytes);
+    streamingOptions.setRunner(DataflowPipelineRunner.class);
+    Pipeline.create(streamingOptions);
+    assertEquals(
+        gcsUploadBufferSizeBytes, streamingOptions.getGcsUploadBufferSizeBytes().intValue());
+  }
+
   /**
    * A fake PTransform for testing.
    */