From 39cba3593843800e4dbd3e52dfb47deed1865ecd Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Wed, 29 Apr 2026 12:23:20 -0700 Subject: [PATCH] only expand update graph if needed --- .../bigquery/StorageApiConvertMessages.java | 206 +++++++++--------- 1 file changed, 107 insertions(+), 99 deletions(-) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiConvertMessages.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiConvertMessages.java index e0713311e2cd..02ef9e9c06a3 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiConvertMessages.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiConvertMessages.java @@ -147,109 +147,117 @@ public PCollectionTuple expand(PCollection> input) { .get(patchTableSchemaTag) .setCoder(KvCoder.of(destinationCoder, ProtoCoder.of(TableSchema.class))); result.get(elementsWaitingForSchemaTag).setCoder(KvCoder.of(destinationCoder, elementCoder)); + if (!hasSchemaUpdateOptions) { + // Don't expand the update graph if it's not needed. + return result; + } else { + final int numShards = + input + .getPipeline() + .getOptions() + .as(BigQueryOptions.class) + .getSchemaUpgradeBufferingShards(); - final int numShards = - input - .getPipeline() - .getOptions() - .as(BigQueryOptions.class) - .getSchemaUpgradeBufferingShards(); + // Throttle the stream to the patch-table function so that only a single update per table per + // two seconds gets processed (to match quotas). The combiner merges incremental schemas, so + // we + // won't miss any updates. + PCollection, ElementT>> tablesPatched = + result + .get(patchTableSchemaTag) + .apply( + "rewindow", + Window.>configure() + .triggering( + Repeatedly.forever( + AfterProcessingTime.pastFirstElementInPane() + .plusDelayOf(Duration.standardSeconds(2)))) + .discardingFiredPanes()) + .apply("merge schemas", Combine.fewKeys(new MergeSchemaCombineFn())) + .setCoder(KvCoder.of(destinationCoder, ProtoCoder.of(TableSchema.class))) + .apply( + "Patch table schema", + ParDo.of( + new PatchTableSchemaDoFn<>(operationName, bqServices, dynamicDestinations))) + .setCoder(KvCoder.of(destinationCoder, NullableCoder.of(elementCoder))) + // We need to make sure that all shards of the buffering transform are notified. + .apply( + "fanout to all shards", + FlatMapElements.via( + new SimpleFunction< + KV, + Iterable, ElementT>>>() { + @Override + public Iterable, ElementT>> apply( + KV elem) { + return IntStream.range(0, numShards) + .mapToObj( + i -> + KV.of( + StorageApiConvertMessages.AssignShardFn.getShardedKey( + elem.getKey(), i, numShards), + elem.getValue())) + .collect(Collectors.toList()); + } + })) + .setCoder( + KvCoder.of(ShardedKey.Coder.of(destinationCoder), NullableCoder.of(elementCoder))) + .apply( + Window., ElementT>>configure() + .triggering(DefaultTrigger.of())); - // Throttle the stream to the patch-table function so that only a single update per table per - // two seconds gets processed (to match quotas). The combiner merges incremental schemas, so we - // won't miss any updates. - PCollection, ElementT>> tablesPatched = - result - .get(patchTableSchemaTag) - .apply( - "rewindow", - Window.>configure() - .triggering( - Repeatedly.forever( - AfterProcessingTime.pastFirstElementInPane() - .plusDelayOf(Duration.standardSeconds(2)))) - .discardingFiredPanes()) - .apply("merge schemas", Combine.fewKeys(new MergeSchemaCombineFn())) - .setCoder(KvCoder.of(destinationCoder, ProtoCoder.of(TableSchema.class))) - .apply( - "Patch table schema", - ParDo.of( - new PatchTableSchemaDoFn<>(operationName, bqServices, dynamicDestinations))) - .setCoder(KvCoder.of(destinationCoder, NullableCoder.of(elementCoder))) - // We need to make sure that all shards of the buffering transform are notified. - .apply( - "fanout to all shards", - FlatMapElements.via( - new SimpleFunction< - KV, - Iterable, ElementT>>>() { - @Override - public Iterable, ElementT>> apply( - KV elem) { - return IntStream.range(0, numShards) - .mapToObj( - i -> - KV.of( - StorageApiConvertMessages.AssignShardFn.getShardedKey( - elem.getKey(), i, numShards), - elem.getValue())) - .collect(Collectors.toList()); - } - })) - .setCoder( - KvCoder.of(ShardedKey.Coder.of(destinationCoder), NullableCoder.of(elementCoder))) - .apply( - Window., ElementT>>configure() - .triggering(DefaultTrigger.of())); + // Any elements that are waiting for a schema update are sent to this stateful DoFn to be + // buffered. + // Note: we currently do not provide the DynamicDestinations object access to the side input + // in + // this path. + // This is because side inputs are not currently available from timer callbacks. Since side + // inputs are generally + // used for getSchema and in this case we read the schema from the table, this is unlikely to + // be + // a problem. + PCollection, ElementT>> shardedWaitingElements = + result + .get(elementsWaitingForSchemaTag) + // TODO: Consider using GroupIntoBatchs.withShardingKey to get auto sharding here + // instead of fixed sharding. + .apply("assignShard", ParDo.of(new AssignShardFn<>(numShards))) + .setCoder( + KvCoder.of( + ShardedKey.Coder.of(destinationCoder), NullableCoder.of(elementCoder))); - // Any elements that are waiting for a schema update are sent to this stateful DoFn to be - // buffered. - // Note: we currently do not provide the DynamicDestinations object access to the side input in - // this path. - // This is because side inputs are not currently available from timer callbacks. Since side - // inputs are generally - // used for getSchema and in this case we read the schema from the table, this is unlikely to be - // a problem. - PCollection, ElementT>> shardedWaitingElements = - result - .get(elementsWaitingForSchemaTag) - // TODO: Consider using GroupIntoBatchs.withShardingKey to get auto sharding here - // instead of fixed sharding. - .apply("assignShard", ParDo.of(new AssignShardFn<>(numShards))) - .setCoder( - KvCoder.of(ShardedKey.Coder.of(destinationCoder), NullableCoder.of(elementCoder))); + PCollectionList, ElementT>> waitingElementsList = + PCollectionList.of(shardedWaitingElements).and(tablesPatched); + PCollectionTuple retryResult = + waitingElementsList + .apply("Buffered flatten", Flatten.pCollections()) + .apply( + "bufferElements", + ParDo.of(new SchemaUpdateHoldingFn<>(elementCoder, convertMessagesDoFn)) + .withOutputTags( + successfulWritesTag, + TupleTagList.of(ImmutableList.of(failedWritesTag, BAD_RECORD_TAG)))); + retryResult.get(successfulWritesTag).setCoder(successCoder); + retryResult.get(failedWritesTag).setCoder(errorCoder); + retryResult.get(BAD_RECORD_TAG).setCoder(BadRecord.getCoder(input.getPipeline())); - PCollectionList, ElementT>> waitingElementsList = - PCollectionList.of(shardedWaitingElements).and(tablesPatched); - PCollectionTuple retryResult = - waitingElementsList - .apply("Buffered flatten", Flatten.pCollections()) - .apply( - "bufferElements", - ParDo.of(new SchemaUpdateHoldingFn<>(elementCoder, convertMessagesDoFn)) - .withOutputTags( - successfulWritesTag, - TupleTagList.of(ImmutableList.of(failedWritesTag, BAD_RECORD_TAG)))); - retryResult.get(successfulWritesTag).setCoder(successCoder); - retryResult.get(failedWritesTag).setCoder(errorCoder); - retryResult.get(BAD_RECORD_TAG).setCoder(BadRecord.getCoder(input.getPipeline())); - - // Flatten successes and failures from both the regular transform and the retry transform. - PCollection> allSuccesses = - PCollectionList.of(result.get(successfulWritesTag)) - .and(retryResult.get(successfulWritesTag)) - .apply("flattenSuccesses", Flatten.pCollections()); - PCollection allFailures = - PCollectionList.of(result.get(failedWritesTag)) - .and(retryResult.get(failedWritesTag)) - .apply("flattenFailures", Flatten.pCollections()); - PCollection allBadRecords = - PCollectionList.of(result.get(BAD_RECORD_TAG)) - .and(retryResult.get(BAD_RECORD_TAG)) - .apply("flattenBadRecords", Flatten.pCollections()); - return PCollectionTuple.of(successfulWritesTag, allSuccesses) - .and(failedWritesTag, allFailures) - .and(BAD_RECORD_TAG, allBadRecords); + // Flatten successes and failures from both the regular transform and the retry transform. + PCollection> allSuccesses = + PCollectionList.of(result.get(successfulWritesTag)) + .and(retryResult.get(successfulWritesTag)) + .apply("flattenSuccesses", Flatten.pCollections()); + PCollection allFailures = + PCollectionList.of(result.get(failedWritesTag)) + .and(retryResult.get(failedWritesTag)) + .apply("flattenFailures", Flatten.pCollections()); + PCollection allBadRecords = + PCollectionList.of(result.get(BAD_RECORD_TAG)) + .and(retryResult.get(BAD_RECORD_TAG)) + .apply("flattenBadRecords", Flatten.pCollections()); + return PCollectionTuple.of(successfulWritesTag, allSuccesses) + .and(failedWritesTag, allFailures) + .and(BAD_RECORD_TAG, allBadRecords); + } } static class AssignShardFn extends DoFn, KV, V>> {