[Java] Honor withMaxRetryJobs for bounded BigQueryIO file loads

CHANGES.md

@@ -125,6 +125,7 @@
 * (Go) Fixed pubsubio importing a `google.golang.org/genproto` package removed in recent releases, which broke builds of Go modules depending on a current `genproto` version ([#40018](https://github.com/apache/beam/issues/40018)).
 * (Python) `VertexAIFeatureStoreEnrichmentHandler` now logs a warning instead of failing when the caller lacks permission to read the feature store metadata ([#40442](https://github.com/apache/beam/issues/40442)).
 * (Java) BigQueryIO now treats a 404 when deleting a temporary table or dataset as success, so a replayed work item whose earlier attempt already deleted it no longer retries forever ([#24997](https://github.com/apache/beam/issues/24997)).
+* (Java) `BigQueryIO.Write.withMaxRetryJobs` is now honored for bounded (batch) pipelines using `FILE_LOADS`, which previously always retried failed load jobs 3 times. Pipelines that never call it keep their existing defaults ([#28281](https://github.com/apache/beam/issues/28281)).
 * (Java) IcebergIO now writes rows containing `EnumerationType` (proto enum) fields as strings, instead of throwing `Unsupported Beam logical type Enum` ([#40299](https://github.com/apache/beam/issues/40299)).
 * (Python) Fixed stateful DoFns with side inputs sometimes taking the timer key coder from a side input instead of the main input, which could make the worker fail to decode timer keys with `Unknown type tag` ([#40374](https://github.com/apache/beam/issues/40374)).
 * (Python) Fixed `AsyncWrapper` finalizing bundles before generator processing, which could drop buffered writes, and nesting iterable outputs ([#40000](https://github.com/apache/beam/pull/40000)).

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

@@ -135,7 +135,11 @@ class BatchLoads<DestinationT, ElementT>
   // It sets to {@code Integer.MAX_VALUE} to block until the BigQuery job finishes.
   static final int LOAD_JOB_POLL_MAX_RETRIES = Integer.MAX_VALUE;
 
+  // The number of times a failed load or copy job is retried in place before the bundle is failed.
+  // A bounded pipeline keeps this low because the runner can just rerun the failed bundle; an
+  // unbounded pipeline retries far more, because failing a bundle in streaming is expensive.
   static final int DEFAULT_MAX_RETRY_JOBS = 3;
+  static final int DEFAULT_MAX_RETRY_JOBS_UNBOUNDED = 1000;
 
   private BigQueryServices bigQueryServices;
   private final WriteDisposition writeDisposition;

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

@@ -2775,7 +2775,6 @@ public static <T> Write<T> write() {
         .setAutoSchemaUpdate(false)
         .setAutoSchemaUpdateStrictTimeout(null)
         .setDeterministicRecordIdFn(null)
-        .setMaxRetryJobs(1000)
         .setPropagateSuccessfulStorageApiWrites(false)
         .setPropagateSuccessfulStorageApiWritesPredicate(Predicates.alwaysTrue())
         .setDirectWriteProtos(true)
@@ -3038,7 +3037,7 @@ public enum Method {
 
     abstract boolean getIgnoreInsertIds();
 
-    abstract int getMaxRetryJobs();
+    abstract @Nullable Integer getMaxRetryJobs();
 
     abstract @Nullable String getKmsKey();
 
@@ -3172,7 +3171,7 @@ abstract Builder<T> setDefaultMissingValueInterpretation(
 
       abstract Builder<T> setAutoSharding(boolean autoSharding);
 
-      abstract Builder<T> setMaxRetryJobs(int maxRetryJobs);
+      abstract Builder<T> setMaxRetryJobs(@Nullable Integer maxRetryJobs);
 
       abstract Builder<T> setPropagateSuccessful(boolean propagateSuccessful);
 
@@ -3825,7 +3824,17 @@ public Write<T> withAutoSharding() {
       return toBuilder().setAutoSharding(true).build();
     }
 
-    /** If set, this will set the max number of retry of batch load jobs. */
+    /**
+     * Sets the maximum number of times a failed BigQuery load or copy job is retried before the
+     * write fails.
+     *
+     * <p>Only applies when the write method is {@link Method#FILE_LOADS}. The streaming insert and
+     * Storage Write API methods retry at the row level and ignore this setting.
+     *
+     * <p>If this is not called, a bounded (batch) pipeline retries 3 times and an unbounded
+     * (streaming) pipeline retries 1000 times. Streaming defaults higher because failing a bundle
+     * in streaming is far more expensive than retrying the load job.
+     */
     public Write<T> withMaxRetryJobs(int maxRetryJobs) {
       return toBuilder().setMaxRetryJobs(maxRetryJobs).build();
     }
@@ -4549,10 +4558,15 @@ private <DestinationT> WriteResult continueExpandTyped(
         batchLoads.setMaxFilesPerPartition(getMaxFilesPerPartition());
         batchLoads.setMaxBytesPerPartition(getMaxBytesPerPartition());
 
-        // When running in streaming (unbounded mode) we want to retry failed load jobs
-        // indefinitely. Failing the bundle is expensive, so we set a fairly high limit on retries.
-        if (IsBounded.UNBOUNDED.equals(input.isBounded())) {
-          batchLoads.setMaxRetryJobs(getMaxRetryJobs());
+        // an explicit withMaxRetryJobs applies to batch and streaming alike. left unset, streaming
+        // retries a failed load job far more often than batch does: failing the bundle in streaming
+        // is expensive, so we would rather keep retrying the job than hand the work back to the
+        // runner. batch leaves BatchLoads on its own lower default
+        Integer maxRetryJobs = getMaxRetryJobs();
+        if (maxRetryJobs != null) {
+          batchLoads.setMaxRetryJobs(maxRetryJobs);
+        } else if (IsBounded.UNBOUNDED.equals(input.isBounded())) {
+          batchLoads.setMaxRetryJobs(BatchLoads.DEFAULT_MAX_RETRY_JOBS_UNBOUNDED);
         }
         batchLoads.setTriggeringFrequency(getTriggeringFrequency());
         if (getAutoSharding()) {

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

@@ -624,7 +624,10 @@ public Row toConfigRow(Write<?> transform) {
       fieldValues.put("skip_invalid_rows", transform.getSkipInvalidRows());
       fieldValues.put("ignore_unknown_values", transform.getIgnoreUnknownValues());
       fieldValues.put("ignore_insert_ids", transform.getIgnoreInsertIds());
-      fieldValues.put("max_retry_jobs", transform.getMaxRetryJobs());
+      Integer maxRetryJobs = transform.getMaxRetryJobs();
+      if (maxRetryJobs != null) {
+        fieldValues.put("max_retry_jobs", maxRetryJobs);
+      }
       fieldValues.put("propagate_successful", transform.getPropagateSuccessful());
       String kmsKey = transform.getKmsKey();
       if (kmsKey != null) {
@@ -898,8 +901,18 @@ public Write<?> fromConfigRow(Row configRow, PipelineOptions options) {
         if (ignoreInsertIds != null) {
           builder = builder.setIgnoreInsertIds(ignoreInsertIds);
         }
+        // before 2.78.0 BigQueryIO.write() always seeded maxRetryJobs with 1000, whether or not the
+        // pipeline ever called withMaxRetryJobs, and a bounded write ignored the value and used
+        // BatchLoads' own default of 3. so a row from one of those versions that holds exactly 1000
+        // tells us nothing about what the user asked for. leaving it unset makes both modes fall
+        // back to the same numbers they used before the upgrade, 3 for bounded and 1000 for
+        // unbounded. any other value was chosen deliberately and is carried over
         Integer maxRetryJobs = configRow.getInt32("max_retry_jobs");
-        if (maxRetryJobs != null) {
+        boolean isPreservedLegacyDefault =
+            TransformUpgrader.compareVersions(updateCompatibilityBeamVersion, "2.78.0") < 0
+                && Integer.valueOf(BatchLoads.DEFAULT_MAX_RETRY_JOBS_UNBOUNDED)
+                    .equals(maxRetryJobs);
+        if (maxRetryJobs != null && !isPreservedLegacyDefault) {
           builder = builder.setMaxRetryJobs(maxRetryJobs);
         }
         String kmsKey = configRow.getString("kms_key");

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

@@ -20,6 +20,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
 import com.google.api.services.bigquery.model.Clustering;
@@ -278,6 +279,70 @@ public void testReCreateWriteTransformFromRowTable() {
         writeTransformFromRow.getJsonClustering().get());
   }
 
+  @Test
+  public void testReCreateWriteTransformDropsLegacyMaxRetryJobsDefault() {
+    // an SDK older than 2.78.0 put 1000 in every config row, so this is what a pipeline that never
+    // called withMaxRetryJobs looks like once one of those versions has serialized it
+    BigQueryIO.Write<?> writeTransform =
+        BigQueryIO.write()
+            .to("dummyproject:dummydataset.dummytable")
+            .withMaxRetryJobs(BatchLoads.DEFAULT_MAX_RETRY_JOBS_UNBOUNDED);
+
+    BigQueryIOTranslation.BigQueryIOWriteTranslator translator =
+        new BigQueryIOTranslation.BigQueryIOWriteTranslator();
+    Row row = translator.toConfigRow(writeTransform);
+
+    PipelineOptions options = PipelineOptionsFactory.create();
+    options.as(StreamingOptions.class).setUpdateCompatibilityVersion("2.77.0");
+    BigQueryIO.Write<?> writeTransformFromRow =
+        (BigQueryIO.Write<?>) translator.fromConfigRow(row, options);
+
+    // unset, so a bounded write falls back to BatchLoads' default of 3 as it did before the
+    // upgrade, rather than jumping to 1000
+    assertNull(writeTransformFromRow.getMaxRetryJobs());
+  }
+
+  @Test
+  public void testReCreateWriteTransformKeepsLegacyMaxRetryJobsOtherThanDefault() {
+    BigQueryIO.Write<?> writeTransform =
+        BigQueryIO.write().to("dummyproject:dummydataset.dummytable").withMaxRetryJobs(7);
+
+    BigQueryIOTranslation.BigQueryIOWriteTranslator translator =
+        new BigQueryIOTranslation.BigQueryIOWriteTranslator();
+    Row row = translator.toConfigRow(writeTransform);
+
+    PipelineOptions options = PipelineOptionsFactory.create();
+    options.as(StreamingOptions.class).setUpdateCompatibilityVersion("2.77.0");
+    BigQueryIO.Write<?> writeTransformFromRow =
+        (BigQueryIO.Write<?>) translator.fromConfigRow(row, options);
+
+    // 7 was never a default in any version, so the pipeline must have asked for it
+    assertEquals(Integer.valueOf(7), writeTransformFromRow.getMaxRetryJobs());
+  }
+
+  @Test
+  public void testReCreateWriteTransformKeepsMaxRetryJobsFromCurrentVersion() {
+    BigQueryIO.Write<?> writeTransform =
+        BigQueryIO.write()
+            .to("dummyproject:dummydataset.dummytable")
+            .withMaxRetryJobs(BatchLoads.DEFAULT_MAX_RETRY_JOBS_UNBOUNDED);
+
+    BigQueryIOTranslation.BigQueryIOWriteTranslator translator =
+        new BigQueryIOTranslation.BigQueryIOWriteTranslator();
+    Row row = translator.toConfigRow(writeTransform);
+
+    PipelineOptions options = PipelineOptionsFactory.create();
+    options.as(StreamingOptions.class).setUpdateCompatibilityVersion("2.78.0");
+    BigQueryIO.Write<?> writeTransformFromRow =
+        (BigQueryIO.Write<?>) translator.fromConfigRow(row, options);
+
+    // 2.78.0 and later only write this field when the pipeline set it, so 1000 is a real choice
+    // here
+    assertEquals(
+        Integer.valueOf(BatchLoads.DEFAULT_MAX_RETRY_JOBS_UNBOUNDED),
+        writeTransformFromRow.getMaxRetryJobs());
+  }
+
   @Test
   public void testWriteTransformRowIncludesAllFields() {
     // These fields do not represent properties of the transform.

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

@@ -2208,12 +2208,40 @@ public void testWriteFailedJobs() throws Exception {
 
     thrown.expect(RuntimeException.class);
     thrown.expectMessage("Failed to create job with prefix");
-    thrown.expectMessage("reached max retries");
+    // a bounded write that never calls withMaxRetryJobs keeps BatchLoads' own default of 3
+    thrown.expectMessage("reached max retries: 3");
     thrown.expectMessage("last failed job");
 
     p.run();
   }
 
+  @Test
+  public void testWriteFailedJobsRespectsMaxRetryJobsWhenBounded() throws Exception {
+    assumeTrue(!useStorageApi);
+    assumeTrue(!useStreaming);
+    p.apply(
+            Create.of(
+                    new TableRow().set("name", "a").set("number", 1),
+                    new TableRow().set("name", "b").set("number", 2),
+                    new TableRow().set("name", "c").set("number", 3))
+                .withCoder(TableRowJsonCoder.of()))
+        .apply(
+            BigQueryIO.writeTableRows()
+                .to("dataset-id.table-id")
+                .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_NEVER)
+                .withMaxRetryJobs(1)
+                .withTestServices(fakeBqServices)
+                .withoutValidation());
+
+    thrown.expect(RuntimeException.class);
+    // the load job fails on every attempt because the destination table does not exist, and the
+    // failure message reports the retry limit that was actually applied. a bounded pipeline used to
+    // drop withMaxRetryJobs on the floor, so this used to read "reached max retries: 3"
+    thrown.expectMessage("reached max retries: 1");
+
+    p.run();
+  }
+
   @Test
   public void testWriteWithMissingSchemaFromView() throws Exception {
     // Because no messages
@@ -3114,7 +3142,7 @@ public void testMaxRetryJobs() {
             .withSchemaUpdateOptions(
                 EnumSet.of(BigQueryIO.Write.SchemaUpdateOption.ALLOW_FIELD_ADDITION))
             .withMaxRetryJobs(500);
-    assertEquals(500, write.getMaxRetryJobs());
+    assertEquals(Integer.valueOf(500), write.getMaxRetryJobs());
   }
 
   @Test