From 71cb92403c08a6045e80330485b7276e1dafa1ab Mon Sep 17 00:00:00 2001 From: bvolpato Date: Thu, 3 Sep 2026 19:08:46 +0000 Subject: [PATCH 1/5] [Java] Preserve BigQuery failed rows across retries --- .../StorageApiWriteUnshardedRecords.java | 3 + .../StorageApiWritesShardedRecords.java | 3 + .../bigquery/StorageApiWriteRetryTest.java | 235 ++++++++++++++++++ 3 files changed, 241 insertions(+) create mode 100644 sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRetryTest.java diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java index 0f52900e9e0b..58fadb3a3fc2 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java @@ -738,15 +738,18 @@ long flush( // rows. ProtoRows.Builder retryRows = ProtoRows.newBuilder(); List retryTimestamps = Lists.newArrayList(); + List<@Nullable TableRow> retryFailsafeTableRows = Lists.newArrayList(); for (int i = 0; i < failedContext.protoRows.getSerializedRowsCount(); ++i) { if (!failedRowIndices.contains(i)) { ByteString rowBytes = failedContext.protoRows.getSerializedRows(i); retryRows.addSerializedRows(rowBytes); retryTimestamps.add(failedContext.timestamps.get(i)); + retryFailsafeTableRows.add(failedContext.failsafeTableRows.get(i)); } } failedContext.protoRows = retryRows.build(); failedContext.timestamps = retryTimestamps; + failedContext.failsafeTableRows = retryFailsafeTableRows; int numRowsRetried = failedContext.protoRows.getSerializedRowsCount(); BigQuerySinkMetrics.appendRowsRowStatusCounter( BigQuerySinkMetrics.RowStatus.RETRIED, errorCode, shortTableUrn) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java index 51be28fedf87..ed9955450c2f 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java @@ -608,15 +608,18 @@ private void handleAppendFailure( // rows. ProtoRows.Builder retryRows = ProtoRows.newBuilder(); @Nullable List timestamps = Lists.newArrayList(); + List<@Nullable TableRow> failsafeTableRows = Lists.newArrayList(); for (int i = 0; i < failedContext.protoRows.getSerializedRowsCount(); ++i) { if (!failedRowIndices.contains(i)) { ByteString rowBytes = failedContext.protoRows.getSerializedRows(i); retryRows.addSerializedRows(rowBytes); timestamps.add(failedContext.timestamps.get(i)); + failsafeTableRows.add(failedContext.failsafeTableRows.get(i)); } } failedContext.protoRows = retryRows.build(); failedContext.timestamps = timestamps; + failedContext.failsafeTableRows = failsafeTableRows; int retriedRows = failedContext.protoRows.getSerializedRowsCount(); BigQuerySinkMetrics.appendRowsRowStatusCounter( BigQuerySinkMetrics.RowStatus.RETRIED, errorCode, shortTableId) diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRetryTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRetryTest.java new file mode 100644 index 000000000000..35a7c5e0c5ae --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRetryTest.java @@ -0,0 +1,235 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.io.gcp.bigquery; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsInAnyOrder; +import static org.hamcrest.Matchers.hasItems; + +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; +import com.google.api.services.bigquery.model.Table; +import com.google.api.services.bigquery.model.TableFieldSchema; +import com.google.api.services.bigquery.model.TableReference; +import com.google.api.services.bigquery.model.TableRow; +import com.google.api.services.bigquery.model.TableSchema; +import com.google.cloud.bigquery.storage.v1.AppendRowsRequest; +import com.google.cloud.bigquery.storage.v1.AppendRowsResponse; +import com.google.cloud.bigquery.storage.v1.Exceptions; +import com.google.cloud.bigquery.storage.v1.ProtoRows; +import com.google.protobuf.DescriptorProtos; +import com.google.protobuf.Descriptors.Descriptor; +import com.google.protobuf.DynamicMessage; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.beam.runners.direct.DirectOptions; +import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.CreateDisposition; +import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.Method; +import org.apache.beam.sdk.io.gcp.testing.FakeBigQueryServices; +import org.apache.beam.sdk.io.gcp.testing.FakeDatasetService; +import org.apache.beam.sdk.io.gcp.testing.FakeJobService; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.testing.TestStream; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.FlatMapElements; +import org.apache.beam.sdk.transforms.MapElements; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.TypeDescriptor; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameter; +import org.junit.runners.Parameterized.Parameters; + +/** Regression tests for row bookkeeping across Storage Write API retries. */ +@RunWith(Parameterized.class) +public class StorageApiWriteRetryTest { + @Parameters(name = "method={0}, streaming={1}") + public static Iterable parameters() { + return ImmutableList.of( + new Object[] {Method.STORAGE_WRITE_API, false}, + new Object[] {Method.STORAGE_WRITE_API, true}, + new Object[] {Method.STORAGE_API_AT_LEAST_ONCE, false}); + } + + @Parameter public Method method; + + @Parameter(1) + public boolean streaming; + + @Rule public final transient TestPipeline pipeline = TestPipeline.create(); + + private static final List APPEND_SIZES = Collections.synchronizedList(new ArrayList<>()); + + @Before + public void setUp() { + FakeDatasetService.setUp(); + APPEND_SIZES.clear(); + pipeline.getOptions().as(BigQueryOptions.class).setProject("project-id"); + pipeline.getOptions().as(DirectOptions.class).setTargetParallelism(1); + } + + @After + public void tearDown() throws Exception { + BigQueryIO.clearStaticCaches(); + } + + @Test + public void testConsecutiveSerializationErrorsPreserveOriginalRows() throws Exception { + TableReference table = + new TableReference() + .setProjectId("project-id") + .setDatasetId("dataset-id") + .setTableId("table-id"); + TableSchema schema = + new TableSchema() + .setFields(ImmutableList.of(new TableFieldSchema().setName("id").setType("STRING"))); + FakeDatasetService dataset = new RejectOneRowPerAppend(); + dataset.createDataset("project-id", "dataset-id", "", "", null); + dataset.createTable(new Table().setTableReference(table).setSchema(schema)); + + PCollection rows; + if (streaming) { + rows = + pipeline.apply( + TestStream.create(TableRowJsonCoder.of()) + .addElements(row("A"), row("B"), row("C")) + .advanceProcessingTime(Duration.standardSeconds(2)) + .advanceWatermarkToInfinity()); + } else { + // Emit all rows from one element so the unsharded writer receives one bundle. + rows = + pipeline + .apply(Create.of(0)) + .apply( + FlatMapElements.into(TypeDescriptor.of(TableRow.class)) + .via(ignored -> ImmutableList.of(row("A"), row("B"), row("C")))) + .setCoder(TableRowJsonCoder.of()); + } + + BigQueryIO.Write write = + BigQueryIO.writeTableRows() + .to(table) + .withMethod(method) + .withCreateDisposition(CreateDisposition.CREATE_NEVER) + .withFormatRecordOnFailureFunction( + input -> new TableRow().set("original_id", input.get("id"))) + .withPropagateSuccessfulStorageApiWrites(true) + .withTestServices( + new FakeBigQueryServices() + .withDatasetService(dataset) + .withJobService(new FakeJobService())) + .withoutValidation(); + if (streaming) { + write = + write + .withNumStorageWriteApiStreams(1) + .withTriggeringFrequency(Duration.standardSeconds(1)); + } else { + write = write.withNumStorageWriteApiStreams(0); + } + WriteResult result = rows.apply(write); + PAssert.that( + result + .getFailedStorageApiInserts() + .apply( + MapElements.into(TypeDescriptor.of(TableRow.class)) + .via(BigQueryStorageApiInsertError::getRow))) + .containsInAnyOrder( + new TableRow().set("original_id", "A"), new TableRow().set("original_id", "B")); + PAssert.that(result.getSuccessfulStorageApiInserts()).containsInAnyOrder(row("C")); + + pipeline.run().waitUntilFinish(); + + assertThat( + dataset.getAllRows("project-id", "dataset-id", "table-id"), containsInAnyOrder(row("C"))); + // Ensure this exercised successive retries of the same multirow batch, not singleton writes. + assertThat(APPEND_SIZES, hasItems(3, 2, 1)); + } + + private static TableRow row(String id) { + return new TableRow().set("id", id); + } + + private static class RejectOneRowPerAppend extends FakeDatasetService { + @Override + public BigQueryServices.StreamAppendClient getStreamAppendClient( + String streamName, + DescriptorProtos.DescriptorProto descriptor, + boolean useConnectionPool, + AppendRowsRequest.MissingValueInterpretation missingValueInterpretation) + throws Exception { + BigQueryServices.StreamAppendClient delegate = + super.getStreamAppendClient( + streamName, descriptor, useConnectionPool, missingValueInterpretation); + Descriptor protoDescriptor = TableRowToStorageApiProto.wrapDescriptorProto(descriptor); + return new BigQueryServices.StreamAppendClient() { + @Override + public ApiFuture appendRows(long offset, ProtoRows rows) + throws Exception { + APPEND_SIZES.add(rows.getSerializedRowsCount()); + try { + for (int i = 0; i < rows.getSerializedRowsCount(); i++) { + String id = + (String) + DynamicMessage.parseFrom(protoDescriptor, rows.getSerializedRows(i)) + .getField(protoDescriptor.findFieldByName("id")); + if (!id.equals("C")) { + return ApiFutures.immediateFailedFuture( + new Exceptions.AppendSerializationError( + 400, "Invalid row", streamName, ImmutableMap.of(i, "Invalid id: " + id))); + } + } + } catch (Exception e) { + return ApiFutures.immediateFailedFuture(e); + } + return delegate.appendRows(offset, rows); + } + + @Override + public com.google.cloud.bigquery.storage.v1.@Nullable TableSchema getUpdatedSchema() { + return delegate.getUpdatedSchema(); + } + + @Override + public void pin() { + delegate.pin(); + } + + @Override + public void unpin() { + delegate.unpin(); + } + + @Override + public void close() throws Exception { + delegate.close(); + } + }; + } + } +} From d803998d98d4607a72dd07ffee4a4452a8a72afc Mon Sep 17 00:00:00 2001 From: bvolpato Date: Thu, 3 Sep 2026 19:14:37 +0000 Subject: [PATCH 2/5] Reset BigQuery caches before retry regression tests --- .../beam/sdk/io/gcp/bigquery/StorageApiWriteRetryTest.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRetryTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRetryTest.java index 35a7c5e0c5ae..c37cf2eb274a 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRetryTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRetryTest.java @@ -86,8 +86,9 @@ public static Iterable parameters() { private static final List APPEND_SIZES = Collections.synchronizedList(new ArrayList<>()); @Before - public void setUp() { + public void setUp() throws Exception { FakeDatasetService.setUp(); + BigQueryIO.clearStaticCaches(); APPEND_SIZES.clear(); pipeline.getOptions().as(BigQueryOptions.class).setProject("project-id"); pipeline.getOptions().as(DirectOptions.class).setTargetParallelism(1); From 3e7359fd2839bd8bb3e885accd41a231d493fb7a Mon Sep 17 00:00:00 2001 From: bvolpato Date: Thu, 3 Sep 2026 19:16:33 +0000 Subject: [PATCH 3/5] Document BigQuery failed-row retry fix --- CHANGES.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGES.md b/CHANGES.md index f0d5d06b9d9f..37a214dedf46 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -92,6 +92,7 @@ ## Bugfixes +* (Java) Fixed BigQuery Storage Write API failure outputs returning the wrong original row after consecutive serialization errors ([#39999](https://github.com/apache/beam/pull/39999)). * (Java) Fixed the Spark runner firing processing-time timers in reverse timestamp order ([#39824](https://github.com/apache/beam/issues/39824)). * (Python) Fixed incorrect profiler options handling on portable runners ([#39613](https://github.com/apache/beam/issues/39613)). * (Java) KafkaIO dynamic reads no longer require the obsolete `beam_fn_api` experiment ([#29998](https://github.com/apache/beam/issues/29998)). From 9d9595be407efa386a0b2b26f562c78c8309603d Mon Sep 17 00:00:00 2001 From: bvolpato Date: Mon, 19 Oct 2026 22:28:59 -0400 Subject: [PATCH 4/5] [CI] Retry BigQuery write PR checks --- .github/trigger_files/beam_PreCommit_Java_GCP_IO_Direct.json | 4 ++++ 1 file changed, 4 insertions(+) create mode 100644 .github/trigger_files/beam_PreCommit_Java_GCP_IO_Direct.json diff --git a/.github/trigger_files/beam_PreCommit_Java_GCP_IO_Direct.json b/.github/trigger_files/beam_PreCommit_Java_GCP_IO_Direct.json new file mode 100644 index 000000000000..5abe02fc09c7 --- /dev/null +++ b/.github/trigger_files/beam_PreCommit_Java_GCP_IO_Direct.json @@ -0,0 +1,4 @@ +{ + "comment": "Modify this file in a trivial way to cause this test suite to run.", + "modification": 1 +} From 165612a46eada1a87bbe020f8c79132de8cfed7e Mon Sep 17 00:00:00 2001 From: bvolpato Date: Fri, 23 Oct 2026 16:31:52 -0400 Subject: [PATCH 5/5] [CI] Retry BigQuery integration after master recovery --- .github/trigger_files/beam_PreCommit_Java_GCP_IO_Direct.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/trigger_files/beam_PreCommit_Java_GCP_IO_Direct.json b/.github/trigger_files/beam_PreCommit_Java_GCP_IO_Direct.json index 5abe02fc09c7..3a009261f4f9 100644 --- a/.github/trigger_files/beam_PreCommit_Java_GCP_IO_Direct.json +++ b/.github/trigger_files/beam_PreCommit_Java_GCP_IO_Direct.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run.", - "modification": 1 + "modification": 2 }