gemini-code-assist[bot] commented on code in PR #38854:
URL: https://github.com/apache/beam/pull/38854#discussion_r3567111686
##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWritesShardedRecords.java:
##########
@@ -1138,11 +1423,103 @@ public void onTimer(
@OnWindowExpiration
public void onWindowExpiration(
+ PipelineOptions pipelineOptions,
@Key ShardedKey<DestinationT> key,
+ @Timestamp org.joda.time.Instant elementTs,
+ @StateId("updatedSchema") ValueState<TableSchema> updatedSchema,
@AlwaysFetched @StateId("streamName") ValueState<String> streamName,
@AlwaysFetched @StateId("streamOffset") ValueState<Long> streamOffset,
+ @AlwaysFetched @StateId("mismatchedRows")
+ BagState<StoragePayloadWithDeadline> 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<Iterable<StoragePayloadWithDeadline>> 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);
+ }
Review Comment:
This `while` loop tightly spins and calls
`messageConverter.updateSchemaFromTable()` (which performs a synchronous
BigQuery API RPC) on every iteration. Without a backoff sleep, this will
immediately exhaust your BigQuery API quota and hammer the service. Please add
a backoff sleep inside the loop.
```java
AtomicReference<Iterable<StoragePayloadWithDeadline>> mismatchedRows =
new AtomicReference<>(mismatchedRowsBag.read());
Instant start = Instant.now();
org.apache.beam.sdk.util.BackOff backoff =
org.apache.beam.sdk.util.FluentBackoff.DEFAULT
.withInitialBackoff(Duration.standardSeconds(1))
.withMaxBackoff(Duration.standardMinutes(1))
.backoff();
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 (!Iterables.isEmpty(mismatchedRows.get())) {
Thread.sleep(backoff.nextBackOffMillis());
}
}
```
##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BufferMismatchedRows.java:
##########
@@ -0,0 +1,355 @@
+/*
+ * 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.nio.ByteBuffer;
+import java.time.Instant;
+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;
+
+class BufferMismatchedRows<DestinationT extends @NonNull Object, ElementT>
+ extends PTransform<
+ PCollection<KV<DestinationT, StoragePayloadWithDeadline>>,
PCollectionTuple> {
+ private final Coder<BigQueryStorageApiInsertError> failedRowsCoder;
+ private final Coder<TableRow> successfulRowsCoder;
+ private final Coder<DestinationT> destinationCoder;
+ private final StorageApiDynamicDestinations<ElementT, DestinationT>
dynamicDestinations;
+ private final
StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl<DestinationT, ElementT>
+ writeDoFn;
+ private final TupleTag<BigQueryStorageApiInsertError> failedRowsTag;
+ private final @Nullable TupleTag<TableRow> successfulRowsTag;
+ // This output is effectively ignored, since we only support this code path
for
+ // StorageApiWriteRecordsInconsistent.
+ private final TupleTag<KV<String, String>> finalizeTag = new
TupleTag<>("finalizeTag");
+ private static final int NUM_DEFAULT_SHARDS = 20;
+
+ public BufferMismatchedRows(
+ Coder<BigQueryStorageApiInsertError> failedRowsCoder,
+ Coder<TableRow> successfulRowsCoder,
+ Coder<DestinationT> destinationCoder,
+ StorageApiDynamicDestinations<ElementT, DestinationT>
dynamicDestinations,
+ StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl<DestinationT,
ElementT> writeDoFn,
+ TupleTag<BigQueryStorageApiInsertError> failedRowsTag,
+ @Nullable TupleTag<TableRow> 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<KV<DestinationT,
StoragePayloadWithDeadline>> 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<DestinationT, StoragePayloadWithDeadline>,
+ KV<ShardedKey<DestinationT>,
StoragePayloadWithDeadline>>() {
+ int shardNumber;
+
+ @Setup
+ public void setup() {
+ shardNumber =
ThreadLocalRandom.current().nextInt(NUM_DEFAULT_SHARDS);
+ }
+
+ @ProcessElement
+ public void process(
+ @Element KV<DestinationT,
StoragePayloadWithDeadline> element,
+ OutputReceiver<KV<ShardedKey<DestinationT>,
StoragePayloadWithDeadline>>
+ 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),
StoragePayloadWithDeadline.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<KV<ShardedKey<DestinationT>, StoragePayloadWithDeadline>,
KV<String, String>> {
+ private final
StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl<DestinationT, ElementT>
+ writeDoFn;
+
+ @StateId("mismatchedRows")
+ private final StateSpec<BagState<StoragePayloadWithDeadline>>
mismatchedRowsSpec =
+ StateSpecs.bag(StoragePayloadWithDeadline.Coder.of());
+
+ @TimerId("retryMismatchedRowsTimer")
+ private final TimerSpec mismatchedRowsTimerSpec =
TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ @StateId("currentMismatchedRowTimerValue")
+ private final StateSpec<ValueState<Long>>
currentMismatchedRowTimerValueSpec =
+ StateSpecs.value();
+
+ @StateId("minPendingTimestamp")
+ private final StateSpec<ValueState<Long>> minPendingTimestampSpec =
StateSpecs.value();
+
+ private final Counter rowsSentToFailedRowsCollection =
+ Metrics.counter(BufferMismatchedRows.BufferingDoFn.class,
"rowsSentToFailedRowsCollection");
+
+ public BufferingDoFn(
+ StorageApiWriteUnshardedRecords.WriteRecordsDoFnImpl<DestinationT,
ElementT> writeDoFn) {
+ this.writeDoFn = writeDoFn;
+ }
+
+ @ProcessElement
+ public void process(
+ PipelineOptions pipelineOptions,
+ ProcessContext processContext,
+ @Element KV<ShardedKey<DestinationT>, StoragePayloadWithDeadline>
element,
+ @StateId("mismatchedRows") BagState<StoragePayloadWithDeadline>
mismatchedRowsBag,
+ @TimerId("retryMismatchedRowsTimer") Timer retryTimer,
+ @StateId("currentMismatchedRowTimerValue") ValueState<Long>
currentTimerValue,
+ @StateId("minPendingTimestamp") ValueState<Long> minPendingTimestamp,
+ MultiOutputReceiver o)
+ throws Exception {
+
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,
+ retryTimer,
+ currentTimerValue,
+ minPendingTimestamp,
+ tableDestination,
+ o.get(failedRowsTag),
+ null,
+ rowsSentToFailedRowsCollection,
+ timerRetryDuration);
+ }
+
+ @Override
+ public Duration getAllowedTimestampSkew() {
+ return Duration.millis(Long.MAX_VALUE);
+ }
+
+ @OnTimer("retryMismatchedRowsTimer")
+ public void onTimer(
+ OnTimerContext context,
+ @Key ShardedKey<DestinationT> shardedDestination,
+ @StateId("mismatchedRows") BagState<StoragePayloadWithDeadline>
mismatchedRowsBag,
+ @StateId("currentMismatchedRowTimerValue") ValueState<Long>
currentTimerValue,
+ @StateId("minPendingTimestamp") ValueState<Long> minPendingTimestamp,
+ @TimerId("retryMismatchedRowsTimer") Timer retryTimer,
+ PipelineOptions pipelineOptions,
+ 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<Iterable<KV<DestinationT, StoragePayloadWithDeadline>>>
mismatchedRowsList =
+ Lists.newArrayList();
+ for (StoragePayloadWithDeadline row : mismatchedRowsBag.read()) {
+ Iterable<KV<DestinationT, StoragePayloadWithDeadline>> 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<KV<DestinationT, StoragePayloadWithDeadline>>
mismatchedDestRows =
+ writeDoFn.finishBundle(
+ o.get(failedRowsTag),
+ successfulRowsTag != null ? o.get(successfulRowsTag) : null,
+ o.get(finalizeTag),
+ null);
+ if (!Iterables.isEmpty(mismatchedDestRows)) {
+ mismatchedRowsList.add(mismatchedDestRows);
+ }
+
+ mismatchedRowsBag.clear();
+ currentTimerValue.clear();
+ minPendingTimestamp.clear();
+ if (!mismatchedRowsList.isEmpty()) {
+ AppendClientInfo appendClientInfo =
+ AppendClientInfo.of(
+ messageConverter.getTableSchema(),
+ messageConverter.getDescriptor(writeDoFn.usesCdc),
+ AutoCloseable::close);
+
+ Iterable<StoragePayloadWithDeadline> mismatchedRows =
+ () ->
+
StreamSupport.stream(Iterables.concat(mismatchedRowsList).spliterator(), false)
+ .map(KV::getValue)
+ .iterator();
+
+ Duration timerRetryDuration =
+ Duration.millis(
+
pipelineOptions.as(BigQueryOptions.class).getStorageApiMismatchRetryTimeMilliSec());
+ SchemaChangeDetectorHelper.bufferMismatchedRows(
+ mismatchedRows,
+ mismatchedRowsBag,
+ retryTimer,
+ currentTimerValue,
+ minPendingTimestamp,
+ tableDestination,
+ o.get(failedRowsTag),
+ appendClientInfo,
+ rowsSentToFailedRowsCollection,
+ timerRetryDuration);
+ }
+ }
+
+ @OnWindowExpiration
+ public void onWindowExpiration(
+ OnWindowExpirationContext context,
+ @Key ShardedKey<DestinationT> shardedDestination,
+ @Timestamp org.joda.time.Instant elementTs,
+ @StateId("mismatchedRows") BagState<StoragePayloadWithDeadline>
mismatchedRowsBag,
+ PipelineOptions pipelineOptions,
+ MultiOutputReceiver o)
+ throws Exception {
+
dynamicDestinations.setSideInputAccessorFromOnWindowExpirationContext(context);
+
+ 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<StoragePayloadWithDeadline> bufferedRows =
mismatchedRowsBag.read();
+ Instant start = Instant.now();
+ while (!Iterables.isEmpty(bufferedRows) &&
start.plus(waitTime).isAfter(Instant.now())) {
+ writeDoFn.startBundle();
+ List<Iterable<KV<DestinationT, StoragePayloadWithDeadline>>>
mismatchedRowsList =
+ Lists.newArrayList();
+ for (StoragePayloadWithDeadline row : bufferedRows) {
+ Iterable<KV<DestinationT, StoragePayloadWithDeadline>>
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<KV<DestinationT, StoragePayloadWithDeadline>>
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();
+ }
Review Comment:
This `while` loop tightly spins without any backoff or sleep. If the schema
update takes time, this loop will continuously call `writeDoFn.processElement`
and `finishBundle` as fast as possible, consuming 100% CPU and hammering the
BigQuery Storage API. Please introduce a backoff sleep (e.g., using
`FluentBackoff`) inside the loop to prevent tight spinning.
```java
Iterable<StoragePayloadWithDeadline> bufferedRows =
mismatchedRowsBag.read();
Instant start = Instant.now();
org.apache.beam.sdk.util.BackOff backoff =
org.apache.beam.sdk.util.FluentBackoff.DEFAULT
.withInitialBackoff(Duration.standardSeconds(1))
.withMaxBackoff(Duration.standardMinutes(1))
.backoff();
while (!Iterables.isEmpty(bufferedRows) &&
start.plus(waitTime).isAfter(Instant.now())) {
writeDoFn.startBundle();
List<Iterable<KV<DestinationT, StoragePayloadWithDeadline>>>
mismatchedRowsList =
Lists.newArrayList();
for (StoragePayloadWithDeadline row : bufferedRows) {
Iterable<KV<DestinationT, StoragePayloadWithDeadline>>
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<KV<DestinationT, StoragePayloadWithDeadline>>
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();
if (!Iterables.isEmpty(bufferedRows)) {
Thread.sleep(backoff.nextBackOffMillis());
}
}
```
##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java:
##########
@@ -0,0 +1,272 @@
+/*
+ * 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.Function;
+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<Instant> getTimestamps();
+
+ abstract List<@Nullable TableRow> getFailsafeTableRows();
+
+ abstract BitSet getSchemaMismatchedRows();
+
+ abstract Map<Integer, byte[]> getSchemaHashes();
+
+ abstract Map<Integer, TableRow> getUnknownFields();
+
+ abstract Map<Integer, byte[]> getOriginalPayloads();
+
+ // Retry deadlines for handling schema updates
+ abstract List<Instant> getDeadlines();
+
+ private static final BitSet EMPTY_BIT_SET = new BitSet(0);
+
+ static AppendRowsPacket fromStorageApiWritePayload(
+ PeekingIterator<StoragePayloadWithDeadline> underlyingIterator,
+ long maxByteSize,
+ SchemaChangeDetectorHelper schemaChangeDetectorHelper,
+ Instant elementTimestamp,
+ Supplier<AppendClientInfo> appendClientInfoSupplier,
+ Function<TimestampedValue<BigQueryStorageApiInsertError>, Boolean>
failedRowsHandler,
+ ThrowingSupplier<byte[]> getCurrentTableSchemaHash,
+ ThrowingSupplier<Descriptors.Descriptor>
getCurrentTableSchemaDescriptor) {
+ List<Instant> timestamps = Lists.newArrayList();
+ List<@Nullable TableRow> failsafeRows = Lists.newArrayList();
+ Map<Integer, byte[]> schemaHashes = Maps.newHashMap();
+ Map<Integer, TableRow> unknownFields = Maps.newHashMap();
+ Map<Integer, byte[]> originalPayloads = Maps.newHashMap();
+ List<Instant> deadlines = Lists.newArrayList();
+ 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().getStoragePayload().getPayload().length
+ > maxByteSize)
+ && inserts.getSerializedRowsCount() > 0) {
+ break;
+ }
+ StoragePayloadWithDeadline payload = underlyingIterator.next();
+ StorageApiWritePayload storagePayload = payload.getStoragePayload();
+
+ @Nullable TableRow failsafeTableRow = null;
+ try {
+ failsafeTableRow = storagePayload.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.
+ ByteString byteString;
+ SchemaChangeDetectorHelper.MergePayloadResult mergeResult =
+ schemaChangeDetectorHelper.getMergedPayload(
+ storagePayload, elementTimestamp, failsafeTableRow,
appendClientInfoSupplier.get());
+ if (mergeResult.getKind() ==
SchemaChangeDetectorHelper.MergePayloadResult.Kind.FAILED) {
+ 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();
+ }
+
+ if (schemaChangeDetectorHelper.isPayloadSchemaOutOfDate(
+ storagePayload,
+ byteString.toByteArray(),
+ getCurrentTableSchemaHash,
+ getCurrentTableSchemaDescriptor)) {
+ mismatchedRows.set(inserts.getSerializedRowsCount());
+ }
+
+ int currentIndex = inserts.getSerializedRowsCount();
+ inserts.addSerializedRows(byteString);
+ Instant timestamp = storagePayload.getTimestamp();
+ if (timestamp == null) {
+ timestamp = elementTimestamp;
+ }
+ timestamps.add(timestamp);
+ failsafeRows.add(failsafeTableRow);
+ if (storagePayload.getSchemaHash() != null) {
+ schemaHashes.put(
+ currentIndex,
Preconditions.checkStateNotNull(storagePayload.getSchemaHash()));
+ }
+ if (storagePayload.getUnknownFields() != null) {
+ unknownFields.put(
+ currentIndex,
Preconditions.checkStateNotNull(storagePayload.getUnknownFields()));
+ originalPayloads.put(currentIndex, storagePayload.getPayload());
+ }
+ deadlines.add(payload.getDeadline());
+ bytesSize += byteString.size();
+ }
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+
+ return new AutoValue_AppendRowsPacket(
+ inserts.build(),
+ timestamps,
+ failsafeRows,
+ mismatchedRows,
+ schemaHashes,
+ unknownFields,
+ originalPayloads,
+ deadlines);
+ }
+
+ AppendRowsPacket getSchemaMismatchedRowsOnly() {
+ ProtoRows.Builder inserts = ProtoRows.newBuilder();
+ List<Instant> timestamps = Lists.newArrayList();
+ List<@Nullable TableRow> failsafeTableRows = Lists.newArrayList();
+ Map<Integer, byte[]> schemaHashes = Maps.newHashMap();
+ Map<Integer, TableRow> unknownFields = Maps.newHashMap();
+ Map<Integer, byte[]> originalPayloads = Maps.newHashMap();
+ List<Instant> deadlines = Lists.newArrayList();
+ 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)));
+ }
+ deadlines.add(getDeadlines().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,
+ deadlines);
+ }
+
+ AppendRowsPacket getSchemaMatchedRowsOnly() {
+ if (getSchemaMismatchedRows().isEmpty()) {
+ return this;
+ }
+
+ ProtoRows.Builder inserts = ProtoRows.newBuilder();
+ List<Instant> timestamps = Lists.newArrayList();
+ Map<Integer, byte[]> schemaHashes = Maps.newHashMap();
+ Map<Integer, TableRow> unknownFields = Maps.newHashMap();
+ Map<Integer, byte[]> originalPayloads = Maps.newHashMap();
+ List<Instant> deadlines = Lists.newArrayList();
+ 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)));
+ }
+ deadlines.add(getDeadlines().get(i));
Review Comment:
Avoid double map lookups (`containsKey` followed by `get`). Perform a single
`get` and check for null instead.
```suggestion
byte[] schemaHash = getSchemaHashes().get(i);
if (schemaHash != null) {
schemaHashes.put(newIndex, schemaHash);
}
TableRow unknownField = getUnknownFields().get(i);
if (unknownField != null) {
unknownFields.put(newIndex, unknownField);
}
byte[] originalPayload = getOriginalPayloads().get(i);
if (originalPayload != null) {
originalPayloads.put(newIndex, originalPayload);
}
```
##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java:
##########
@@ -4090,6 +4131,14 @@ private <DestinationT> WriteResult continueExpandTyped(
DynamicDestinations<T, DestinationT> dynamicDestinations,
RowWriterFactory<T, DestinationT> 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");
+ }
Review Comment:
The validation checks for `getAutoSchemaUpdateStrictTimeout()` are executed
even if `getAutoSchemaUpdate()` is `false`. If a user disables auto schema
updates but has a strict timeout set (or if it's left as default/configured
elsewhere), this will throw an `IllegalArgumentException` unnecessarily. The
check should only run if `getAutoSchemaUpdate()` is enabled.
```suggestion
if (getAutoSchemaUpdate() && 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");
}
```
##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendRowsPacket.java:
##########
@@ -0,0 +1,272 @@
+/*
+ * 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.Function;
+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<Instant> getTimestamps();
+
+ abstract List<@Nullable TableRow> getFailsafeTableRows();
+
+ abstract BitSet getSchemaMismatchedRows();
+
+ abstract Map<Integer, byte[]> getSchemaHashes();
+
+ abstract Map<Integer, TableRow> getUnknownFields();
+
+ abstract Map<Integer, byte[]> getOriginalPayloads();
+
+ // Retry deadlines for handling schema updates
+ abstract List<Instant> getDeadlines();
+
+ private static final BitSet EMPTY_BIT_SET = new BitSet(0);
+
+ static AppendRowsPacket fromStorageApiWritePayload(
+ PeekingIterator<StoragePayloadWithDeadline> underlyingIterator,
+ long maxByteSize,
+ SchemaChangeDetectorHelper schemaChangeDetectorHelper,
+ Instant elementTimestamp,
+ Supplier<AppendClientInfo> appendClientInfoSupplier,
+ Function<TimestampedValue<BigQueryStorageApiInsertError>, Boolean>
failedRowsHandler,
+ ThrowingSupplier<byte[]> getCurrentTableSchemaHash,
+ ThrowingSupplier<Descriptors.Descriptor>
getCurrentTableSchemaDescriptor) {
+ List<Instant> timestamps = Lists.newArrayList();
+ List<@Nullable TableRow> failsafeRows = Lists.newArrayList();
+ Map<Integer, byte[]> schemaHashes = Maps.newHashMap();
+ Map<Integer, TableRow> unknownFields = Maps.newHashMap();
+ Map<Integer, byte[]> originalPayloads = Maps.newHashMap();
+ List<Instant> deadlines = Lists.newArrayList();
+ 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().getStoragePayload().getPayload().length
+ > maxByteSize)
+ && inserts.getSerializedRowsCount() > 0) {
+ break;
+ }
+ StoragePayloadWithDeadline payload = underlyingIterator.next();
+ StorageApiWritePayload storagePayload = payload.getStoragePayload();
+
+ @Nullable TableRow failsafeTableRow = null;
+ try {
+ failsafeTableRow = storagePayload.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.
+ ByteString byteString;
+ SchemaChangeDetectorHelper.MergePayloadResult mergeResult =
+ schemaChangeDetectorHelper.getMergedPayload(
+ storagePayload, elementTimestamp, failsafeTableRow,
appendClientInfoSupplier.get());
+ if (mergeResult.getKind() ==
SchemaChangeDetectorHelper.MergePayloadResult.Kind.FAILED) {
+ 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();
+ }
+
+ if (schemaChangeDetectorHelper.isPayloadSchemaOutOfDate(
+ storagePayload,
+ byteString.toByteArray(),
+ getCurrentTableSchemaHash,
+ getCurrentTableSchemaDescriptor)) {
+ mismatchedRows.set(inserts.getSerializedRowsCount());
+ }
+
+ int currentIndex = inserts.getSerializedRowsCount();
+ inserts.addSerializedRows(byteString);
+ Instant timestamp = storagePayload.getTimestamp();
+ if (timestamp == null) {
+ timestamp = elementTimestamp;
+ }
+ timestamps.add(timestamp);
+ failsafeRows.add(failsafeTableRow);
+ if (storagePayload.getSchemaHash() != null) {
+ schemaHashes.put(
+ currentIndex,
Preconditions.checkStateNotNull(storagePayload.getSchemaHash()));
+ }
+ if (storagePayload.getUnknownFields() != null) {
+ unknownFields.put(
+ currentIndex,
Preconditions.checkStateNotNull(storagePayload.getUnknownFields()));
+ originalPayloads.put(currentIndex, storagePayload.getPayload());
+ }
+ deadlines.add(payload.getDeadline());
+ bytesSize += byteString.size();
+ }
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+
+ return new AutoValue_AppendRowsPacket(
+ inserts.build(),
+ timestamps,
+ failsafeRows,
+ mismatchedRows,
+ schemaHashes,
+ unknownFields,
+ originalPayloads,
+ deadlines);
+ }
+
+ AppendRowsPacket getSchemaMismatchedRowsOnly() {
+ ProtoRows.Builder inserts = ProtoRows.newBuilder();
+ List<Instant> timestamps = Lists.newArrayList();
+ List<@Nullable TableRow> failsafeTableRows = Lists.newArrayList();
+ Map<Integer, byte[]> schemaHashes = Maps.newHashMap();
+ Map<Integer, TableRow> unknownFields = Maps.newHashMap();
+ Map<Integer, byte[]> originalPayloads = Maps.newHashMap();
+ List<Instant> deadlines = Lists.newArrayList();
+ 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)));
+ }
+ deadlines.add(getDeadlines().get(i));
+ if (getSchemaHashes().containsKey(i)) {
+ schemaHashes.put(newIndex,
Preconditions.checkStateNotNull(getSchemaHashes().get(i)));
+ }
Review Comment:
Avoid double map lookups (`containsKey` followed by `get`). Instead, perform
a single `get` and check if the returned value is non-null. This is cleaner and
more efficient.
```java
TableRow unknownField = getUnknownFields().get(i);
if (unknownField != null) {
unknownFields.put(newIndex, unknownField);
}
byte[] originalPayload = getOriginalPayloads().get(i);
if (originalPayload != null) {
originalPayloads.put(newIndex, originalPayload);
}
deadlines.add(getDeadlines().get(i));
byte[] schemaHash = getSchemaHashes().get(i);
if (schemaHash != null) {
schemaHashes.put(newIndex, schemaHash);
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
