Improve sharding for bounded pcolleciton, to better control concurrent connections. this will keep elements for same destination close to each other and shard them. For single table write it's same behaviour, for dynamic destination it will improve reduce amount of connections used

Author: stankiewicz

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiLoads.java

@@ -23,6 +23,7 @@
 import com.google.cloud.bigquery.storage.v1.AppendRowsRequest;
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
 import java.util.Map;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.function.Predicate;
@@ -38,6 +39,8 @@
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.transforms.Redistribute;
 import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.transforms.Values;
+import org.apache.beam.sdk.transforms.WithKeys;
 import org.apache.beam.sdk.transforms.errorhandling.BadRecord;
 import org.apache.beam.sdk.transforms.errorhandling.BadRecordRouter;
 import org.apache.beam.sdk.transforms.errorhandling.BadRecordRouter.ThrowingBadRecordRouter;
@@ -50,6 +53,7 @@
 import org.apache.beam.sdk.values.PCollectionList;
 import org.apache.beam.sdk.values.PCollectionTuple;
 import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing;
 import org.joda.time.Duration;
 
 /** This {@link PTransform} manages loads into BigQuery using the Storage API. */
@@ -379,12 +383,32 @@ public WriteResult expandUntriggered(
     PCollection<KV<DestinationT, StorageApiWritePayload>> successfulConvertedRows =
         convertMessagesResult.get(successfulConvertedRowsTag);
 
-    if (numShards > 0) {
+    if (numShards > 0 && input.isBounded() == PCollection.IsBounded.UNBOUNDED) {
       successfulConvertedRows =
           successfulConvertedRows.apply(
               "ResdistibuteNumShards",
               Redistribute.<KV<DestinationT, StorageApiWritePayload>>arbitrarily()
                   .withNumBuckets(numShards));
+    } else if (numShards > 0 && input.isBounded() == PCollection.IsBounded.BOUNDED) {
+      successfulConvertedRows =
+          successfulConvertedRows
+              .apply(
+                  "Add shard",
+                  WithKeys.of(
+                      (SerializableFunction<KV<DestinationT, StorageApiWritePayload>, Integer>)
+                          kv ->
+                              Math.abs(
+                                      Hashing.murmur3_32_fixed()
+                                              .hashString(
+                                                  dynamicDestinations
+                                                      .getTable(kv.getKey())
+                                                      .getShortTableUrn(),
+                                                  StandardCharsets.UTF_8)
+                                              .asInt()
+                                          ^ ThreadLocalRandom.current().nextInt(numShards))
+                                  % numShards))
+              .apply("RedistributeNumShards", Redistribute.byKey())
+              .apply("Remove shard", Values.create());
     }
 
     PCollectionTuple writeRecordsResult =