split write rename and temp table deletion into to dofns with a shuff…

…le (#30023) * split write rename and temp table deletion into to dofns with a shuffle * fix tests * spotless + checkstyle * remove unserializable data being passed between write-rename and table cleanup * Update sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteRename.java Co-authored-by: Ahmed Abualsaud <[email protected]> * address comments * spotkess --------- Co-authored-by: Ahmed Abualsaud <[email protected]>
apache · Jan 22, 2024 · 173d834 · 173d834
1 parent d014a98
commit 173d834
Show file tree

Hide file tree

Showing 4 changed files with 392 additions and 254 deletions.
diff --git a/...o/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java b/...o/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java
@@ -419,16 +419,15 @@ private WriteResult expandTriggered(PCollection<KV<DestinationT, ElementT>> inpu
                     .withSideInputs(sideInputsForUpdateSchema))
             .apply(
                 "WriteRenameTriggered",
-                ParDo.of(
-                        new WriteRename(
-                            bigQueryServices,
-                            copyJobIdPrefixView,
-                            writeDisposition,
-                            createDisposition,
-                            maxRetryJobs,
-                            kmsKey,
-                            loadJobProjectId))
-                    .withSideInputs(copyJobIdPrefixView));
+                new WriteRename(
+                    bigQueryServices,
+                    copyJobIdPrefixView,
+                    writeDisposition,
+                    createDisposition,
+                    maxRetryJobs,
+                    kmsKey,
+                    loadJobProjectId,
+                    copyJobIdPrefixView));
 
     PCollection<TableDestination> successfulSinglePartitionWrites =
         writeSinglePartition(partitions.get(singlePartitionTag), loadJobIdPrefixView)
@@ -517,16 +516,15 @@ public WriteResult expandUntriggered(PCollection<KV<DestinationT, ElementT>> inp
                     .withSideInputs(sideInputsForUpdateSchema))
             .apply(
                 "WriteRenameUntriggered",
-                ParDo.of(
-                        new WriteRename(
-                            bigQueryServices,
-                            copyJobIdPrefixView,
-                            writeDisposition,
-                            createDisposition,
-                            maxRetryJobs,
-                            kmsKey,
-                            loadJobProjectId))
-                    .withSideInputs(copyJobIdPrefixView))
+                new WriteRename(
+                    bigQueryServices,
+                    copyJobIdPrefixView,
+                    writeDisposition,
+                    createDisposition,
+                    maxRetryJobs,
+                    kmsKey,
+                    loadJobProjectId,
+                    copyJobIdPrefixView))
             .setCoder(tableDestinationCoder);
 
     PCollectionList<TableDestination> allSuccessfulWrites =

diff --git a/...gle-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpers.java b/...gle-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpers.java
@@ -32,6 +32,7 @@
 import com.google.api.services.bigquery.model.TimePartitioning;
 import com.google.cloud.hadoop.util.ApiErrorExtractor;
 import java.io.IOException;
+import java.io.Serializable;
 import java.math.BigInteger;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -168,7 +169,7 @@ private static boolean nextBackOff(Sleeper sleeper, BackOff backOff)
     }
   }
 
-  static class PendingJob {
+  static class PendingJob implements Serializable {
     private final SerializableFunction<RetryJobId, Void> executeJob;
     private final SerializableFunction<RetryJobId, Job> pollJob;
     private final SerializableFunction<RetryJobId, Job> lookupJob;
@@ -275,7 +276,7 @@ boolean shouldRetry() {
     }
   }
 
-  static class RetryJobId {
+  static class RetryJobId implements Serializable {
     private final String jobIdPrefix;
     private final int retryIndex;