Skip to content

[BigQueryIO] Improve sharding for batch pipelines #38783

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
stankiewicz wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[BigQueryIO] Improve sharding for batch pipelines #38783

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 60 additions & 1 deletion ...gle-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiLoads.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.CreateDisposition;
import org.apache.beam.sdk.options.StreamingOptions;
import org.apache.beam.sdk.schemas.NoSuchSchemaException;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.Flatten;
Expand All @@ -38,6 +39,7 @@
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.Redistribute;
import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.transforms.Values;
import org.apache.beam.sdk.transforms.errorhandling.BadRecord;
import org.apache.beam.sdk.transforms.errorhandling.BadRecordRouter;
import org.apache.beam.sdk.transforms.errorhandling.BadRecordRouter.ThrowingBadRecordRouter;
Expand Down Expand Up @@ -379,12 +381,19 @@ public WriteResult expandUntriggered(
PCollection<KV<DestinationT, StorageApiWritePayload>> successfulConvertedRows =
convertMessagesResult.get(successfulConvertedRowsTag);

if (numShards > 0) {
boolean streaming = input.getPipeline().getOptions().as(StreamingOptions.class).isStreaming();

@reuvenlax reuvenlax Jun 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FYI while this check is fine, it's currently impossible to hit this codepath in streaming. STORAGE_API_WRITES in streaming usually must have a triggering frequency set (except for the weird case of a bounded PCollection in streaming). STORAGE_API_WRITES_AT_LEAST_ONCE uses a different code path

@stankiewicz stankiewicz Jun 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That is the weird case i wanted to protect from. happy to remove the check.

if (numShards > 0 && streaming) {
successfulConvertedRows =
successfulConvertedRows.apply(
"ResdistibuteNumShards",
Redistribute.<KV<DestinationT, StorageApiWritePayload>>arbitrarily()
.withNumBuckets(numShards));
} else if (numShards > 0 && !streaming) {
successfulConvertedRows =
successfulConvertedRows
.apply("AddKeyWithSideInputs", ParDo.of(new AddShardKeyFn<>(numShards)))

@reuvenlax reuvenlax Jun 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

worth noting that this change potentially greatly increases the number of shards you group by. Previously the whole PCollection was sharded into numShards shards. Now it will be numDestinations * numShards.

@stankiewicz stankiewicz Jun 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are right. Parallelism will be in line with amount of connections used. If backend allows, it will scale and split the work up to this number.
Also we are not talking about many keys. Most of the users will have single/few tables.

@reuvenlax reuvenlax Jun 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Have you tested this?

The number of streams created by the transform fundamentally depends on the number of bundles, not on locality within a worker. For each bundle and destination, a stream will be created and destroyed. This PR might be effective and increasing bundle locality of destinations, but this is something we'd have to test (e.g. in streaming it would have the opposite effect!)

@reuvenlax reuvenlax Jun 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FYI We have definitely seen pipelines with 1000+ tables! Much of the code in this sink is there to deal with this and the scalability issues it creates.

@stankiewicz stankiewicz Jun 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tested 700 destination this with dataflow batch which have long bundles. Also this change is isolated for batch.

Baseline - today to sw api in batch with that amount of tables and num shards set you can't make it work as you will reach a concurrent connection quota within seconds. Workaround would have be to write without shards and partition data prior to incoming bq transform.

.apply("RedistributeNumShards", Redistribute.byKey())
.apply("Remove shard", Values.create());
}

PCollectionTuple writeRecordsResult =
Expand Down Expand Up @@ -457,6 +466,56 @@ private void addErrorCollections(
}
}

/**
* A {@link DoFn} that applies a composite sharding key to incoming records to optimize BigQuery
* Storage API throughput.
*
* <p>This transform manages the balance between connection count (resource overhead) and
* processing parallelism by distributing data across {@code numShards} buckets:
*
* <ul>
* <li><b>Data Affinity:</b> By using a composite key {@code KV<DestT, Integer>}, this transform

@reuvenlax reuvenlax Jun 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't understand this comment. <dest, 1> and <dest, 2> are two different keys and can end up on different workers. The fact that they have the same prefix generally doesn't mean anything for most runners.

@stankiewicz stankiewicz Jun 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reuven, I've updated the description as it wasn't accurate. Goal is to group by destination+shard, not by prefix.

* (along with GBK downstream) ensures that records for a specific sharded destination
* (table, shard) are grouped together. This allows downstream transforms to maintain stable
* {@code StreamConnection} sessions for each destination.
* <li><b>Parallel Throughput:</b> By appending a pseudo-random integer shard index, this
* transform allows the runner to distribute the records for a single destination across up
* to {@code numShards} parallel streams, parallelizing the write throughput of "hot"
* (high-volume) destinations.
* <li><b>Concurrency control:</b> The {@code numShards} parameter acts as the parallelism
* multiplier per destination. The total potential concurrency across the pipeline is {@code
* numShards * total_destinations}.
* </ul>
*
* <p>The output structure is {@code KV<KV<DestT, Integer>, KV<DestT, Payload>>}. Downstream,
* {@link Redistribute#byKey()} uses this composite key to partition the data, ensuring the runner
* effectively balances load while respecting the per-destination parallelism limits configured
* here.
*/
private static class AddShardKeyFn<DestT, ElemT>
extends DoFn<
KV<DestT, StorageApiWritePayload>,
KV<KV<DestT, Integer>, KV<DestT, StorageApiWritePayload>>> {
private final int numShards;
private int shardNumber = 0;

public AddShardKeyFn(int numShards) {
this.numShards = Math.max(1, numShards);
}

@Setup
public void setup() {
shardNumber = ThreadLocalRandom.current().nextInt(numShards);
}

@ProcessElement
public void processElement(
@Element KV<DestT, StorageApiWritePayload> element,
OutputReceiver<KV<KV<DestT, Integer>, KV<DestT, StorageApiWritePayload>>> outputReceiver) {
outputReceiver.output(KV.of(KV.of(element.getKey(), ++shardNumber % numShards), element));
}
}
Comment thread
stankiewicz marked this conversation as resolved.

private static class ConvertInsertErrorToBadRecord
extends DoFn<BigQueryStorageApiInsertError, BadRecord> {

Expand Down
Loading