Merge pull request #15500 from [BEAM-6721] Set numShards dynamically for TextIO.write()

pabloem · web-flow · commit 84e24ea01f4a · 2021-10-26T11:01:41.000-07:00
[BEAM-6721] Set numShards dynamically for TextIO.write()
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
@@ -268,7 +268,6 @@ public static <UserT> TypedWrite<UserT, Void> writeCustomType() {
         .setDelimiter(new char[] {'\n'})
         .setWritableByteChannelFactory(FileBasedSink.CompressionType.UNCOMPRESSED)
         .setWindowedWrites(false)
-        .setNumShards(0)
         .setNoSpilling(false)
         .build();
   }
@@ -623,7 +622,7 @@ public abstract static class TypedWrite<UserT, DestinationT>
     abstract @Nullable String getFooter();
 
     /** Requested number of shards. 0 for automatic. */
-    abstract int getNumShards();
+    abstract @Nullable ValueProvider<Integer> getNumShards();
 
     /** The shard template of each file written, combined with prefix and suffix. */
     abstract @Nullable String getShardTemplate();
@@ -689,7 +688,8 @@ abstract Builder<UserT, DestinationT> setDestinationFunction(
       abstract Builder<UserT, DestinationT> setFormatFunction(
           @Nullable SerializableFunction<UserT, String> formatFunction);
 
-      abstract Builder<UserT, DestinationT> setNumShards(int numShards);
+      abstract Builder<UserT, DestinationT> setNumShards(
+          @Nullable ValueProvider<Integer> numShards);
 
       abstract Builder<UserT, DestinationT> setWindowedWrites(boolean windowedWrites);
 
@@ -846,6 +846,14 @@ public TypedWrite<UserT, DestinationT> withSuffix(String filenameSuffix) {
      */
     public TypedWrite<UserT, DestinationT> withNumShards(int numShards) {
       checkArgument(numShards >= 0);
+      return withNumShards(StaticValueProvider.of(numShards));
+    }
+
+    /**
+     * Like {@link #withNumShards(int)}. Specifying {@code null} means runner-determined sharding.
+     */
+    public TypedWrite<UserT, DestinationT> withNumShards(
+        @Nullable ValueProvider<Integer> numShards) {
       return toBuilder().setNumShards(numShards).build();
     }
 
@@ -1002,7 +1010,7 @@ public WriteFilesResult<DestinationT> expand(PCollection<UserT> input) {
                   getHeader(),
                   getFooter(),
                   getWritableByteChannelFactory()));
-      if (getNumShards() > 0) {
+      if (getNumShards() != null) {
         write = write.withNumShards(getNumShards());
       }
       if (getWindowedWrites()) {
@@ -1020,8 +1028,8 @@ public void populateDisplayData(DisplayData.Builder builder) {
 
       resolveDynamicDestinations().populateDisplayData(builder);
       builder
-          .addIfNotDefault(
-              DisplayData.item("numShards", getNumShards()).withLabel("Maximum Output Shards"), 0)
+          .addIfNotNull(
+              DisplayData.item("numShards", getNumShards()).withLabel("Maximum Output Shards"))
           .addIfNotNull(
               DisplayData.item("tempDirectory", getTempDirectory())
                   .withLabel("Directory for temporary files"))
@@ -1139,6 +1147,11 @@ public Write withNumShards(int numShards) {
       return new Write(inner.withNumShards(numShards));
     }
 
+    /** See {@link TypedWrite#withNumShards(ValueProvider)}. */
+    public Write withNumShards(@Nullable ValueProvider<Integer> numShards) {
+      return new Write(inner.withNumShards(numShards));
+    }
+
     /** See {@link TypedWrite#withoutSharding()}. */
     public Write withoutSharding() {
       return new Write(inner.withoutSharding());
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java
@@ -884,7 +884,8 @@ public void processElement(ProcessContext context) throws Exception {
         shardCount = context.sideInput(numShardsView);
       } else {
         checkNotNull(getNumShardsProvider());
-        shardCount = getNumShardsProvider().get();
+        shardCount =
+            checkNotNull(getNumShardsProvider().get(), "Must have non-null number of shards.");
       }
       checkArgument(
           shardCount > 0,
