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

* Fix BigQuery Storage Write API stream count for bounded writes

Author: Akshatsharma2205

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProvider.java

@@ -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();
 
       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
@@ -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();
 

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java

@@ -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,
@@ -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();
 

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProviderTest.java

@@ -122,7 +122,10 @@ public void testInvalidConfig() {
         Arrays.asList(
             BigQueryWriteConfiguration.builder()
                 .setTable("project:dataset.table")
-                .setCreateDisposition("INVALID_DISPOSITION"));
+                .setCreateDisposition("INVALID_DISPOSITION"),
+            BigQueryWriteConfiguration.builder()
+                .setTable("project:dataset.table")
+                .setNumStreams(-1));
 
     for (BigQueryWriteConfiguration.Builder config : invalidConfigs) {
       assertThrows(

sdks/python/apache_beam/io/gcp/bigquery.py

@@ -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

website/www/site/content/en/documentation/io/built-in/google-bigquery.md

@@ -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