Skip to content

Fix BigQuery Storage Write API stream count for bounded writes #38776

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
Akshatsharma2205 wants to merge 2 commits into
base: master
Choose a base branch
from
+14 −9
Open

Fix BigQuery Storage Write API stream count for bounded writes #38776

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
400b41d
Fix BigQuery Storage Write API stream count for bounded writes
Akshatsharma2205 Jun 2, 2026
07868d1
Address numStreams validation feedback
Akshatsharma2205 Jun 2, 2026
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
9 changes: 5 additions & 4 deletions ...he/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -179,11 +179,11 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) {
PCollection<Row> inputRows = input.getSinglePCollection();

BigQueryIO.Write<Row> write = createStorageWriteApiTransform(inputRows.getSchema());
int numStreams = configuration.getNumStreams() == null ? 0 : configuration.getNumStreams();
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Jun 2, 2026

Choose a reason for hiding this comment

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

medium

If configuration.getNumStreams() is configured with a negative value, it will be silently ignored and treated as 0 (which enables auto-sharding if unbounded, or default behavior if bounded). We should validate that numStreams is non-negative and throw an IllegalArgumentException if it is negative to prevent silent configuration errors.

      int numStreams = configuration.getNumStreams() == null ? 0 : configuration.getNumStreams();
      if (numStreams < 0) {
        throw new IllegalArgumentException("numStreams must be non-negative, but was: " + numStreams);
      }

Copy link
Copy Markdown
Author

@Akshatsharma2205 Akshatsharma2205 Jun 2, 2026

Choose a reason for hiding this comment

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

Thanks for the suggestion. I added validation for negative numStreams in BigQueryWriteConfiguration.validate() using the existing checkArgument style.


if (inputRows.isBounded() == IsBounded.UNBOUNDED) {
Long triggeringFrequency = configuration.getTriggeringFrequencySeconds();
Boolean autoSharding = configuration.getAutoSharding();
int numStreams = configuration.getNumStreams() == null ? 0 : configuration.getNumStreams();

boolean useAtLeastOnceSemantics =
configuration.getUseAtLeastOnceSemantics() != null
Expand All @@ -197,12 +197,13 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) {
: Duration.standardSeconds(triggeringFrequency));
}
// set num streams if specified, otherwise default to autoSharding
if (numStreams > 0) {
write = write.withNumStorageWriteApiStreams(numStreams);
} else if (autoSharding == null || autoSharding) {
if (numStreams == 0 && (autoSharding == null || autoSharding)) {
write = write.withAutoSharding();
}
}
if (numStreams > 0) {
write = write.withNumStorageWriteApiStreams(numStreams);
}

Schema inputSchema = inputRows.getSchema();

Expand Down
9 changes: 7 additions & 2 deletions ...c/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,12 @@ public void validate() {

Boolean autoSharding = getAutoSharding();
Integer numStreams = getNumStreams();
if (numStreams != null) {
checkArgument(
numStreams >= 0,
invalidConfigMessage + "numStreams must be non-negative, but was: %s",
numStreams);
}
if (autoSharding != null && autoSharding && numStreams != null) {
checkArgument(
numStreams == 0,
Expand Down Expand Up @@ -152,8 +158,7 @@ public static Builder builder() {
public abstract Boolean getAutoSharding();

@SchemaFieldDescription(
"Specifies the number of write streams that the Storage API sink will use. "
+ "This parameter is only applicable when writing unbounded data.")
"Specifies the number of write streams that the Storage API sink will use.")
@Nullable
public abstract Integer getNumStreams();

Expand Down
3 changes: 1 addition & 2 deletions sdks/python/apache_beam/io/gcp/bigquery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2155,8 +2155,7 @@ def __init__(
all of FILE_LOADS, STREAMING_INSERTS, and STORAGE_WRITE_API. Only
applicable to unbounded input.
num_storage_api_streams: Specifies the number of write streams that the
Storage API sink will use. This parameter is only applicable when
writing unbounded data.
Storage API sink will use.
ignore_unknown_columns: Accept rows that contain values that do not match
the schema. The unknown values are ignored. Default is False,
which treats unknown values as errors. This option is only valid for
Expand Down
2 changes: 1 addition & 1 deletion website/www/site/content/en/documentation/io/built-in/google-bigquery.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -855,7 +855,7 @@ pipeline uses. You can set it explicitly on the transform via
[`withNumStorageWriteApiStreams`](https://beam.apache.org/releases/javadoc/current/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.Write.html#withNumStorageWriteApiStreams-int-)
or provide the `numStorageWriteApiStreams` option to the pipeline as defined in
[`BigQueryOptions`](https://beam.apache.org/releases/javadoc/current/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.html).
Please note this is only supported for streaming pipelines.
Fixed stream counts can be used with both batch and streaming pipelines.

Triggering frequency determines how soon the data is visible for querying in
BigQuery. You can explicitly set it via
Expand Down
Loading