From 8ccf744dd0dd5df0a3de63940d4789fe7dfc0449 Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Tue, 2 Jun 2026 16:20:32 -0700 Subject: [PATCH 01/10] Strongly-consistent schemna detection in BigQueryIO. --- .../transforms/reflect/DoFnSignatures.java | 1 + .../sdk/io/gcp/bigquery/AppendRowsPacket.java | 249 ++++++++++ .../beam/sdk/io/gcp/bigquery/BigQueryIO.java | 58 ++- .../sdk/io/gcp/bigquery/BigQueryOptions.java | 7 + .../io/gcp/bigquery/BigQuerySinkMetrics.java | 1 + .../io/gcp/bigquery/BufferMismatchedRows.java | 280 ++++++++++++ .../io/gcp/bigquery/DynamicDestinations.java | 38 ++ .../bigquery/DynamicDestinationsHelpers.java | 13 + .../sdk/io/gcp/bigquery/MismatchedRow.java | 86 ++++ .../bigquery/SchemaChangeDetectorHelper.java | 217 +++++++++ .../gcp/bigquery/SchemaUpdateHoldingFn.java | 10 + .../io/gcp/bigquery/SplittingIterable.java | 160 ++----- .../StorageApiDynamicDestinations.java | 18 +- ...StorageApiDynamicDestinationsTableRow.java | 17 +- .../sdk/io/gcp/bigquery/StorageApiLoads.java | 7 + .../gcp/bigquery/StorageApiWritePayload.java | 16 + .../StorageApiWriteRecordsInconsistent.java | 112 ++++- .../StorageApiWriteUnshardedRecords.java | 426 ++++++++++-------- .../StorageApiWritesShardedRecords.java | 379 +++++++++++++--- .../bigquery/TableRowToStorageApiProto.java | 1 + .../io/gcp/bigquery/UpgradeTableSchema.java | 107 ++++- .../io/gcp/bigquery/BigQueryIOWriteTest.java | 56 ++- .../SchemaChangeDetectorHelperTest.java | 305 +++++++++++++ .../gcp/bigquery/SplittingIterableTest.java | 244 ++++++++++ .../gcp/bigquery/UpgradeTableSchemaTest.java | 125 ++++- 25 files changed, 2509 insertions(+), 424 deletions(-) create mode 100644 sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java create mode 100644 sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java create mode 100644 sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/MismatchedRow.java create mode 100644 sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelper.java create mode 100644 sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelperTest.java create mode 100644 sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterableTest.java diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java index 9f3491bca7b9..4a60ef4e5b5b 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java @@ -222,6 +222,7 @@ private DoFnSignatures() {} private static final Collection> ALLOWED_ON_WINDOW_EXPIRATION_PARAMETERS = ImmutableList.of( + Parameter.OnWindowExpirationContextParameter.class, Parameter.WindowParameter.class, Parameter.PipelineOptionsParameter.class, Parameter.OutputReceiverParameter.class, diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java new file mode 100644 index 000000000000..445857abe43a --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java @@ -0,0 +1,249 @@ +/* + * 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 com.google.api.services.bigquery.model.TableRow; +import com.google.auto.value.AutoValue; +import com.google.cloud.bigquery.storage.v1.ProtoRows; +import com.google.protobuf.ByteString; +import com.google.protobuf.Descriptors; +import java.io.IOException; +import java.util.BitSet; +import java.util.List; +import java.util.Map; +import java.util.function.Consumer; +import java.util.function.Supplier; +import java.util.stream.IntStream; +import java.util.stream.Stream; +import org.apache.beam.sdk.util.Preconditions; +import org.apache.beam.sdk.util.ThrowingSupplier; +import org.apache.beam.sdk.values.TimestampedValue; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.PeekingIterator; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +@AutoValue +abstract class AppendRowsPacket { + abstract ProtoRows getProtoRows(); + + abstract List getTimestamps(); + + abstract List<@Nullable TableRow> getFailsafeTableRows(); + + abstract BitSet getSchemaMismatchedRows(); + + abstract Map getSchemaHashes(); + + abstract Map getUnknownFields(); + + abstract Map getOriginalPayloads(); + + private static final BitSet EMPTY_BIT_SET = new BitSet(0); + + static AppendRowsPacket fromStorageApiWritePayload( + PeekingIterator underlyingIterator, + long maxByteSize, + SchemaChangeDetectorHelper schemaChangeDetectorHelper, + Instant elementTimestamp, + Supplier appendClientInfoSupplier, + Consumer> failedRowsConsumer, + ThrowingSupplier getCurrentTableSchemaHash, + ThrowingSupplier getCurrentTableSchemaDescriptor) { + List timestamps = Lists.newArrayList(); + List<@Nullable TableRow> failsafeRows = Lists.newArrayList(); + Map schemaHashes = Maps.newHashMap(); + Map unknownFields = Maps.newHashMap(); + Map originalPayloads = Maps.newHashMap(); + ProtoRows.Builder inserts = ProtoRows.newBuilder(); + long bytesSize = 0; + BitSet mismatchedRows = new BitSet(); + try { + while (underlyingIterator.hasNext()) { + // Make sure that we don't exceed the maxByteSize over multiple elements. A single + // element can exceed + // the split threshold, but in that case it should be the only element returned. + if ((bytesSize + underlyingIterator.peek().getPayload().length > maxByteSize) + && inserts.getSerializedRowsCount() > 0) { + break; + } + StorageApiWritePayload payload = underlyingIterator.next(); + + @Nullable TableRow failsafeTableRow = null; + try { + failsafeTableRow = payload.getFailsafeTableRow(); + } catch (IOException e) { + // Do nothing, table row will be generated later from row bytes + } + + // If autoUpdateSchema is set, we try to automatically merge in unknown fields. + SchemaChangeDetectorHelper.MergePayloadResult mergeResult = + schemaChangeDetectorHelper.getMergedPayload( + payload, elementTimestamp, failsafeTableRow, appendClientInfoSupplier.get()); + if (mergeResult.getKind() == SchemaChangeDetectorHelper.MergePayloadResult.Kind.FAILED) { + failedRowsConsumer.accept(mergeResult.getFailed()); + continue; + } + ByteString byteString = mergeResult.getMerged(); + + if (schemaChangeDetectorHelper.isPayloadSchemaOutOfDate( + payload, + byteString.toByteArray(), + getCurrentTableSchemaHash, + getCurrentTableSchemaDescriptor)) { + mismatchedRows.set(inserts.getSerializedRowsCount()); + } + + int currentIndex = inserts.getSerializedRowsCount(); + inserts.addSerializedRows(byteString); + Instant timestamp = payload.getTimestamp(); + if (timestamp == null) { + timestamp = elementTimestamp; + } + timestamps.add(timestamp); + failsafeRows.add(failsafeTableRow); + if (payload.getSchemaHash() != null) { + schemaHashes.put(currentIndex, Preconditions.checkStateNotNull(payload.getSchemaHash())); + } + if (payload.getUnknownFields() != null) { + unknownFields.put( + currentIndex, Preconditions.checkStateNotNull(payload.getUnknownFields())); + originalPayloads.put(currentIndex, payload.getPayload()); + } + bytesSize += byteString.size(); + } + } catch (Exception e) { + throw new RuntimeException(e); + } + + return new AutoValue_AppendRowsPacket( + inserts.build(), + timestamps, + failsafeRows, + mismatchedRows, + schemaHashes, + unknownFields, + originalPayloads); + } + + AppendRowsPacket getSchemaMismatchedRowsOnly() { + ProtoRows.Builder inserts = ProtoRows.newBuilder(); + List timestamps = Lists.newArrayList(); + List<@Nullable TableRow> failsafeTableRows = Lists.newArrayList(); + Map schemaHashes = Maps.newHashMap(); + Map unknownFields = Maps.newHashMap(); + Map originalPayloads = Maps.newHashMap(); + if (!getSchemaMismatchedRows().isEmpty()) { + int newIndex = 0; + for (int i = 0; i < getProtoRows().getSerializedRowsCount(); i++) { + if (getSchemaMismatchedRows().get(i)) { + inserts.addSerializedRows(getProtoRows().getSerializedRows(i)); + timestamps.add(getTimestamps().get(i)); + failsafeTableRows.add(getFailsafeTableRows().get(i)); + if (getUnknownFields().containsKey(i)) { + unknownFields.put(newIndex, Preconditions.checkStateNotNull(getUnknownFields().get(i))); + } + if (getOriginalPayloads().containsKey(i)) { + originalPayloads.put( + newIndex, Preconditions.checkStateNotNull(getOriginalPayloads().get(i))); + } + if (getSchemaHashes().containsKey(i)) { + schemaHashes.put(newIndex, Preconditions.checkStateNotNull(getSchemaHashes().get(i))); + } + newIndex++; + } + } + } + BitSet allBits = new BitSet(inserts.getSerializedRowsCount()); + allBits.set(0, inserts.getSerializedRowsCount()); + return new AutoValue_AppendRowsPacket( + inserts.build(), + timestamps, + failsafeTableRows, + allBits, + schemaHashes, + unknownFields, + originalPayloads); + } + + AppendRowsPacket getSchemaMatchedRowsOnly() { + if (getSchemaMismatchedRows().isEmpty()) { + return this; + } + + ProtoRows.Builder inserts = ProtoRows.newBuilder(); + List timestamps = Lists.newArrayList(); + Map schemaHashes = Maps.newHashMap(); + Map unknownFields = Maps.newHashMap(); + Map originalPayloads = Maps.newHashMap(); + List<@Nullable TableRow> failsafeTableRows = Lists.newArrayList(); + int newIndex = 0; + for (int i = 0; i < getProtoRows().getSerializedRowsCount(); i++) { + if (!getSchemaMismatchedRows().get(i)) { + inserts.addSerializedRows(getProtoRows().getSerializedRows(i)); + timestamps.add(getTimestamps().get(i)); + failsafeTableRows.add(getFailsafeTableRows().get(i)); + if (getSchemaHashes().containsKey(i)) { + schemaHashes.put(newIndex, Preconditions.checkStateNotNull(getSchemaHashes().get(i))); + } + if (getUnknownFields().containsKey(i)) { + unknownFields.put(newIndex, Preconditions.checkStateNotNull(getUnknownFields().get(i))); + } + if (getOriginalPayloads().containsKey(i)) { + originalPayloads.put( + newIndex, Preconditions.checkStateNotNull(getOriginalPayloads().get(i))); + } + newIndex++; + } + } + return new AutoValue_AppendRowsPacket( + inserts.build(), + timestamps, + failsafeTableRows, + EMPTY_BIT_SET, + schemaHashes, + unknownFields, + originalPayloads); + } + + Stream toPayloadStream() { + return IntStream.range(0, getProtoRows().getSerializedRowsCount()) + .mapToObj( + i -> { + try { + byte @Nullable [] originalBytes = getOriginalPayloads().get(i); + StorageApiWritePayload payload = + StorageApiWritePayload.of( + originalBytes != null + ? originalBytes + : getProtoRows().getSerializedRows(i).toByteArray(), + getUnknownFields().get(i), + getFailsafeTableRows().get(i)) + .withTimestamp(getTimestamps().get(i)); + byte @Nullable [] hash = getSchemaHashes().get(i); + if (hash != null) { + payload = payload.withSchemaHash(hash); + } + return payload; + } catch (IOException e) { + throw new RuntimeException(e); + } + }); + } +} diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java index b222b358f547..1360f58c6df5 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java @@ -108,6 +108,7 @@ import org.apache.beam.sdk.io.gcp.bigquery.PassThroughThenCleanup.ContextContainer; import org.apache.beam.sdk.io.gcp.bigquery.RowWriterFactory.OutputType; import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.StreamingOptions; import org.apache.beam.sdk.options.ValueProvider; import org.apache.beam.sdk.options.ValueProvider.NestedValueProvider; import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider; @@ -2496,6 +2497,7 @@ public static Write write() { .setAutoSharding(false) .setPropagateSuccessful(true) .setAutoSchemaUpdate(false) + .setAutoSchemaUpdateStrictTimeout(null) .setDeterministicRecordIdFn(null) .setMaxRetryJobs(1000) .setPropagateSuccessfulStorageApiWrites(false) @@ -2777,6 +2779,8 @@ public enum Method { abstract boolean getAutoSchemaUpdate(); + abstract @Nullable Duration getAutoSchemaUpdateStrictTimeout(); + abstract @Nullable Class getWriteProtosClass(); abstract boolean getDirectWriteProtos(); @@ -2897,6 +2901,8 @@ abstract Builder setDefaultMissingValueInterpretation( abstract Builder setAutoSchemaUpdate(boolean autoSchemaUpdate); + abstract Builder setAutoSchemaUpdateStrictTimeout(@Nullable Duration timeout); + abstract Builder setWriteProtosClass(@Nullable Class clazz); abstract Builder setDirectWriteProtos(boolean direct); @@ -3562,11 +3568,46 @@ public Write withSuccessfulInsertsPropagation(boolean propagateSuccessful) { * If true, enables automatically detecting BigQuery table schema updates. Table schema updates * are usually noticed within several minutes. Only supported when using one of the STORAGE_API * insert methods. + * + *

Rows that contain new columns will only have the new columns sent to BigQuery after the + * new table schema has been observed. Until then, only the known columns will be sent to + * BigQuery. + * + *

Note that this option detects table schema updates performed on the table, usually by an + * external process. If you want Beam to update the table schema for you, please see {@link + * #withSchemaUpdateOptions}. */ public Write withAutoSchemaUpdate(boolean autoSchemaUpdate) { return toBuilder().setAutoSchemaUpdate(autoSchemaUpdate).build(); } + /** + * If true, enables automatically detecting BigQuery table schema updates. Table schema updates + * are usually noticed within several minutes. Only supported when using one of the STORAGE_API + * insert methods. + * + *

This mode ensures consistent writes - rows with new schemas will not be sent to BigQuery + * until we have observed the new table schema. This comes with a few caveats: - Detecting + * unknown columns requires extra parsing. Some increase in CPU usage may be noticed. - Rows + * with unknown columns will be retried until a new schema is observed. This may temporarily + * block inserts of rows into BigQuery whenever a schema update happens. + * + *

The timeout parameter specifies how long to wait until we see the new table schema. If + * more than the specified time goes by before a matching table schema is seen, the row will be + * sent to the failedRows output collection. + * + *

Note that this option detects table schema updates performed on the table, usually by an + * external process. If you want Beam to update the table schema for you, please see {@link + * #withSchemaUpdateOptions}. + */ + public Write withAutoSchemaUpdateConsistent( + boolean autoSchemaUpdate, Duration waitForSchemaTimeout) { + return toBuilder() + .setAutoSchemaUpdate(autoSchemaUpdate) + .setAutoSchemaUpdateStrictTimeout(waitForSchemaTimeout) + .build(); + } + /* * Provides a function which can serve as a source of deterministic unique ids for each record * to be written, replacing the unique ids generated with the default scheme. When used with @@ -4090,6 +4131,14 @@ private WriteResult continueExpandTyped( DynamicDestinations dynamicDestinations, RowWriterFactory rowWriterFactory, Write.Method method) { + if (getAutoSchemaUpdateStrictTimeout() != null) { + checkArgument( + method == Method.STORAGE_API_AT_LEAST_ONCE || method == Method.STORAGE_WRITE_API, + "Auto update schema only supported when using storage write API"); + checkArgument( + input.getPipeline().getOptions().as(StreamingOptions.class).isStreaming(), + "auto update schema only supported on streaming pipelines"); + } if (method == Write.Method.STREAMING_INSERTS) { checkArgument( getWriteDisposition() != WriteDisposition.WRITE_TRUNCATE, @@ -4294,6 +4343,10 @@ private WriteResult continueExpandTyped( RowWriterFactory.TableRowWriterFactory tableRowWriterFactory = (RowWriterFactory.TableRowWriterFactory) rowWriterFactory; // Fallback behavior: convert to JSON TableRows and convert those into Beam TableRows. + @Nullable Set schemaUpdateOptions = getSchemaUpdateOptions(); + boolean useSchemaUpdatingTableRow = + (schemaUpdateOptions != null && !schemaUpdateOptions.isEmpty()) + || (getAutoSchemaUpdate() && getAutoSchemaUpdateStrictTimeout() != null); storageApiDynamicDestinations = new StorageApiDynamicDestinationsTableRow<>( dynamicDestinations, @@ -4303,9 +4356,7 @@ private WriteResult continueExpandTyped( getCreateDisposition(), getIgnoreUnknownValues(), getAutoSchemaUpdate(), - getSchemaUpdateOptions() == null - ? Collections.emptySet() - : getSchemaUpdateOptions()); + useSchemaUpdatingTableRow); } int numShards = getStorageApiNumStreams(bqOptions); @@ -4328,6 +4379,7 @@ private WriteResult continueExpandTyped( method == Method.STORAGE_API_AT_LEAST_ONCE, enableAutoSharding, getAutoSchemaUpdate(), + getAutoSchemaUpdateStrictTimeout(), getIgnoreUnknownValues(), getPropagateSuccessfulStorageApiWrites(), getPropagateSuccessfulStorageApiWritesPredicate(), diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java index face2ef5841a..bc1c8fcaf1ab 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java @@ -260,4 +260,11 @@ public interface BigQueryOptions Integer getSchemaUpgradeBufferingShards(); void setSchemaUpgradeBufferingShards(Integer value); + + @Hidden + @Description("The initial retry time in milliseconds when a schema mismatch is detected.") + @Default.Integer(5000) + Integer getStorageApiInitialMismatchRetryTimeMilliSec(); + + void setStorageApiInitialMismatchRetryTimeMilliSec(Integer value); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQuerySinkMetrics.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQuerySinkMetrics.java index aec011cebb54..acb0bd5f6946 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQuerySinkMetrics.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQuerySinkMetrics.java @@ -51,6 +51,7 @@ public class BigQuerySinkMetrics { public static final String OK = Status.Code.OK.toString(); static final String INTERNAL = "INTERNAL"; public static final String PAYLOAD_TOO_LARGE = "PayloadTooLarge"; + public static final String SCHEMA_MISMATCHED = "SchemaMismatched"; // Base Metric names private static final String RPC_REQUESTS = "RpcRequestsCount"; diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java new file mode 100644 index 000000000000..03c0dbfaa393 --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java @@ -0,0 +1,280 @@ +/* + * 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 com.google.api.services.bigquery.model.TableRow; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.ThreadLocalRandom; +import java.util.stream.StreamSupport; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.metrics.Counter; +import org.apache.beam.sdk.metrics.Metrics; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.state.BagState; +import org.apache.beam.sdk.state.StateSpec; +import org.apache.beam.sdk.state.StateSpecs; +import org.apache.beam.sdk.state.TimeDomain; +import org.apache.beam.sdk.state.Timer; +import org.apache.beam.sdk.state.TimerSpec; +import org.apache.beam.sdk.state.TimerSpecs; +import org.apache.beam.sdk.state.ValueState; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.util.ShardedKey; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.TupleTagList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.joda.time.Instant; + +class BufferMismatchedRows + extends PTransform>, PCollectionTuple> { + private final Coder failedRowsCoder; + private final Coder successfulRowsCoder; + private final Coder destinationCoder; + private final StorageApiDynamicDestinations dynamicDestinations; + private final StorageApiWriteUnshardedRecords.WriteRecordsDoFn writeDoFn; + private final TupleTag failedRowsTag; + private final @Nullable TupleTag successfulRowsTag; + // This output is effectively ignored, since we only support this code path for + // StorageApiWriteRecordsInconsistent. + private final TupleTag> finalizeTag = new TupleTag<>("finalizeTag"); + private static final int NUM_DEFAULT_SHARDS = 20; + private static final Duration RETRY_MISMATCHED_ROWS_PERIOD = Duration.standardMinutes(1); + + public BufferMismatchedRows( + Coder failedRowsCoder, + Coder successfulRowsCoder, + Coder destinationCoder, + StorageApiDynamicDestinations dynamicDestinations, + StorageApiWriteUnshardedRecords.WriteRecordsDoFn writeDoFn, + TupleTag failedRowsTag, + @Nullable TupleTag successfulRowsTag) { + this.failedRowsCoder = failedRowsCoder; + this.successfulRowsCoder = successfulRowsCoder; + this.destinationCoder = destinationCoder; + this.dynamicDestinations = dynamicDestinations; + this.writeDoFn = writeDoFn; + this.failedRowsTag = failedRowsTag; + this.successfulRowsTag = successfulRowsTag; + } + + @Override + public PCollectionTuple expand(PCollection> input) { + // Append records to the Storage API streams. + TupleTagList tupleTagList = TupleTagList.of(failedRowsTag); + if (successfulRowsTag != null) { + tupleTagList = tupleTagList.and(successfulRowsTag); + } + + PCollectionTuple result = + input + .apply( + "addShard", + ParDo.of( + new DoFn< + KV, + KV, MismatchedRow>>() { + int shardNumber; + + @Setup + public void setup() { + shardNumber = ThreadLocalRandom.current().nextInt(NUM_DEFAULT_SHARDS); + } + + @ProcessElement + public void process( + @Element KV element, + OutputReceiver, MismatchedRow>> o) { + ByteBuffer buffer = ByteBuffer.allocate(Integer.BYTES); + buffer.putInt(++shardNumber % NUM_DEFAULT_SHARDS); + o.output( + KV.of( + ShardedKey.of(element.getKey(), buffer.array()), + element.getValue())); + } + })) + .setCoder(KvCoder.of(ShardedKey.Coder.of(destinationCoder), MismatchedRow.Coder.of())) + .apply( + "bufferMismatchedRows", + ParDo.of(new BufferingDoFn(writeDoFn)) + .withOutputTags(finalizeTag, tupleTagList) + .withSideInputs(dynamicDestinations.getSideInputs())); + + result.get(failedRowsTag).setCoder(failedRowsCoder); + if (successfulRowsTag != null) { + result.get(successfulRowsTag).setCoder(successfulRowsCoder); + } + return result; + } + + class BufferingDoFn + extends DoFn, MismatchedRow>, KV> { + private final StorageApiWriteUnshardedRecords.WriteRecordsDoFn + writeDoFn; + + @StateId("mismatchedRows") + private final StateSpec> mismatchedRowsSpec = + StateSpecs.bag(MismatchedRow.Coder.of()); + + @TimerId("retryMismatchedRowsTimer") + private final TimerSpec mismatchedRowsTimerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME); + + @StateId("currentMismatchedRowTimerValue") + private final StateSpec> currentMismatchedRowTimerValueSpec = + StateSpecs.value(); + + @StateId("minPendingTimestamp") + private final StateSpec> minPendingTimestampSpec = StateSpecs.value(); + + private final Counter rowsSentToFailedRowsCollection = + Metrics.counter(BufferMismatchedRows.BufferingDoFn.class, "rowsSentToFailedRowsCollection"); + + public BufferingDoFn( + StorageApiWriteUnshardedRecords.WriteRecordsDoFn writeDoFn) { + this.writeDoFn = writeDoFn; + } + + @StartBundle + public void startBundle() throws IOException { + writeDoFn.startBundle(); + } + + @ProcessElement + public void process( + ProcessContext processContext, + @Element KV, MismatchedRow> element, + @StateId("mismatchedRows") BagState mismatchedRowsBag, + @TimerId("retryMismatchedRowsTimer") Timer retryTimer, + @StateId("currentMismatchedRowTimerValue") ValueState currentTimerValue, + @StateId("minPendingTimestamp") ValueState minPendingTimestamp, + MultiOutputReceiver o) + throws Exception { + dynamicDestinations.setSideInputAccessorFromProcessContext(processContext); + TableDestination tableDestination = dynamicDestinations.getTable(element.getKey().getKey()); + + SchemaChangeDetectorHelper.bufferMismatchedRows( + Collections.singleton(element.getValue()), + mismatchedRowsBag, + retryTimer, + currentTimerValue, + minPendingTimestamp, + tableDestination, + o.get(failedRowsTag), + null, + rowsSentToFailedRowsCollection, + RETRY_MISMATCHED_ROWS_PERIOD); + } + + @OnTimer("retryMismatchedRowsTimer") + public void onTimer( + OnTimerContext context, + @Key ShardedKey shardedDestination, + @Timestamp Instant timestamp, + @StateId("mismatchedRows") BagState mismatchedRowsBag, + @StateId("currentMismatchedRowTimerValue") ValueState currentTimerValue, + @StateId("minPendingTimestamp") ValueState minPendingTimestamp, + @TimerId("retryMismatchedRowsTimer") Timer retryTimer, + PipelineOptions pipelineOptions, + MultiOutputReceiver o) + throws Exception { + dynamicDestinations.setSideInputAccessorFromOnTimerContext(context); + + mismatchedRowsBag.readLater(); + currentTimerValue.readLater(); + minPendingTimestamp.readLater(); + + List>> mismatchedRowsList = Lists.newArrayList(); + for (MismatchedRow row : mismatchedRowsBag.read()) { + Iterable> mismatchedRows = + writeDoFn.processElement( + pipelineOptions, + KV.of(shardedDestination.getKey(), row.getPayload()), + timestamp, + o); + if (!Iterables.isEmpty(mismatchedRows)) { + mismatchedRowsList.add(mismatchedRows); + } + } + // Once we're done, delegate to finishBundle to finish things. + Iterable> mismatchedDestRows = + writeDoFn.finishBundle( + o.get(failedRowsTag), + successfulRowsTag != null ? o.get(successfulRowsTag) : null, + o.get(finalizeTag)); + if (!Iterables.isEmpty(mismatchedDestRows)) { + mismatchedRowsList.add(mismatchedDestRows); + } + + mismatchedRowsBag.clear(); + if (!mismatchedRowsList.isEmpty()) { + TableDestination tableDestination = + dynamicDestinations.getTable(shardedDestination.getKey()); + StorageApiDynamicDestinations.MessageConverter messageConverter = + writeDoFn.messageConverters.get( + shardedDestination.getKey(), + dynamicDestinations, + pipelineOptions, + writeDoFn.getDatasetService(pipelineOptions), + writeDoFn.getWriteStreamService(pipelineOptions)); + messageConverter.updateSchemaFromTable(); + + AppendClientInfo appendClientInfo = + AppendClientInfo.of( + messageConverter.getTableSchema(), + messageConverter.getDescriptor(writeDoFn.usesCdc), + AutoCloseable::close); + + Iterable mismatchedRows = + () -> + StreamSupport.stream(Iterables.concat(mismatchedRowsList).spliterator(), false) + .map(KV::getValue) + .iterator(); + + currentTimerValue.clear(); + SchemaChangeDetectorHelper.bufferMismatchedRows( + mismatchedRows, + mismatchedRowsBag, + retryTimer, + currentTimerValue, + minPendingTimestamp, + tableDestination, + o.get(failedRowsTag), + appendClientInfo, + rowsSentToFailedRowsCollection, + RETRY_MISMATCHED_ROWS_PERIOD); + } + } + + @Teardown + public void onTeardown() { + writeDoFn.teardown(); + } + } +} diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java index 105da60c75b1..f40918391200 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinations.java @@ -82,6 +82,33 @@ interface SideInputAccessor { private transient @Nullable SideInputAccessor sideInputAccessor; private transient @Nullable PipelineOptions options; + static class SideInputAccessorViaOnTimerContext implements SideInputAccessor { + private DoFn.OnTimerContext onTimerContext; + + public SideInputAccessorViaOnTimerContext(DoFn.OnTimerContext onTimerContext) { + this.onTimerContext = onTimerContext; + } + + @Override + public SideInputT sideInput(PCollectionView view) { + return onTimerContext.sideInput(view); + } + } + + static class SideInputAccessorViaOnWindowExpirationContext implements SideInputAccessor { + private DoFn.OnWindowExpirationContext onWindowExpirationContext; + + public SideInputAccessorViaOnWindowExpirationContext( + DoFn.OnWindowExpirationContext onWindowExpirationContext) { + this.onWindowExpirationContext = onWindowExpirationContext; + } + + @Override + public SideInputT sideInput(PCollectionView view) { + return onWindowExpirationContext.sideInput(view); + } + } + static class SideInputAccessorViaProcessContext implements SideInputAccessor { private DoFn.ProcessContext processContext; @@ -129,6 +156,17 @@ void setSideInputAccessorFromProcessContext(DoFn.ProcessContext context) { this.options = context.getPipelineOptions(); } + void setSideInputAccessorFromOnTimerContext(DoFn.OnTimerContext context) { + this.sideInputAccessor = new SideInputAccessorViaOnTimerContext(context); + this.options = context.getPipelineOptions(); + } + + void setSideInputAccessorFromOnWindowExpirationContext( + DoFn.OnWindowExpirationContext context) { + this.sideInputAccessor = new SideInputAccessorViaOnWindowExpirationContext(context); + this.options = context.getPipelineOptions(); + } + /** * Returns an object that represents at a high level which table is being written to. May not * return null. diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java index 6370244c268c..14a7ef3de660 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java @@ -215,6 +215,19 @@ void setSideInputAccessorFromProcessContext(DoFn.ProcessContext context) { inner.setSideInputAccessorFromProcessContext(context); } + @Override + void setSideInputAccessorFromOnTimerContext(DoFn.OnTimerContext context) { + super.setSideInputAccessorFromOnTimerContext(context); + inner.setSideInputAccessorFromOnTimerContext(context); + } + + @Override + void setSideInputAccessorFromOnWindowExpirationContext( + DoFn.OnWindowExpirationContext context) { + super.setSideInputAccessorFromOnWindowExpirationContext(context); + inner.setSideInputAccessorFromOnWindowExpirationContext(context); + } + @Override public String toString() { return MoreObjects.toStringHelper(this).add("inner", inner).toString(); diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/MismatchedRow.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/MismatchedRow.java new file mode 100644 index 000000000000..168e4cc6705f --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/MismatchedRow.java @@ -0,0 +1,86 @@ +/* + * 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 com.google.auto.value.AutoValue; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import org.apache.beam.sdk.coders.ByteArrayCoder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.CustomCoder; +import org.apache.beam.sdk.coders.InstantCoder; +import org.apache.beam.sdk.coders.NullableCoder; +import org.apache.beam.sdk.schemas.AutoValueSchema; +import org.apache.beam.sdk.schemas.annotations.DefaultSchema; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +@DefaultSchema(AutoValueSchema.class) +@AutoValue +abstract class MismatchedRow { + abstract StorageApiWritePayload getPayload(); + + abstract org.joda.time.Instant getDeadline(); + + static MismatchedRow of(StorageApiWritePayload payload, org.joda.time.Instant deadline) { + return new AutoValue_MismatchedRow(payload, deadline); + } + + // Schemas give us a coder, however there are still some limitations to storing schema objects + // inside of state + // variables (mostly involving the Dataflow runner and update). Therefore we use a custom coder. + static class Coder extends CustomCoder { + static final ByteArrayCoder BYTE_ARRAY_CODER = ByteArrayCoder.of(); + static final InstantCoder INSTANT_CODER = InstantCoder.of(); + static final NullableCoder NULLABLE_BYTE_ARRAY_CODER = + NullableCoder.of(BYTE_ARRAY_CODER); + static final NullableCoder NULLABLE_INSTANT_CODER = + NullableCoder.of(InstantCoder.of()); + + static MismatchedRow.Coder of() { + return new Coder(); + } + + @Override + public void encode(MismatchedRow value, OutputStream outStream) + throws CoderException, IOException { + BYTE_ARRAY_CODER.encode(value.getPayload().getPayload(), outStream); + NULLABLE_BYTE_ARRAY_CODER.encode(value.getPayload().getUnknownFieldsPayload(), outStream); + NULLABLE_INSTANT_CODER.encode(value.getPayload().getTimestamp(), outStream); + NULLABLE_BYTE_ARRAY_CODER.encode(value.getPayload().getFailsafeTableRowPayload(), outStream); + NULLABLE_BYTE_ARRAY_CODER.encode(value.getPayload().getSchemaHash(), outStream); + INSTANT_CODER.encode(value.getDeadline(), outStream); + } + + @Override + public MismatchedRow decode(InputStream inStream) throws CoderException, IOException { + byte[] innerPayload = BYTE_ARRAY_CODER.decode(inStream); + byte @Nullable [] unknownFieldsPayload = NULLABLE_BYTE_ARRAY_CODER.decode(inStream); + @Nullable Instant timestamp = NULLABLE_INSTANT_CODER.decode(inStream); + byte @Nullable [] failsafeTableRowPayload = NULLABLE_BYTE_ARRAY_CODER.decode(inStream); + byte @Nullable [] schemaHash = NULLABLE_BYTE_ARRAY_CODER.decode(inStream); + Instant deadline = INSTANT_CODER.decode(inStream); + + StorageApiWritePayload payload = + StorageApiWritePayload.of( + innerPayload, timestamp, unknownFieldsPayload, failsafeTableRowPayload, schemaHash); + return new AutoValue_MismatchedRow(payload, deadline); + } + } +} diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelper.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelper.java new file mode 100644 index 000000000000..3defd8315be7 --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelper.java @@ -0,0 +1,217 @@ +/* + * 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 com.google.api.services.bigquery.model.TableReference; +import com.google.api.services.bigquery.model.TableRow; +import com.google.auto.value.AutoOneOf; +import com.google.cloud.bigquery.storage.v1.TableSchema; +import com.google.protobuf.ByteString; +import com.google.protobuf.Descriptors; +import java.io.IOException; +import java.util.Optional; +import org.apache.beam.sdk.metrics.Counter; +import org.apache.beam.sdk.state.BagState; +import org.apache.beam.sdk.state.Timer; +import org.apache.beam.sdk.state.ValueState; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.Preconditions; +import org.apache.beam.sdk.util.ThrowingSupplier; +import org.apache.beam.sdk.values.TimestampedValue; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Predicates; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.joda.time.Instant; + +class SchemaChangeDetectorHelper { + @AutoOneOf(MergePayloadResult.Kind.class) + abstract static class MergePayloadResult { + enum Kind { + MERGED, + FAILED + }; + + abstract Kind getKind(); + + abstract ByteString getMerged(); + + abstract TimestampedValue getFailed(); + + static MergePayloadResult ofMerged(ByteString merged) { + return AutoOneOf_SchemaChangeDetectorHelper_MergePayloadResult.merged(merged); + } + + static MergePayloadResult ofFailed(TimestampedValue failed) { + return AutoOneOf_SchemaChangeDetectorHelper_MergePayloadResult.failed(failed); + } + } + + private final boolean autoUpdateSchema; + private final boolean ignoreUnknownValues; + private final TableReference tableReference; + private final boolean ignoreSchemaHashes; + + public SchemaChangeDetectorHelper( + boolean autoUpdateSchema, + boolean ignoreUnknownValues, + TableReference tableReference, + boolean ignoreSchemaHashes) { + this.autoUpdateSchema = autoUpdateSchema; + this.ignoreUnknownValues = ignoreUnknownValues; + this.tableReference = tableReference; + this.ignoreSchemaHashes = ignoreSchemaHashes; + } + + Optional checkUpdatedSchema( + TableSchema currentSchema, + String streamName, + BigQueryServices.WriteStreamService writeStreamService) { + if (autoUpdateSchema) { + @Nullable TableSchema streamSchema = writeStreamService.getWriteStreamSchema(streamName); + if (streamSchema != null) { + return TableSchemaUpdateUtils.getUpdatedSchema(currentSchema, streamSchema); + } + } + return Optional.empty(); + } + + Optional checkResponseForUpdatedSchema( + TableSchema currentSchema, BigQueryServices.StreamAppendClient streamAppendClient) { + @Nullable TableSchema updatedSchemaReturned = streamAppendClient.getUpdatedSchema(); + + // Update the table schema and clear the append client. + if (updatedSchemaReturned != null) { + return TableSchemaUpdateUtils.getUpdatedSchema(currentSchema, updatedSchemaReturned); + } + return Optional.empty(); + } + + MergePayloadResult getMergedPayload( + StorageApiWritePayload payload, + Instant elementTimestamp, + @Nullable TableRow failsafeTableRow, + AppendClientInfo appendClientInfo) + throws IOException { + ByteString byteString = ByteString.copyFrom(payload.getPayload()); + + if (autoUpdateSchema) { + @Nullable TableRow unknownFields = payload.getUnknownFields(); + + if (unknownFields != null && !unknownFields.isEmpty()) { + try { + // Protocol buffer serialization format supports concatenation. We serialize any new + // "known" fields into a proto and concatenate to the existing proto. + byteString = + Preconditions.checkStateNotNull(appendClientInfo) + .mergeNewFields(byteString, unknownFields, ignoreUnknownValues); + } catch (TableRowToStorageApiProto.SchemaConversionException e) { + // This generally implies that ignoreUnknownValues=false and there were still + // unknown values here. + // Reconstitute the TableRow and send it to the failed-rows consumer. + TableRow tableRow = + failsafeTableRow != null + ? failsafeTableRow + : appendClientInfo.toTableRow(byteString, Predicates.alwaysTrue()); + BigQueryStorageApiInsertError error = + new BigQueryStorageApiInsertError(tableRow, e.toString(), tableReference); + Instant timestamp = MoreObjects.firstNonNull(payload.getTimestamp(), elementTimestamp); + return MergePayloadResult.ofFailed(TimestampedValue.of(error, timestamp)); + } + } + } + return MergePayloadResult.ofMerged(byteString); + } + + boolean isPayloadSchemaOutOfDate( + StorageApiWritePayload payload, + byte[] mergedPayload, + ThrowingSupplier schemaHash, + ThrowingSupplier schemaDescriptor) + throws Exception { + if (ignoreSchemaHashes) { + + if (payload.getUnknownFields() == null) { + return false; + } + // TODO: CHECK FOR REQUIRED FIELDS AS WELL. + return UpgradeTableSchema.missingUnknownField( + Preconditions.checkStateNotNull(payload.getUnknownFields()), schemaDescriptor); + } else { + return UpgradeTableSchema.isPayloadSchemaOutOfDate( + payload.getSchemaHash(), () -> mergedPayload, schemaHash, schemaDescriptor); + } + } + + static void bufferMismatchedRows( + Iterable rows, + BagState bufferedBag, + Timer retryRowsTimers, + ValueState currentTimerValue, + ValueState minPendingTimestamp, + TableDestination tableDestination, + DoFn.OutputReceiver failedRowsReceiver, + @Nullable AppendClientInfo appendClientInfo, + Counter rowsSentToFailedRowsCollection, + Duration retryPeriod) + throws IOException { + org.joda.time.Instant minTimestamp = + org.joda.time.Instant.ofEpochMilli( + MoreObjects.firstNonNull( + minPendingTimestamp.read(), BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis())); + + for (MismatchedRow row : rows) { + if (appendClientInfo == null + || row.getDeadline().isAfter(retryRowsTimers.getCurrentRelativeTime())) { + bufferedBag.add(row); + org.joda.time.@Nullable Instant timestamp = row.getPayload().getTimestamp(); + if (timestamp != null && timestamp.isBefore(minTimestamp)) { + minTimestamp = timestamp; + } + } else { + @Nullable TableRow failedRow = row.getPayload().getFailsafeTableRow(); + if (failedRow == null) { + ByteString rowBytes = ByteString.copyFrom(row.getPayload().getPayload()); + failedRow = appendClientInfo.toTableRow(rowBytes, Predicates.alwaysTrue()); + } + BigQueryStorageApiInsertError error = + new BigQueryStorageApiInsertError( + failedRow, "Mismatched schema", tableDestination.getTableReference()); + failedRowsReceiver.outputWithTimestamp( + error, Preconditions.checkStateNotNull(row.getPayload().getTimestamp())); + rowsSentToFailedRowsCollection.inc(); + BigQuerySinkMetrics.appendRowsRowStatusCounter( + BigQuerySinkMetrics.RowStatus.FAILED, + BigQuerySinkMetrics.SCHEMA_MISMATCHED, + tableDestination.getShortTableUrn()) + .inc(1); + } + } + minPendingTimestamp.write(minTimestamp.getMillis()); + + long targetTs = + MoreObjects.firstNonNull( + currentTimerValue.read(), + retryRowsTimers.getCurrentRelativeTime().getMillis() + retryPeriod.getMillis()); + retryRowsTimers + .withOutputTimestamp(minTimestamp) + .set(org.joda.time.Instant.ofEpochMilli(targetTs)); + currentTimerValue.write(targetTs); + } +} diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaUpdateHoldingFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaUpdateHoldingFn.java index b65b11ef3194..fa37327ac843 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaUpdateHoldingFn.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaUpdateHoldingFn.java @@ -165,6 +165,7 @@ public Duration getAllowedTimestampSkew() { @OnTimer("pollTimer") public void onPollTimer( + OnTimerContext context, @Key ShardedKey key, PipelineOptions pipelineOptions, @StateId("bufferedElements") BagState> bag, @@ -174,6 +175,8 @@ public void onPollTimer( BoundedWindow window, MultiOutputReceiver o) throws Exception { + convertMessagesDoFn.getDynamicDestinations().setSideInputAccessorFromOnTimerContext(context); + if (tryFlushBuffer(key.getKey(), pipelineOptions, bag, minBufferedTimestamp, o)) { timerTs.clear(); } else { @@ -188,12 +191,19 @@ public void onPollTimer( @OnWindowExpiration public void onWindowExpiration( + // OnWindowExpirationContext context, @Key ShardedKey key, PipelineOptions pipelineOptions, @StateId("bufferedElements") BagState> bag, @StateId("minBufferedTimestamp") CombiningState minBufferedTimestamp, MultiOutputReceiver o) throws Exception { + // TODO: Beam doesn't current support OnWindowExpirationContext. Reenable after this is fixed. + // https://github.com/apache/beam/issues/38875 + // convertMessagesDoFn + // .getDynamicDestinations() + // .setSideInputAccessorFromOnWindowExpirationContext(context); + // This can happen on test completion or drain. We can't set any more timers in window // expiration, so we just have to loop until the schema is updated. BackOff backoff = diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterable.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterable.java index 26eb6d55d208..3ef7309d8aee 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterable.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterable.java @@ -17,23 +17,15 @@ */ package org.apache.beam.sdk.io.gcp.bigquery; -import static org.apache.beam.sdk.io.gcp.bigquery.UpgradeTableSchema.isPayloadSchemaOutOfDate; - -import com.google.api.services.bigquery.model.TableRow; -import com.google.auto.value.AutoValue; -import com.google.cloud.bigquery.storage.v1.ProtoRows; -import com.google.protobuf.ByteString; import com.google.protobuf.Descriptors; -import java.io.IOException; import java.util.Iterator; -import java.util.List; import java.util.NoSuchElementException; -import java.util.function.BiConsumer; -import java.util.function.Function; +import java.util.function.Consumer; +import java.util.function.Supplier; +import org.apache.beam.sdk.util.Preconditions; import org.apache.beam.sdk.util.ThrowingSupplier; import org.apache.beam.sdk.values.TimestampedValue; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterators; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.PeekingIterator; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Instant; @@ -43,147 +35,73 @@ * parameter controls how many rows are batched into a single ProtoRows object before we move on to * the next one. */ -class SplittingIterable implements Iterable { - @AutoValue - abstract static class Value { - abstract ProtoRows getProtoRows(); - - abstract List getTimestamps(); - - abstract List<@Nullable TableRow> getFailsafeTableRows(); - - abstract boolean getSchemaMismatchSeen(); - } - - interface ConcatFields { - ByteString concat(ByteString bytes, TableRow tableRows) - throws TableRowToStorageApiProto.SchemaConversionException; - } - +class SplittingIterable implements Iterable { private final Iterable underlying; private final long splitSize; - private final ConcatFields concatProtoAndTableRow; - private final Function protoToTableRow; - private final BiConsumer, String> failedRowsConsumer; + private final Consumer> failedRowsConsumer; private final ThrowingSupplier getCurrentTableSchemaHash; private final ThrowingSupplier getCurrentTableSchemaDescriptor; - private final boolean autoUpdateSchema; - private final Instant elementsTimestamp; + private final Instant elementTimestamp; + private final Supplier appendClientSupplier; + private final SchemaChangeDetectorHelper schemaChangeDetectorHelper; public SplittingIterable( Iterable underlying, long splitSize, - ConcatFields concatProtoAndTableRow, - Function protoToTableRow, - BiConsumer, String> failedRowsConsumer, + Consumer> failedRowsConsumer, ThrowingSupplier getCurrentTableSchemaHash, ThrowingSupplier getCurrentTableSchemaDescriptor, - boolean autoUpdateSchema, - Instant elementsTimestamp) { + Instant elementTimestamp, + Supplier appendClientSupplier, + SchemaChangeDetectorHelper schemaChangeDetectorHelper) { this.underlying = underlying; this.splitSize = splitSize; - this.concatProtoAndTableRow = concatProtoAndTableRow; - this.protoToTableRow = protoToTableRow; this.failedRowsConsumer = failedRowsConsumer; this.getCurrentTableSchemaHash = getCurrentTableSchemaHash; this.getCurrentTableSchemaDescriptor = getCurrentTableSchemaDescriptor; - this.autoUpdateSchema = autoUpdateSchema; - this.elementsTimestamp = elementsTimestamp; + this.elementTimestamp = elementTimestamp; + this.appendClientSupplier = appendClientSupplier; + this.schemaChangeDetectorHelper = schemaChangeDetectorHelper; } @Override - public Iterator iterator() { - return new Iterator() { + public Iterator iterator() { + return new Iterator() { final PeekingIterator underlyingIterator = Iterators.peekingIterator(underlying.iterator()); + @Nullable AppendRowsPacket cachedNext = null; @Override public boolean hasNext() { - return underlyingIterator.hasNext(); + if (cachedNext == null && underlyingIterator.hasNext()) { + cachedNext = calculateNext(); + } + return cachedNext != null; } @Override - public Value next() { + public AppendRowsPacket next() { if (!hasNext()) { throw new NoSuchElementException(); } + AppendRowsPacket result = Preconditions.checkStateNotNull(cachedNext); + cachedNext = null; + return result; + } - List timestamps = Lists.newArrayList(); - List<@Nullable TableRow> failsafeRows = Lists.newArrayList(); - ProtoRows.Builder inserts = ProtoRows.newBuilder(); - long bytesSize = 0; - - boolean schemaMismatchSeen = false; - try { - while (underlyingIterator.hasNext()) { - // Make sure that we don't exceed the split-size length over multiple elements. A single - // element can exceed - // the split threshold, but in that case it should be the only element returned. - if ((bytesSize + underlyingIterator.peek().getPayload().length > splitSize) - && inserts.getSerializedRowsCount() > 0) { - break; - } - StorageApiWritePayload payload = underlyingIterator.next(); - schemaMismatchSeen = - schemaMismatchSeen - || isPayloadSchemaOutOfDate( - payload, getCurrentTableSchemaHash, getCurrentTableSchemaDescriptor); - - ByteString byteString = ByteString.copyFrom(payload.getPayload()); - @Nullable TableRow failsafeTableRow = null; - try { - failsafeTableRow = payload.getFailsafeTableRow(); - } catch (IOException e) { - // Do nothing, table row will be generated later from row bytes - } - if (autoUpdateSchema) { - @Nullable TableRow unknownFields = payload.getUnknownFields(); - if (unknownFields != null && !unknownFields.isEmpty()) { - // Protocol buffer serialization format supports concatenation. We serialize any new - // "known" fields - // into a proto and concatenate to the existing proto. - - try { - byteString = concatProtoAndTableRow.concat(byteString, unknownFields); - } catch (TableRowToStorageApiProto.SchemaConversionException e) { - // This generally implies that ignoreUnknownValues=false and there were still - // unknown values here. - // Reconstitute the TableRow and send it to the failed-rows consumer. - TableRow tableRow = - failsafeTableRow != null - ? failsafeTableRow - : protoToTableRow.apply(byteString); - // TODO(24926, reuvenlax): We need to merge the unknown fields in! Currently we - // only execute this - // codepath when ignoreUnknownFields==true, so we should never hit this codepath. - // However once - // 24926 is fixed, we need to merge the unknownFields back into the main row - // before outputting to the - // failed-rows consumer. - Instant timestamp = payload.getTimestamp(); - if (timestamp == null) { - timestamp = elementsTimestamp; - } - failedRowsConsumer.accept(TimestampedValue.of(tableRow, timestamp), e.toString()); - continue; - } - } - } - inserts.addSerializedRows(byteString); - Instant timestamp = payload.getTimestamp(); - if (timestamp == null) { - timestamp = elementsTimestamp; - } - timestamps.add(timestamp); - failsafeRows.add(failsafeTableRow); - bytesSize += byteString.size(); - } - } catch (Exception e) { - throw new RuntimeException(e); - } - return new AutoValue_SplittingIterable_Value( - inserts.build(), timestamps, failsafeRows, schemaMismatchSeen); + public @Nullable AppendRowsPacket calculateNext() { + AppendRowsPacket value = + AppendRowsPacket.fromStorageApiWritePayload( + underlyingIterator, + splitSize, + schemaChangeDetectorHelper, + elementTimestamp, + appendClientSupplier, + failedRowsConsumer, + getCurrentTableSchemaHash, + getCurrentTableSchemaDescriptor); + return value.getProtoRows().getSerializedRowsCount() == 0 ? null : value; } }; } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinations.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinations.java index d6981858a4b6..1c21260b864c 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinations.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinations.java @@ -18,6 +18,7 @@ package org.apache.beam.sdk.io.gcp.bigquery; import com.google.api.services.bigquery.model.TableRow; +import com.google.cloud.bigquery.storage.v1.TableSchema; import com.google.protobuf.DescriptorProtos; import java.io.IOException; import javax.annotation.Nullable; @@ -29,7 +30,7 @@ abstract class StorageApiDynamicDestinations extends DynamicDestinationsHelpers.DelegatingDynamicDestinations { public interface MessageConverter { - com.google.cloud.bigquery.storage.v1.TableSchema getTableSchema(); + TableSchema getTableSchema(); DescriptorProtos.DescriptorProto getDescriptor(boolean includeCdcColumns) throws Exception; @@ -42,6 +43,8 @@ StorageApiWritePayload toMessage( TableRow toFailsafeTableRow(T element); void updateSchemaFromTable() throws IOException, InterruptedException; + + default void updateSchema(TableSchema schema) {} } StorageApiDynamicDestinations(DynamicDestinations inner) { @@ -60,4 +63,17 @@ void setSideInputAccessorFromProcessContext(DoFn.ProcessContext context) { super.setSideInputAccessorFromProcessContext(context); inner.setSideInputAccessorFromProcessContext(context); } + + @Override + void setSideInputAccessorFromOnTimerContext(DoFn.OnTimerContext context) { + super.setSideInputAccessorFromOnTimerContext(context); + inner.setSideInputAccessorFromOnTimerContext(context); + } + + @Override + void setSideInputAccessorFromOnWindowExpirationContext( + DoFn.OnWindowExpirationContext context) { + super.setSideInputAccessorFromOnWindowExpirationContext(context); + inner.setSideInputAccessorFromOnWindowExpirationContext(context); + } } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsTableRow.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsTableRow.java index 1710d32689c9..ac6aaa2846fb 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsTableRow.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiDynamicDestinationsTableRow.java @@ -24,7 +24,6 @@ import com.google.protobuf.Descriptors.Descriptor; import com.google.protobuf.Message; import java.io.IOException; -import java.util.Set; import java.util.concurrent.ExecutionException; import java.util.concurrent.atomic.AtomicReference; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.CreateDisposition; @@ -48,7 +47,7 @@ public class StorageApiDynamicDestinationsTableRow schemaUpdateOptions; + private final boolean useSchemaUpdatingTableRow; private static final TableSchemaCache SCHEMA_CACHE = new TableSchemaCache(Duration.standardSeconds(1)); @@ -64,7 +63,7 @@ public class StorageApiDynamicDestinationsTableRow schemaUpdateOptions) { + boolean useSchemaUpdatingTableRow) { super(inner); this.formatFunction = formatFunction; this.formatRecordOnFailureFunction = formatRecordOnFailureFunction; @@ -72,7 +71,7 @@ public class StorageApiDynamicDestinationsTableRow getMessageConverter( } }; - return schemaUpdateOptions.isEmpty() + return !useSchemaUpdatingTableRow ? getConverter.apply(getSchema(destination)) : new SchemaUpgradingTableRowConverter( getConverter, options, datasetService, writeStreamService); @@ -143,6 +142,7 @@ public StorageApiWritePayload toMessage( converter.toMessage(element, rowMutationInformation, collectedExceptions); // Set the schema hash on the payload so the next transform knows whether it has an // out-of-date schema. + // TODO: Don't set this for autoUpdateSchema, as it's ignored. payload = payload.toBuilder().setSchemaHash(converter.getSchemaHash()).build(); return payload; @@ -153,9 +153,16 @@ public void updateSchemaFromTable() throws IOException, InterruptedException { SCHEMA_CACHE.refreshSchema( delegate.get().tableReference, datasetService, writeStreamService, bigQueryOptions); // Recycle the internal MessageConverter so that we pick up the new schema from the cache. + // TODO: only do this if the schema has changed. this.delegate.set(getConverter.apply(null)); } + @Override + public void updateSchema(com.google.cloud.bigquery.storage.v1.TableSchema schema) { + TableSchema modelTableSchema = TableRowToStorageApiProto.protoSchemaToTableSchema(schema); + this.delegate.set(getConverter.apply(modelTableSchema)); + } + @Override public TableRow toFailsafeTableRow(T element) { return delegate.get().toFailsafeTableRow(element); diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiLoads.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiLoads.java index 007bba5c6cdf..d301c33ff529 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiLoads.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiLoads.java @@ -75,6 +75,7 @@ public class StorageApiLoads private final boolean allowInconsistentWrites; private final boolean allowAutosharding; private final boolean autoUpdateSchema; + private final @Nullable Duration autoUpdateSchemaStrictTimeout; private final boolean ignoreUnknownValues; private final boolean usesCdc; @@ -99,6 +100,7 @@ public StorageApiLoads( boolean allowInconsistentWrites, boolean allowAutosharding, boolean autoUpdateSchema, + @Nullable Duration autoUpdateSchemaStrictTimeout, boolean ignoreUnknownValues, boolean propagateSuccessfulStorageApiWrites, Predicate propagateSuccessfulStorageApiWritesPredicate, @@ -120,6 +122,7 @@ public StorageApiLoads( this.allowInconsistentWrites = allowInconsistentWrites; this.allowAutosharding = allowAutosharding; this.autoUpdateSchema = autoUpdateSchema; + this.autoUpdateSchemaStrictTimeout = autoUpdateSchemaStrictTimeout; this.ignoreUnknownValues = ignoreUnknownValues; if (propagateSuccessfulStorageApiWrites) { this.successfulWrittenRowsTag = new TupleTag<>("successfulPublishedRowsTag"); @@ -195,7 +198,9 @@ public WriteResult expandInconsistent( successfulRowsPredicate, BigQueryStorageApiInsertErrorCoder.of(), TableRowJsonCoder.of(), + destinationCoder, autoUpdateSchema, + autoUpdateSchemaStrictTimeout, ignoreUnknownValues, createDisposition, kmsKey, @@ -296,6 +301,7 @@ public WriteResult expandTriggered( successfulWrittenRowsTag, successfulRowsPredicate, autoUpdateSchema, + autoUpdateSchemaStrictTimeout, ignoreUnknownValues, defaultMissingValueInterpretation, bigLakeConfiguration)); @@ -399,6 +405,7 @@ public WriteResult expandUntriggered( BigQueryStorageApiInsertErrorCoder.of(), TableRowJsonCoder.of(), autoUpdateSchema, + autoUpdateSchemaStrictTimeout, ignoreUnknownValues, createDisposition, kmsKey, diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritePayload.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritePayload.java index 3f00ed67a8a5..280906d08975 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritePayload.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritePayload.java @@ -83,6 +83,22 @@ static StorageApiWritePayload of( .build(); } + static StorageApiWritePayload of( + byte[] payload, + @Nullable Instant timestamp, + @Nullable byte[] unknownFieldsPayload, + @Nullable byte[] failsafeTableRowPayload, + @Nullable byte[] schemaHash) + throws IOException { + return new AutoValue_StorageApiWritePayload.Builder() + .setPayload(payload) + .setTimestamp(timestamp) + .setUnknownFieldsPayload(unknownFieldsPayload) + .setFailsafeTableRowPayload(failsafeTableRowPayload) + .setSchemaHash(schemaHash) + .build(); + } + public StorageApiWritePayload withTimestamp(Instant instant) { return toBuilder().setTimestamp(instant).build(); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRecordsInconsistent.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRecordsInconsistent.java index 58bbed8ba5a9..79643f0653f4 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRecordsInconsistent.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRecordsInconsistent.java @@ -23,13 +23,18 @@ import java.util.function.Predicate; import javax.annotation.Nullable; import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.transforms.Flatten; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; +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.sdk.values.TupleTagList; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.joda.time.Duration; /** * A transform to write sharded records to BigQuery using the Storage API. This transform uses the @@ -38,7 +43,7 @@ * writes, use {@link StorageApiWritesShardedRecords} or {@link StorageApiWriteUnshardedRecords}. */ @SuppressWarnings("FutureReturnValueIgnored") -public class StorageApiWriteRecordsInconsistent +public class StorageApiWriteRecordsInconsistent extends PTransform>, PCollectionTuple> { private final StorageApiDynamicDestinations dynamicDestinations; private final BigQueryServices bqServices; @@ -47,10 +52,15 @@ public class StorageApiWriteRecordsInconsistent private final Predicate successfulRowsPredicate; + // This output is effectively ignored, since this code path uses the default stream which is never + // finalized. private final TupleTag> finalizeTag = new TupleTag<>("finalizeTag"); private final Coder failedRowsCoder; private final Coder successfulRowsCoder; + private final Coder destinationCoder; private final boolean autoUpdateSchema; + private final @Nullable Duration autoUpdateSchemaStrictTimeout; + private final @Nullable TupleTag> mismatchedRowsTag; private final boolean ignoreUnknownValues; private final BigQueryIO.Write.CreateDisposition createDisposition; private final @Nullable String kmsKey; @@ -66,7 +76,9 @@ public StorageApiWriteRecordsInconsistent( Predicate successfulRowsPredicate, Coder failedRowsCoder, Coder successfulRowsCoder, + Coder destinationCoder, boolean autoUpdateSchema, + @Nullable Duration autoUpdateSchemaStrictTimeout, boolean ignoreUnknownValues, BigQueryIO.Write.CreateDisposition createDisposition, @Nullable String kmsKey, @@ -78,9 +90,16 @@ public StorageApiWriteRecordsInconsistent( this.failedRowsTag = failedRowsTag; this.failedRowsCoder = failedRowsCoder; this.successfulRowsCoder = successfulRowsCoder; + this.destinationCoder = destinationCoder; this.successfulRowsTag = successfulRowsTag; this.successfulRowsPredicate = successfulRowsPredicate; this.autoUpdateSchema = autoUpdateSchema; + this.autoUpdateSchemaStrictTimeout = autoUpdateSchemaStrictTimeout; + if (autoUpdateSchemaStrictTimeout != null) { + this.mismatchedRowsTag = new TupleTag<>("mismatchedRowsTag"); + } else { + this.mismatchedRowsTag = null; + } this.ignoreUnknownValues = ignoreUnknownValues; this.createDisposition = createDisposition; this.kmsKey = kmsKey; @@ -98,36 +117,83 @@ public PCollectionTuple expand(PCollection fn = + new StorageApiWriteUnshardedRecords.WriteRecordsDoFn<>( + operationName, + dynamicDestinations, + bqServices, + true, + bigQueryOptions.getStorageApiAppendThresholdBytes(), + bigQueryOptions.getStorageApiAppendThresholdRecordCount(), + bigQueryOptions.getNumStorageWriteApiStreamAppendClients(), + finalizeTag, + failedRowsTag, + successfulRowsTag, + successfulRowsPredicate, + autoUpdateSchema, + autoUpdateSchemaStrictTimeout, + mismatchedRowsTag, + ignoreUnknownValues, + createDisposition, + kmsKey, + usesCdc, + defaultMissingValueInterpretation, + bigQueryOptions.getStorageWriteApiMaxRetries(), + bigLakeConfiguration, + bigQueryOptions.getStorageApiInitialMismatchRetryTimeMilliSec()); PCollectionTuple result = input.apply( "Write Records", - ParDo.of( - new StorageApiWriteUnshardedRecords.WriteRecordsDoFn<>( - operationName, - dynamicDestinations, - bqServices, - true, - bigQueryOptions.getStorageApiAppendThresholdBytes(), - bigQueryOptions.getStorageApiAppendThresholdRecordCount(), - bigQueryOptions.getNumStorageWriteApiStreamAppendClients(), - finalizeTag, - failedRowsTag, - successfulRowsTag, - successfulRowsPredicate, - autoUpdateSchema, - ignoreUnknownValues, - createDisposition, - kmsKey, - usesCdc, - defaultMissingValueInterpretation, - bigQueryOptions.getStorageWriteApiMaxRetries(), - bigLakeConfiguration)) + ParDo.of(fn) .withOutputTags(finalizeTag, tupleTagList) .withSideInputs(dynamicDestinations.getSideInputs())); result.get(failedRowsTag).setCoder(failedRowsCoder); if (successfulRowsTag != null) { result.get(successfulRowsTag).setCoder(successfulRowsCoder); } - return result; + + @Nullable PCollectionTuple mismatchedResult = null; + if (mismatchedRowsTag != null) { + PCollection> mismatchedRows = result.get(mismatchedRowsTag); + mismatchedRows.setCoder(KvCoder.of(destinationCoder, MismatchedRow.Coder.of())); + mismatchedResult = + mismatchedRows.apply( + "bufferMismatched", + new BufferMismatchedRows<>( + failedRowsCoder, + successfulRowsCoder, + destinationCoder, + dynamicDestinations, + fn, + failedRowsTag, + successfulRowsTag)); + mismatchedResult.get(failedRowsTag).setCoder(failedRowsCoder); + if (successfulRowsTag != null) { + mismatchedResult.get(successfulRowsTag).setCoder(successfulRowsCoder); + } + } + + if (mismatchedResult != null) { + // We need to merge in any results from BufferMismatchRows. + PCollection flattenedErrors = + PCollectionList.of(result.get(failedRowsTag)) + .and(mismatchedResult.get(failedRowsTag)) + .apply("flattenErrors", Flatten.pCollections()); + + PCollectionTuple flattenedResult = PCollectionTuple.of(failedRowsTag, flattenedErrors); + if (successfulRowsTag != null) { + PCollection flattenedSuccesses = + PCollectionList.of(result.get(successfulRowsTag)) + .and(mismatchedResult.get(successfulRowsTag)) + .apply("flattenSucesses", Flatten.pCollections()); + flattenedResult = flattenedResult.and(successfulRowsTag, flattenedSuccesses); + } + return flattenedResult; + } else { + return result; + } } } 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 2dfc8b2f1c00..e031bc00fe00 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 @@ -17,9 +17,7 @@ */ package org.apache.beam.sdk.io.gcp.bigquery; -import static org.apache.beam.sdk.io.gcp.bigquery.UpgradeTableSchema.isPayloadSchemaOutOfDate; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; -import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import com.google.api.core.ApiFuture; @@ -35,11 +33,11 @@ import com.google.protobuf.DescriptorProtos; import com.google.protobuf.Descriptors.Descriptor; import com.google.protobuf.Descriptors.DescriptorValidationException; -import com.google.protobuf.DynamicMessage; import io.grpc.Status; import java.io.IOException; import java.time.Instant; import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Objects; @@ -49,7 +47,10 @@ import java.util.concurrent.Callable; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BiConsumer; +import java.util.function.Consumer; import java.util.function.Predicate; +import java.util.function.Supplier; import java.util.stream.Collectors; import java.util.stream.StreamSupport; import org.apache.beam.sdk.coders.Coder; @@ -81,14 +82,18 @@ import org.apache.beam.sdk.values.OutputBuilder; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.TimestampedValue; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.TupleTagList; import org.apache.beam.sdk.values.WindowedValues; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Predicates; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterators; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.PeekingIterator; import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; @@ -138,6 +143,7 @@ public StorageApiWriteUnshardedRecords( Coder failedRowsCoder, Coder successfulRowsCoder, boolean autoUpdateSchema, + @Nullable Duration autoUpdateSchemaStrictTimeout, boolean ignoreUnknownValues, BigQueryIO.Write.CreateDisposition createDisposition, @Nullable String kmsKey, @@ -152,6 +158,7 @@ public StorageApiWriteUnshardedRecords( this.failedRowsCoder = failedRowsCoder; this.successfulRowsCoder = successfulRowsCoder; this.autoUpdateSchema = autoUpdateSchema; + checkState(autoUpdateSchemaStrictTimeout == null, "Not supported in this configuration"); this.ignoreUnknownValues = ignoreUnknownValues; this.createDisposition = createDisposition; this.kmsKey = kmsKey; @@ -188,13 +195,16 @@ public PCollectionTuple expand(PCollection private final Predicate successfulRowsPredicate; private final boolean autoUpdateSchema; + private final @Nullable Duration autoUpdateSchemaStrictTimeout; + private final @Nullable TupleTag> mismatchedRowsTag; private final boolean ignoreUnknownValues; private final BigQueryIO.Write.CreateDisposition createDisposition; private final @Nullable String kmsKey; - private final boolean usesCdc; + final boolean usesCdc; private final AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation; static class AppendRowsContext extends RetryManager.Operation.Context { @@ -250,6 +262,7 @@ public AppendRowsContext( } class DestinationState { + private final DestinationT destination; private final TableDestination tableDestination; private final String tableUrn; private final String shortTableUrn; @@ -257,7 +270,6 @@ class DestinationState { private @Nullable AppendClientInfo appendClientInfo = null; private long currentOffset = 0; private List pendingMessages; - private List pendingTimestamps; private transient @Nullable WriteStreamService maybeWriteStreamService; private final Counter recordsAppended = Metrics.counter(WriteRecordsDoFn.class, "recordsAppended"); @@ -279,8 +291,10 @@ class DestinationState { private final long maxRequestSize; private final boolean includeCdcColumns; + private final SchemaChangeDetectorHelper schemaChangeDetectorHelper; public DestinationState( + DestinationT destination, TableDestination tableDestination, String tableUrn, String shortTableUrn, @@ -293,11 +307,11 @@ public DestinationState( Callable tryCreateTable, boolean includeCdcColumns) throws Exception { + this.destination = destination; this.tableDestination = tableDestination; this.tableUrn = tableUrn; this.shortTableUrn = shortTableUrn; this.pendingMessages = Lists.newArrayList(); - this.pendingTimestamps = Lists.newArrayList(); this.maybeWriteStreamService = writeStreamService; this.useDefaultStream = useDefaultStream; this.messageConverter = messageConverter; @@ -309,6 +323,12 @@ public DestinationState( if (includeCdcColumns) { checkState(useDefaultStream); } + this.schemaChangeDetectorHelper = + new SchemaChangeDetectorHelper( + autoUpdateSchema, + ignoreUnknownValues, + tableDestination.getTableReference(), + autoUpdateSchemaStrictTimeout != null); } public TableDestination getTableDestination() { @@ -409,22 +429,15 @@ SchemaAndDescriptor getCurrentTableSchema(String stream, @Nullable TableSchema u AtomicBoolean updated = new AtomicBoolean(); CreateTableHelpers.createTableWrapper( () -> { - if (autoUpdateSchema) { - @Nullable - TableSchema streamSchema = - Preconditions.checkStateNotNull(maybeWriteStreamService) - .getWriteStreamSchema(streamName); - if (streamSchema != null) { - Optional newSchema = - TableSchemaUpdateUtils.getUpdatedSchema( - messageConverter.getTableSchema(), streamSchema); - if (newSchema.isPresent()) { - currentSchema.set(newSchema.get()); - updated.set(true); - LOG.debug( - "Fetched updated schema for table {}:\n\t{}", tableUrn, newSchema.get()); - } - } + Optional newSchema = + schemaChangeDetectorHelper.checkUpdatedSchema( + messageConverter.getTableSchema(), + streamName, + Preconditions.checkStateNotNull(maybeWriteStreamService)); + if (newSchema.isPresent()) { + currentSchema.set(newSchema.get()); + updated.set(true); + LOG.debug("Fetched updated schema for table {}:\n\t{}", tableUrn, newSchema.get()); } return null; }, @@ -484,76 +497,21 @@ void invalidateAppendClient(boolean invalidateCache) { } } - void addMessage( - StorageApiWritePayload payload, - org.joda.time.Instant elementTs, - OutputReceiver failedRowsReceiver) - throws Exception { + void addMessage(StorageApiWritePayload payload) throws Exception { maybeTickleCache(); - - @Nullable ByteString mergedPayloadBytes = null; - if (autoUpdateSchema) { - if (appendClientInfo == null) { - appendClientInfo = getAppendClientInfo(true, null); - } - @Nullable TableRow unknownFields = payload.getUnknownFields(); - if (unknownFields != null && !unknownFields.isEmpty()) { - // check if unknownFields contains repeated struct, merge - // otherwise use concat - mergedPayloadBytes = ByteString.copyFrom(payload.getPayload()); - try { - mergedPayloadBytes = - Preconditions.checkStateNotNull(appendClientInfo) - .mergeNewFields(mergedPayloadBytes, unknownFields, ignoreUnknownValues); - } catch (TableRowToStorageApiProto.SchemaConversionException e) { - @Nullable TableRow tableRow = payload.getFailsafeTableRow(); - if (tableRow == null) { - tableRow = - checkNotNull(appendClientInfo) - .toTableRow(mergedPayloadBytes, Predicates.alwaysTrue()); - } - // TODO(24926, reuvenlax): We need to merge the unknown fields in! Currently we only - // execute this - // codepath when ignoreUnknownFields==true, so we should never hit this codepath. - // However once - // 24926 is fixed, we need to merge the unknownFields back into the main row before - // outputting to the - // failed-rows consumer. - org.joda.time.Instant timestamp = payload.getTimestamp(); - rowsSentToFailedRowsCollection.inc(); - failedRowsReceiver.outputWithTimestamp( - new BigQueryStorageApiInsertError( - tableRow, e.toString(), tableDestination.getTableReference()), - timestamp != null ? timestamp : elementTs); - return; - } - } - } - byte[] byteArray = - mergedPayloadBytes != null ? mergedPayloadBytes.toByteArray() : payload.getPayload(); - StorageApiWritePayload pending = - StorageApiWritePayload.of(byteArray, null, payload.getFailsafeTableRow()); - byte[] schemaHash = payload.getSchemaHash(); - if (schemaHash != null) { - pending = pending.withSchemaHash(schemaHash); - } - pendingMessages.add(pending); - - org.joda.time.Instant timestamp = payload.getTimestamp(); - pendingTimestamps.add(timestamp != null ? timestamp : elementTs); + pendingMessages.add(payload); } long flush( RetryManager retryManager, OutputReceiver failedRowsReceiver, - @Nullable OutputReceiver successfulRowsReceiver) + @Nullable OutputReceiver successfulRowsReceiver, + BiConsumer bufferMismatchedRow) throws Exception { if (pendingMessages.isEmpty()) { return 0; } - final ProtoRows.Builder insertsBuilder = ProtoRows.newBuilder(); - List<@Nullable TableRow> failsafeTableRows = Lists.newArrayList(); boolean schemaMismatchSeen; BackOff backoff = FluentBackoff.DEFAULT @@ -564,26 +522,46 @@ long flush( BigQuerySinkMetrics.throttledTimeCounter( BigQuerySinkMetrics.RpcMethod.OPEN_WRITE_STREAM)) .backoff(); + + org.joda.time.Instant startMismatchTime = org.joda.time.Instant.now(); + final Duration initialMismatchRetryTime = + Duration.millis(WriteRecordsDoFn.this.initialMismatchRetryTimeMilliSec); + + AppendRowsPacket rowsToProcess; + Iterable payloadsToIterate = pendingMessages; do { - insertsBuilder.clear(); - failsafeTableRows.clear(); - schemaMismatchSeen = false; - for (StorageApiWritePayload payload : pendingMessages) { - schemaMismatchSeen = - isPayloadSchemaOutOfDate( - payload, - () -> getAppendClientInfo(true, null).getTableSchemaHash(), - () -> - TableRowToStorageApiProto.wrapDescriptorProto( - messageConverter.getDescriptor(includeCdcColumns))); - if (schemaMismatchSeen) { - break; - } + final AppendRowsPacket processingRows = + buildInserts(payloadsToIterate, failedRowsReceiver); + rowsToProcess = processingRows; - insertsBuilder.addSerializedRows(ByteString.copyFrom(payload.getPayload())); - failsafeTableRows.add(payload.getFailsafeTableRow()); - } + schemaMismatchSeen = !processingRows.getSchemaMismatchedRows().isEmpty(); if (schemaMismatchSeen) { + if (autoUpdateSchemaStrictTimeout != null) { + if (startMismatchTime + .plus(initialMismatchRetryTime) + .isBefore(org.joda.time.Instant.now())) { + // Local retry time has expired! Pull out the mismatched rows so that we can buffer + // them for later retrying. + org.joda.time.Instant targetExpirationTime = + startMismatchTime.plus(autoUpdateSchemaStrictTimeout); + + // Buffer those rows. + processingRows + .getSchemaMismatchedRowsOnly() + .toPayloadStream() + .map(payload -> MismatchedRow.of(payload, targetExpirationTime)) + .forEach(row -> bufferMismatchedRow.accept(destination, row)); + + // Process the good rows! + Iterable goodRows = + () -> processingRows.getSchemaMatchedRowsOnly().toPayloadStream().iterator(); + + rowsToProcess = buildInserts(goodRows, failedRowsReceiver); + checkState(rowsToProcess.getSchemaMismatchedRows().isEmpty()); + break; + } + } + LOG.info("Schema out of date: refreshing table schema for {}.", tableUrn); // Refresh our view of the schema and try again.. this.messageConverter.updateSchemaFromTable(); @@ -592,16 +570,25 @@ long flush( Preconditions.checkStateNotNull( getAppendClientInfo( false, null /* read updated schema from messageConverter */)); + // Convert back to an Iterable to try again. We can't reuse the + // existing + // payloadsToIterate as if there are any failures in it, that would cause us to output + // them again to + // the dead-letter collection, resulting in duplicate outputs to dead letter. + payloadsToIterate = () -> processingRows.toPayloadStream().iterator(); } } while (schemaMismatchSeen && BackOffUtils.next(Sleeper.DEFAULT, backoff)); pendingMessages.clear(); - final ProtoRows inserts = insertsBuilder.build(); - List insertTimestamps = pendingTimestamps; + + if (rowsToProcess.getProtoRows().getSerializedRowsCount() == 0) { + // All rows either failed or were buffered with schema mismatches. + return 0; + } // Handle the case where the request is too large. - if (inserts.getSerializedSize() >= maxRequestSize) { - if (inserts.getSerializedRowsCount() > 1) { + if (rowsToProcess.getProtoRows().getSerializedSize() >= maxRequestSize) { + if (rowsToProcess.getProtoRows().getSerializedRowsCount() > 1) { // TODO(reuvenlax): Is it worth trying to handle this case by splitting the protoRows? // Given that we split // the ProtoRows iterable at 2MB and the max request size is 10MB, this scenario seems @@ -611,21 +598,14 @@ long flush( + "This is unexpected. All rows in the request will be sent to the failed-rows PCollection.", maxRequestSize); } - for (int i = 0; i < inserts.getSerializedRowsCount(); ++i) { - @Nullable TableRow failedRow = failsafeTableRows.get(i); + for (int i = 0; i < rowsToProcess.getProtoRows().getSerializedRowsCount(); ++i) { + @Nullable TableRow failedRow = rowsToProcess.getFailsafeTableRows().get(i); if (failedRow == null) { - ByteString rowBytes = inserts.getSerializedRows(i); + ByteString rowBytes = rowsToProcess.getProtoRows().getSerializedRows(i); AppendClientInfo aci = getAppendClientInfo(true, null); - failedRow = - TableRowToStorageApiProto.tableRowFromMessage( - aci.getSchemaInformation(), - DynamicMessage.parseFrom( - TableRowToStorageApiProto.wrapDescriptorProto(aci.getDescriptor()), - rowBytes), - true, - successfulRowsPredicate); + failedRow = aci.toTableRow(rowBytes, successfulRowsPredicate); } - org.joda.time.Instant timestamp = insertTimestamps.get(i); + org.joda.time.Instant timestamp = rowsToProcess.getTimestamps().get(i); failedRowsReceiver.outputWithTimestamp( new BigQueryStorageApiInsertError( failedRow, @@ -633,7 +613,7 @@ long flush( tableDestination.getTableReference()), timestamp); } - int numRowsFailed = inserts.getSerializedRowsCount(); + int numRowsFailed = rowsToProcess.getProtoRows().getSerializedRowsCount(); BigQuerySinkMetrics.appendRowsRowStatusCounter( BigQuerySinkMetrics.RowStatus.FAILED, BigQuerySinkMetrics.PAYLOAD_TOO_LARGE, @@ -647,10 +627,14 @@ long flush( if (!this.useDefaultStream) { getOrCreateStreamName(); // Force creation of the stream before we get offsets. offset = this.currentOffset; - this.currentOffset += inserts.getSerializedRowsCount(); + this.currentOffset += rowsToProcess.getProtoRows().getSerializedRowsCount(); } AppendRowsContext appendRowsContext = - new AppendRowsContext(offset, inserts, insertTimestamps, failsafeTableRows); + new AppendRowsContext( + offset, + rowsToProcess.getProtoRows(), + rowsToProcess.getTimestamps(), + rowsToProcess.getFailsafeTableRows()); retryManager.addOperation( c -> { @@ -703,15 +687,7 @@ long flush( ByteString protoBytes = failedContext.protoRows.getSerializedRows(failedIndex); AppendClientInfo aci = Preconditions.checkStateNotNull(this.appendClientInfo); - failedRow = - TableRowToStorageApiProto.tableRowFromMessage( - aci.getSchemaInformation(), - DynamicMessage.parseFrom( - TableRowToStorageApiProto.wrapDescriptorProto( - aci.getDescriptor()), - protoBytes), - true, - Predicates.alwaysTrue()); + failedRow = aci.toTableRow(protoBytes, Predicates.alwaysTrue()); } element = new BigQueryStorageApiInsertError( @@ -894,12 +870,8 @@ long flush( ByteString rowBytes = c.protoRows.getSerializedRowsList().get(i); try { TableRow row = - TableRowToStorageApiProto.tableRowFromMessage( - Preconditions.checkStateNotNull(appendClientInfo) - .getSchemaInformation(), - DynamicMessage.parseFrom(descriptor, rowBytes), - true, - successfulRowsPredicate); + Preconditions.checkStateNotNull(appendClientInfo) + .toTableRow(rowBytes, successfulRowsPredicate); org.joda.time.Instant timestamp = c.timestamps.get(i); successfulRowsReceiver.outputWithTimestamp(row, timestamp); } catch (Exception e) { @@ -911,7 +883,39 @@ long flush( }, appendRowsContext); maybeTickleCache(); - return inserts.getSerializedRowsCount(); + return rowsToProcess.getProtoRows().getSerializedRowsCount(); + } + + private AppendRowsPacket buildInserts( + Iterable payloads, + OutputReceiver failedRowsReceiver) { + Supplier appendClientInfoSupplier = + () -> { + if (appendClientInfo == null) { + appendClientInfo = getAppendClientInfo(true, null); + } + return appendClientInfo; + }; + + Consumer> failedRowsConsumer = + tv -> { + rowsSentToFailedRowsCollection.inc(); + failedRowsReceiver.outputWithTimestamp(tv.getValue(), tv.getTimestamp()); + }; + + PeekingIterator peekingIterator = + Iterators.peekingIterator(payloads.iterator()); + return AppendRowsPacket.fromStorageApiWritePayload( + peekingIterator, + Long.MAX_VALUE, + schemaChangeDetectorHelper, + BoundedWindow.TIMESTAMP_MAX_VALUE, + appendClientInfoSupplier, + failedRowsConsumer, + () -> getAppendClientInfo(true, null).getTableSchemaHash(), + () -> + TableRowToStorageApiProto.wrapDescriptorProto( + messageConverter.getDescriptor(includeCdcColumns))); } String retrieveErrorDetails(Iterable failedContext) { @@ -930,16 +934,15 @@ String retrieveErrorDetails(Iterable failedContext) { void postFlush() { // If we got a response indicating an updated schema, recreate the client. - if (this.appendClientInfo != null && autoUpdateSchema) { + if (this.appendClientInfo != null + && autoUpdateSchema + && autoUpdateSchemaStrictTimeout == null) { @Nullable StreamAppendClient streamAppendClient = appendClientInfo.getStreamAppendClient(); - @Nullable - TableSchema updatedTableSchemaReturned = - (streamAppendClient != null) ? streamAppendClient.getUpdatedSchema() : null; - if (updatedTableSchemaReturned != null) { + if (streamAppendClient != null) { Optional updatedTableSchema = - TableSchemaUpdateUtils.getUpdatedSchema( - this.messageConverter.getTableSchema(), updatedTableSchemaReturned); + schemaChangeDetectorHelper.checkResponseForUpdatedSchema( + this.messageConverter.getTableSchema(), streamAppendClient); if (updatedTableSchema.isPresent()) { invalidateAppendClient(false); // TODO: This overwrites whatever is in the cache which can cause races between @@ -956,7 +959,7 @@ void postFlush() { } private @Nullable Map destinations = Maps.newHashMap(); - private final TwoLevelMessageConverterCache messageConverters; + final TwoLevelMessageConverterCache messageConverters; private transient @Nullable DatasetService maybeDatasetService; private transient @Nullable WriteStreamService maybeWriteStreamService; private int numPendingRecords = 0; @@ -964,11 +967,12 @@ void postFlush() { private final int flushThresholdBytes; private final int flushThresholdCount; private final int maxRetries; - private final StorageApiDynamicDestinations dynamicDestinations; + final StorageApiDynamicDestinations dynamicDestinations; private final BigQueryServices bqServices; private final boolean useDefaultStream; private int streamAppendClientCount; private final @Nullable Map bigLakeConfiguration; + private final int initialMismatchRetryTimeMilliSec; WriteRecordsDoFn( String operationName, @@ -983,13 +987,16 @@ void postFlush() { @Nullable TupleTag successfulRowsTag, Predicate successfulRowsPredicate, boolean autoUpdateSchema, + @Nullable Duration autoUpdateSchemaStrictTimeout, + @Nullable TupleTag> mismatchedRowsTag, boolean ignoreUnknownValues, BigQueryIO.Write.CreateDisposition createDisposition, @Nullable String kmsKey, boolean usesCdc, AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation, int maxRetries, - @Nullable Map bigLakeConfiguration) { + @Nullable Map bigLakeConfiguration, + int initialMismatchRetryTimeMilliSec) { this.messageConverters = new TwoLevelMessageConverterCache<>(operationName); this.dynamicDestinations = dynamicDestinations; this.bqServices = bqServices; @@ -1002,6 +1009,8 @@ void postFlush() { this.successfulRowsTag = successfulRowsTag; this.successfulRowsPredicate = successfulRowsPredicate; this.autoUpdateSchema = autoUpdateSchema; + this.autoUpdateSchemaStrictTimeout = autoUpdateSchemaStrictTimeout; + this.mismatchedRowsTag = mismatchedRowsTag; this.ignoreUnknownValues = ignoreUnknownValues; this.createDisposition = createDisposition; this.kmsKey = kmsKey; @@ -1009,6 +1018,7 @@ void postFlush() { this.defaultMissingValueInterpretation = defaultMissingValueInterpretation; this.maxRetries = maxRetries; this.bigLakeConfiguration = bigLakeConfiguration; + this.initialMismatchRetryTimeMilliSec = initialMismatchRetryTimeMilliSec; } boolean shouldFlush(int recordBytes) { @@ -1017,7 +1027,7 @@ boolean shouldFlush(int recordBytes) { && numPendingRecords > 0); } - void flushIfNecessary( + Iterable> flushIfNecessary( OutputReceiver failedRowsReceiver, @Nullable OutputReceiver successfulRowsReceiver, int recordBytes) @@ -1027,14 +1037,16 @@ void flushIfNecessary( // Too much memory being used. Flush the state and wait for it to drain out. // TODO(reuvenlax): Consider waiting for memory usage to drop instead of waiting for all the // appends to finish. - flushAll(failedRowsReceiver, successfulRowsReceiver); + return flushAll(failedRowsReceiver, successfulRowsReceiver); } + return Collections.emptyList(); } - void flushAll( + Iterable> flushAll( OutputReceiver failedRowsReceiver, @Nullable OutputReceiver successfulRowsReceiver) throws Exception { + List> mismatchedRows = Lists.newArrayList(); List> retryManagers = Lists.newArrayListWithCapacity(Preconditions.checkStateNotNull(destinations).size()); long numRowsWritten = 0; @@ -1049,7 +1061,11 @@ void flushAll( BigQuerySinkMetrics.RpcMethod.APPEND_ROWS)); retryManagers.add(retryManager); numRowsWritten += - destinationState.flush(retryManager, failedRowsReceiver, successfulRowsReceiver); + destinationState.flush( + retryManager, + failedRowsReceiver, + successfulRowsReceiver, + (d, m) -> mismatchedRows.add(KV.of(d, m))); retryManager.run(false); } if (numRowsWritten > 0) { @@ -1067,9 +1083,10 @@ void flushAll( } numPendingRecords = 0; numPendingRecordBytes = 0; + return mismatchedRows; } - private DatasetService initializeDatasetService(PipelineOptions pipelineOptions) { + DatasetService getDatasetService(PipelineOptions pipelineOptions) { if (maybeDatasetService == null) { maybeDatasetService = bqServices.getDatasetService(pipelineOptions.as(BigQueryOptions.class)); @@ -1077,7 +1094,7 @@ private DatasetService initializeDatasetService(PipelineOptions pipelineOptions) return maybeDatasetService; } - private WriteStreamService initializeWriteStreamService(PipelineOptions pipelineOptions) { + WriteStreamService getWriteStreamService(PipelineOptions pipelineOptions) { if (maybeWriteStreamService == null) { maybeWriteStreamService = bqServices.getWriteStreamService(pipelineOptions.as(BigQueryOptions.class)); @@ -1093,7 +1110,7 @@ public void startBundle() throws IOException { } DestinationState createDestinationState( - ProcessContext c, + PipelineOptions options, DestinationT destination, boolean useCdc, DatasetService datasetService, @@ -1110,7 +1127,7 @@ DestinationState createDestinationState( Callable tryCreateTable = () -> { CreateTableHelpers.possiblyCreateTable( - c.getPipelineOptions().as(BigQueryOptions.class), + options.as(BigQueryOptions.class), tableDestination1, () -> dynamicDestinations.getSchema(destination), () -> dynamicDestinations.getTableConstraints(destination), @@ -1126,12 +1143,9 @@ DestinationState createDestinationState( try { messageConverter = messageConverters.get( - destination, - dynamicDestinations, - c.getPipelineOptions(), - datasetService, - writeStreamService); + destination, dynamicDestinations, options, datasetService, writeStreamService); return new DestinationState( + destination, tableDestination1, tableDestination1.getTableUrn(bigQueryOptions), tableDestination1.getShortTableUrn(), @@ -1156,17 +1170,36 @@ public void process( @Timestamp org.joda.time.Instant elementTs, MultiOutputReceiver o) throws Exception { - DatasetService initializedDatasetService = initializeDatasetService(pipelineOptions); - WriteStreamService initializedWriteStreamService = - initializeWriteStreamService(pipelineOptions); dynamicDestinations.setSideInputAccessorFromProcessContext(c); + Iterable> mismatchedRows = + processElement(pipelineOptions, element, elementTs, o); + if (!Iterables.isEmpty(mismatchedRows)) { + final OutputReceiver> mismatchedReceiver = + o.get(Preconditions.checkStateNotNull(mismatchedRowsTag)); + mismatchedRows.forEach( + m -> { + org.joda.time.Instant ts = + MoreObjects.firstNonNull(m.getValue().getPayload().getTimestamp(), elementTs); + mismatchedReceiver.outputWithTimestamp(m, ts); + }); + } + } + + Iterable> processElement( + PipelineOptions pipelineOptions, + KV element, + org.joda.time.Instant elementTs, + MultiOutputReceiver o) + throws Exception { + DatasetService initializedDatasetService = getDatasetService(pipelineOptions); + WriteStreamService initializedWriteStreamService = getWriteStreamService(pipelineOptions); DestinationState state = Preconditions.checkStateNotNull(destinations) .computeIfAbsent( element.getKey(), destination -> createDestinationState( - c, + pipelineOptions, destination, usesCdc, initializedDatasetService, @@ -1185,13 +1218,16 @@ public void process( (successfulRowsTag != null) ? o.get(successfulRowsTag) : null; int recordBytes = element.getValue().getPayload().length; - flushIfNecessary(failedRowsReceiver, successfulRowsReceiver, recordBytes); - state.addMessage(element.getValue(), elementTs, failedRowsReceiver); + Iterable> mismatchedRows = + flushIfNecessary(failedRowsReceiver, successfulRowsReceiver, recordBytes); + state.addMessage(element.getValue()); ++numPendingRecords; numPendingRecordBytes += recordBytes; + + return mismatchedRows; } - private OutputReceiver makeSuccessfulRowsreceiver( + private OutputReceiver makeSuccessfulRowsReceiver( FinishBundleContext context, TupleTag successfulRowsTag) { return new OutputReceiver() { @Override @@ -1216,14 +1252,10 @@ public OutputBuilder builder(TableRow value) { } @FinishBundle - public void finishBundle(FinishBundleContext context) throws Exception { - + public void finishBundle(final FinishBundleContext context) throws Exception { OutputReceiver failedRowsReceiver = - new OutputReceiver() { - @Override - public OutputBuilder builder( - BigQueryStorageApiInsertError value) { - return WindowedValues.builder() + value -> + WindowedValues.builder() .setValue(value) .setTimestamp(GlobalWindow.INSTANCE.maxTimestamp()) .setWindow(GlobalWindow.INSTANCE) @@ -1238,30 +1270,64 @@ public OutputBuilder builder( window); } }); - } - }; @Nullable OutputReceiver successfulRowsReceiver = null; if (successfulRowsTag != null) { - successfulRowsReceiver = makeSuccessfulRowsreceiver(context, successfulRowsTag); + successfulRowsReceiver = makeSuccessfulRowsReceiver(context, successfulRowsTag); + } + + OutputReceiver> finalizeOutputRecever = + value -> + WindowedValues.>builder() + .setValue(value) + .setTimestamp(GlobalWindow.INSTANCE.maxTimestamp()) + .setWindow(GlobalWindow.INSTANCE) + .setPaneInfo(PaneInfo.NO_FIRING) + .setReceiver( + windowedValue -> { + for (BoundedWindow window : windowedValue.getWindows()) { + context.output( + finalizeTag, + windowedValue.getValue(), + windowedValue.getTimestamp(), + window); + } + }); + Iterable> mismatchedRows = + finishBundle(failedRowsReceiver, successfulRowsReceiver, finalizeOutputRecever); + if (!Iterables.isEmpty(mismatchedRows)) { + mismatchedRows.forEach( + m -> { + org.joda.time.Instant ts = + MoreObjects.firstNonNull( + m.getValue().getPayload().getTimestamp(), + GlobalWindow.INSTANCE.maxTimestamp()); + context.output( + Preconditions.checkStateNotNull(mismatchedRowsTag), m, ts, GlobalWindow.INSTANCE); + }); } + } + + public Iterable> finishBundle( + OutputReceiver failedRowsReceiver, + @Nullable OutputReceiver successfulRowsReceiver, + OutputReceiver> finalizeOutputReceiver) + throws Exception { - flushAll(failedRowsReceiver, successfulRowsReceiver); + Iterable> mismatchedRows = + flushAll(failedRowsReceiver, successfulRowsReceiver); final Map destinations = Preconditions.checkStateNotNull(this.destinations); for (DestinationState state : destinations.values()) { if (!useDefaultStream && !Strings.isNullOrEmpty(state.streamName)) { - context.output( - finalizeTag, - KV.of(state.tableUrn, state.streamName), - GlobalWindow.INSTANCE.maxTimestamp(), - GlobalWindow.INSTANCE); + finalizeOutputReceiver.output(KV.of(state.tableUrn, state.streamName)); } state.teardown(); } destinations.clear(); this.destinations = null; + return mismatchedRows; } @Teardown 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 b644d7aa752c..295a67d1247c 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 @@ -18,6 +18,7 @@ package org.apache.beam.sdk.io.gcp.bigquery; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; @@ -48,10 +49,13 @@ import java.util.function.Function; import java.util.function.Predicate; import java.util.function.Supplier; +import java.util.stream.Collectors; +import java.util.stream.StreamSupport; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.extensions.protobuf.ProtoCoder; +import org.apache.beam.sdk.function.ThrowingConsumer; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.CreateDisposition; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.DatasetService; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.StreamAppendClient; @@ -66,6 +70,7 @@ import org.apache.beam.sdk.schemas.NoSuchSchemaException; import org.apache.beam.sdk.schemas.SchemaCoder; import org.apache.beam.sdk.schemas.SchemaRegistry; +import org.apache.beam.sdk.state.BagState; import org.apache.beam.sdk.state.StateSpec; import org.apache.beam.sdk.state.StateSpecs; import org.apache.beam.sdk.state.TimeDomain; @@ -120,6 +125,7 @@ public class StorageApiWritesShardedRecords { private static final Logger LOG = LoggerFactory.getLogger(StorageApiWritesShardedRecords.class); private static final Duration DEFAULT_STREAM_IDLE_TIME = Duration.standardHours(1); + private static final Duration RETRY_MISMATCHED_ROWS_PERIOD = Duration.standardMinutes(1); private final StorageApiDynamicDestinations dynamicDestinations; private final CreateDisposition createDisposition; @@ -128,6 +134,7 @@ public class StorageApiWritesShardedRecords destinationCoder; private final Coder failedRowsCoder; private final boolean autoUpdateSchema; + private final @Nullable Duration autoUpdateSchemaStrictTimeout; private final boolean ignoreUnknownValues; private final AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation; private final @Nullable Map bigLakeConfiguration; @@ -160,6 +167,7 @@ public StorageApiWritesShardedRecords( @Nullable TupleTag successfulRowsTag, Predicate successfulRowsPredicate, boolean autoUpdateSchema, + @Nullable Duration autoUpdateSchemaStrictTimeout, boolean ignoreUnknownValues, AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation, @Nullable Map bigLakeConfiguration) { @@ -174,6 +182,7 @@ public StorageApiWritesShardedRecords( this.successfulRowsPredicate = successfulRowsPredicate; this.succussfulRowsCoder = successfulRowsCoder; this.autoUpdateSchema = autoUpdateSchema; + this.autoUpdateSchemaStrictTimeout = autoUpdateSchemaStrictTimeout; this.ignoreUnknownValues = ignoreUnknownValues; this.defaultMissingValueInterpretation = defaultMissingValueInterpretation; this.bigLakeConfiguration = bigLakeConfiguration; @@ -344,11 +353,25 @@ class WriteRecordsDoFn private final StateSpec> streamOffsetSpec = StateSpecs.value(); @StateId("updatedSchema") - private final StateSpec> updatedSchema = + private final StateSpec> updatedSchemaSpe = StateSpecs.value(ProtoCoder.of(TableSchema.class)); + @StateId("mismatchedRows") + private final StateSpec> mismatchedRowsSpec = + StateSpecs.bag(MismatchedRow.Coder.of()); + + @TimerId("retryMismatchedRowsTimer") + private final TimerSpec mismatchedRowsTimerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME); + + @StateId("currentMismatchedRowTimerValue") + private final StateSpec> currentMismatchedRowTimerValueSpec = + StateSpecs.value(); + + @StateId("minPendingTimestamp") + private final StateSpec> minPendingTimestampSpec = StateSpecs.value(); + @TimerId("idleTimer") - private final TimerSpec idleTimer = TimerSpecs.timer(TimeDomain.PROCESSING_TIME); + private final TimerSpec idleTimerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME); private final Duration streamIdleTime; private final long splitSize; @@ -475,8 +498,7 @@ public void close() { } AppendClientInfo get() { - org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState( - valid); + checkState(valid); return appendClientInfo; } @@ -487,7 +509,7 @@ StreamAppendClient getStreamAppendClient() { private CreateRetryManagerResult createRetryManager( ShardedKey key, - Iterable messages, + Iterable messages, Function, ApiFuture> runOperation, Function>, RetryType> onError, Consumer> onSuccess, @@ -504,8 +526,8 @@ private CreateRetryManagerResult createRetryManager( List> failedRows = Lists.newArrayList(); int recordsAppended = 0; List histogramUpdates = Lists.newArrayList(); - for (SplittingIterable.Value splitValue : messages) { - if (splitValue.getSchemaMismatchSeen()) { + for (AppendRowsPacket splitValue : messages) { + if (!splitValue.getSchemaMismatchedRows().isEmpty()) { return CreateRetryManagerResult.schemaMismatch(); } // Handle the case of a row that is too large. @@ -767,18 +789,16 @@ public void process( final PipelineOptions pipelineOptions, @Element KV, Iterable> element, @Timestamp org.joda.time.Instant elementTs, - final @AlwaysFetched @StateId("streamName") ValueState streamName, - final @AlwaysFetched @StateId("streamOffset") ValueState streamOffset, - final @StateId("updatedSchema") ValueState updatedSchema, + @AlwaysFetched @StateId("streamName") ValueState streamName, + @AlwaysFetched @StateId("streamOffset") ValueState streamOffset, + @StateId("updatedSchema") ValueState updatedSchema, @TimerId("idleTimer") Timer idleTimer, + @StateId("mismatchedRows") BagState mismatchedRowsBag, + @TimerId("retryMismatchedRowsTimer") Timer mismatchedRowsRetryTimer, + @StateId("currentMismatchedRowTimerValue") ValueState mismatchedRowsRetryTimerValue, + @StateId("minPendingTimestamp") ValueState minPendingTimestamp, final MultiOutputReceiver o) throws Exception { - BigQueryOptions bigQueryOptions = pipelineOptions.as(BigQueryOptions.class); - - if (autoUpdateSchema) { - updatedSchema.readLater(); - } - dynamicDestinations.setSideInputAccessorFromProcessContext(c); TableDestination tableDestination = destinations.computeIfAbsent( @@ -793,11 +813,85 @@ public void process( dest); return tableDestination1; }); + + StorageApiDynamicDestinations.MessageConverter messageConverter = + messageConverters.get( + element.getKey().getKey(), + dynamicDestinations, + pipelineOptions, + getDatasetService(pipelineOptions), + getWriteStreamService(pipelineOptions)); + + ThrowingConsumer> processMismatchedRows = + mismatchedRows -> { + if (!Iterables.isEmpty(mismatchedRows)) { + AppendClientInfo info = + AppendClientInfo.of( + Preconditions.checkStateNotNull(messageConverter.getTableSchema()), + messageConverter.getDescriptor(false), + AutoCloseable::close); + + SchemaChangeDetectorHelper.bufferMismatchedRows( + mismatchedRows, + mismatchedRowsBag, + mismatchedRowsRetryTimer, + mismatchedRowsRetryTimerValue, + minPendingTimestamp, + tableDestination, + o.get(failedRowsTag), + info, + rowsSentToFailedRowsCollection, + RETRY_MISMATCHED_ROWS_PERIOD); + } + }; + processPayloads( + pipelineOptions, + element.getKey(), + tableDestination, + messageConverter, + element.getValue(), + elementTs, + streamName, + streamOffset, + updatedSchema, + idleTimer, + o, + processMismatchedRows); + } + + private void processPayloads( + final PipelineOptions pipelineOptions, + ShardedKey shardedDestination, + TableDestination tableDestination, + StorageApiDynamicDestinations.MessageConverter messageConverter, + Iterable element, + org.joda.time.Instant elementTs, + final ValueState streamName, + final ValueState streamOffset, + final ValueState updatedSchema, + Timer idleTimer, + final MultiOutputReceiver o, + ThrowingConsumer> processMismatchedRows) + throws Exception { + BigQueryOptions bigQueryOptions = pipelineOptions.as(BigQueryOptions.class); + + if (autoUpdateSchema) { + updatedSchema.readLater(); + } + + final DestinationT destination = shardedDestination.getKey(); + final String tableId = tableDestination.getTableUrn(bigQueryOptions); final String shortTableId = tableDestination.getShortTableUrn(); final TableReference tableReference = tableDestination.getTableReference(); final DatasetService datasetService = getDatasetService(pipelineOptions); final WriteStreamService writeStreamService = getWriteStreamService(pipelineOptions); + SchemaChangeDetectorHelper schemaChangeDetectorHelper = + new SchemaChangeDetectorHelper( + autoUpdateSchema, + ignoreUnknownValues, + tableReference, + autoUpdateSchemaStrictTimeout != null); Lineage.getSinks() .add( @@ -808,12 +902,11 @@ public void process( Coder destinationCoder = dynamicDestinations.getDestinationCoder(); Callable tryCreateTable = () -> { - DestinationT dest = element.getKey().getKey(); CreateTableHelpers.possiblyCreateTable( - c.getPipelineOptions().as(BigQueryOptions.class), + bigQueryOptions, tableDestination, - () -> dynamicDestinations.getSchema(dest), - () -> dynamicDestinations.getTableConstraints(dest), + () -> dynamicDestinations.getSchema(destination), + () -> dynamicDestinations.getTableConstraints(destination), createDisposition, destinationCoder, kmsKey, @@ -827,20 +920,16 @@ public void process( getOrCreateStream( tableId, streamName, streamOffset, idleTimer, writeStreamService, tryCreateTable); - StorageApiDynamicDestinations.MessageConverter messageConverter = - messageConverters.get( - element.getKey().getKey(), - dynamicDestinations, - pipelineOptions, - datasetService, - writeStreamService); Callable getAppendClientInfo = () -> { @Nullable TableSchema tableSchema; DescriptorProtos.DescriptorProto descriptor; + TableSchema updatedSchemaValue = autoUpdateSchema ? updatedSchema.read() : null; - if (autoUpdateSchema && updatedSchemaValue != null) { + if (autoUpdateSchema + && updatedSchemaValue != null + && autoUpdateSchemaStrictTimeout == null) { // This means that Vortex has told us in the past that the table schema has been // updated. We should use // this updated schema instead of the initial schema from the messageConverter. @@ -849,23 +938,23 @@ public void process( TableRowToStorageApiProto.descriptorSchemaFromTableSchema( tableSchema, true, false); } else { - // Start off with the base schema. As we get notified of schema updates, we - // will update the descriptor. + if (autoUpdateSchemaStrictTimeout != null && updatedSchemaValue != null) { + Optional updated = + TableSchemaUpdateUtils.getUpdatedSchema( + messageConverter.getTableSchema(), updatedSchemaValue); + updated.ifPresent(messageConverter::updateSchema); + } + tableSchema = messageConverter.getTableSchema(); descriptor = messageConverter.getDescriptor(false); - if (autoUpdateSchema) { + if (autoUpdateSchema && autoUpdateSchemaStrictTimeout == null) { // A StreamWriter ignores table schema updates that happen prior to its creation. // So before creating a StreamWriter below, we fetch the table schema to check if we // missed an update. If so, use the new schema instead of the base schema. - // TODO: There's still a race here! - @Nullable - TableSchema streamSchema = - MoreObjects.firstNonNull( - writeStreamService.getWriteStreamSchema(getOrCreateStream.get()), - TableSchema.getDefaultInstance()); Optional newSchema = - TableSchemaUpdateUtils.getUpdatedSchema(tableSchema, streamSchema); + schemaChangeDetectorHelper.checkUpdatedSchema( + tableSchema, getOrCreateStream.get(), writeStreamService); if (newSchema.isPresent()) { tableSchema = newSchema.get(); @@ -895,8 +984,9 @@ public void process( // cache could in // theory evict the object during execution and we want a pin held throughout the execution of // this function. + Iterable mismatchedRows = Collections.emptyList(); try (AppendClientHolder appendClientHolder = - new AppendClientHolder(element.getKey(), getAppendClientInfo)) { + new AppendClientHolder(shardedDestination, getAppendClientInfo)) { String currentStream = getOrCreateStream.get(); if (!currentStream.equals(appendClientHolder.get().getStreamName())) { // Cached append client is inconsistent with persisted state. Throw away cached item and @@ -981,28 +1071,26 @@ public void process( BigQuerySinkMetrics.throttledTimeCounter( BigQuerySinkMetrics.RpcMethod.OPEN_WRITE_STREAM)) .backoff(); + CreateRetryManagerResult createRetryManagerResult; + org.joda.time.Instant startMismatchTime = org.joda.time.Instant.now(); + final Duration initialMismatchRetryTime = + Duration.millis( + bigQueryOptions.getStorageApiInitialMismatchRetryTimeMilliSec()); // Retry locally + Iterable payloadsToIterate = element; do { // Each ProtoRows object contains at most 1MB of rows. // TODO: Push messageFromTableRow up to top level. That we we cans skip TableRow entirely // if // already proto or already schema. - Iterable messages = + final Iterable messages = new SplittingIterable( - element.getValue(), + payloadsToIterate, splitSize, - // Unknown field merger - (bytes, tableRow) -> - appendClientHolder.get().mergeNewFields(bytes, tableRow, ignoreUnknownValues), - // Convert back to TableRow - bytes -> appendClientHolder.get().toTableRow(bytes, Predicates.alwaysTrue()), // Failed rows consumer - (failedRow, errorMessage) -> { + (TimestampedValue error) -> { o.get(failedRowsTag) - .outputWithTimestamp( - new BigQueryStorageApiInsertError( - failedRow.getValue(), errorMessage, tableReference), - failedRow.getTimestamp()); + .outputWithTimestamp(error.getValue(), error.getTimestamp()); rowsSentToFailedRowsCollection.inc(); BigQuerySinkMetrics.appendRowsRowStatusCounter( BigQuerySinkMetrics.RowStatus.FAILED, @@ -1015,19 +1103,62 @@ public void process( () -> TableRowToStorageApiProto.wrapDescriptorProto( messageConverter.getDescriptor(false)), - autoUpdateSchema, - elementTs); + elementTs, + appendClientHolder::get, + schemaChangeDetectorHelper); + Iterable messagesToProcess = messages; createRetryManagerResult = createRetryManager( - element.getKey(), - messages, + shardedDestination, + messagesToProcess, runOperation, onError, onSuccess, appendClientHolder.get(), tableReference); + if (createRetryManagerResult.getSchemaMismatchSeen()) { + if (autoUpdateSchemaStrictTimeout != null) { + if (startMismatchTime + .plus(initialMismatchRetryTime) + .isBefore(org.joda.time.Instant.now())) { + // Local retry time has expired! Pull out the mismatched rows so that we can buffer + // them for later + // retrying. + org.joda.time.Instant targetExpirationTime = + startMismatchTime.plus(autoUpdateSchemaStrictTimeout); + + mismatchedRows = + () -> + StreamSupport.stream(messages.spliterator(), false) + .map(AppendRowsPacket::getSchemaMismatchedRowsOnly) + .flatMap(AppendRowsPacket::toPayloadStream) + .map(payload -> MismatchedRow.of(payload, targetExpirationTime)) + .map(MismatchedRow.class::cast) + .iterator(); + + // Continue processing only the messages that matched the schema. + messagesToProcess = + () -> + StreamSupport.stream(messages.spliterator(), false) + .map(AppendRowsPacket::getSchemaMatchedRowsOnly) + .iterator(); + + createRetryManagerResult = + createRetryManager( + shardedDestination, + messagesToProcess, + runOperation, + onError, + onSuccess, + appendClientHolder.get(), + tableReference); + checkState(!createRetryManagerResult.getSchemaMismatchSeen()); + // Break out of the loop and start processing. + break; + } + } // TODO: The call to updateSchemaFromTable will throttle the DoFn (both because of the // RPC // call and because @@ -1035,9 +1166,23 @@ public void process( LOG.info("Schema out of date: refreshing table schema for {}", tableId); // Force the message converter to get the schema again from the table. messageConverter.updateSchemaFromTable(); + if (autoUpdateSchemaStrictTimeout != null) { + updatedSchema.write(messageConverter.getTableSchema()); + } // Close all RPC clients that were opened with the old descriptor. Clear the cache, // forcing us to create a new append client with the updated descriptor. appendClientHolder.invalidateAndReset(); + + // Convert back to an Iterable to try again. We can't reuse the + // existing + // payloadsToIterate as if there are any failures in it, that would cause us to output + // them again to + // the dead-letter collection, resulting in duplicate outputs to dead letter. + payloadsToIterate = + () -> + StreamSupport.stream(messages.spliterator(), false) + .flatMap(AppendRowsPacket::toPayloadStream) + .iterator(); } } while (createRetryManagerResult.getSchemaMismatchSeen() && BackOffUtils.next(Sleeper.DEFAULT, backoff)); @@ -1070,21 +1215,19 @@ public void process( appendSplitDistribution.update(numAppends); if (autoUpdateSchema) { - @Nullable StreamAppendClient streamAppendClient = appendClientHolder.getStreamAppendClient(); TableSchema originalSchema = appendClientHolder.get().getTableSchema(); - @Nullable - TableSchema updatedSchemaReturned = - (streamAppendClient != null) ? streamAppendClient.getUpdatedSchema() : null; - // Update the table schema and clear the append client. - if (updatedSchemaReturned != null) { - Optional newSchema = - TableSchemaUpdateUtils.getUpdatedSchema(originalSchema, updatedSchemaReturned); - if (newSchema.isPresent()) { - APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(element.getKey())); - LOG.debug( - "Fetched updated schema for table {}:\n\t{}", tableId, updatedSchemaReturned); + Optional newSchema = + schemaChangeDetectorHelper.checkResponseForUpdatedSchema( + originalSchema, streamAppendClient); + if (newSchema.isPresent()) { + APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(shardedDestination)); + LOG.debug("Fetched updated schema for table {}:\n\t{}", tableId, newSchema.get()); + if (autoUpdateSchemaStrictTimeout != null) { + messageConverter.updateSchemaFromTable(); + updatedSchema.write(messageConverter.getTableSchema()); + } else { updatedSchema.write(newSchema.get()); } } @@ -1093,10 +1236,112 @@ public void process( java.time.Duration timeElapsed = java.time.Duration.between(now, Instant.now()); appendLatencyDistribution.update(timeElapsed.toMillis()); } + processMismatchedRows.accept(mismatchedRows); } idleTimer.offset(streamIdleTime).withNoOutputTimestamp().setRelative(); } + @OnTimer("retryMismatchedRowsTimer") + public void onMismatchedRowsTimer( + OnTimerContext context, + PipelineOptions pipelineOptions, + @Key ShardedKey shardedDestination, + @Timestamp org.joda.time.Instant elementTs, + @AlwaysFetched @StateId("streamName") ValueState streamName, + @AlwaysFetched @StateId("streamOffset") ValueState streamOffset, + @StateId("updatedSchema") ValueState updatedSchema, + @TimerId("idleTimer") Timer idleTimer, + MultiOutputReceiver o, + @TimerId("retryMismatchedRowsTimer") Timer retryRowsTimer, + @StateId("mismatchedRows") BagState mismatchedRowsBag, + @StateId("currentMismatchedRowTimerValue") ValueState currentTimerValue, + @StateId("minPendingTimestamp") ValueState minPendingTimestamp) + throws Exception { + dynamicDestinations.setSideInputAccessorFromOnTimerContext(context); + + mismatchedRowsBag.readLater(); + currentTimerValue.readLater(); + minPendingTimestamp.readLater(); + + TableDestination tableDestination = + destinations.computeIfAbsent( + shardedDestination.getKey(), + dest -> { + TableDestination tableDestination1 = dynamicDestinations.getTable(dest); + checkArgument( + tableDestination1 != null, + "DynamicDestinations.getTable() may not return null, " + + "but %s returned null for destination %s", + dynamicDestinations, + dest); + return tableDestination1; + }); + + Iterable payloads = + StreamSupport.stream(mismatchedRowsBag.read().spliterator(), false) + .map(MismatchedRow::getPayload) + .collect(Collectors.toList()); + + StorageApiDynamicDestinations.MessageConverter messageConverter = + messageConverters.get( + shardedDestination.getKey(), + dynamicDestinations, + pipelineOptions, + getDatasetService(pipelineOptions), + getWriteStreamService(pipelineOptions)); + LOG.info( + "Schema out of date: refreshing table schema for {}", + tableDestination.getShortTableUrn()); + // Force the message converter to get the schema again from the table. + messageConverter.updateSchemaFromTable(); + updatedSchema.write(messageConverter.getTableSchema()); + + // Close all RPC clients that were opened with the old descriptor. Clear the cache, + // forcing us to create a new append client with the updated descriptor. + APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(shardedDestination)); + + // Try to reprocess all of these rows. + ThrowingConsumer> processMismatchedRows = + mismatchedRows -> { + // Rebuffer the ones that are still not succeeding. + mismatchedRowsBag.clear(); + if (!Iterables.isEmpty(mismatchedRows)) { + AppendClientInfo info = + AppendClientInfo.of( + Preconditions.checkStateNotNull(messageConverter.getTableSchema()), + messageConverter.getDescriptor(false), + AutoCloseable::close); + + SchemaChangeDetectorHelper.bufferMismatchedRows( + mismatchedRows, + mismatchedRowsBag, + retryRowsTimer, + currentTimerValue, + minPendingTimestamp, + tableDestination, + o.get(failedRowsTag), + info, + rowsSentToFailedRowsCollection, + RETRY_MISMATCHED_ROWS_PERIOD); + } + }; + + currentTimerValue.clear(); + processPayloads( + pipelineOptions, + shardedDestination, + tableDestination, + messageConverter, + payloads, + elementTs, + streamName, + streamOffset, + updatedSchema, + idleTimer, + o, + processMismatchedRows); + } + // called by the idleTimer and window-expiration handlers. private void finalizeStream( @AlwaysFetched @StateId("streamName") ValueState streamName, diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java index ba72bb8682fd..898c486baa7f 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java @@ -1193,6 +1193,7 @@ public static Descriptor wrapDescriptorProto(DescriptorProto descriptorProto) if (!collectedExceptions.isEmpty()) { return null; } + try { return builder.build(); } catch (Exception e) { diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpgradeTableSchema.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpgradeTableSchema.java index 6c3c028b6d0d..a6cc916b7802 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpgradeTableSchema.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpgradeTableSchema.java @@ -17,18 +17,22 @@ */ package org.apache.beam.sdk.io.gcp.bigquery; +import com.google.api.services.bigquery.model.TableCell; +import com.google.cloud.bigquery.storage.v1.BigQuerySchemaUtil; import com.google.cloud.bigquery.storage.v1.TableFieldSchema; import com.google.cloud.bigquery.storage.v1.TableSchema; import com.google.protobuf.Descriptors; import com.google.protobuf.DynamicMessage; import com.google.protobuf.Message; import com.google.protobuf.UnknownFieldSet; +import java.util.AbstractMap; import java.util.Arrays; import java.util.Collections; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Set; +import java.util.function.Supplier; import org.apache.beam.sdk.util.ThrowingSupplier; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; @@ -247,26 +251,31 @@ private static TableFieldSchema mergeField(TableFieldSchema f1, TableFieldSchema } public static boolean isPayloadSchemaOutOfDate( - StorageApiWritePayload payload, + byte @Nullable [] payloadSchemaHash, + Supplier payloadSupplier, ThrowingSupplier schemaHash, ThrowingSupplier schemaDescriptor) throws Exception { - byte @Nullable [] payloadSchemaHash = payload.getSchemaHash(); if (payloadSchemaHash != null) { // Schema hash is only included in the payload if schema update options are set. HashCode lhs = HashCode.fromBytes(payloadSchemaHash); HashCode rhs = HashCode.fromBytes(schemaHash.get()); if (!lhs.equals(rhs)) { - DynamicMessage msg = - DynamicMessage.newBuilder(schemaDescriptor.get()) - .mergeFrom(payload.getPayload()) - .buildPartial(); - return !msg.isInitialized() || hasUnknownFields(msg); + return isPayloadSchemaOutOfDate(payloadSupplier, schemaDescriptor); } } return false; } + public static boolean isPayloadSchemaOutOfDate( + Supplier payloadSupplier, ThrowingSupplier schemaDescriptor) + throws Exception { + Descriptors.Descriptor descriptor = schemaDescriptor.get(); + byte[] payload = payloadSupplier.get(); + DynamicMessage msg = DynamicMessage.newBuilder(descriptor).mergeFrom(payload).buildPartial(); + return !msg.isInitialized() || hasUnknownFields(msg); + } + private static boolean hasUnknownFields(Message message) { if (message == null) { return false; @@ -305,4 +314,88 @@ private static boolean hasUnknownFields(Message message) { // If we reach here, neither this message nor its descendants have unknown fields return false; } + + public static boolean missingUnknownField( + AbstractMap unknownFields, + ThrowingSupplier schemaDescriptor) + throws Exception { + @Nullable Object fValue = unknownFields.get("f"); + if (fValue instanceof List) { + List cells = (List) fValue; + return missingUnknownField(cells, schemaDescriptor.get()); + } else { + return missingUnknownField(unknownFields, schemaDescriptor.get()); + } + } + + public static boolean missingUnknownField( + List unknownFields, Descriptors.Descriptor schemaDescriptor) { + for (int i = 0; i < unknownFields.size(); i++) { + Object cell = unknownFields.get(i); + Object value; + if (cell instanceof TableCell) { + value = ((TableCell) cell).getV(); + } else if (cell instanceof Map) { + value = ((Map) cell).get("v"); + } else { + value = cell; + } + if (i >= schemaDescriptor.getFields().size()) { + return true; + } + Descriptors.FieldDescriptor fieldDescriptor = schemaDescriptor.getFields().get(i); + if (missingUnknownFieldObject(value, fieldDescriptor)) { + return true; + } + } + return false; + } + + public static boolean missingUnknownField( + AbstractMap unknownFields, Descriptors.Descriptor schemaDescriptor) { + for (Map.Entry entry : unknownFields.entrySet()) { + String key = entry.getKey().toLowerCase(); + String protoFieldName = + BigQuerySchemaUtil.isProtoCompatible(key) + ? key + : BigQuerySchemaUtil.generatePlaceholderFieldName(key); + + Descriptors.FieldDescriptor fieldDescriptor = + schemaDescriptor.findFieldByName(protoFieldName); + if (fieldDescriptor == null) { + return true; + } + Object value = entry.getValue(); + if (value == null) { + continue; + } + if (missingUnknownFieldObject(value, fieldDescriptor)) { + return true; + } + } + return false; + } + + private static boolean missingUnknownFieldObject( + @Nullable Object value, Descriptors.FieldDescriptor fieldDescriptor) { + if (value == null) { + return false; + } + + if (fieldDescriptor.getType() != Descriptors.FieldDescriptor.Type.MESSAGE) { + return false; + } + if (value instanceof List) { + for (Object element : (List) value) { + if (missingUnknownFieldObject(element, fieldDescriptor)) { + return true; + } + } + return false; + } else if (value instanceof AbstractMap) { + return missingUnknownField( + (AbstractMap) value, fieldDescriptor.getMessageType()); + } + return false; + } } diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java index 601ed71473ed..3e0ac4463152 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java @@ -136,6 +136,7 @@ import org.apache.beam.sdk.metrics.Lineage; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.options.StreamingOptions; import org.apache.beam.sdk.options.ValueProvider; import org.apache.beam.sdk.schemas.JavaFieldSchema; import org.apache.beam.sdk.schemas.Schema; @@ -2301,12 +2302,22 @@ public TableRow apply(Long input) { @Test public void testUpdateTableSchemaUseSet() throws Exception { - updateTableSchemaTest(true); + updateTableSchemaTest(true, false); } @Test public void testUpdateTableSchemaUseSetF() throws Exception { - updateTableSchemaTest(false); + updateTableSchemaTest(false, false); + } + + @Test + public void testUpdateTableSchemaConsistentUseSet() throws Exception { + updateTableSchemaTest(true, true); + } + + @Test + public void testUpdateTableSchemaConsistentUseSetF() throws Exception { + updateTableSchemaTest(false, true); } @Test @@ -2396,13 +2407,18 @@ public void onTimer(@Key String tableSpec) throws IOException { } } - public void updateTableSchemaTest(boolean useSet) throws Exception { + public void updateTableSchemaTest(boolean useSet, boolean withConsistentSchemaUpdate) + throws Exception { assumeTrue(useStreaming); assumeTrue(useStorageApi); + if (withConsistentSchemaUpdate) { + p.getOptions().as(StreamingOptions.class).setStreaming(true); + } // Make sure that GroupIntoBatches does not buffer data. p.getOptions().as(BigQueryOptions.class).setStorageApiAppendThresholdBytes(1); p.getOptions().as(BigQueryOptions.class).setNumStorageWriteApiStreams(1); + p.getOptions().as(BigQueryOptions.class).setStorageApiInitialMismatchRetryTimeMilliSec(0); BigQueryIO.Write.Method method = useStorageApiApproximate ? Method.STORAGE_API_AT_LEAST_ONCE : Method.STORAGE_WRITE_API; @@ -2417,15 +2433,12 @@ public void updateTableSchemaTest(boolean useSet) throws Exception { new TableFieldSchema().setName("name").setType("STRING"), new TableFieldSchema().setName("req").setType("STRING").setMode("REQUIRED"))); - // Add new fields to the update schema. Also reorder some existing fields to validate that we - // handle update - // field reordering correctly. TableSchema tableSchemaUpdated = new TableSchema() .setFields( ImmutableList.of( - new TableFieldSchema().setName("name").setType("STRING"), new TableFieldSchema().setName("number").setType("INTEGER"), + new TableFieldSchema().setName("name").setType("STRING"), new TableFieldSchema().setName("req").setType("STRING"), new TableFieldSchema().setName("double_number").setType("INTEGER"), new TableFieldSchema().setName("12_special_name").setType("STRING"))); @@ -2476,6 +2489,8 @@ public void updateTableSchemaTest(boolean useSet) throws Exception { for (long i = 6; i < 10; i++) { testStream = testStream.addElements(i); } + // Expire the buffering timer so the elements get processed + testStream = testStream.advanceProcessingTime(Duration.standardMinutes(2)); PCollection tableRows = p.apply(testStream.advanceWatermarkToInfinity()) @@ -2488,25 +2503,41 @@ public void updateTableSchemaTest(boolean useSet) throws Exception { Duration.standardSeconds(5), tableSchemaUpdated, fakeDatasetService))) .setCoder(TableRowJsonCoder.of()); - tableRows.apply( + BigQueryIO.Write write = BigQueryIO.writeTableRows() .to(tableRef) .withMethod(method) .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_NEVER) .ignoreUnknownValues() - .withAutoSchemaUpdate(true) .withTestServices(fakeBqServices) - .withoutValidation()); + .withoutValidation(); + if (withConsistentSchemaUpdate) { + write = write.withAutoSchemaUpdateConsistent(true, Duration.standardHours(1)); + } else { + write = write.withAutoSchemaUpdate(true); + } + tableRows.apply(write); p.run(); Iterable expectedDroppedValues = - LongStream.range(0, 6) + (withConsistentSchemaUpdate ? LongStream.empty() : LongStream.range(0, 6)) .mapToObj(getRowSet) .map(tr -> filterUnknownValues(tr, tableSchema.getFields())) .collect(Collectors.toList()); Iterable expectedFullValues = - LongStream.range(6, 10).mapToObj(getRowSet).collect(Collectors.toList()); + (withConsistentSchemaUpdate ? LongStream.range(0, 10) : LongStream.range(6, 10)) + .mapToObj(getRowSet) + .collect(Collectors.toList()); + System.err.println( + "GOT " + + fakeDatasetService.getAllRows( + tableRef.getProjectId(), tableRef.getDatasetId(), tableRef.getTableId())); + System.err.println( + "WANT " + + Arrays.toString( + Iterables.toArray( + Iterables.concat(expectedDroppedValues, expectedFullValues), TableRow.class))); assertThat( fakeDatasetService.getAllRows( tableRef.getProjectId(), tableRef.getDatasetId(), tableRef.getTableId()), @@ -2524,6 +2555,7 @@ public void testAutoPatchTableSchemaTest() throws Exception { p.getOptions().as(BigQueryOptions.class).setStorageApiAppendThresholdBytes(1); p.getOptions().as(BigQueryOptions.class).setNumStorageWriteApiStreams(1); p.getOptions().as(BigQueryOptions.class).setSchemaUpgradeBufferingShards(2); + p.getOptions().as(BigQueryOptions.class).setStorageApiInitialMismatchRetryTimeMilliSec(0); BigQueryIO.Write.Method method = useStorageApiApproximate ? Method.STORAGE_API_AT_LEAST_ONCE : Method.STORAGE_WRITE_API; diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelperTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelperTest.java new file mode 100644 index 000000000000..8ba90ad01d6d --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelperTest.java @@ -0,0 +1,305 @@ +/* + * 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.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.api.services.bigquery.model.TableReference; +import com.google.api.services.bigquery.model.TableRow; +import com.google.cloud.bigquery.storage.v1.TableFieldSchema; +import com.google.cloud.bigquery.storage.v1.TableSchema; +import com.google.protobuf.ByteString; +import com.google.protobuf.DescriptorProtos; +import com.google.protobuf.Descriptors; +import java.nio.charset.StandardCharsets; +import java.util.Optional; +import org.apache.beam.sdk.values.TimestampedValue; +import org.joda.time.Instant; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class SchemaChangeDetectorHelperTest { + + private TableReference tableReference; + private BigQueryServices.WriteStreamService mockWriteStreamService; + private BigQueryServices.StreamAppendClient mockStreamAppendClient; + private AppendClientInfo mockAppendClientInfo; + + @Before + public void setUp() { + tableReference = new TableReference().setProjectId("p").setDatasetId("d").setTableId("t"); + mockWriteStreamService = mock(BigQueryServices.WriteStreamService.class); + mockStreamAppendClient = mock(BigQueryServices.StreamAppendClient.class); + mockAppendClientInfo = mock(AppendClientInfo.class); + } + + @Test + public void testCheckUpdatedSchema_autoUpdateTrue_hasUpdate() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(true, false, tableReference, false); + TableSchema currentSchema = TableSchema.newBuilder().build(); + TableSchema newSchema = + TableSchema.newBuilder() + .addFields(TableFieldSchema.newBuilder().setName("new_field").build()) + .build(); + + when(mockWriteStreamService.getWriteStreamSchema("test_stream")).thenReturn(newSchema); + + Optional updated = + helper.checkUpdatedSchema(currentSchema, "test_stream", mockWriteStreamService); + + assertTrue(updated.isPresent()); + assertEquals(1, updated.get().getFieldsCount()); + assertEquals("new_field", updated.get().getFields(0).getName()); + } + + @Test + public void testCheckUpdatedSchema_autoUpdateFalse() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(false, false, tableReference, false); + TableSchema currentSchema = TableSchema.newBuilder().build(); + + Optional updated = + helper.checkUpdatedSchema(currentSchema, "test_stream", mockWriteStreamService); + + assertFalse(updated.isPresent()); + } + + @Test + public void testCheckUpdatedSchema_autoUpdateTrue_noUpdate() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(true, false, tableReference, false); + TableSchema currentSchema = TableSchema.newBuilder().build(); + + when(mockWriteStreamService.getWriteStreamSchema("test_stream")).thenReturn(null); + + Optional updated = + helper.checkUpdatedSchema(currentSchema, "test_stream", mockWriteStreamService); + + assertFalse(updated.isPresent()); + } + + @Test + public void testCheckResponseForUpdatedSchema_hasUpdate() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(true, false, tableReference, false); + TableSchema currentSchema = TableSchema.newBuilder().build(); + TableSchema newSchema = + TableSchema.newBuilder() + .addFields(TableFieldSchema.newBuilder().setName("new_field").build()) + .build(); + + when(mockStreamAppendClient.getUpdatedSchema()).thenReturn(newSchema); + + Optional updated = + helper.checkResponseForUpdatedSchema(currentSchema, mockStreamAppendClient); + + assertTrue(updated.isPresent()); + assertEquals(1, updated.get().getFieldsCount()); + assertEquals("new_field", updated.get().getFields(0).getName()); + } + + @Test + public void testCheckResponseForUpdatedSchema_noUpdate() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(true, false, tableReference, false); + TableSchema currentSchema = TableSchema.newBuilder().build(); + + when(mockStreamAppendClient.getUpdatedSchema()).thenReturn(null); + + Optional updated = + helper.checkResponseForUpdatedSchema(currentSchema, mockStreamAppendClient); + + assertFalse(updated.isPresent()); + } + + @Test + public void testGetMergedPayload_autoUpdateFalse() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(false, false, tableReference, false); + StorageApiWritePayload payload = + StorageApiWritePayload.of(new byte[] {1, 2, 3}, new TableRow().set("foo", "bar"), null); + + SchemaChangeDetectorHelper.MergePayloadResult result = + helper.getMergedPayload(payload, Instant.now(), null, mockAppendClientInfo); + + assertEquals(SchemaChangeDetectorHelper.MergePayloadResult.Kind.MERGED, result.getKind()); + assertArrayEquals(new byte[] {1, 2, 3}, result.getMerged().toByteArray()); + } + + @Test + public void testGetMergedPayload_autoUpdateTrue_emptyUnknownFields() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(true, false, tableReference, false); + StorageApiWritePayload payload = StorageApiWritePayload.of(new byte[] {1, 2, 3}, null, null); + + SchemaChangeDetectorHelper.MergePayloadResult result = + helper.getMergedPayload(payload, Instant.now(), null, mockAppendClientInfo); + + assertEquals(SchemaChangeDetectorHelper.MergePayloadResult.Kind.MERGED, result.getKind()); + assertArrayEquals(new byte[] {1, 2, 3}, result.getMerged().toByteArray()); + } + + @Test + public void testGetMergedPayload_autoUpdateTrue_mergeSuccess() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(true, false, tableReference, false); + TableRow unknownFields = new TableRow().set("foo", "bar"); + StorageApiWritePayload payload = + StorageApiWritePayload.of(new byte[] {1, 2, 3}, unknownFields, null); + + ByteString mergedBytes = ByteString.copyFrom(new byte[] {4, 5, 6}); + when(mockAppendClientInfo.mergeNewFields(any(ByteString.class), eq(unknownFields), eq(false))) + .thenReturn(mergedBytes); + + SchemaChangeDetectorHelper.MergePayloadResult result = + helper.getMergedPayload(payload, Instant.now(), null, mockAppendClientInfo); + + assertEquals(SchemaChangeDetectorHelper.MergePayloadResult.Kind.MERGED, result.getKind()); + assertArrayEquals(new byte[] {4, 5, 6}, result.getMerged().toByteArray()); + } + + @Test + public void testGetMergedPayload_autoUpdateTrue_mergeFailure() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(true, false, tableReference, false); + TableRow unknownFields = new TableRow().set("foo", "bar"); + StorageApiWritePayload payload = + StorageApiWritePayload.of(new byte[] {1, 2, 3}, unknownFields, null); + + when(mockAppendClientInfo.mergeNewFields(any(ByteString.class), eq(unknownFields), eq(false))) + .thenThrow(new TableRowToStorageApiProto.SchemaDoesntMatchException("conversion error")); + TableRow expectedFailsafe = new TableRow().set("failsafe", "true"); + + SchemaChangeDetectorHelper.MergePayloadResult result = + helper.getMergedPayload(payload, Instant.now(), expectedFailsafe, mockAppendClientInfo); + + assertEquals(SchemaChangeDetectorHelper.MergePayloadResult.Kind.FAILED, result.getKind()); + TimestampedValue failed = result.getFailed(); + assertEquals( + "org.apache.beam.sdk.io.gcp.bigquery.TableRowToStorageApiProto$SchemaDoesntMatchException: conversion error", + failed.getValue().getErrorMessage()); + assertEquals(expectedFailsafe, failed.getValue().getRow()); + } + + @Test + public void testIsPayloadSchemaOutOfDate_ignoreSchemaHashesTrue_hasMissingUnknownField() + throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(false, false, tableReference, true); + TableRow unknownFields = new TableRow().set("foo", "bar"); + StorageApiWritePayload payload = + StorageApiWritePayload.of(new byte[] {1, 2, 3}, unknownFields, null); + + DescriptorProtos.DescriptorProto descriptorProto = + DescriptorProtos.DescriptorProto.newBuilder().setName("test").build(); + Descriptors.Descriptor descriptor = + TableRowToStorageApiProto.wrapDescriptorProto(descriptorProto); + + boolean outOfDate = + helper.isPayloadSchemaOutOfDate(payload, new byte[0], () -> new byte[0], () -> descriptor); + + assertTrue(outOfDate); + } + + @Test + public void testIsPayloadSchemaOutOfDate_ignoreSchemaHashesTrue_noUnknownFields() + throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(false, false, tableReference, true); + StorageApiWritePayload payload = StorageApiWritePayload.of(new byte[] {1, 2, 3}, null, null); + + boolean outOfDate = + helper.isPayloadSchemaOutOfDate(payload, new byte[0], () -> new byte[0], () -> null); + + assertFalse(outOfDate); + } + + @Test + public void testIsPayloadSchemaOutOfDate_ignoreSchemaHashesFalse_hashMatch() throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(false, false, tableReference, false); + byte[] hash = "hash".getBytes(StandardCharsets.UTF_8); + StorageApiWritePayload payload = + StorageApiWritePayload.of(new byte[] {1, 2, 3}, null, null).withSchemaHash(hash); + + boolean outOfDate = + helper.isPayloadSchemaOutOfDate(payload, new byte[0], () -> hash, () -> null); + + assertFalse(outOfDate); + } + + @Test + public void testIsPayloadSchemaOutOfDate_ignoreSchemaHashesFalse_hashMismatch_noUnknownFields() + throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(false, false, tableReference, false); + byte[] payloadHash = "hash1".getBytes(StandardCharsets.UTF_8); + byte[] currentHash = "hash2".getBytes(StandardCharsets.UTF_8); + StorageApiWritePayload payload = + StorageApiWritePayload.of(new byte[0], null, null).withSchemaHash(payloadHash); + + DescriptorProtos.DescriptorProto descriptorProto = + DescriptorProtos.DescriptorProto.newBuilder().setName("test").build(); + Descriptors.Descriptor descriptor = + TableRowToStorageApiProto.wrapDescriptorProto(descriptorProto); + + // DynamicMessage will parse new byte[0] cleanly, and we have no unknown fields. + boolean outOfDate = + helper.isPayloadSchemaOutOfDate(payload, new byte[0], () -> currentHash, () -> descriptor); + + assertFalse(outOfDate); + } + + @Test + public void testIsPayloadSchemaOutOfDate_ignoreSchemaHashesFalse_hashMismatch_hasUnknownFields() + throws Exception { + SchemaChangeDetectorHelper helper = + new SchemaChangeDetectorHelper(false, false, tableReference, false); + byte[] payloadHash = "hash1".getBytes(StandardCharsets.UTF_8); + byte[] currentHash = "hash2".getBytes(StandardCharsets.UTF_8); + StorageApiWritePayload payload = + StorageApiWritePayload.of(new byte[0], null, null).withSchemaHash(payloadHash); + + DescriptorProtos.DescriptorProto descriptorProto = + DescriptorProtos.DescriptorProto.newBuilder().setName("test").build(); + Descriptors.Descriptor descriptor = + TableRowToStorageApiProto.wrapDescriptorProto(descriptorProto); + + // To have unknown fields, we provide some valid protobuf bytes that contain a tag not in + // descriptor. + // Tag 1 (varint) = 1 << 3 | 0 = 8. Value 1. + byte[] payloadWithUnknown = new byte[] {8, 1}; + + boolean outOfDate = + helper.isPayloadSchemaOutOfDate( + payload, payloadWithUnknown, () -> currentHash, () -> descriptor); + + assertTrue(outOfDate); + } +} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterableTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterableTest.java new file mode 100644 index 000000000000..c5e4fb0aa012 --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterableTest.java @@ -0,0 +1,244 @@ +/* + * 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.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; + +import com.google.api.services.bigquery.model.TableReference; +import com.google.api.services.bigquery.model.TableRow; +import com.google.cloud.bigquery.storage.v1.TableSchema; +import com.google.protobuf.DescriptorProtos; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.stream.Collectors; +import org.apache.beam.sdk.values.TimestampedValue; +import org.joda.time.Instant; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class SplittingIterableTest { + + private AppendClientInfo createAppendClientInfo() throws Exception { + TableSchema tableSchema = TableSchema.newBuilder().build(); + DescriptorProtos.DescriptorProto descriptor = + DescriptorProtos.DescriptorProto.newBuilder().setName("test").build(); + return AppendClientInfo.of(tableSchema, descriptor, client -> {}); + } + + @Test + public void testBatchingBySplitSize() throws Exception { + List payloads = new ArrayList<>(); + // Payload of 10 bytes each + for (int i = 0; i < 5; i++) { + payloads.add( + StorageApiWritePayload.of(new byte[10], null, null) + .withTimestamp(Instant.ofEpochMilli(i))); + } + + List> failedRows = new ArrayList<>(); + SchemaChangeDetectorHelper schemaChangeDetectorHelper = + new SchemaChangeDetectorHelper(false, false, new TableReference(), false); + AppendClientInfo appendClientInfo = createAppendClientInfo(); + + // Split size 25 means 2 payloads (20 bytes) per batch, then 2, then 1. + SplittingIterable iterable = + new SplittingIterable( + payloads, + 25, + failedRows::add, + () -> new byte[0], + () -> TableRowToStorageApiProto.wrapDescriptorProto(appendClientInfo.getDescriptor()), + Instant.now(), + () -> appendClientInfo, + schemaChangeDetectorHelper); + + Iterator it = iterable.iterator(); + + assertTrue(it.hasNext()); + AppendRowsPacket batch1 = it.next(); + assertEquals(2, batch1.getProtoRows().getSerializedRowsCount()); + assertEquals(Instant.ofEpochMilli(0), batch1.getTimestamps().get(0)); + assertEquals(Instant.ofEpochMilli(1), batch1.getTimestamps().get(1)); + + assertTrue(it.hasNext()); + AppendRowsPacket batch2 = it.next(); + assertEquals(2, batch2.getProtoRows().getSerializedRowsCount()); + + assertTrue(it.hasNext()); + AppendRowsPacket batch3 = it.next(); + assertEquals(1, batch3.getProtoRows().getSerializedRowsCount()); + + assertFalse(it.hasNext()); + assertTrue(failedRows.isEmpty()); + } + + @Test + public void testLargeElementExceedingSplitSize() throws Exception { + List payloads = new ArrayList<>(); + // Payload of 10 bytes, 100 bytes, 10 bytes + payloads.add(StorageApiWritePayload.of(new byte[10], null, null)); + payloads.add(StorageApiWritePayload.of(new byte[100], null, null)); + payloads.add(StorageApiWritePayload.of(new byte[10], null, null)); + + List> failedRows = new ArrayList<>(); + SchemaChangeDetectorHelper schemaChangeDetectorHelper = + new SchemaChangeDetectorHelper(false, false, new TableReference(), false); + AppendClientInfo appendClientInfo = createAppendClientInfo(); + + // Split size 25 + SplittingIterable iterable = + new SplittingIterable( + payloads, + 25, + failedRows::add, + () -> new byte[0], + () -> TableRowToStorageApiProto.wrapDescriptorProto(appendClientInfo.getDescriptor()), + Instant.now(), + () -> appendClientInfo, + schemaChangeDetectorHelper); + + Iterator it = iterable.iterator(); + + assertTrue(it.hasNext()); + AppendRowsPacket batch1 = it.next(); + assertEquals(1, batch1.getProtoRows().getSerializedRowsCount()); + + assertTrue(it.hasNext()); + AppendRowsPacket batch2 = it.next(); + assertEquals(1, batch2.getProtoRows().getSerializedRowsCount()); + assertEquals(100, batch2.getProtoRows().getSerializedRows(0).size()); + + assertTrue(it.hasNext()); + AppendRowsPacket batch3 = it.next(); + assertEquals(1, batch3.getProtoRows().getSerializedRowsCount()); + + assertFalse(it.hasNext()); + } + + @Test + public void testSchemaMismatchedAndMatchedRows() throws Exception { + List payloads = new ArrayList<>(); + byte[] hash1 = "currentHash".getBytes(StandardCharsets.UTF_8); + byte[] hash2 = "oldHash".getBytes(StandardCharsets.UTF_8); + + TableRow unknownFieldsRow = new TableRow().set("foo", "bar"); + + payloads.add(StorageApiWritePayload.of(new byte[0], null, null).withSchemaHash(hash1)); + payloads.add( + StorageApiWritePayload.of(new byte[0], unknownFieldsRow, null).withSchemaHash(hash2)); + payloads.add(StorageApiWritePayload.of(new byte[0], null, null).withSchemaHash(hash1)); + + List> failedRows = new ArrayList<>(); + SchemaChangeDetectorHelper schemaChangeDetectorHelper = + new SchemaChangeDetectorHelper(false, true, new TableReference(), true); + AppendClientInfo appendClientInfo = createAppendClientInfo(); + + SplittingIterable iterable = + new SplittingIterable( + payloads, + 100, + failedRows::add, + () -> "currentHash".getBytes(StandardCharsets.UTF_8), // Matches hash1 + () -> TableRowToStorageApiProto.wrapDescriptorProto(appendClientInfo.getDescriptor()), + Instant.now(), + () -> appendClientInfo, + schemaChangeDetectorHelper); + + Iterator it = iterable.iterator(); + assertTrue(it.hasNext()); + AppendRowsPacket batch = it.next(); + + assertEquals(3, batch.getProtoRows().getSerializedRowsCount()); + + // Check getting only mismatched rows + AppendRowsPacket mismatched = batch.getSchemaMismatchedRowsOnly(); + assertEquals(1, mismatched.getProtoRows().getSerializedRowsCount()); + assertArrayEquals(new byte[0], mismatched.getProtoRows().getSerializedRows(0).toByteArray()); + assertEquals(unknownFieldsRow, mismatched.getUnknownFields().get(0)); + assertArrayEquals(new byte[0], mismatched.getOriginalPayloads().get(0)); + assertArrayEquals(hash2, mismatched.getSchemaHashes().get(0)); + + // Check getting only matched rows + AppendRowsPacket matched = batch.getSchemaMatchedRowsOnly(); + assertEquals(2, matched.getProtoRows().getSerializedRowsCount()); + assertArrayEquals(new byte[0], matched.getProtoRows().getSerializedRows(0).toByteArray()); + assertArrayEquals(new byte[0], matched.getProtoRows().getSerializedRows(1).toByteArray()); + assertArrayEquals(hash1, matched.getSchemaHashes().get(0)); + assertArrayEquals(hash1, matched.getSchemaHashes().get(1)); + assertNull(matched.getUnknownFields().get(0)); + assertNull(matched.getOriginalPayloads().get(0)); + + // Reconstruct stream + List reconstructed = + batch.toPayloadStream().collect(Collectors.toList()); + assertEquals(3, reconstructed.size()); + assertArrayEquals(new byte[0], reconstructed.get(0).getPayload()); + assertArrayEquals(new byte[0], reconstructed.get(1).getPayload()); + assertArrayEquals(new byte[0], reconstructed.get(2).getPayload()); + } + + @Test + public void testFailedRows() throws Exception { + List payloads = new ArrayList<>(); + payloads.add(StorageApiWritePayload.of(new byte[0], null, null)); + // Provide unknown fields, so auto update schema tries to merge and fails + TableRow unknownFieldsRow = new TableRow().set("foo", "bar"); + payloads.add(StorageApiWritePayload.of(new byte[0], unknownFieldsRow, null)); + + List> failedRows = new ArrayList<>(); + SchemaChangeDetectorHelper schemaChangeDetectorHelper = + new SchemaChangeDetectorHelper( + true, // autoUpdateSchema + false, // ignoreUnknownValues (this will trigger SchemaConversionException) + new TableReference().setTableId("test_table"), + false); + AppendClientInfo appendClientInfo = createAppendClientInfo(); + + SplittingIterable iterable = + new SplittingIterable( + payloads, + 100, + failedRows::add, + () -> new byte[0], + () -> TableRowToStorageApiProto.wrapDescriptorProto(appendClientInfo.getDescriptor()), + Instant.now(), + () -> appendClientInfo, + schemaChangeDetectorHelper); + + Iterator it = iterable.iterator(); + assertTrue(it.hasNext()); + AppendRowsPacket batch = it.next(); + + // Only 1 row successfully added to batch + assertEquals(1, batch.getProtoRows().getSerializedRowsCount()); + assertArrayEquals(new byte[0], batch.getProtoRows().getSerializedRows(0).toByteArray()); + + // 1 row failed + assertEquals(1, failedRows.size()); + BigQueryStorageApiInsertError error = failedRows.get(0).getValue(); + assertEquals("test_table", error.getTable().getTableId()); + } +} diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/UpgradeTableSchemaTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/UpgradeTableSchemaTest.java index 13724ab83e61..2e1da866b260 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/UpgradeTableSchemaTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/UpgradeTableSchemaTest.java @@ -22,6 +22,8 @@ import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; +import com.google.api.services.bigquery.model.TableCell; +import com.google.api.services.bigquery.model.TableRow; import com.google.cloud.bigquery.storage.v1.TableFieldSchema; import com.google.cloud.bigquery.storage.v1.TableSchema; import com.google.protobuf.DescriptorProtos; @@ -29,6 +31,7 @@ import com.google.protobuf.DynamicMessage; import com.google.protobuf.UnknownFieldSet; import java.nio.charset.StandardCharsets; +import java.util.Arrays; import org.apache.beam.sdk.util.Preconditions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing; @@ -325,14 +328,19 @@ public void testIsPayloadSchemaOutOfDate() throws Exception { // 1. Missing hash in payload StorageApiWritePayload payloadNoHash = StorageApiWritePayload.of(new byte[0], null, null); - assertFalse(UpgradeTableSchema.isPayloadSchemaOutOfDate(payloadNoHash, () -> hash1, null)); + assertFalse( + UpgradeTableSchema.isPayloadSchemaOutOfDate( + payloadNoHash.getSchemaHash(), payloadNoHash::getPayload, () -> hash1, null)); // 2. Equal hash - assertFalse(UpgradeTableSchema.isPayloadSchemaOutOfDate(payload, () -> hash1, null)); + assertFalse( + UpgradeTableSchema.isPayloadSchemaOutOfDate( + payload.getSchemaHash(), payload::getPayload, () -> hash1, null)); // 3. Different hash, but message doesn't have unknown fields (initialized and empty) assertFalse( - UpgradeTableSchema.isPayloadSchemaOutOfDate(payload, () -> hash2, () -> descriptor)); + UpgradeTableSchema.isPayloadSchemaOutOfDate( + payload.getSchemaHash(), payload::getPayload, () -> hash2, () -> descriptor)); // 4. Different hash with unknown fields. DynamicMessage unknownFieldSet = @@ -348,7 +356,10 @@ public void testIsPayloadSchemaOutOfDate() throws Exception { StorageApiWritePayload.of(unknownFieldSet.toByteArray(), null, null).withSchemaHash(hash1); assertTrue( UpgradeTableSchema.isPayloadSchemaOutOfDate( - payloadWithUnknown, () -> hash2, () -> descriptor)); + payloadWithUnknown.getSchemaHash(), + payloadWithUnknown::getPayload, + () -> hash2, + () -> descriptor)); // 5. Different hash with missing required fields. DynamicMessage missingField = DynamicMessage.newBuilder(descriptor).buildPartial(); @@ -356,6 +367,110 @@ public void testIsPayloadSchemaOutOfDate() throws Exception { StorageApiWritePayload.of(missingField.toByteArray(), null, null).withSchemaHash(hash1); assertTrue( UpgradeTableSchema.isPayloadSchemaOutOfDate( - payloadMissingField, () -> hash2, () -> descriptor)); + payloadMissingField.getSchemaHash(), + payloadMissingField::getPayload, + () -> hash2, + () -> descriptor)); + } + + @Test + public void testMissingUnknownField_TableRowWithF() throws Exception { + DescriptorProtos.DescriptorProto descriptorProto = + DescriptorProtos.DescriptorProto.newBuilder() + .setName("TestMessage") + .addField( + DescriptorProtos.FieldDescriptorProto.newBuilder() + .setName("field1") + .setNumber(1) + .setType(DescriptorProtos.FieldDescriptorProto.Type.TYPE_STRING) + .build()) + .build(); + Descriptors.Descriptor descriptor = + TableRowToStorageApiProto.wrapDescriptorProto(descriptorProto); + + TableRow row = new TableRow(); + TableCell cell1 = new TableCell(); + cell1.setV("value1"); + row.set("f", Arrays.asList(cell1)); + + assertFalse(UpgradeTableSchema.missingUnknownField(row, () -> descriptor)); + + TableCell cell2 = new TableCell(); + cell2.setV("value2"); + row.set("f", Arrays.asList(cell1, cell2)); + + assertTrue(UpgradeTableSchema.missingUnknownField(row, () -> descriptor)); + } + + @Test + public void testMissingUnknownField_TableRow() throws Exception { + DescriptorProtos.DescriptorProto descriptorProto = + DescriptorProtos.DescriptorProto.newBuilder() + .setName("TestMessage") + .addField( + DescriptorProtos.FieldDescriptorProto.newBuilder() + .setName("field1") + .setNumber(1) + .setType(DescriptorProtos.FieldDescriptorProto.Type.TYPE_STRING) + .build()) + .build(); + Descriptors.Descriptor descriptor = + TableRowToStorageApiProto.wrapDescriptorProto(descriptorProto); + + TableRow row = new TableRow(); + row.set("field1", "value1"); + + assertFalse(UpgradeTableSchema.missingUnknownField(row, () -> descriptor)); + + row.set("field2", "value2"); + + assertTrue(UpgradeTableSchema.missingUnknownField(row, () -> descriptor)); + } + + @Test + public void testMissingUnknownField_NestedObject() throws Exception { + DescriptorProtos.DescriptorProto nestedProto = + DescriptorProtos.DescriptorProto.newBuilder() + .setName("NestedMessage") + .addField( + DescriptorProtos.FieldDescriptorProto.newBuilder() + .setName("nested1") + .setNumber(1) + .setType(DescriptorProtos.FieldDescriptorProto.Type.TYPE_STRING) + .build()) + .build(); + + DescriptorProtos.DescriptorProto parentProto = + DescriptorProtos.DescriptorProto.newBuilder() + .setName("TestMessage") + .addField( + DescriptorProtos.FieldDescriptorProto.newBuilder() + .setName("field1") + .setNumber(1) + .setType(DescriptorProtos.FieldDescriptorProto.Type.TYPE_MESSAGE) + .setTypeName("NestedMessage") + .build()) + .build(); + + DescriptorProtos.FileDescriptorProto fileProto = + DescriptorProtos.FileDescriptorProto.newBuilder() + .addMessageType(nestedProto) + .addMessageType(parentProto) + .build(); + + Descriptors.FileDescriptor fileDescriptor = + Descriptors.FileDescriptor.buildFrom(fileProto, new Descriptors.FileDescriptor[0]); + Descriptors.Descriptor descriptor = fileDescriptor.findMessageTypeByName("TestMessage"); + + TableRow nestedRow = new TableRow(); + nestedRow.set("nested1", "value1"); + + TableRow parentRow = new TableRow(); + parentRow.set("field1", nestedRow); + + assertFalse(UpgradeTableSchema.missingUnknownField(parentRow, () -> descriptor)); + + nestedRow.set("nested2", "value2"); + assertTrue(UpgradeTableSchema.missingUnknownField(parentRow, () -> descriptor)); } } From 4ac6c182ce3317667f38d4e0f6888b53d98374c6 Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Tue, 9 Jun 2026 15:21:47 -0700 Subject: [PATCH 02/10] reenable window expiration context --- .../beam/sdk/transforms/reflect/DoFnSignatures.java | 1 - .../sdk/io/gcp/bigquery/SchemaUpdateHoldingFn.java | 10 ++++------ 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java index 4a60ef4e5b5b..9f3491bca7b9 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/reflect/DoFnSignatures.java @@ -222,7 +222,6 @@ private DoFnSignatures() {} private static final Collection> ALLOWED_ON_WINDOW_EXPIRATION_PARAMETERS = ImmutableList.of( - Parameter.OnWindowExpirationContextParameter.class, Parameter.WindowParameter.class, Parameter.PipelineOptionsParameter.class, Parameter.OutputReceiverParameter.class, diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaUpdateHoldingFn.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaUpdateHoldingFn.java index fa37327ac843..456b557d61c4 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaUpdateHoldingFn.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaUpdateHoldingFn.java @@ -191,18 +191,16 @@ public void onPollTimer( @OnWindowExpiration public void onWindowExpiration( - // OnWindowExpirationContext context, + OnWindowExpirationContext context, @Key ShardedKey key, PipelineOptions pipelineOptions, @StateId("bufferedElements") BagState> bag, @StateId("minBufferedTimestamp") CombiningState minBufferedTimestamp, MultiOutputReceiver o) throws Exception { - // TODO: Beam doesn't current support OnWindowExpirationContext. Reenable after this is fixed. - // https://github.com/apache/beam/issues/38875 - // convertMessagesDoFn - // .getDynamicDestinations() - // .setSideInputAccessorFromOnWindowExpirationContext(context); + convertMessagesDoFn + .getDynamicDestinations() + .setSideInputAccessorFromOnWindowExpirationContext(context); // This can happen on test completion or drain. We can't set any more timers in window // expiration, so we just have to loop until the schema is updated. From b91e8aab56f5202ee0da4b3b107e3e2b020019ee Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Tue, 9 Jun 2026 19:24:50 -0700 Subject: [PATCH 03/10] fix test failure --- .../sdk/io/gcp/bigquery/BigQueryIOTranslation.java | 13 +++++++++++++ .../io/gcp/bigquery/BigQueryIOTranslationTest.java | 2 ++ 2 files changed, 15 insertions(+) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java index dd59939726bf..7a7f2300a459 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java @@ -438,6 +438,7 @@ static class BigQueryIOWriteTranslator implements TransformPayloadTranslator transform) { fieldValues.put("use_beam_schema", transform.getUseBeamSchema()); fieldValues.put("auto_sharding", transform.getAutoSharding()); fieldValues.put("auto_schema_update", transform.getAutoSchemaUpdate()); + fieldValues.put( + "auto_schema_update_strict_timeout_ms", + transform.getAutoSchemaUpdateStrictTimeout().getMillis()); if (transform.getWriteProtosClass() != null) { fieldValues.put("write_protos_class", toByteArray(transform.getWriteProtosClass())); } @@ -870,6 +874,15 @@ public Write fromConfigRow(Row configRow, PipelineOptions options) { if (autoSchemaUpdate != null) { builder = builder.setAutoSchemaUpdate(autoSchemaUpdate); } + Long autoSchemaUpdateStrictTimeoutMillis = + configRow.getInt64("auto_schema_update_strict_timeout_ms"); + if (autoSchemaUpdateStrictTimeoutMillis != null) { + builder = + builder + .setAutoSchemaUpdate(true) + .setAutoSchemaUpdateStrictTimeout( + org.joda.time.Duration.millis(autoSchemaUpdateStrictTimeoutMillis)); + } byte[] writeProtosClasses = configRow.getBytes("write_protos_class"); if (writeProtosClasses != null) { builder = diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslationTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslationTest.java index de63120c93cc..4391a66e27ee 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslationTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslationTest.java @@ -132,6 +132,8 @@ public class BigQueryIOTranslationTest { WRITE_TRANSFORM_SCHEMA_MAPPING.put("getAutoSharding", "auto_sharding"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getPropagateSuccessful", "propagate_successful"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getAutoSchemaUpdate", "auto_schema_update"); + WRITE_TRANSFORM_SCHEMA_MAPPING.put( + "getAutoSchemaUpdateStrictTimeout", "auto_schema_update_strict_timeout_ms"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getWriteProtosClass", "write_protos_class"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getDirectWriteProtos", "direct_write_protos"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getDeterministicRecordIdFn", "deterministic_record_id_fn"); From 668631977957b002f3402fe7fefde545d046c05f Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Tue, 9 Jun 2026 22:22:18 -0700 Subject: [PATCH 04/10] fix error --- .../beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java index 7a7f2300a459..ee490b227242 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java @@ -583,9 +583,11 @@ public Row toConfigRow(Write transform) { fieldValues.put("use_beam_schema", transform.getUseBeamSchema()); fieldValues.put("auto_sharding", transform.getAutoSharding()); fieldValues.put("auto_schema_update", transform.getAutoSchemaUpdate()); - fieldValues.put( - "auto_schema_update_strict_timeout_ms", - transform.getAutoSchemaUpdateStrictTimeout().getMillis()); + if (transform.getAutoSchemaUpdateStrictTimeout() != null) { + fieldValues.put( + "auto_schema_update_strict_timeout_ms", + transform.getAutoSchemaUpdateStrictTimeout().getMillis()); + } if (transform.getWriteProtosClass() != null) { fieldValues.put("write_protos_class", toByteArray(transform.getWriteProtosClass())); } From e73fd702daf3c5d605820742664e4c2e7efe9eda Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Tue, 16 Jun 2026 10:40:38 -0700 Subject: [PATCH 05/10] foo --- .../io/gcp/bigquery/BufferMismatchedRows.java | 8 +++- .../StorageApiWritesShardedRecords.java | 3 +- .../StorageApiSinkSchemaUpdateIT.java | 37 +++++++++++++++---- 3 files changed, 38 insertions(+), 10 deletions(-) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java index 03c0dbfaa393..633c3295acb2 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java @@ -192,6 +192,11 @@ public void process( RETRY_MISMATCHED_ROWS_PERIOD); } + @Override + public Duration getAllowedTimestampSkew() { + return Duration.millis(Long.MAX_VALUE); + } + @OnTimer("retryMismatchedRowsTimer") public void onTimer( OnTimerContext context, @@ -233,6 +238,8 @@ public void onTimer( } mismatchedRowsBag.clear(); + currentTimerValue.clear(); + minPendingTimestamp.clear(); if (!mismatchedRowsList.isEmpty()) { TableDestination tableDestination = dynamicDestinations.getTable(shardedDestination.getKey()); @@ -257,7 +264,6 @@ public void onTimer( .map(KV::getValue) .iterator(); - currentTimerValue.clear(); SchemaChangeDetectorHelper.bufferMismatchedRows( mismatchedRows, mismatchedRowsBag, 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 295a67d1247c..be28713a8652 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 @@ -1305,6 +1305,8 @@ public void onMismatchedRowsTimer( mismatchedRows -> { // Rebuffer the ones that are still not succeeding. mismatchedRowsBag.clear(); + currentTimerValue.clear(); + minPendingTimestamp.clear(); if (!Iterables.isEmpty(mismatchedRows)) { AppendClientInfo info = AppendClientInfo.of( @@ -1326,7 +1328,6 @@ public void onMismatchedRowsTimer( } }; - currentTimerValue.clear(); processPayloads( pipelineOptions, shardedDestination, diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java index 3cb9897ada52..ba2f74f7378a 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java @@ -49,6 +49,7 @@ import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.WriteDisposition; import org.apache.beam.sdk.io.gcp.testing.BigqueryClient; import org.apache.beam.sdk.options.ExperimentalOptions; +import org.apache.beam.sdk.options.StreamingOptions; import org.apache.beam.sdk.state.StateSpec; import org.apache.beam.sdk.state.StateSpecs; import org.apache.beam.sdk.state.ValueState; @@ -358,7 +359,10 @@ private static TableSchema makeTableSchemaFromTypes( } private void runStreamingPipelineWithSchemaChange( - Write.Method method, boolean useAutoSchemaUpdate, boolean useIgnoreUnknownValues) + Write.Method method, + boolean useAutoSchemaUpdate, + boolean consistentAutoUpdate, + boolean useIgnoreUnknownValues) throws Exception { Pipeline p = Pipeline.create(TestPipeline.testingPipelineOptions()); // Set threshold bytes to 0 so that the stream attempts to fetch an updated schema after each @@ -367,6 +371,8 @@ private void runStreamingPipelineWithSchemaChange( // Limit parallelism so that all streams recognize the new schema in an expected short amount // of time (before we start writing rows with updated schema) p.getOptions().as(BigQueryOptions.class).setNumStorageWriteApiStreams(TOTAL_NUM_STREAMS); + p.getOptions().as(StreamingOptions.class).setStreaming(true); + // Need to manually enable streaming engine for legacy dataflow runner ExperimentalOptions.addExperiment( p.getOptions().as(ExperimentalOptions.class), GcpOptions.STREAMING_ENGINE_EXPERIMENT); @@ -374,7 +380,7 @@ private void runStreamingPipelineWithSchemaChange( if (p.getOptions().getRunner().getName().contains("DataflowRunner")) { assumeTrue( "Skipping in favor of more relevant test case and to avoid timing issues", - !changeTableSchema && useInputSchema && useAutoSchemaUpdate); + consistentAutoUpdate || (!changeTableSchema && useInputSchema && useAutoSchemaUpdate)); } List fieldNamesOrigin = new ArrayList(Arrays.asList(FIELDS)); @@ -402,6 +408,7 @@ private void runStreamingPipelineWithSchemaChange( BigQueryIO.writeTableRows() .to(tableSpec) .withAutoSchemaUpdate(useAutoSchemaUpdate) + .withAutoSchemaUpdateConsistent(consistentAutoUpdate, Duration.standardMinutes(5)) .withMethod(method) .withCreateDisposition(CreateDisposition.CREATE_NEVER) .withWriteDisposition(WriteDisposition.WRITE_APPEND); @@ -413,7 +420,8 @@ private void runStreamingPipelineWithSchemaChange( } // We give a healthy waiting period between each element to give Storage API streams a chance to // recognize the new schema. Apply on relevant tests. - boolean waitLonger = changeTableSchema && (useAutoSchemaUpdate || !useInputSchema); + boolean waitLonger = + changeTableSchema && (useAutoSchemaUpdate || !useInputSchema) && !consistentAutoUpdate; if (method == Write.Method.STORAGE_WRITE_API) { write = write.withTriggeringFrequency( @@ -579,33 +587,46 @@ public void testExactlyOnce() throws Exception { Write.Method.STORAGE_WRITE_API, /** autoSchemaUpdate */ false, + false, /** ignoreUnknownvalues */ false); } @Test public void testExactlyOnceWithIgnoreUnknownValues() throws Exception { - runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_WRITE_API, false, true); + runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_WRITE_API, false, false, true); } @Test public void testExactlyOnceWithAutoSchemaUpdate() throws Exception { - runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_WRITE_API, true, true); + runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_WRITE_API, true, false, true); + } + + @Test + public void testExactlyOnceWithAutoSchemaUpdateConsistent() throws Exception { + runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_WRITE_API, true, true, true); } @Test public void testAtLeastOnce() throws Exception { - runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_API_AT_LEAST_ONCE, false, false); + runStreamingPipelineWithSchemaChange( + Write.Method.STORAGE_API_AT_LEAST_ONCE, false, false, false); } @Test public void testAtLeastOnceWithIgnoreUnknownValues() throws Exception { - runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_API_AT_LEAST_ONCE, false, true); + runStreamingPipelineWithSchemaChange( + Write.Method.STORAGE_API_AT_LEAST_ONCE, false, false, true); } @Test public void testAtLeastOnceWithAutoSchemaUpdate() throws Exception { - runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_API_AT_LEAST_ONCE, true, true); + runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_API_AT_LEAST_ONCE, true, false, true); + } + + @Test + public void testAtLeastOnceWithAutoSchemaUpdateConsistent() throws Exception { + runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_API_AT_LEAST_ONCE, true, true, true); } public void runDynamicDestinationsWithAutoSchemaUpdate(boolean useAtLeastOnce) throws Exception { From dab71a70eaba70db9250181877e35bb92f2741ef Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Thu, 25 Jun 2026 15:53:05 -0700 Subject: [PATCH 06/10] review --- .../sdk/io/gcp/bigquery/BigQueryOptions.java | 13 +- .../io/gcp/bigquery/BufferMismatchedRows.java | 21 +- .../StorageApiWriteRecordsInconsistent.java | 12 +- .../StorageApiWriteUnshardedRecords.java | 322 ++++++++++-------- .../StorageApiWritesShardedRecords.java | 17 +- .../io/gcp/bigquery/BigQueryIOWriteTest.java | 4 +- 6 files changed, 220 insertions(+), 169 deletions(-) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java index bc1c8fcaf1ab..41413cdfa930 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java @@ -262,9 +262,16 @@ public interface BigQueryOptions void setSchemaUpgradeBufferingShards(Integer value); @Hidden - @Description("The initial retry time in milliseconds when a schema mismatch is detected.") + @Description("How long to retry locally before buffering when a schema mismatch is detected.") @Default.Integer(5000) - Integer getStorageApiInitialMismatchRetryTimeMilliSec(); + Integer getStorageApiMismatchLocalRetryTimeMilliSec(); - void setStorageApiInitialMismatchRetryTimeMilliSec(Integer value); + void setStorageApiMismatchLocalRetryTimeMilliSec(Integer value); + + @Hidden + @Description("The retry time in milliseconds when a schema mismatch is detected.") + @Default.Integer(60000) + Integer getStorageApiMismatchRetryTimeMilliSec(); + + void setStorageApiMismatchRetryTimeMilliSec(Integer value); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java index 633c3295acb2..a3cbcf5c16f6 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java @@ -59,21 +59,21 @@ class BufferMismatchedRows private final Coder successfulRowsCoder; private final Coder destinationCoder; private final StorageApiDynamicDestinations dynamicDestinations; - private final StorageApiWriteUnshardedRecords.WriteRecordsDoFn writeDoFn; + private final StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl + writeDoFn; private final TupleTag failedRowsTag; private final @Nullable TupleTag successfulRowsTag; // This output is effectively ignored, since we only support this code path for // StorageApiWriteRecordsInconsistent. private final TupleTag> finalizeTag = new TupleTag<>("finalizeTag"); private static final int NUM_DEFAULT_SHARDS = 20; - private static final Duration RETRY_MISMATCHED_ROWS_PERIOD = Duration.standardMinutes(1); public BufferMismatchedRows( Coder failedRowsCoder, Coder successfulRowsCoder, Coder destinationCoder, StorageApiDynamicDestinations dynamicDestinations, - StorageApiWriteUnshardedRecords.WriteRecordsDoFn writeDoFn, + StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl writeDoFn, TupleTag failedRowsTag, @Nullable TupleTag successfulRowsTag) { this.failedRowsCoder = failedRowsCoder; @@ -136,7 +136,7 @@ public void process( class BufferingDoFn extends DoFn, MismatchedRow>, KV> { - private final StorageApiWriteUnshardedRecords.WriteRecordsDoFn + private final StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl writeDoFn; @StateId("mismatchedRows") @@ -157,7 +157,7 @@ class BufferingDoFn Metrics.counter(BufferMismatchedRows.BufferingDoFn.class, "rowsSentToFailedRowsCollection"); public BufferingDoFn( - StorageApiWriteUnshardedRecords.WriteRecordsDoFn writeDoFn) { + StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl writeDoFn) { this.writeDoFn = writeDoFn; } @@ -168,6 +168,7 @@ public void startBundle() throws IOException { @ProcessElement public void process( + PipelineOptions pipelineOptions, ProcessContext processContext, @Element KV, MismatchedRow> element, @StateId("mismatchedRows") BagState mismatchedRowsBag, @@ -179,6 +180,9 @@ public void process( dynamicDestinations.setSideInputAccessorFromProcessContext(processContext); TableDestination tableDestination = dynamicDestinations.getTable(element.getKey().getKey()); + Duration timerRetryDuration = + Duration.millis( + pipelineOptions.as(BigQueryOptions.class).getStorageApiMismatchRetryTimeMilliSec()); SchemaChangeDetectorHelper.bufferMismatchedRows( Collections.singleton(element.getValue()), mismatchedRowsBag, @@ -189,7 +193,7 @@ public void process( o.get(failedRowsTag), null, rowsSentToFailedRowsCollection, - RETRY_MISMATCHED_ROWS_PERIOD); + timerRetryDuration); } @Override @@ -264,6 +268,9 @@ public void onTimer( .map(KV::getValue) .iterator(); + Duration timerRetryDuration = + Duration.millis( + pipelineOptions.as(BigQueryOptions.class).getStorageApiMismatchRetryTimeMilliSec()); SchemaChangeDetectorHelper.bufferMismatchedRows( mismatchedRows, mismatchedRowsBag, @@ -274,7 +281,7 @@ public void onTimer( o.get(failedRowsTag), appendClientInfo, rowsSentToFailedRowsCollection, - RETRY_MISMATCHED_ROWS_PERIOD); + timerRetryDuration); } } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRecordsInconsistent.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRecordsInconsistent.java index 79643f0653f4..68d4e38a5f90 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRecordsInconsistent.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRecordsInconsistent.java @@ -120,8 +120,9 @@ public PCollectionTuple expand(PCollection fn = - new StorageApiWriteUnshardedRecords.WriteRecordsDoFn<>( + + StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl fnImpl = + new StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl<>( operationName, dynamicDestinations, bqServices, @@ -143,7 +144,10 @@ public PCollectionTuple expand(PCollection fn = + new StorageApiWriteUnshardedRecords.WriteRecordsDoFn<>(fnImpl); + ; PCollectionTuple result = input.apply( "Write Records", @@ -167,7 +171,7 @@ public PCollectionTuple expand(PCollection( - operationName, - dynamicDestinations, - bqServices, - false, - options.getStorageApiAppendThresholdBytes(), - options.getStorageApiAppendThresholdRecordCount(), - options.getNumStorageWriteApiStreamAppendClients(), - finalizeTag, - failedRowsTag, - successfulRowsTag, - successfulRowsPredicate, - autoUpdateSchema, - null, - null, - ignoreUnknownValues, - createDisposition, - kmsKey, - usesCdc, - defaultMissingValueInterpretation, - options.getStorageWriteApiMaxRetries(), - bigLakeConfiguration, - options.getStorageApiInitialMismatchRetryTimeMilliSec())) + new WriteRecordsDoFnImpl<>( + operationName, + dynamicDestinations, + bqServices, + false, + options.getStorageApiAppendThresholdBytes(), + options.getStorageApiAppendThresholdRecordCount(), + options.getNumStorageWriteApiStreamAppendClients(), + finalizeTag, + failedRowsTag, + successfulRowsTag, + successfulRowsPredicate, + autoUpdateSchema, + null, + null, + ignoreUnknownValues, + createDisposition, + kmsKey, + usesCdc, + defaultMissingValueInterpretation, + options.getStorageWriteApiMaxRetries(), + bigLakeConfiguration, + options.getStorageApiMismatchLocalRetryTimeMilliSec()))) .withOutputTags(finalizeTag, tupleTagList) .withSideInputs(dynamicDestinations.getSideInputs())); @@ -225,6 +228,139 @@ public PCollectionTuple expand(PCollection extends DoFn, KV> { + WriteRecordsDoFnImpl writeRecordsDoFnImpl; + + WriteRecordsDoFn(WriteRecordsDoFnImpl writeRecordsDoFnImpl) { + this.writeRecordsDoFnImpl = writeRecordsDoFnImpl; + } + + @StartBundle + public void startBundle() throws IOException { + writeRecordsDoFnImpl.startBundle(); + } + + @FinishBundle + public void finishBundle(final FinishBundleContext context) throws Exception { + OutputReceiver failedRowsReceiver = + value -> + WindowedValues.builder() + .setValue(value) + .setTimestamp(GlobalWindow.INSTANCE.maxTimestamp()) + .setWindow(GlobalWindow.INSTANCE) + .setPaneInfo(PaneInfo.NO_FIRING) + .setReceiver( + windowedValue -> { + for (BoundedWindow window : windowedValue.getWindows()) { + context.output( + writeRecordsDoFnImpl.failedRowsTag, + windowedValue.getValue(), + windowedValue.getTimestamp(), + window); + } + }); + + OutputReceiver successfulRowsReceiver = null; + if (writeRecordsDoFnImpl.successfulRowsTag != null) { + successfulRowsReceiver = + makeSuccessfulRowsReceiver(context, writeRecordsDoFnImpl.successfulRowsTag); + } + + OutputReceiver> finalizeOutputRecever = + value -> + WindowedValues.>builder() + .setValue(value) + .setTimestamp(GlobalWindow.INSTANCE.maxTimestamp()) + .setWindow(GlobalWindow.INSTANCE) + .setPaneInfo(PaneInfo.NO_FIRING) + .setReceiver( + windowedValue -> { + for (BoundedWindow window : windowedValue.getWindows()) { + context.output( + writeRecordsDoFnImpl.finalizeTag, + windowedValue.getValue(), + windowedValue.getTimestamp(), + window); + } + }); + Iterable> mismatchedRows = + writeRecordsDoFnImpl.finishBundle( + failedRowsReceiver, successfulRowsReceiver, finalizeOutputRecever); + if (!Iterables.isEmpty(mismatchedRows)) { + mismatchedRows.forEach( + m -> { + org.joda.time.Instant ts = + MoreObjects.firstNonNull( + m.getValue().getPayload().getTimestamp(), + GlobalWindow.INSTANCE.maxTimestamp()); + context.output( + Preconditions.checkStateNotNull(writeRecordsDoFnImpl.mismatchedRowsTag), + m, + ts, + GlobalWindow.INSTANCE); + }); + } + } + + private OutputReceiver makeSuccessfulRowsReceiver( + FinishBundleContext context, TupleTag successfulRowsTag) { + return new OutputReceiver() { + @Override + public OutputBuilder builder(TableRow value) { + return WindowedValues.builder() + .setValue(value) + .setTimestamp(GlobalWindow.INSTANCE.maxTimestamp()) + .setWindow(GlobalWindow.INSTANCE) + .setPaneInfo(PaneInfo.NO_FIRING) + .setReceiver( + windowedValue -> { + for (BoundedWindow window : windowedValue.getWindows()) { + context.output( + successfulRowsTag, + windowedValue.getValue(), + windowedValue.getTimestamp(), + window); + } + }); + } + }; + } + + @ProcessElement + public void process( + ProcessContext c, + PipelineOptions pipelineOptions, + @Element KV element, + @Timestamp org.joda.time.Instant elementTs, + MultiOutputReceiver o) + throws Exception { + writeRecordsDoFnImpl.dynamicDestinations.setSideInputAccessorFromProcessContext(c); + Iterable> mismatchedRows = + writeRecordsDoFnImpl.processElement(pipelineOptions, element, elementTs, o); + if (!Iterables.isEmpty(mismatchedRows)) { + final OutputReceiver> mismatchedReceiver = + o.get(Preconditions.checkStateNotNull(writeRecordsDoFnImpl.mismatchedRowsTag)); + mismatchedRows.forEach( + m -> { + org.joda.time.Instant ts = + MoreObjects.firstNonNull(m.getValue().getPayload().getTimestamp(), elementTs); + mismatchedReceiver.outputWithTimestamp(m, ts); + }); + } + } + + @Teardown + public void teardown() { + writeRecordsDoFnImpl.teardown(); + } + + @Override + public Duration getAllowedTimestampSkew() { + return Duration.millis(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()); + } + } + + static class WriteRecordsDoFnImpl + implements Serializable { private final Counter forcedFlushes = Metrics.counter(WriteRecordsDoFn.class, "forcedFlushes"); private final TupleTag> finalizeTag; private final TupleTag failedRowsTag; @@ -504,8 +640,8 @@ void addMessage(StorageApiWritePayload payload) throws Exception { long flush( RetryManager retryManager, - OutputReceiver failedRowsReceiver, - @Nullable OutputReceiver successfulRowsReceiver, + DoFn.OutputReceiver failedRowsReceiver, + DoFn.@Nullable OutputReceiver successfulRowsReceiver, BiConsumer bufferMismatchedRow) throws Exception { if (pendingMessages.isEmpty()) { @@ -525,7 +661,7 @@ long flush( org.joda.time.Instant startMismatchTime = org.joda.time.Instant.now(); final Duration initialMismatchRetryTime = - Duration.millis(WriteRecordsDoFn.this.initialMismatchRetryTimeMilliSec); + Duration.millis(WriteRecordsDoFnImpl.this.initialMismatchRetryTimeMilliSec); AppendRowsPacket rowsToProcess; Iterable payloadsToIterate = pendingMessages; @@ -888,7 +1024,7 @@ long flush( private AppendRowsPacket buildInserts( Iterable payloads, - OutputReceiver failedRowsReceiver) { + DoFn.OutputReceiver failedRowsReceiver) { Supplier appendClientInfoSupplier = () -> { if (appendClientInfo == null) { @@ -974,7 +1110,7 @@ void postFlush() { private final @Nullable Map bigLakeConfiguration; private final int initialMismatchRetryTimeMilliSec; - WriteRecordsDoFn( + WriteRecordsDoFnImpl( String operationName, StorageApiDynamicDestinations dynamicDestinations, BigQueryServices bqServices, @@ -1028,8 +1164,8 @@ boolean shouldFlush(int recordBytes) { } Iterable> flushIfNecessary( - OutputReceiver failedRowsReceiver, - @Nullable OutputReceiver successfulRowsReceiver, + DoFn.OutputReceiver failedRowsReceiver, + DoFn.@Nullable OutputReceiver successfulRowsReceiver, int recordBytes) throws Exception { if (shouldFlush(recordBytes)) { @@ -1043,8 +1179,8 @@ Iterable> flushIfNecessary( } Iterable> flushAll( - OutputReceiver failedRowsReceiver, - @Nullable OutputReceiver successfulRowsReceiver) + DoFn.OutputReceiver failedRowsReceiver, + DoFn.@Nullable OutputReceiver successfulRowsReceiver) throws Exception { List> mismatchedRows = Lists.newArrayList(); List> retryManagers = @@ -1102,7 +1238,6 @@ WriteStreamService getWriteStreamService(PipelineOptions pipelineOptions) { return maybeWriteStreamService; } - @StartBundle public void startBundle() throws IOException { destinations = Maps.newHashMap(); numPendingRecords = 0; @@ -1162,34 +1297,11 @@ DestinationState createDestinationState( } } - @ProcessElement - public void process( - ProcessContext c, - PipelineOptions pipelineOptions, - @Element KV element, - @Timestamp org.joda.time.Instant elementTs, - MultiOutputReceiver o) - throws Exception { - dynamicDestinations.setSideInputAccessorFromProcessContext(c); - Iterable> mismatchedRows = - processElement(pipelineOptions, element, elementTs, o); - if (!Iterables.isEmpty(mismatchedRows)) { - final OutputReceiver> mismatchedReceiver = - o.get(Preconditions.checkStateNotNull(mismatchedRowsTag)); - mismatchedRows.forEach( - m -> { - org.joda.time.Instant ts = - MoreObjects.firstNonNull(m.getValue().getPayload().getTimestamp(), elementTs); - mismatchedReceiver.outputWithTimestamp(m, ts); - }); - } - } - Iterable> processElement( PipelineOptions pipelineOptions, KV element, org.joda.time.Instant elementTs, - MultiOutputReceiver o) + DoFn.MultiOutputReceiver o) throws Exception { DatasetService initializedDatasetService = getDatasetService(pipelineOptions); WriteStreamService initializedWriteStreamService = getWriteStreamService(pipelineOptions); @@ -1212,9 +1324,8 @@ Iterable> processElement( state.getTableDestination().getTableReference(), pipelineOptions.as(BigQueryOptions.class))); - OutputReceiver failedRowsReceiver = o.get(failedRowsTag); - @Nullable - OutputReceiver successfulRowsReceiver = + DoFn.OutputReceiver failedRowsReceiver = o.get(failedRowsTag); + DoFn.@Nullable OutputReceiver successfulRowsReceiver = (successfulRowsTag != null) ? o.get(successfulRowsTag) : null; int recordBytes = element.getValue().getPayload().length; @@ -1227,91 +1338,10 @@ Iterable> processElement( return mismatchedRows; } - private OutputReceiver makeSuccessfulRowsReceiver( - FinishBundleContext context, TupleTag successfulRowsTag) { - return new OutputReceiver() { - @Override - public OutputBuilder builder(TableRow value) { - return WindowedValues.builder() - .setValue(value) - .setTimestamp(GlobalWindow.INSTANCE.maxTimestamp()) - .setWindow(GlobalWindow.INSTANCE) - .setPaneInfo(PaneInfo.NO_FIRING) - .setReceiver( - windowedValue -> { - for (BoundedWindow window : windowedValue.getWindows()) { - context.output( - successfulRowsTag, - windowedValue.getValue(), - windowedValue.getTimestamp(), - window); - } - }); - } - }; - } - - @FinishBundle - public void finishBundle(final FinishBundleContext context) throws Exception { - OutputReceiver failedRowsReceiver = - value -> - WindowedValues.builder() - .setValue(value) - .setTimestamp(GlobalWindow.INSTANCE.maxTimestamp()) - .setWindow(GlobalWindow.INSTANCE) - .setPaneInfo(PaneInfo.NO_FIRING) - .setReceiver( - windowedValue -> { - for (BoundedWindow window : windowedValue.getWindows()) { - context.output( - failedRowsTag, - windowedValue.getValue(), - windowedValue.getTimestamp(), - window); - } - }); - - @Nullable OutputReceiver successfulRowsReceiver = null; - if (successfulRowsTag != null) { - successfulRowsReceiver = makeSuccessfulRowsReceiver(context, successfulRowsTag); - } - - OutputReceiver> finalizeOutputRecever = - value -> - WindowedValues.>builder() - .setValue(value) - .setTimestamp(GlobalWindow.INSTANCE.maxTimestamp()) - .setWindow(GlobalWindow.INSTANCE) - .setPaneInfo(PaneInfo.NO_FIRING) - .setReceiver( - windowedValue -> { - for (BoundedWindow window : windowedValue.getWindows()) { - context.output( - finalizeTag, - windowedValue.getValue(), - windowedValue.getTimestamp(), - window); - } - }); - Iterable> mismatchedRows = - finishBundle(failedRowsReceiver, successfulRowsReceiver, finalizeOutputRecever); - if (!Iterables.isEmpty(mismatchedRows)) { - mismatchedRows.forEach( - m -> { - org.joda.time.Instant ts = - MoreObjects.firstNonNull( - m.getValue().getPayload().getTimestamp(), - GlobalWindow.INSTANCE.maxTimestamp()); - context.output( - Preconditions.checkStateNotNull(mismatchedRowsTag), m, ts, GlobalWindow.INSTANCE); - }); - } - } - public Iterable> finishBundle( - OutputReceiver failedRowsReceiver, - @Nullable OutputReceiver successfulRowsReceiver, - OutputReceiver> finalizeOutputReceiver) + DoFn.OutputReceiver failedRowsReceiver, + DoFn.@Nullable OutputReceiver successfulRowsReceiver, + DoFn.OutputReceiver> finalizeOutputReceiver) throws Exception { Iterable> mismatchedRows = @@ -1330,7 +1360,6 @@ public Iterable> finishBundle( return mismatchedRows; } - @Teardown public void teardown() { destinations = null; try { @@ -1346,10 +1375,5 @@ public void teardown() { throw new RuntimeException(e); } } - - @Override - public Duration getAllowedTimestampSkew() { - return Duration.millis(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis()); - } } } 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 be28713a8652..970bbfd94ce4 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 @@ -125,7 +125,6 @@ public class StorageApiWritesShardedRecords { private static final Logger LOG = LoggerFactory.getLogger(StorageApiWritesShardedRecords.class); private static final Duration DEFAULT_STREAM_IDLE_TIME = Duration.standardHours(1); - private static final Duration RETRY_MISMATCHED_ROWS_PERIOD = Duration.standardMinutes(1); private final StorageApiDynamicDestinations dynamicDestinations; private final CreateDisposition createDisposition; @@ -831,6 +830,11 @@ public void process( messageConverter.getDescriptor(false), AutoCloseable::close); + Duration timerRetryDuration = + Duration.millis( + pipelineOptions + .as(BigQueryOptions.class) + .getStorageApiMismatchRetryTimeMilliSec()); SchemaChangeDetectorHelper.bufferMismatchedRows( mismatchedRows, mismatchedRowsBag, @@ -841,7 +845,7 @@ public void process( o.get(failedRowsTag), info, rowsSentToFailedRowsCollection, - RETRY_MISMATCHED_ROWS_PERIOD); + timerRetryDuration); } }; processPayloads( @@ -1076,7 +1080,7 @@ private void processPayloads( org.joda.time.Instant startMismatchTime = org.joda.time.Instant.now(); final Duration initialMismatchRetryTime = Duration.millis( - bigQueryOptions.getStorageApiInitialMismatchRetryTimeMilliSec()); // Retry locally + bigQueryOptions.getStorageApiMismatchLocalRetryTimeMilliSec()); // Retry locally Iterable payloadsToIterate = element; do { // Each ProtoRows object contains at most 1MB of rows. @@ -1314,6 +1318,11 @@ public void onMismatchedRowsTimer( messageConverter.getDescriptor(false), AutoCloseable::close); + Duration timerRetryDuration = + Duration.millis( + pipelineOptions + .as(BigQueryOptions.class) + .getStorageApiMismatchRetryTimeMilliSec()); SchemaChangeDetectorHelper.bufferMismatchedRows( mismatchedRows, mismatchedRowsBag, @@ -1324,7 +1333,7 @@ public void onMismatchedRowsTimer( o.get(failedRowsTag), info, rowsSentToFailedRowsCollection, - RETRY_MISMATCHED_ROWS_PERIOD); + timerRetryDuration); } }; diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java index 3e0ac4463152..c1cd2a010c52 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java @@ -2418,7 +2418,7 @@ public void updateTableSchemaTest(boolean useSet, boolean withConsistentSchemaUp // Make sure that GroupIntoBatches does not buffer data. p.getOptions().as(BigQueryOptions.class).setStorageApiAppendThresholdBytes(1); p.getOptions().as(BigQueryOptions.class).setNumStorageWriteApiStreams(1); - p.getOptions().as(BigQueryOptions.class).setStorageApiInitialMismatchRetryTimeMilliSec(0); + p.getOptions().as(BigQueryOptions.class).setStorageApiMismatchLocalRetryTimeMilliSec(0); BigQueryIO.Write.Method method = useStorageApiApproximate ? Method.STORAGE_API_AT_LEAST_ONCE : Method.STORAGE_WRITE_API; @@ -2555,7 +2555,7 @@ public void testAutoPatchTableSchemaTest() throws Exception { p.getOptions().as(BigQueryOptions.class).setStorageApiAppendThresholdBytes(1); p.getOptions().as(BigQueryOptions.class).setNumStorageWriteApiStreams(1); p.getOptions().as(BigQueryOptions.class).setSchemaUpgradeBufferingShards(2); - p.getOptions().as(BigQueryOptions.class).setStorageApiInitialMismatchRetryTimeMilliSec(0); + p.getOptions().as(BigQueryOptions.class).setStorageApiMismatchLocalRetryTimeMilliSec(0); BigQueryIO.Write.Method method = useStorageApiApproximate ? Method.STORAGE_API_AT_LEAST_ONCE : Method.STORAGE_WRITE_API; From 749569718204e8333555b099da5f6e5e0441db0f Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Mon, 29 Jun 2026 22:30:47 -0700 Subject: [PATCH 07/10] fixes --- .../sdk/io/gcp/bigquery/AppendRowsPacket.java | 49 ++++++---- .../io/gcp/bigquery/BufferMismatchedRows.java | 52 +++++----- .../bigquery/SchemaChangeDetectorHelper.java | 14 +-- .../io/gcp/bigquery/SplittingIterable.java | 6 +- .../StorageApiWriteRecordsInconsistent.java | 7 +- .../StorageApiWriteUnshardedRecords.java | 97 +++++++++++-------- .../StorageApiWritesShardedRecords.java | 55 ++++++----- ...w.java => StoragePayloadWithDeadline.java} | 32 +++--- .../gcp/bigquery/SplittingIterableTest.java | 46 ++++++--- .../StorageApiSinkSchemaUpdateIT.java | 7 +- 10 files changed, 217 insertions(+), 148 deletions(-) rename sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/{MismatchedRow.java => StoragePayloadWithDeadline.java} (71%) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java index 445857abe43a..16c6d37ed7f2 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java @@ -55,10 +55,13 @@ abstract class AppendRowsPacket { abstract Map getOriginalPayloads(); + // Retry deadlines for handling schema updates + abstract List getDeadlines(); + private static final BitSet EMPTY_BIT_SET = new BitSet(0); static AppendRowsPacket fromStorageApiWritePayload( - PeekingIterator underlyingIterator, + PeekingIterator underlyingIterator, long maxByteSize, SchemaChangeDetectorHelper schemaChangeDetectorHelper, Instant elementTimestamp, @@ -71,6 +74,7 @@ static AppendRowsPacket fromStorageApiWritePayload( Map schemaHashes = Maps.newHashMap(); Map unknownFields = Maps.newHashMap(); Map originalPayloads = Maps.newHashMap(); + List deadlines = Lists.newArrayList(); ProtoRows.Builder inserts = ProtoRows.newBuilder(); long bytesSize = 0; BitSet mismatchedRows = new BitSet(); @@ -79,15 +83,17 @@ static AppendRowsPacket fromStorageApiWritePayload( // Make sure that we don't exceed the maxByteSize over multiple elements. A single // element can exceed // the split threshold, but in that case it should be the only element returned. - if ((bytesSize + underlyingIterator.peek().getPayload().length > maxByteSize) + if ((bytesSize + underlyingIterator.peek().getStoragePayload().getPayload().length + > maxByteSize) && inserts.getSerializedRowsCount() > 0) { break; } - StorageApiWritePayload payload = underlyingIterator.next(); + StoragePayloadWithDeadline payload = underlyingIterator.next(); + StorageApiWritePayload storagePayload = payload.getStoragePayload(); @Nullable TableRow failsafeTableRow = null; try { - failsafeTableRow = payload.getFailsafeTableRow(); + failsafeTableRow = storagePayload.getFailsafeTableRow(); } catch (IOException e) { // Do nothing, table row will be generated later from row bytes } @@ -95,7 +101,7 @@ static AppendRowsPacket fromStorageApiWritePayload( // If autoUpdateSchema is set, we try to automatically merge in unknown fields. SchemaChangeDetectorHelper.MergePayloadResult mergeResult = schemaChangeDetectorHelper.getMergedPayload( - payload, elementTimestamp, failsafeTableRow, appendClientInfoSupplier.get()); + storagePayload, elementTimestamp, failsafeTableRow, appendClientInfoSupplier.get()); if (mergeResult.getKind() == SchemaChangeDetectorHelper.MergePayloadResult.Kind.FAILED) { failedRowsConsumer.accept(mergeResult.getFailed()); continue; @@ -103,7 +109,7 @@ static AppendRowsPacket fromStorageApiWritePayload( ByteString byteString = mergeResult.getMerged(); if (schemaChangeDetectorHelper.isPayloadSchemaOutOfDate( - payload, + storagePayload, byteString.toByteArray(), getCurrentTableSchemaHash, getCurrentTableSchemaDescriptor)) { @@ -112,20 +118,22 @@ static AppendRowsPacket fromStorageApiWritePayload( int currentIndex = inserts.getSerializedRowsCount(); inserts.addSerializedRows(byteString); - Instant timestamp = payload.getTimestamp(); + Instant timestamp = storagePayload.getTimestamp(); if (timestamp == null) { timestamp = elementTimestamp; } timestamps.add(timestamp); failsafeRows.add(failsafeTableRow); - if (payload.getSchemaHash() != null) { - schemaHashes.put(currentIndex, Preconditions.checkStateNotNull(payload.getSchemaHash())); + if (storagePayload.getSchemaHash() != null) { + schemaHashes.put( + currentIndex, Preconditions.checkStateNotNull(storagePayload.getSchemaHash())); } - if (payload.getUnknownFields() != null) { + if (storagePayload.getUnknownFields() != null) { unknownFields.put( - currentIndex, Preconditions.checkStateNotNull(payload.getUnknownFields())); - originalPayloads.put(currentIndex, payload.getPayload()); + currentIndex, Preconditions.checkStateNotNull(storagePayload.getUnknownFields())); + originalPayloads.put(currentIndex, storagePayload.getPayload()); } + deadlines.add(payload.getDeadline()); bytesSize += byteString.size(); } } catch (Exception e) { @@ -139,7 +147,8 @@ static AppendRowsPacket fromStorageApiWritePayload( mismatchedRows, schemaHashes, unknownFields, - originalPayloads); + originalPayloads, + deadlines); } AppendRowsPacket getSchemaMismatchedRowsOnly() { @@ -149,6 +158,7 @@ AppendRowsPacket getSchemaMismatchedRowsOnly() { Map schemaHashes = Maps.newHashMap(); Map unknownFields = Maps.newHashMap(); Map originalPayloads = Maps.newHashMap(); + List deadlines = Lists.newArrayList(); if (!getSchemaMismatchedRows().isEmpty()) { int newIndex = 0; for (int i = 0; i < getProtoRows().getSerializedRowsCount(); i++) { @@ -163,6 +173,7 @@ AppendRowsPacket getSchemaMismatchedRowsOnly() { originalPayloads.put( newIndex, Preconditions.checkStateNotNull(getOriginalPayloads().get(i))); } + deadlines.add(getDeadlines().get(i)); if (getSchemaHashes().containsKey(i)) { schemaHashes.put(newIndex, Preconditions.checkStateNotNull(getSchemaHashes().get(i))); } @@ -179,7 +190,8 @@ AppendRowsPacket getSchemaMismatchedRowsOnly() { allBits, schemaHashes, unknownFields, - originalPayloads); + originalPayloads, + deadlines); } AppendRowsPacket getSchemaMatchedRowsOnly() { @@ -192,6 +204,7 @@ AppendRowsPacket getSchemaMatchedRowsOnly() { Map schemaHashes = Maps.newHashMap(); Map unknownFields = Maps.newHashMap(); Map originalPayloads = Maps.newHashMap(); + List deadlines = Lists.newArrayList(); List<@Nullable TableRow> failsafeTableRows = Lists.newArrayList(); int newIndex = 0; for (int i = 0; i < getProtoRows().getSerializedRowsCount(); i++) { @@ -209,6 +222,7 @@ AppendRowsPacket getSchemaMatchedRowsOnly() { originalPayloads.put( newIndex, Preconditions.checkStateNotNull(getOriginalPayloads().get(i))); } + deadlines.add(getDeadlines().get(i)); newIndex++; } } @@ -219,10 +233,11 @@ AppendRowsPacket getSchemaMatchedRowsOnly() { EMPTY_BIT_SET, schemaHashes, unknownFields, - originalPayloads); + originalPayloads, + deadlines); } - Stream toPayloadStream() { + Stream toPayloadStream() { return IntStream.range(0, getProtoRows().getSerializedRowsCount()) .mapToObj( i -> { @@ -240,7 +255,7 @@ Stream toPayloadStream() { if (hash != null) { payload = payload.withSchemaHash(hash); } - return payload; + return StoragePayloadWithDeadline.of(payload, getDeadlines().get(i)); } catch (IOException e) { throw new RuntimeException(e); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java index a3cbcf5c16f6..0140ae3b7f36 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java @@ -51,10 +51,10 @@ import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; -import org.joda.time.Instant; class BufferMismatchedRows - extends PTransform>, PCollectionTuple> { + extends PTransform< + PCollection>, PCollectionTuple> { private final Coder failedRowsCoder; private final Coder successfulRowsCoder; private final Coder destinationCoder; @@ -86,7 +86,7 @@ public BufferMismatchedRows( } @Override - public PCollectionTuple expand(PCollection> input) { + public PCollectionTuple expand(PCollection> input) { // Append records to the Storage API streams. TupleTagList tupleTagList = TupleTagList.of(failedRowsTag); if (successfulRowsTag != null) { @@ -99,8 +99,8 @@ public PCollectionTuple expand(PCollection> inpu "addShard", ParDo.of( new DoFn< - KV, - KV, MismatchedRow>>() { + KV, + KV, StoragePayloadWithDeadline>>() { int shardNumber; @Setup @@ -110,8 +110,9 @@ public void setup() { @ProcessElement public void process( - @Element KV element, - OutputReceiver, MismatchedRow>> o) { + @Element KV element, + OutputReceiver, StoragePayloadWithDeadline>> + o) { ByteBuffer buffer = ByteBuffer.allocate(Integer.BYTES); buffer.putInt(++shardNumber % NUM_DEFAULT_SHARDS); o.output( @@ -120,7 +121,9 @@ public void process( element.getValue())); } })) - .setCoder(KvCoder.of(ShardedKey.Coder.of(destinationCoder), MismatchedRow.Coder.of())) + .setCoder( + KvCoder.of( + ShardedKey.Coder.of(destinationCoder), StoragePayloadWithDeadline.Coder.of())) .apply( "bufferMismatchedRows", ParDo.of(new BufferingDoFn(writeDoFn)) @@ -135,13 +138,13 @@ public void process( } class BufferingDoFn - extends DoFn, MismatchedRow>, KV> { + extends DoFn, StoragePayloadWithDeadline>, KV> { private final StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl writeDoFn; @StateId("mismatchedRows") - private final StateSpec> mismatchedRowsSpec = - StateSpecs.bag(MismatchedRow.Coder.of()); + private final StateSpec> mismatchedRowsSpec = + StateSpecs.bag(StoragePayloadWithDeadline.Coder.of()); @TimerId("retryMismatchedRowsTimer") private final TimerSpec mismatchedRowsTimerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME); @@ -170,8 +173,8 @@ public void startBundle() throws IOException { public void process( PipelineOptions pipelineOptions, ProcessContext processContext, - @Element KV, MismatchedRow> element, - @StateId("mismatchedRows") BagState mismatchedRowsBag, + @Element KV, StoragePayloadWithDeadline> element, + @StateId("mismatchedRows") BagState mismatchedRowsBag, @TimerId("retryMismatchedRowsTimer") Timer retryTimer, @StateId("currentMismatchedRowTimerValue") ValueState currentTimerValue, @StateId("minPendingTimestamp") ValueState minPendingTimestamp, @@ -205,8 +208,7 @@ public Duration getAllowedTimestampSkew() { public void onTimer( OnTimerContext context, @Key ShardedKey shardedDestination, - @Timestamp Instant timestamp, - @StateId("mismatchedRows") BagState mismatchedRowsBag, + @StateId("mismatchedRows") BagState mismatchedRowsBag, @StateId("currentMismatchedRowTimerValue") ValueState currentTimerValue, @StateId("minPendingTimestamp") ValueState minPendingTimestamp, @TimerId("retryMismatchedRowsTimer") Timer retryTimer, @@ -219,24 +221,24 @@ public void onTimer( currentTimerValue.readLater(); minPendingTimestamp.readLater(); - List>> mismatchedRowsList = Lists.newArrayList(); - for (MismatchedRow row : mismatchedRowsBag.read()) { - Iterable> mismatchedRows = + List>> mismatchedRowsList = + Lists.newArrayList(); + for (StoragePayloadWithDeadline row : mismatchedRowsBag.read()) { + // TODO: This will reset the target expiration time! + Iterable> mismatchedRows = writeDoFn.processElement( - pipelineOptions, - KV.of(shardedDestination.getKey(), row.getPayload()), - timestamp, - o); + pipelineOptions, KV.of(shardedDestination.getKey(), row), null, o); if (!Iterables.isEmpty(mismatchedRows)) { mismatchedRowsList.add(mismatchedRows); } } // Once we're done, delegate to finishBundle to finish things. - Iterable> mismatchedDestRows = + Iterable> mismatchedDestRows = writeDoFn.finishBundle( o.get(failedRowsTag), successfulRowsTag != null ? o.get(successfulRowsTag) : null, - o.get(finalizeTag)); + o.get(finalizeTag), + null); if (!Iterables.isEmpty(mismatchedDestRows)) { mismatchedRowsList.add(mismatchedDestRows); } @@ -262,7 +264,7 @@ public void onTimer( messageConverter.getDescriptor(writeDoFn.usesCdc), AutoCloseable::close); - Iterable mismatchedRows = + Iterable mismatchedRows = () -> StreamSupport.stream(Iterables.concat(mismatchedRowsList).spliterator(), false) .map(KV::getValue) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelper.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelper.java index 3defd8315be7..36146246d476 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelper.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelper.java @@ -160,8 +160,8 @@ boolean isPayloadSchemaOutOfDate( } static void bufferMismatchedRows( - Iterable rows, - BagState bufferedBag, + Iterable rows, + BagState bufferedBag, Timer retryRowsTimers, ValueState currentTimerValue, ValueState minPendingTimestamp, @@ -176,25 +176,25 @@ static void bufferMismatchedRows( MoreObjects.firstNonNull( minPendingTimestamp.read(), BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis())); - for (MismatchedRow row : rows) { + for (StoragePayloadWithDeadline row : rows) { if (appendClientInfo == null || row.getDeadline().isAfter(retryRowsTimers.getCurrentRelativeTime())) { bufferedBag.add(row); - org.joda.time.@Nullable Instant timestamp = row.getPayload().getTimestamp(); + org.joda.time.@Nullable Instant timestamp = row.getStoragePayload().getTimestamp(); if (timestamp != null && timestamp.isBefore(minTimestamp)) { minTimestamp = timestamp; } } else { - @Nullable TableRow failedRow = row.getPayload().getFailsafeTableRow(); + @Nullable TableRow failedRow = row.getStoragePayload().getFailsafeTableRow(); if (failedRow == null) { - ByteString rowBytes = ByteString.copyFrom(row.getPayload().getPayload()); + ByteString rowBytes = ByteString.copyFrom(row.getStoragePayload().getPayload()); failedRow = appendClientInfo.toTableRow(rowBytes, Predicates.alwaysTrue()); } BigQueryStorageApiInsertError error = new BigQueryStorageApiInsertError( failedRow, "Mismatched schema", tableDestination.getTableReference()); failedRowsReceiver.outputWithTimestamp( - error, Preconditions.checkStateNotNull(row.getPayload().getTimestamp())); + error, Preconditions.checkStateNotNull(row.getStoragePayload().getTimestamp())); rowsSentToFailedRowsCollection.inc(); BigQuerySinkMetrics.appendRowsRowStatusCounter( BigQuerySinkMetrics.RowStatus.FAILED, diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterable.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterable.java index 3ef7309d8aee..ca41d49a0334 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterable.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterable.java @@ -36,7 +36,7 @@ * the next one. */ class SplittingIterable implements Iterable { - private final Iterable underlying; + private final Iterable underlying; private final long splitSize; private final Consumer> failedRowsConsumer; @@ -47,7 +47,7 @@ class SplittingIterable implements Iterable { private final SchemaChangeDetectorHelper schemaChangeDetectorHelper; public SplittingIterable( - Iterable underlying, + Iterable underlying, long splitSize, Consumer> failedRowsConsumer, ThrowingSupplier getCurrentTableSchemaHash, @@ -68,7 +68,7 @@ public SplittingIterable( @Override public Iterator iterator() { return new Iterator() { - final PeekingIterator underlyingIterator = + final PeekingIterator underlyingIterator = Iterators.peekingIterator(underlying.iterator()); @Nullable AppendRowsPacket cachedNext = null; diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRecordsInconsistent.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRecordsInconsistent.java index 68d4e38a5f90..3e32249b0e4b 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRecordsInconsistent.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRecordsInconsistent.java @@ -60,7 +60,7 @@ public class StorageApiWriteRecordsInconsistent destinationCoder; private final boolean autoUpdateSchema; private final @Nullable Duration autoUpdateSchemaStrictTimeout; - private final @Nullable TupleTag> mismatchedRowsTag; + private final @Nullable TupleTag> mismatchedRowsTag; private final boolean ignoreUnknownValues; private final BigQueryIO.Write.CreateDisposition createDisposition; private final @Nullable String kmsKey; @@ -161,8 +161,9 @@ public PCollectionTuple expand(PCollection> mismatchedRows = result.get(mismatchedRowsTag); - mismatchedRows.setCoder(KvCoder.of(destinationCoder, MismatchedRow.Coder.of())); + PCollection> mismatchedRows = + result.get(mismatchedRowsTag); + mismatchedRows.setCoder(KvCoder.of(destinationCoder, StoragePayloadWithDeadline.Coder.of())); mismatchedResult = mismatchedRows.apply( "bufferMismatched", 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 0a8466dad772..13417f5b3b13 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 @@ -282,15 +282,18 @@ public void finishBundle(final FinishBundleContext context) throws Exception { window); } }); - Iterable> mismatchedRows = + Iterable> mismatchedRows = writeRecordsDoFnImpl.finishBundle( - failedRowsReceiver, successfulRowsReceiver, finalizeOutputRecever); + failedRowsReceiver, + successfulRowsReceiver, + finalizeOutputRecever, + org.joda.time.Instant.now()); if (!Iterables.isEmpty(mismatchedRows)) { mismatchedRows.forEach( m -> { org.joda.time.Instant ts = MoreObjects.firstNonNull( - m.getValue().getPayload().getTimestamp(), + m.getValue().getStoragePayload().getTimestamp(), GlobalWindow.INSTANCE.maxTimestamp()); context.output( Preconditions.checkStateNotNull(writeRecordsDoFnImpl.mismatchedRowsTag), @@ -334,15 +337,28 @@ public void process( MultiOutputReceiver o) throws Exception { writeRecordsDoFnImpl.dynamicDestinations.setSideInputAccessorFromProcessContext(c); - Iterable> mismatchedRows = - writeRecordsDoFnImpl.processElement(pipelineOptions, element, elementTs, o); + org.joda.time.Instant startTimeForLocalRetry = org.joda.time.Instant.now(); + org.joda.time.Instant targetDeadline = + (writeRecordsDoFnImpl.autoUpdateSchemaStrictTimeout != null) + ? startTimeForLocalRetry.plus(writeRecordsDoFnImpl.autoUpdateSchemaStrictTimeout) + : BoundedWindow.TIMESTAMP_MAX_VALUE; + + Iterable> mismatchedRows = + writeRecordsDoFnImpl.processElement( + pipelineOptions, + KV.of( + element.getKey(), + StoragePayloadWithDeadline.of(element.getValue(), targetDeadline)), + startTimeForLocalRetry, + o); if (!Iterables.isEmpty(mismatchedRows)) { - final OutputReceiver> mismatchedReceiver = + final OutputReceiver> mismatchedReceiver = o.get(Preconditions.checkStateNotNull(writeRecordsDoFnImpl.mismatchedRowsTag)); mismatchedRows.forEach( m -> { org.joda.time.Instant ts = - MoreObjects.firstNonNull(m.getValue().getPayload().getTimestamp(), elementTs); + MoreObjects.firstNonNull( + m.getValue().getStoragePayload().getTimestamp(), elementTs); mismatchedReceiver.outputWithTimestamp(m, ts); }); } @@ -369,7 +385,8 @@ static class WriteRecordsDoFnImpl successfulRowsPredicate; private final boolean autoUpdateSchema; private final @Nullable Duration autoUpdateSchemaStrictTimeout; - private final @Nullable TupleTag> mismatchedRowsTag; + private final @Nullable TupleTag> + mismatchedRowsTag; private final boolean ignoreUnknownValues; private final BigQueryIO.Write.CreateDisposition createDisposition; private final @Nullable String kmsKey; @@ -405,7 +422,7 @@ class DestinationState { private String streamName = ""; private @Nullable AppendClientInfo appendClientInfo = null; private long currentOffset = 0; - private List pendingMessages; + private List pendingMessages; private transient @Nullable WriteStreamService maybeWriteStreamService; private final Counter recordsAppended = Metrics.counter(WriteRecordsDoFn.class, "recordsAppended"); @@ -633,16 +650,17 @@ void invalidateAppendClient(boolean invalidateCache) { } } - void addMessage(StorageApiWritePayload payload) throws Exception { + void addMessage(StoragePayloadWithDeadline payload) throws Exception { maybeTickleCache(); pendingMessages.add(payload); } long flush( RetryManager retryManager, + org.joda.time.@Nullable Instant startTimeForLocalRetry, DoFn.OutputReceiver failedRowsReceiver, DoFn.@Nullable OutputReceiver successfulRowsReceiver, - BiConsumer bufferMismatchedRow) + BiConsumer bufferMismatchedRow) throws Exception { if (pendingMessages.isEmpty()) { return 0; @@ -659,12 +677,11 @@ long flush( BigQuerySinkMetrics.RpcMethod.OPEN_WRITE_STREAM)) .backoff(); - org.joda.time.Instant startMismatchTime = org.joda.time.Instant.now(); final Duration initialMismatchRetryTime = Duration.millis(WriteRecordsDoFnImpl.this.initialMismatchRetryTimeMilliSec); AppendRowsPacket rowsToProcess; - Iterable payloadsToIterate = pendingMessages; + Iterable payloadsToIterate = pendingMessages; do { final AppendRowsPacket processingRows = buildInserts(payloadsToIterate, failedRowsReceiver); @@ -673,23 +690,20 @@ long flush( schemaMismatchSeen = !processingRows.getSchemaMismatchedRows().isEmpty(); if (schemaMismatchSeen) { if (autoUpdateSchemaStrictTimeout != null) { - if (startMismatchTime - .plus(initialMismatchRetryTime) - .isBefore(org.joda.time.Instant.now())) { + if (startTimeForLocalRetry == null + || startTimeForLocalRetry + .plus(initialMismatchRetryTime) + .isBefore(org.joda.time.Instant.now())) { // Local retry time has expired! Pull out the mismatched rows so that we can buffer // them for later retrying. - org.joda.time.Instant targetExpirationTime = - startMismatchTime.plus(autoUpdateSchemaStrictTimeout); - // Buffer those rows. processingRows .getSchemaMismatchedRowsOnly() .toPayloadStream() - .map(payload -> MismatchedRow.of(payload, targetExpirationTime)) .forEach(row -> bufferMismatchedRow.accept(destination, row)); // Process the good rows! - Iterable goodRows = + Iterable goodRows = () -> processingRows.getSchemaMatchedRowsOnly().toPayloadStream().iterator(); rowsToProcess = buildInserts(goodRows, failedRowsReceiver); @@ -1023,7 +1037,7 @@ long flush( } private AppendRowsPacket buildInserts( - Iterable payloads, + Iterable payloads, DoFn.OutputReceiver failedRowsReceiver) { Supplier appendClientInfoSupplier = () -> { @@ -1039,7 +1053,7 @@ private AppendRowsPacket buildInserts( failedRowsReceiver.outputWithTimestamp(tv.getValue(), tv.getTimestamp()); }; - PeekingIterator peekingIterator = + PeekingIterator peekingIterator = Iterators.peekingIterator(payloads.iterator()); return AppendRowsPacket.fromStorageApiWritePayload( peekingIterator, @@ -1124,7 +1138,7 @@ void postFlush() { Predicate successfulRowsPredicate, boolean autoUpdateSchema, @Nullable Duration autoUpdateSchemaStrictTimeout, - @Nullable TupleTag> mismatchedRowsTag, + @Nullable TupleTag> mismatchedRowsTag, boolean ignoreUnknownValues, BigQueryIO.Write.CreateDisposition createDisposition, @Nullable String kmsKey, @@ -1163,9 +1177,10 @@ boolean shouldFlush(int recordBytes) { && numPendingRecords > 0); } - Iterable> flushIfNecessary( + Iterable> flushIfNecessary( DoFn.OutputReceiver failedRowsReceiver, DoFn.@Nullable OutputReceiver successfulRowsReceiver, + org.joda.time.@Nullable Instant startTimeForLocalRetry, int recordBytes) throws Exception { if (shouldFlush(recordBytes)) { @@ -1173,16 +1188,17 @@ Iterable> flushIfNecessary( // Too much memory being used. Flush the state and wait for it to drain out. // TODO(reuvenlax): Consider waiting for memory usage to drop instead of waiting for all the // appends to finish. - return flushAll(failedRowsReceiver, successfulRowsReceiver); + return flushAll(failedRowsReceiver, successfulRowsReceiver, startTimeForLocalRetry); } return Collections.emptyList(); } - Iterable> flushAll( + Iterable> flushAll( DoFn.OutputReceiver failedRowsReceiver, - DoFn.@Nullable OutputReceiver successfulRowsReceiver) + DoFn.@Nullable OutputReceiver successfulRowsReceiver, + org.joda.time.@Nullable Instant startTimeForLocalRetry) throws Exception { - List> mismatchedRows = Lists.newArrayList(); + List> mismatchedRows = Lists.newArrayList(); List> retryManagers = Lists.newArrayListWithCapacity(Preconditions.checkStateNotNull(destinations).size()); long numRowsWritten = 0; @@ -1199,6 +1215,7 @@ Iterable> flushAll( numRowsWritten += destinationState.flush( retryManager, + startTimeForLocalRetry, failedRowsReceiver, successfulRowsReceiver, (d, m) -> mismatchedRows.add(KV.of(d, m))); @@ -1297,10 +1314,10 @@ DestinationState createDestinationState( } } - Iterable> processElement( + Iterable> processElement( PipelineOptions pipelineOptions, - KV element, - org.joda.time.Instant elementTs, + KV element, + org.joda.time.@Nullable Instant startTimeForLocalRetry, DoFn.MultiOutputReceiver o) throws Exception { DatasetService initializedDatasetService = getDatasetService(pipelineOptions); @@ -1328,9 +1345,10 @@ Iterable> processElement( DoFn.@Nullable OutputReceiver successfulRowsReceiver = (successfulRowsTag != null) ? o.get(successfulRowsTag) : null; - int recordBytes = element.getValue().getPayload().length; - Iterable> mismatchedRows = - flushIfNecessary(failedRowsReceiver, successfulRowsReceiver, recordBytes); + int recordBytes = element.getValue().getStoragePayload().getPayload().length; + Iterable> mismatchedRows = + flushIfNecessary( + failedRowsReceiver, successfulRowsReceiver, startTimeForLocalRetry, recordBytes); state.addMessage(element.getValue()); ++numPendingRecords; numPendingRecordBytes += recordBytes; @@ -1338,14 +1356,15 @@ Iterable> processElement( return mismatchedRows; } - public Iterable> finishBundle( + public Iterable> finishBundle( DoFn.OutputReceiver failedRowsReceiver, DoFn.@Nullable OutputReceiver successfulRowsReceiver, - DoFn.OutputReceiver> finalizeOutputReceiver) + DoFn.OutputReceiver> finalizeOutputReceiver, + org.joda.time.@Nullable Instant startTimeForLocalRetry) throws Exception { - Iterable> mismatchedRows = - flushAll(failedRowsReceiver, successfulRowsReceiver); + Iterable> mismatchedRows = + flushAll(failedRowsReceiver, successfulRowsReceiver, startTimeForLocalRetry); final Map destinations = Preconditions.checkStateNotNull(this.destinations); 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 970bbfd94ce4..97b2a87d090c 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 @@ -49,7 +49,6 @@ import java.util.function.Function; import java.util.function.Predicate; import java.util.function.Supplier; -import java.util.stream.Collectors; import java.util.stream.StreamSupport; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.KvCoder; @@ -356,8 +355,8 @@ class WriteRecordsDoFn StateSpecs.value(ProtoCoder.of(TableSchema.class)); @StateId("mismatchedRows") - private final StateSpec> mismatchedRowsSpec = - StateSpecs.bag(MismatchedRow.Coder.of()); + private final StateSpec> mismatchedRowsSpec = + StateSpecs.bag(StoragePayloadWithDeadline.Coder.of()); @TimerId("retryMismatchedRowsTimer") private final TimerSpec mismatchedRowsTimerSpec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME); @@ -792,7 +791,7 @@ public void process( @AlwaysFetched @StateId("streamOffset") ValueState streamOffset, @StateId("updatedSchema") ValueState updatedSchema, @TimerId("idleTimer") Timer idleTimer, - @StateId("mismatchedRows") BagState mismatchedRowsBag, + @StateId("mismatchedRows") BagState mismatchedRowsBag, @TimerId("retryMismatchedRowsTimer") Timer mismatchedRowsRetryTimer, @StateId("currentMismatchedRowTimerValue") ValueState mismatchedRowsRetryTimerValue, @StateId("minPendingTimestamp") ValueState minPendingTimestamp, @@ -821,7 +820,7 @@ public void process( getDatasetService(pipelineOptions), getWriteStreamService(pipelineOptions)); - ThrowingConsumer> processMismatchedRows = + ThrowingConsumer> processMismatchedRows = mismatchedRows -> { if (!Iterables.isEmpty(mismatchedRows)) { AppendClientInfo info = @@ -848,12 +847,25 @@ public void process( timerRetryDuration); } }; + + org.joda.time.Instant now = org.joda.time.Instant.now(); + org.joda.time.Instant targetDeadline = + (autoUpdateSchemaStrictTimeout != null) + ? now.plus(autoUpdateSchemaStrictTimeout) + : BoundedWindow.TIMESTAMP_MAX_VALUE; + + Iterable withDeadlines = + () -> + StreamSupport.stream(element.getValue().spliterator(), false) + .map(p -> StoragePayloadWithDeadline.of(p, targetDeadline)) + .iterator(); processPayloads( pipelineOptions, element.getKey(), tableDestination, messageConverter, - element.getValue(), + withDeadlines, + now, elementTs, streamName, streamOffset, @@ -868,14 +880,15 @@ private void processPayloads( ShardedKey shardedDestination, TableDestination tableDestination, StorageApiDynamicDestinations.MessageConverter messageConverter, - Iterable element, + Iterable element, + org.joda.time.@Nullable Instant startTimeForLocalRetry, org.joda.time.Instant elementTs, final ValueState streamName, final ValueState streamOffset, final ValueState updatedSchema, Timer idleTimer, final MultiOutputReceiver o, - ThrowingConsumer> processMismatchedRows) + ThrowingConsumer> processMismatchedRows) throws Exception { BigQueryOptions bigQueryOptions = pipelineOptions.as(BigQueryOptions.class); @@ -988,7 +1001,7 @@ private void processPayloads( // cache could in // theory evict the object during execution and we want a pin held throughout the execution of // this function. - Iterable mismatchedRows = Collections.emptyList(); + Iterable mismatchedRows = Collections.emptyList(); try (AppendClientHolder appendClientHolder = new AppendClientHolder(shardedDestination, getAppendClientInfo)) { String currentStream = getOrCreateStream.get(); @@ -1077,11 +1090,10 @@ private void processPayloads( .backoff(); CreateRetryManagerResult createRetryManagerResult; - org.joda.time.Instant startMismatchTime = org.joda.time.Instant.now(); final Duration initialMismatchRetryTime = Duration.millis( bigQueryOptions.getStorageApiMismatchLocalRetryTimeMilliSec()); // Retry locally - Iterable payloadsToIterate = element; + Iterable payloadsToIterate = element; do { // Each ProtoRows object contains at most 1MB of rows. // TODO: Push messageFromTableRow up to top level. That we we cans skip TableRow entirely @@ -1124,22 +1136,19 @@ private void processPayloads( if (createRetryManagerResult.getSchemaMismatchSeen()) { if (autoUpdateSchemaStrictTimeout != null) { - if (startMismatchTime - .plus(initialMismatchRetryTime) - .isBefore(org.joda.time.Instant.now())) { + if (startTimeForLocalRetry == null + || startTimeForLocalRetry + .plus(initialMismatchRetryTime) + .isBefore(org.joda.time.Instant.now())) { // Local retry time has expired! Pull out the mismatched rows so that we can buffer // them for later // retrying. - org.joda.time.Instant targetExpirationTime = - startMismatchTime.plus(autoUpdateSchemaStrictTimeout); mismatchedRows = () -> StreamSupport.stream(messages.spliterator(), false) .map(AppendRowsPacket::getSchemaMismatchedRowsOnly) .flatMap(AppendRowsPacket::toPayloadStream) - .map(payload -> MismatchedRow.of(payload, targetExpirationTime)) - .map(MismatchedRow.class::cast) .iterator(); // Continue processing only the messages that matched the schema. @@ -1257,7 +1266,7 @@ public void onMismatchedRowsTimer( @TimerId("idleTimer") Timer idleTimer, MultiOutputReceiver o, @TimerId("retryMismatchedRowsTimer") Timer retryRowsTimer, - @StateId("mismatchedRows") BagState mismatchedRowsBag, + @StateId("mismatchedRows") BagState mismatchedRowsBag, @StateId("currentMismatchedRowTimerValue") ValueState currentTimerValue, @StateId("minPendingTimestamp") ValueState minPendingTimestamp) throws Exception { @@ -1281,10 +1290,7 @@ public void onMismatchedRowsTimer( return tableDestination1; }); - Iterable payloads = - StreamSupport.stream(mismatchedRowsBag.read().spliterator(), false) - .map(MismatchedRow::getPayload) - .collect(Collectors.toList()); + Iterable payloads = mismatchedRowsBag.read(); StorageApiDynamicDestinations.MessageConverter messageConverter = messageConverters.get( @@ -1305,7 +1311,7 @@ public void onMismatchedRowsTimer( APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(shardedDestination)); // Try to reprocess all of these rows. - ThrowingConsumer> processMismatchedRows = + ThrowingConsumer> processMismatchedRows = mismatchedRows -> { // Rebuffer the ones that are still not succeeding. mismatchedRowsBag.clear(); @@ -1343,6 +1349,7 @@ public void onMismatchedRowsTimer( tableDestination, messageConverter, payloads, + null, elementTs, streamName, streamOffset, diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/MismatchedRow.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StoragePayloadWithDeadline.java similarity index 71% rename from sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/MismatchedRow.java rename to sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StoragePayloadWithDeadline.java index 168e4cc6705f..8000986405fc 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/MismatchedRow.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StoragePayloadWithDeadline.java @@ -33,19 +33,20 @@ @DefaultSchema(AutoValueSchema.class) @AutoValue -abstract class MismatchedRow { - abstract StorageApiWritePayload getPayload(); +abstract class StoragePayloadWithDeadline { + abstract StorageApiWritePayload getStoragePayload(); abstract org.joda.time.Instant getDeadline(); - static MismatchedRow of(StorageApiWritePayload payload, org.joda.time.Instant deadline) { - return new AutoValue_MismatchedRow(payload, deadline); + static StoragePayloadWithDeadline of( + StorageApiWritePayload payload, org.joda.time.Instant deadline) { + return new AutoValue_StoragePayloadWithDeadline(payload, deadline); } // Schemas give us a coder, however there are still some limitations to storing schema objects // inside of state // variables (mostly involving the Dataflow runner and update). Therefore we use a custom coder. - static class Coder extends CustomCoder { + static class Coder extends CustomCoder { static final ByteArrayCoder BYTE_ARRAY_CODER = ByteArrayCoder.of(); static final InstantCoder INSTANT_CODER = InstantCoder.of(); static final NullableCoder NULLABLE_BYTE_ARRAY_CODER = @@ -53,23 +54,26 @@ static class Coder extends CustomCoder { static final NullableCoder NULLABLE_INSTANT_CODER = NullableCoder.of(InstantCoder.of()); - static MismatchedRow.Coder of() { + static StoragePayloadWithDeadline.Coder of() { return new Coder(); } @Override - public void encode(MismatchedRow value, OutputStream outStream) + public void encode(StoragePayloadWithDeadline value, OutputStream outStream) throws CoderException, IOException { - BYTE_ARRAY_CODER.encode(value.getPayload().getPayload(), outStream); - NULLABLE_BYTE_ARRAY_CODER.encode(value.getPayload().getUnknownFieldsPayload(), outStream); - NULLABLE_INSTANT_CODER.encode(value.getPayload().getTimestamp(), outStream); - NULLABLE_BYTE_ARRAY_CODER.encode(value.getPayload().getFailsafeTableRowPayload(), outStream); - NULLABLE_BYTE_ARRAY_CODER.encode(value.getPayload().getSchemaHash(), outStream); + BYTE_ARRAY_CODER.encode(value.getStoragePayload().getPayload(), outStream); + NULLABLE_BYTE_ARRAY_CODER.encode( + value.getStoragePayload().getUnknownFieldsPayload(), outStream); + NULLABLE_INSTANT_CODER.encode(value.getStoragePayload().getTimestamp(), outStream); + NULLABLE_BYTE_ARRAY_CODER.encode( + value.getStoragePayload().getFailsafeTableRowPayload(), outStream); + NULLABLE_BYTE_ARRAY_CODER.encode(value.getStoragePayload().getSchemaHash(), outStream); INSTANT_CODER.encode(value.getDeadline(), outStream); } @Override - public MismatchedRow decode(InputStream inStream) throws CoderException, IOException { + public StoragePayloadWithDeadline decode(InputStream inStream) + throws CoderException, IOException { byte[] innerPayload = BYTE_ARRAY_CODER.decode(inStream); byte @Nullable [] unknownFieldsPayload = NULLABLE_BYTE_ARRAY_CODER.decode(inStream); @Nullable Instant timestamp = NULLABLE_INSTANT_CODER.decode(inStream); @@ -80,7 +84,7 @@ public MismatchedRow decode(InputStream inStream) throws CoderException, IOExcep StorageApiWritePayload payload = StorageApiWritePayload.of( innerPayload, timestamp, unknownFieldsPayload, failsafeTableRowPayload, schemaHash); - return new AutoValue_MismatchedRow(payload, deadline); + return new AutoValue_StoragePayloadWithDeadline(payload, deadline); } } } diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterableTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterableTest.java index c5e4fb0aa012..5bc7aea1b3c1 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterableTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterableTest.java @@ -50,12 +50,13 @@ private AppendClientInfo createAppendClientInfo() throws Exception { @Test public void testBatchingBySplitSize() throws Exception { - List payloads = new ArrayList<>(); + List payloads = new ArrayList<>(); // Payload of 10 bytes each for (int i = 0; i < 5; i++) { - payloads.add( + StorageApiWritePayload payload = StorageApiWritePayload.of(new byte[10], null, null) - .withTimestamp(Instant.ofEpochMilli(i))); + .withTimestamp(Instant.ofEpochMilli(i)); + payloads.add(StoragePayloadWithDeadline.of(payload, null)); } List> failedRows = new ArrayList<>(); @@ -97,11 +98,14 @@ public void testBatchingBySplitSize() throws Exception { @Test public void testLargeElementExceedingSplitSize() throws Exception { - List payloads = new ArrayList<>(); + List payloads = new ArrayList<>(); // Payload of 10 bytes, 100 bytes, 10 bytes - payloads.add(StorageApiWritePayload.of(new byte[10], null, null)); - payloads.add(StorageApiWritePayload.of(new byte[100], null, null)); - payloads.add(StorageApiWritePayload.of(new byte[10], null, null)); + payloads.add( + StoragePayloadWithDeadline.of(StorageApiWritePayload.of(new byte[10], null, null), null)); + payloads.add( + StoragePayloadWithDeadline.of(StorageApiWritePayload.of(new byte[100], null, null), null)); + payloads.add( + StoragePayloadWithDeadline.of(StorageApiWritePayload.of(new byte[10], null, null), null)); List> failedRows = new ArrayList<>(); SchemaChangeDetectorHelper schemaChangeDetectorHelper = @@ -140,16 +144,22 @@ public void testLargeElementExceedingSplitSize() throws Exception { @Test public void testSchemaMismatchedAndMatchedRows() throws Exception { - List payloads = new ArrayList<>(); + List payloads = new ArrayList<>(); byte[] hash1 = "currentHash".getBytes(StandardCharsets.UTF_8); byte[] hash2 = "oldHash".getBytes(StandardCharsets.UTF_8); TableRow unknownFieldsRow = new TableRow().set("foo", "bar"); - payloads.add(StorageApiWritePayload.of(new byte[0], null, null).withSchemaHash(hash1)); payloads.add( - StorageApiWritePayload.of(new byte[0], unknownFieldsRow, null).withSchemaHash(hash2)); - payloads.add(StorageApiWritePayload.of(new byte[0], null, null).withSchemaHash(hash1)); + StoragePayloadWithDeadline.of( + StorageApiWritePayload.of(new byte[0], null, null).withSchemaHash(hash1), null)); + payloads.add( + StoragePayloadWithDeadline.of( + StorageApiWritePayload.of(new byte[0], unknownFieldsRow, null).withSchemaHash(hash2), + null)); + payloads.add( + StoragePayloadWithDeadline.of( + StorageApiWritePayload.of(new byte[0], null, null).withSchemaHash(hash1), null)); List> failedRows = new ArrayList<>(); SchemaChangeDetectorHelper schemaChangeDetectorHelper = @@ -193,7 +203,10 @@ public void testSchemaMismatchedAndMatchedRows() throws Exception { // Reconstruct stream List reconstructed = - batch.toPayloadStream().collect(Collectors.toList()); + batch + .toPayloadStream() + .map(StoragePayloadWithDeadline::getStoragePayload) + .collect(Collectors.toList()); assertEquals(3, reconstructed.size()); assertArrayEquals(new byte[0], reconstructed.get(0).getPayload()); assertArrayEquals(new byte[0], reconstructed.get(1).getPayload()); @@ -202,11 +215,14 @@ public void testSchemaMismatchedAndMatchedRows() throws Exception { @Test public void testFailedRows() throws Exception { - List payloads = new ArrayList<>(); - payloads.add(StorageApiWritePayload.of(new byte[0], null, null)); + List payloads = new ArrayList<>(); + payloads.add( + StoragePayloadWithDeadline.of(StorageApiWritePayload.of(new byte[0], null, null), null)); // Provide unknown fields, so auto update schema tries to merge and fails TableRow unknownFieldsRow = new TableRow().set("foo", "bar"); - payloads.add(StorageApiWritePayload.of(new byte[0], unknownFieldsRow, null)); + payloads.add( + StoragePayloadWithDeadline.of( + StorageApiWritePayload.of(new byte[0], unknownFieldsRow, null), null)); List> failedRows = new ArrayList<>(); SchemaChangeDetectorHelper schemaChangeDetectorHelper = diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java index ba2f74f7378a..b23f961065a0 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java @@ -382,6 +382,10 @@ private void runStreamingPipelineWithSchemaChange( "Skipping in favor of more relevant test case and to avoid timing issues", consistentAutoUpdate || (!changeTableSchema && useInputSchema && useAutoSchemaUpdate)); } + if (consistentAutoUpdate) { + assumeTrue(changeTableSchema); + assumeFalse(useAutoSchemaUpdate); + } List fieldNamesOrigin = new ArrayList(Arrays.asList(FIELDS)); @@ -490,7 +494,8 @@ private void runStreamingPipelineWithSchemaChange( boolean checkNoDuplication = (method == Write.Method.STORAGE_WRITE_API) ? true : false; checkRowCompleteness(tableSpec, expectedCount, checkNoDuplication); if (useIgnoreUnknownValues) { - checkRowsWithUpdatedSchema(tableSpec, extraField, useAutoSchemaUpdate); + checkRowsWithUpdatedSchema( + tableSpec, extraField, useAutoSchemaUpdate || consistentAutoUpdate); } } From f0b6418bab0c3ff11fc24573c949f87fe1661e52 Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Sat, 11 Jul 2026 23:49:57 -0700 Subject: [PATCH 08/10] more fixes --- .../sdk/io/gcp/bigquery/AppendRowsPacket.java | 18 +- .../sdk/io/gcp/bigquery/BigQueryOptions.java | 10 +- .../io/gcp/bigquery/BufferMismatchedRows.java | 97 ++++++++-- .../bigquery/SchemaChangeDetectorHelper.java | 15 +- .../io/gcp/bigquery/SplittingIterable.java | 11 +- .../sdk/io/gcp/bigquery/StorageApiLoads.java | 3 +- .../StorageApiWriteUnshardedRecords.java | 79 ++++++-- .../StorageApiWritesShardedRecords.java | 175 +++++++++++++++--- .../StorageApiSinkSchemaUpdateIT.java | 14 +- 9 files changed, 335 insertions(+), 87 deletions(-) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java index 16c6d37ed7f2..3e7c5d1a5a37 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java @@ -26,7 +26,7 @@ import java.util.BitSet; import java.util.List; import java.util.Map; -import java.util.function.Consumer; +import java.util.function.Function; import java.util.function.Supplier; import java.util.stream.IntStream; import java.util.stream.Stream; @@ -66,7 +66,7 @@ static AppendRowsPacket fromStorageApiWritePayload( SchemaChangeDetectorHelper schemaChangeDetectorHelper, Instant elementTimestamp, Supplier appendClientInfoSupplier, - Consumer> failedRowsConsumer, + Function, Boolean> failedRowsHandler, ThrowingSupplier getCurrentTableSchemaHash, ThrowingSupplier getCurrentTableSchemaDescriptor) { List timestamps = Lists.newArrayList(); @@ -99,14 +99,22 @@ static AppendRowsPacket fromStorageApiWritePayload( } // If autoUpdateSchema is set, we try to automatically merge in unknown fields. + ByteString byteString; SchemaChangeDetectorHelper.MergePayloadResult mergeResult = schemaChangeDetectorHelper.getMergedPayload( storagePayload, elementTimestamp, failsafeTableRow, appendClientInfoSupplier.get()); if (mergeResult.getKind() == SchemaChangeDetectorHelper.MergePayloadResult.Kind.FAILED) { - failedRowsConsumer.accept(mergeResult.getFailed()); - continue; + if (failedRowsHandler.apply(mergeResult.getFailed())) { + continue; + } else { + // This implies that instead of skipping failed rows, we should mark it as a mismatched + // row. + byteString = ByteString.empty(); + mismatchedRows.set(inserts.getSerializedRowsCount()); + } + } else { + byteString = mergeResult.getMerged(); } - ByteString byteString = mergeResult.getMerged(); if (schemaChangeDetectorHelper.isPayloadSchemaOutOfDate( storagePayload, diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java index 41413cdfa930..da8526bd9948 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java @@ -261,17 +261,23 @@ public interface BigQueryOptions void setSchemaUpgradeBufferingShards(Integer value); - @Hidden @Description("How long to retry locally before buffering when a schema mismatch is detected.") @Default.Integer(5000) Integer getStorageApiMismatchLocalRetryTimeMilliSec(); void setStorageApiMismatchLocalRetryTimeMilliSec(Integer value); - @Hidden @Description("The retry time in milliseconds when a schema mismatch is detected.") @Default.Integer(60000) Integer getStorageApiMismatchRetryTimeMilliSec(); void setStorageApiMismatchRetryTimeMilliSec(Integer value); + + @Description( + "If a pipeline is drained while waiting for a BigQuery schema update, we will wait this long for the " + + "schema update before sending the rows to the failed-rows collection.") + @Default.Integer(300000) + Integer getStorageApiMismatchDrainRetryTimeMilliSec(); + + void setStorageApiMismatchDrainRetryTimeMilliSec(Integer value); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java index 0140ae3b7f36..3ac210b5942f 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java @@ -18,8 +18,8 @@ package org.apache.beam.sdk.io.gcp.bigquery; import com.google.api.services.bigquery.model.TableRow; -import java.io.IOException; import java.nio.ByteBuffer; +import java.time.Instant; import java.util.Collections; import java.util.List; import java.util.concurrent.ThreadLocalRandom; @@ -164,11 +164,6 @@ public BufferingDoFn( this.writeDoFn = writeDoFn; } - @StartBundle - public void startBundle() throws IOException { - writeDoFn.startBundle(); - } - @ProcessElement public void process( PipelineOptions pipelineOptions, @@ -216,15 +211,25 @@ public void onTimer( MultiOutputReceiver o) throws Exception { dynamicDestinations.setSideInputAccessorFromOnTimerContext(context); + writeDoFn.startBundle(); mismatchedRowsBag.readLater(); currentTimerValue.readLater(); minPendingTimestamp.readLater(); + TableDestination tableDestination = dynamicDestinations.getTable(shardedDestination.getKey()); + StorageApiDynamicDestinations.MessageConverter messageConverter = + writeDoFn.messageConverters.get( + shardedDestination.getKey(), + dynamicDestinations, + pipelineOptions, + writeDoFn.getDatasetService(pipelineOptions), + writeDoFn.getWriteStreamService(pipelineOptions)); + messageConverter.updateSchemaFromTable(); + List>> mismatchedRowsList = Lists.newArrayList(); for (StoragePayloadWithDeadline row : mismatchedRowsBag.read()) { - // TODO: This will reset the target expiration time! Iterable> mismatchedRows = writeDoFn.processElement( pipelineOptions, KV.of(shardedDestination.getKey(), row), null, o); @@ -247,17 +252,6 @@ public void onTimer( currentTimerValue.clear(); minPendingTimestamp.clear(); if (!mismatchedRowsList.isEmpty()) { - TableDestination tableDestination = - dynamicDestinations.getTable(shardedDestination.getKey()); - StorageApiDynamicDestinations.MessageConverter messageConverter = - writeDoFn.messageConverters.get( - shardedDestination.getKey(), - dynamicDestinations, - pipelineOptions, - writeDoFn.getDatasetService(pipelineOptions), - writeDoFn.getWriteStreamService(pipelineOptions)); - messageConverter.updateSchemaFromTable(); - AppendClientInfo appendClientInfo = AppendClientInfo.of( messageConverter.getTableSchema(), @@ -287,6 +281,73 @@ public void onTimer( } } + @OnWindowExpiration + public void onWindowExpiration( + OnWindowExpirationContext context, + @Key ShardedKey shardedDestination, + @Timestamp org.joda.time.Instant elementTs, + @StateId("mismatchedRows") BagState mismatchedRowsBag, + PipelineOptions pipelineOptions, + MultiOutputReceiver o) + throws Exception { + dynamicDestinations.setSideInputAccessorFromOnWindowExpirationContext(context); + + TableDestination tableDestination = dynamicDestinations.getTable(shardedDestination.getKey()); + StorageApiDynamicDestinations.MessageConverter messageConverter = + writeDoFn.messageConverters.get( + shardedDestination.getKey(), + dynamicDestinations, + pipelineOptions, + writeDoFn.getDatasetService(pipelineOptions), + writeDoFn.getWriteStreamService(pipelineOptions)); + messageConverter.updateSchemaFromTable(); + + java.time.Duration waitTime = + java.time.Duration.ofMillis( + pipelineOptions + .as(BigQueryOptions.class) + .getStorageApiMismatchDrainRetryTimeMilliSec()); + + Iterable bufferedRows = mismatchedRowsBag.read(); + Instant start = Instant.now(); + while (!Iterables.isEmpty(bufferedRows) && start.plus(waitTime).isAfter(Instant.now())) { + writeDoFn.startBundle(); + List>> mismatchedRowsList = + Lists.newArrayList(); + for (StoragePayloadWithDeadline row : bufferedRows) { + Iterable> mismatchedRows = + writeDoFn.processElement( + pipelineOptions, KV.of(shardedDestination.getKey(), row), null, o); + if (!Iterables.isEmpty(mismatchedRows)) { + mismatchedRowsList.add(mismatchedRows); + } + } + // Once we're done, delegate to finishBundle to finish things. + Iterable> mismatchedDestRows = + writeDoFn.finishBundle( + o.get(failedRowsTag), + successfulRowsTag != null ? o.get(successfulRowsTag) : null, + o.get(finalizeTag), + null); + if (!Iterables.isEmpty(mismatchedDestRows)) { + mismatchedRowsList.add(mismatchedDestRows); + } + + bufferedRows = + () -> + StreamSupport.stream(Iterables.concat(mismatchedRowsList).spliterator(), false) + .map(KV::getValue) + .iterator(); + } + writeDoFn.writeFailedRows( + shardedDestination.getKey(), + bufferedRows, + "Timed out waiting for table schema update in OnWindowExpiration", + pipelineOptions, + elementTs, + o.get(failedRowsTag)); + } + @Teardown public void onTeardown() { writeDoFn.teardown(); diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelper.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelper.java index 36146246d476..517b319a8c14 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelper.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SchemaChangeDetectorHelper.java @@ -146,17 +146,20 @@ boolean isPayloadSchemaOutOfDate( ThrowingSupplier schemaDescriptor) throws Exception { if (ignoreSchemaHashes) { - - if (payload.getUnknownFields() == null) { - return false; + if (payload.getUnknownFields() != null + && UpgradeTableSchema.missingUnknownField( + Preconditions.checkStateNotNull(payload.getUnknownFields()), schemaDescriptor)) { + return true; } - // TODO: CHECK FOR REQUIRED FIELDS AS WELL. - return UpgradeTableSchema.missingUnknownField( - Preconditions.checkStateNotNull(payload.getUnknownFields()), schemaDescriptor); + // We currently rely on getting an append failure from Vortex if there are + // missing required + // fields. However + // we should consider explicitly checking here in the future. } else { return UpgradeTableSchema.isPayloadSchemaOutOfDate( payload.getSchemaHash(), () -> mergedPayload, schemaHash, schemaDescriptor); } + return false; } static void bufferMismatchedRows( diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterable.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterable.java index ca41d49a0334..6dc321be2e0a 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterable.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/SplittingIterable.java @@ -20,7 +20,7 @@ import com.google.protobuf.Descriptors; import java.util.Iterator; import java.util.NoSuchElementException; -import java.util.function.Consumer; +import java.util.function.Function; import java.util.function.Supplier; import org.apache.beam.sdk.util.Preconditions; import org.apache.beam.sdk.util.ThrowingSupplier; @@ -39,7 +39,8 @@ class SplittingIterable implements Iterable { private final Iterable underlying; private final long splitSize; - private final Consumer> failedRowsConsumer; + private final Function, Boolean> + failedRowsHandler; private final ThrowingSupplier getCurrentTableSchemaHash; private final ThrowingSupplier getCurrentTableSchemaDescriptor; private final Instant elementTimestamp; @@ -49,7 +50,7 @@ class SplittingIterable implements Iterable { public SplittingIterable( Iterable underlying, long splitSize, - Consumer> failedRowsConsumer, + Function, Boolean> failedRowsHandler, ThrowingSupplier getCurrentTableSchemaHash, ThrowingSupplier getCurrentTableSchemaDescriptor, Instant elementTimestamp, @@ -57,7 +58,7 @@ public SplittingIterable( SchemaChangeDetectorHelper schemaChangeDetectorHelper) { this.underlying = underlying; this.splitSize = splitSize; - this.failedRowsConsumer = failedRowsConsumer; + this.failedRowsHandler = failedRowsHandler; this.getCurrentTableSchemaHash = getCurrentTableSchemaHash; this.getCurrentTableSchemaDescriptor = getCurrentTableSchemaDescriptor; this.elementTimestamp = elementTimestamp; @@ -98,7 +99,7 @@ public AppendRowsPacket next() { schemaChangeDetectorHelper, elementTimestamp, appendClientSupplier, - failedRowsConsumer, + failedRowsHandler, getCurrentTableSchemaHash, getCurrentTableSchemaDescriptor); return value.getProtoRows().getSerializedRowsCount() == 0 ? null : value; diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiLoads.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiLoads.java index d301c33ff529..5bd5a6d0dd38 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiLoads.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiLoads.java @@ -304,7 +304,8 @@ public WriteResult expandTriggered( autoUpdateSchemaStrictTimeout, ignoreUnknownValues, defaultMissingValueInterpretation, - bigLakeConfiguration)); + bigLakeConfiguration, + hasSchemaUpdateOptions)); PCollection insertErrors = PCollectionList.of(convertMessagesResult.get(failedRowsTag)) 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 13417f5b3b13..86573e34f265 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 @@ -49,7 +49,7 @@ import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiConsumer; -import java.util.function.Consumer; +import java.util.function.Function; import java.util.function.Predicate; import java.util.function.Supplier; import java.util.stream.Collectors; @@ -650,6 +650,37 @@ void invalidateAppendClient(boolean invalidateCache) { } } + void writeFailedRows( + Iterable failedRows, + String msg, + org.joda.time.Instant defaultTs, + DoFn.OutputReceiver failedRowsReceiver) + throws IOException { + + for (StoragePayloadWithDeadline mismatchedRow : failedRows) { + TableRow failedRow = mismatchedRow.getStoragePayload().getFailsafeTableRow(); + if (failedRow == null) { + AppendClientInfo aci = getAppendClientInfo(true, null); + failedRow = + aci.toTableRow( + ByteString.copyFrom(mismatchedRow.getStoragePayload().getPayload()), + Predicates.alwaysTrue()); + } + BigQueryStorageApiInsertError error = + new BigQueryStorageApiInsertError( + failedRow, msg, tableDestination.getTableReference()); + org.joda.time.Instant ts = + MoreObjects.firstNonNull(mismatchedRow.getStoragePayload().getTimestamp(), defaultTs); + failedRowsReceiver.outputWithTimestamp(error, ts); + rowsSentToFailedRowsCollection.inc(); + BigQuerySinkMetrics.appendRowsRowStatusCounter( + BigQuerySinkMetrics.RowStatus.FAILED, + BigQuerySinkMetrics.PAYLOAD_TOO_LARGE, + tableDestination.getShortTableUrn()) + .inc(1); + } + } + void addMessage(StoragePayloadWithDeadline payload) throws Exception { maybeTickleCache(); pendingMessages.add(payload); @@ -720,12 +751,6 @@ long flush( Preconditions.checkStateNotNull( getAppendClientInfo( false, null /* read updated schema from messageConverter */)); - // Convert back to an Iterable to try again. We can't reuse the - // existing - // payloadsToIterate as if there are any failures in it, that would cause us to output - // them again to - // the dead-letter collection, resulting in duplicate outputs to dead letter. - payloadsToIterate = () -> processingRows.toPayloadStream().iterator(); } } while (schemaMismatchSeen && BackOffUtils.next(Sleeper.DEFAULT, backoff)); @@ -1047,11 +1072,14 @@ private AppendRowsPacket buildInserts( return appendClientInfo; }; - Consumer> failedRowsConsumer = - tv -> { - rowsSentToFailedRowsCollection.inc(); - failedRowsReceiver.outputWithTimestamp(tv.getValue(), tv.getTimestamp()); - }; + Function, Boolean> failedRowsHandler = + (autoUpdateSchemaStrictTimeout == null) + ? tv -> { + rowsSentToFailedRowsCollection.inc(); + failedRowsReceiver.outputWithTimestamp(tv.getValue(), tv.getTimestamp()); + return true; + } + : tv -> false; PeekingIterator peekingIterator = Iterators.peekingIterator(payloads.iterator()); @@ -1061,7 +1089,7 @@ private AppendRowsPacket buildInserts( schemaChangeDetectorHelper, BoundedWindow.TIMESTAMP_MAX_VALUE, appendClientInfoSupplier, - failedRowsConsumer, + failedRowsHandler, () -> getAppendClientInfo(true, null).getTableSchemaHash(), () -> TableRowToStorageApiProto.wrapDescriptorProto( @@ -1314,6 +1342,31 @@ DestinationState createDestinationState( } } + void writeFailedRows( + DestinationT destination, + Iterable failedRows, + String msg, + PipelineOptions pipelineOptions, + org.joda.time.Instant defaultTs, + DoFn.OutputReceiver failedRowsReceiver) + throws IOException { + DatasetService initializedDatasetService = getDatasetService(pipelineOptions); + WriteStreamService initializedWriteStreamService = getWriteStreamService(pipelineOptions); + DestinationState state = + Preconditions.checkStateNotNull(destinations) + .computeIfAbsent( + destination, + d -> + createDestinationState( + pipelineOptions, + destination, + usesCdc, + initializedDatasetService, + initializedWriteStreamService, + pipelineOptions.as(BigQueryOptions.class))); + state.writeFailedRows(failedRows, msg, defaultTs, failedRowsReceiver); + } + Iterable> processElement( PipelineOptions pipelineOptions, KV element, 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 97b2a87d090c..82c396d76c26 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 @@ -34,6 +34,7 @@ import com.google.cloud.bigquery.storage.v1.WriteStream; import com.google.protobuf.ByteString; import com.google.protobuf.DescriptorProtos; +import com.google.protobuf.Descriptors; import io.grpc.Status; import io.grpc.Status.Code; import java.io.IOException; @@ -49,6 +50,7 @@ import java.util.function.Function; import java.util.function.Predicate; import java.util.function.Supplier; +import java.util.stream.Collectors; import java.util.stream.StreamSupport; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.KvCoder; @@ -145,6 +147,7 @@ public class StorageApiWritesShardedRecords succussfulRowsCoder; private final TupleTag> flushTag = new TupleTag<>("flushTag"); + private final boolean managedSchemaUpdate; private static final AppendClientCache>> APPEND_CLIENTS = new AppendClientCache<>(Duration.standardMinutes(5)); @@ -168,7 +171,8 @@ public StorageApiWritesShardedRecords( @Nullable Duration autoUpdateSchemaStrictTimeout, boolean ignoreUnknownValues, AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation, - @Nullable Map bigLakeConfiguration) { + @Nullable Map bigLakeConfiguration, + boolean managedSchemaUpdate) { this.dynamicDestinations = dynamicDestinations; this.createDisposition = createDisposition; this.kmsKey = kmsKey; @@ -184,6 +188,7 @@ public StorageApiWritesShardedRecords( this.ignoreUnknownValues = ignoreUnknownValues; this.defaultMissingValueInterpretation = defaultMissingValueInterpretation; this.bigLakeConfiguration = bigLakeConfiguration; + this.managedSchemaUpdate = managedSchemaUpdate; } @Override @@ -400,7 +405,7 @@ String getOrCreateStream( String tableId, ValueState streamName, ValueState streamOffset, - Timer streamIdleTimer, + @Nullable Timer streamIdleTimer, WriteStreamService writeStreamService, Callable tryCreateTable) { try { @@ -425,7 +430,9 @@ String getOrCreateStream( stream.set(streamValue); } // Reset the idle timer. - streamIdleTimer.offset(streamIdleTime).withNoOutputTimestamp().setRelative(); + if (streamIdleTimer != null) { + streamIdleTimer.offset(streamIdleTime).withNoOutputTimestamp().setRelative(); + } return stream.get(); } catch (Exception e) { @@ -886,7 +893,7 @@ private void processPayloads( final ValueState streamName, final ValueState streamOffset, final ValueState updatedSchema, - Timer idleTimer, + @Nullable Timer idleTimer, final MultiOutputReceiver o, ThrowingConsumer> processMismatchedRows) throws Exception { @@ -1099,12 +1106,16 @@ private void processPayloads( // TODO: Push messageFromTableRow up to top level. That we we cans skip TableRow entirely // if // already proto or already schema. - final Iterable messages = - new SplittingIterable( - payloadsToIterate, - splitSize, - // Failed rows consumer - (TimestampedValue error) -> { + + // If either managedSchemaUpdate==true (implying that the schema will be updated by + // another branch of this + // same pipeline) or autoUpdateSchemaStrictTimeout != null (implying that we want to + // strictly buffer records + // until the schema is updated or the timeout expires), we want the failedRowsHandler to + // be a noop. + Function, Boolean> failedRowsHandler = + (autoUpdateSchemaStrictTimeout == null) && !managedSchemaUpdate + ? (TimestampedValue error) -> { o.get(failedRowsTag) .outputWithTimestamp(error.getValue(), error.getTimestamp()); rowsSentToFailedRowsCollection.inc(); @@ -1113,12 +1124,20 @@ private void processPayloads( BigQuerySinkMetrics.PAYLOAD_TOO_LARGE, shortTableId) .inc(1); - }, + return true; + } + : e -> false; + final byte[] currentTableSchemaHash = appendClientHolder.get().getTableSchemaHash(); + final Descriptors.Descriptor currentDescriptorProto = + TableRowToStorageApiProto.wrapDescriptorProto(messageConverter.getDescriptor(false)); + final Iterable messages = + new SplittingIterable( + payloadsToIterate, + splitSize, + failedRowsHandler, // Get the currently-known TableSchema hash - () -> appendClientHolder.get().getTableSchemaHash(), - () -> - TableRowToStorageApiProto.wrapDescriptorProto( - messageConverter.getDescriptor(false)), + () -> currentTableSchemaHash, + () -> currentDescriptorProto, elementTs, appendClientHolder::get, schemaChangeDetectorHelper); @@ -1172,11 +1191,17 @@ private void processPayloads( break; } } + // TODO: The call to updateSchemaFromTable will throttle the DoFn (both because of the // RPC // call and because // the cache has a delay on refresh). We should update throttling counters here as well. LOG.info("Schema out of date: refreshing table schema for {}", tableId); + LOG.info( + "DEBUG: updatedSchemaValue={}, autoUpdateSchemaStrictTimeout={}, autoUpdateSchema={}", + updatedSchemaValue, + autoUpdateSchemaStrictTimeout, + autoUpdateSchema); // Force the message converter to get the schema again from the table. messageConverter.updateSchemaFromTable(); if (autoUpdateSchemaStrictTimeout != null) { @@ -1185,20 +1210,12 @@ private void processPayloads( // Close all RPC clients that were opened with the old descriptor. Clear the cache, // forcing us to create a new append client with the updated descriptor. appendClientHolder.invalidateAndReset(); - - // Convert back to an Iterable to try again. We can't reuse the - // existing - // payloadsToIterate as if there are any failures in it, that would cause us to output - // them again to - // the dead-letter collection, resulting in duplicate outputs to dead letter. - payloadsToIterate = - () -> - StreamSupport.stream(messages.spliterator(), false) - .flatMap(AppendRowsPacket::toPayloadStream) - .iterator(); } } while (createRetryManagerResult.getSchemaMismatchSeen() && BackOffUtils.next(Sleeper.DEFAULT, backoff)); + if (createRetryManagerResult.getSchemaMismatchSeen()) { + throw new RuntimeException("Timed out waiting for schema update"); + } // Output any rows that failed along they way. createRetryManagerResult @@ -1251,7 +1268,9 @@ private void processPayloads( } processMismatchedRows.accept(mismatchedRows); } - idleTimer.offset(streamIdleTime).withNoOutputTimestamp().setRelative(); + if (idleTimer != null) { + idleTimer.offset(streamIdleTime).withNoOutputTimestamp().setRelative(); + } } @OnTimer("retryMismatchedRowsTimer") @@ -1313,11 +1332,15 @@ public void onMismatchedRowsTimer( // Try to reprocess all of these rows. ThrowingConsumer> processMismatchedRows = mismatchedRows -> { + // TODO: Try to refactor so that we don't have to materialize this list. + List mismatchedRowsList = + StreamSupport.stream(mismatchedRows.spliterator(), false) + .collect(Collectors.toList()); // Rebuffer the ones that are still not succeeding. mismatchedRowsBag.clear(); currentTimerValue.clear(); minPendingTimestamp.clear(); - if (!Iterables.isEmpty(mismatchedRows)) { + if (!Iterables.isEmpty(mismatchedRowsList)) { AppendClientInfo info = AppendClientInfo.of( Preconditions.checkStateNotNull(messageConverter.getTableSchema()), @@ -1330,7 +1353,7 @@ public void onMismatchedRowsTimer( .as(BigQueryOptions.class) .getStorageApiMismatchRetryTimeMilliSec()); SchemaChangeDetectorHelper.bufferMismatchedRows( - mismatchedRows, + mismatchedRowsList, mismatchedRowsBag, retryRowsTimer, currentTimerValue, @@ -1400,11 +1423,103 @@ public void onTimer( @OnWindowExpiration public void onWindowExpiration( + PipelineOptions pipelineOptions, @Key ShardedKey key, + @Timestamp org.joda.time.Instant elementTs, + @StateId("updatedSchema") ValueState updatedSchema, @AlwaysFetched @StateId("streamName") ValueState streamName, @AlwaysFetched @StateId("streamOffset") ValueState streamOffset, + @AlwaysFetched @StateId("mismatchedRows") + BagState mismatchedRowsBag, MultiOutputReceiver o, - BoundedWindow window) { + BoundedWindow window) + throws Exception { + TableDestination tableDestination = + destinations.computeIfAbsent( + key.getKey(), + dest -> { + TableDestination tableDestination1 = dynamicDestinations.getTable(dest); + checkArgument( + tableDestination1 != null, + "DynamicDestinations.getTable() may not return null, " + + "but %s returned null for destination %s", + dynamicDestinations, + dest); + return tableDestination1; + }); + + StorageApiDynamicDestinations.MessageConverter messageConverter = + messageConverters.get( + key.getKey(), + dynamicDestinations, + pipelineOptions, + getDatasetService(pipelineOptions), + getWriteStreamService(pipelineOptions)); + + java.time.Duration waitTime = + java.time.Duration.ofMillis( + pipelineOptions + .as(BigQueryOptions.class) + .getStorageApiMismatchDrainRetryTimeMilliSec()); + AtomicReference> mismatchedRows = + new AtomicReference<>(mismatchedRowsBag.read()); + Instant start = Instant.now(); + while (!Iterables.isEmpty(mismatchedRows.get()) + && start.plus(waitTime).isAfter(Instant.now())) { + messageConverter.updateSchemaFromTable(); + APPEND_CLIENTS.invalidate(messageConverters.getAppendClientKey(key)); + processPayloads( + pipelineOptions, + key, + tableDestination, + messageConverter, + mismatchedRows.get(), + null, + elementTs, + streamName, + streamOffset, + updatedSchema, + null, + o, + mismatchedRows::set); + } + + if (mismatchedRows.get() != null) { + // At this point, there's no more waiting. Output the remaining elements to the failed-rows + // collection. + AppendClientInfo appendClientInfo = + AppendClientInfo.of( + Preconditions.checkStateNotNull(messageConverter.getTableSchema()), + messageConverter.getDescriptor(false), + AutoCloseable::close); + + for (StoragePayloadWithDeadline mismatchedRow : mismatchedRows.get()) { + TableRow failedRow = mismatchedRow.getStoragePayload().getFailsafeTableRow(); + if (failedRow == null) { + failedRow = + appendClientInfo.toTableRow( + ByteString.copyFrom(mismatchedRow.getStoragePayload().getPayload()), + Predicates.alwaysTrue()); + } + + BigQueryStorageApiInsertError error = + new BigQueryStorageApiInsertError( + failedRow, + "Timed out waiting for table schema update in OnWindowExpiration", + tableDestination.getTableReference()); + org.joda.time.Instant ts = + MoreObjects.firstNonNull(mismatchedRow.getStoragePayload().getTimestamp(), elementTs); + o.get(failedRowsTag).outputWithTimestamp(error, ts); + rowsSentToFailedRowsCollection.inc(); + BigQuerySinkMetrics.appendRowsRowStatusCounter( + BigQuerySinkMetrics.RowStatus.FAILED, + BigQuerySinkMetrics.PAYLOAD_TOO_LARGE, + tableDestination.getShortTableUrn()) + .inc(1); + } + ; + } + // Window is done - usually because the pipeline has been drained. Make sure to clean up // streams so that they are not leaked. finalizeStream(streamName, streamOffset, key, o, window.maxTimestamp()); diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java index b23f961065a0..21e694fa2591 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiSinkSchemaUpdateIT.java @@ -477,7 +477,7 @@ private void runStreamingPipelineWithSchemaChange( PROJECT, BIG_QUERY_DATASET_ID, ImmutableMap.of(tableId, updatedSchema)))); } WriteResult result = rows.apply("Stream to BigQuery", write); - if (useIgnoreUnknownValues) { + if (useIgnoreUnknownValues || consistentAutoUpdate) { // We ignore the extra fields, so no rows should have been sent to DLQ PAssert.that("Check DLQ is empty", result.getFailedStorageApiInserts()).empty(); } else { @@ -490,10 +490,10 @@ private void runStreamingPipelineWithSchemaChange( p.run().waitUntilFinish(); // Check row completeness, non-duplication, and that schema update works as intended. - int expectedCount = useIgnoreUnknownValues ? TOTAL_N : ORIGINAL_N; - boolean checkNoDuplication = (method == Write.Method.STORAGE_WRITE_API) ? true : false; + int expectedCount = (useIgnoreUnknownValues || consistentAutoUpdate) ? TOTAL_N : ORIGINAL_N; + boolean checkNoDuplication = (method == Write.Method.STORAGE_WRITE_API); checkRowCompleteness(tableSpec, expectedCount, checkNoDuplication); - if (useIgnoreUnknownValues) { + if (useIgnoreUnknownValues || consistentAutoUpdate) { checkRowsWithUpdatedSchema( tableSpec, extraField, useAutoSchemaUpdate || consistentAutoUpdate); } @@ -566,7 +566,6 @@ public void checkRowsWithUpdatedSchema( List actualRows = BQ_CLIENT.queryUnflattened( String.format("SELECT * FROM [%s]", tableSpec), PROJECT, true, false, bigQueryLocation); - for (TableRow row : actualRows) { // Rows written to the table should not have the extra field if // 1. The row has original schema @@ -609,7 +608,7 @@ public void testExactlyOnceWithAutoSchemaUpdate() throws Exception { @Test public void testExactlyOnceWithAutoSchemaUpdateConsistent() throws Exception { - runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_WRITE_API, true, true, true); + runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_WRITE_API, false, true, true); } @Test @@ -631,13 +630,14 @@ public void testAtLeastOnceWithAutoSchemaUpdate() throws Exception { @Test public void testAtLeastOnceWithAutoSchemaUpdateConsistent() throws Exception { - runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_API_AT_LEAST_ONCE, true, true, true); + runStreamingPipelineWithSchemaChange(Write.Method.STORAGE_API_AT_LEAST_ONCE, false, true, true); } public void runDynamicDestinationsWithAutoSchemaUpdate(boolean useAtLeastOnce) throws Exception { Pipeline p = Pipeline.create(TestPipeline.testingPipelineOptions()); // 0 threshold so that the stream tries fetching an updated schema after each append p.getOptions().as(BigQueryOptions.class).setStorageApiAppendThresholdBytes(0); + p.getOptions().as(BigQueryOptions.class).setStorageApiMismatchRetryTimeMilliSec(20); // Total streams per destination p.getOptions() .as(BigQueryOptions.class) From acb9b76dc32286980594ca0819b50279dd34acec Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Sun, 12 Jul 2026 12:03:16 -0700 Subject: [PATCH 09/10] remove line --- .../apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java | 1 - 1 file changed, 1 deletion(-) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java index 3ac210b5942f..a8202744642b 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java @@ -292,7 +292,6 @@ public void onWindowExpiration( throws Exception { dynamicDestinations.setSideInputAccessorFromOnWindowExpirationContext(context); - TableDestination tableDestination = dynamicDestinations.getTable(shardedDestination.getKey()); StorageApiDynamicDestinations.MessageConverter messageConverter = writeDoFn.messageConverters.get( shardedDestination.getKey(), From 012623fdec9ff52d18d9013541aecad3e38299a0 Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Sun, 12 Jul 2026 14:08:16 -0700 Subject: [PATCH 10/10] remove double lookup --- .../sdk/io/gcp/bigquery/AppendRowsPacket.java | 32 +++++++++++-------- 1 file changed, 18 insertions(+), 14 deletions(-) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java index 3e7c5d1a5a37..a3bef076bd9c 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java @@ -174,16 +174,18 @@ AppendRowsPacket getSchemaMismatchedRowsOnly() { inserts.addSerializedRows(getProtoRows().getSerializedRows(i)); timestamps.add(getTimestamps().get(i)); failsafeTableRows.add(getFailsafeTableRows().get(i)); - if (getUnknownFields().containsKey(i)) { - unknownFields.put(newIndex, Preconditions.checkStateNotNull(getUnknownFields().get(i))); + TableRow unknown = getUnknownFields().get(i); + if (unknown != null) { + unknownFields.put(newIndex, unknown); } - if (getOriginalPayloads().containsKey(i)) { - originalPayloads.put( - newIndex, Preconditions.checkStateNotNull(getOriginalPayloads().get(i))); + byte[] originalPayload = getOriginalPayloads().get(i); + if (originalPayload != null) { + originalPayloads.put(newIndex, originalPayload); } deadlines.add(getDeadlines().get(i)); - if (getSchemaHashes().containsKey(i)) { - schemaHashes.put(newIndex, Preconditions.checkStateNotNull(getSchemaHashes().get(i))); + byte[] schemaHash = getSchemaHashes().get(i); + if (schemaHash != null) { + schemaHashes.put(newIndex, schemaHash); } newIndex++; } @@ -220,15 +222,17 @@ AppendRowsPacket getSchemaMatchedRowsOnly() { inserts.addSerializedRows(getProtoRows().getSerializedRows(i)); timestamps.add(getTimestamps().get(i)); failsafeTableRows.add(getFailsafeTableRows().get(i)); - if (getSchemaHashes().containsKey(i)) { - schemaHashes.put(newIndex, Preconditions.checkStateNotNull(getSchemaHashes().get(i))); + byte[] schemaHash = getSchemaHashes().get(i); + if (schemaHash != null) { + schemaHashes.put(newIndex, schemaHash); } - if (getUnknownFields().containsKey(i)) { - unknownFields.put(newIndex, Preconditions.checkStateNotNull(getUnknownFields().get(i))); + TableRow unknown = getUnknownFields().get(i); + if (unknown != null) { + unknownFields.put(newIndex, unknown); } - if (getOriginalPayloads().containsKey(i)) { - originalPayloads.put( - newIndex, Preconditions.checkStateNotNull(getOriginalPayloads().get(i))); + byte[] originalPayload = getOriginalPayloads().get(i); + if (originalPayload != null) { + originalPayloads.put(newIndex, originalPayload); } deadlines.add(getDeadlines().get(i)); newIndex++;