[beam] branch master updated: Handle runner-provided shards for TextIO

[pa...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

pabloem pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new 552cca7  Handle runner-provided shards for TextIO
     new b061886  Merge pull request #15831 from [BEAM-13144] Handle runner-provided shards for TextIO
552cca7 is described below

commit 552cca7ecc05bd65585af748af8fefb689d148b8
Author: Pablo Estrada <pa...@apache.org>
AuthorDate: Thu Oct 28 15:38:44 2021 -0700

    Handle runner-provided shards for TextIO
---
 sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

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 8f8a699..6682df5 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
@@ -846,7 +846,14 @@ public class TextIO {
      */
     public TypedWrite<UserT, DestinationT> withNumShards(int numShards) {
       checkArgument(numShards >= 0);
-      return withNumShards(StaticValueProvider.of(numShards));
+      if (numShards == 0) {
+        // If 0 shards are passed, then the user wants runner-determined
+        // sharding to kick in, thus we pass a null StaticValueProvider
+        // so that the runner-determined-sharding path will be activated.
+        return withNumShards(null);
+      } else {
+        return withNumShards(StaticValueProvider.of(numShards));
+      }
     }
 
     /**