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..3a009261f4f9 --- /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": 2 +} diff --git a/CHANGES.md b/CHANGES.md index fda4c9edc218..eb7da6f10f81 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -99,6 +99,7 @@ * (Java) Fixed the declared schema of the error output of the Kafka write SchemaTransform, which wrapped the error schema a second time and did not match the rows it emits ([#39760](https://github.com/apache/beam/issues/39760)). * (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)). * (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) 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) 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) `Duration` built from float seconds now rounds to the nearest microsecond instead of truncating, which could lose a microsecond ([#40263](https://github.com/apache/beam/issues/40263)). 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 8ea8d75ed17e..da51604c5272 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 @@ -885,15 +885,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 336311e3421e..9361ae5d5694 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 @@ -668,15 +668,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..c37cf2eb274a --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRetryTest.java @@ -0,0 +1,236 @@ +/* + * 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() throws Exception { + FakeDatasetService.setUp(); + BigQueryIO.clearStaticCaches(); + 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(); + } + }; + } + } +}