add WRITE_TRUNCATE_DATA write_disposition to the bigquery python sdk

CHANGES.md

@@ -66,7 +66,7 @@
 
 * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
 * (Java) FirestoreIO: added `RpcQosOptions.Builder.withRampupThrottlingDisabled()`, mirroring the existing DatastoreIO option, to opt out of the 500/50/5 write ramp-up ([#27170](https://github.com/apache/beam/issues/27170)).
-
+* (Python) The BigQuery IO writer now allows the `WRITE_TRUNCATE_DATA` write_disposition ([#40622](https://github.com/apache/beam/pull/40622)).
 
 ## New Features / Improvements
 

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

@@ -3038,6 +3038,17 @@
        */
       WRITE_TRUNCATE,
 
+      /**
+       * Specifies that write should replace a table, but not the table schema.
+       *
+       * <p>The replacement may occur in multiple steps - for instance by first removing the
+       * existing table, then creating a replacement, then filling it in. This is not an atomic
+       * operation, and external programs may see the table in any of these intermediate steps.
+       *
+       * <p>Note: This write disposition is only supported for the FILE_LOADS write method.
+       */
+      WRITE_TRUNCATE_DATA,
+
       /** Specifies that rows may be appended to an existing table. */
       WRITE_APPEND,
 
@@ -3973,8 +3984,9 @@
             getBigLakeConfiguration() == null,
             "bigLakeConfiguration is only supported when using STORAGE_WRITE_API or STORAGE_API_AT_LEAST_ONCE.");
       } else {
-        if (getWriteDisposition() == WriteDisposition.WRITE_TRUNCATE) {
-          LOG.error("The Storage API sink does not support the WRITE_TRUNCATE write disposition.");
+        WriteDisposition disposition = getWriteDisposition() 
+        if (disposition == WriteDisposition.WRITE_TRUNCATE || disposition == WriteDisposition.WRITE_TRUNCATE_DATA) {
+          LOG.error("The Storage API sink does not support the "+disposition.toString()+" write disposition.");
         }
         Map<String, String> bigLakeConfiguration = getBigLakeConfiguration();
         if (bigLakeConfiguration != null) {
@@ -4229,6 +4241,9 @@
         checkArgument(
             getWriteDisposition() != WriteDisposition.WRITE_TRUNCATE,
             "WriteDisposition.WRITE_TRUNCATE is not supported for an unbounded PCollection.");
+        checkArgument(
+            getWriteDisposition() != WriteDisposition.WRITE_TRUNCATE_DATA,
+            "WriteDisposition.WRITE_TRUNCATE_DATA is not supported for an unbounded PCollection.");
         InsertRetryPolicy retryPolicy =
             MoreObjects.firstNonNull(getFailedInsertRetryPolicy(), InsertRetryPolicy.alwaysRetry());
 

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

@@ -257,9 +257,12 @@
     if (!loadConfig
             .getWriteDisposition()
             .equals(BigQueryIO.Write.WriteDisposition.WRITE_TRUNCATE.toString())
+        && !loadConfig
+            .getWriteDisposition()
+            .equals(BigQueryIO.Write.WriteDisposition.WRITE_TRUNCATE_DATA.toString()
         && !loadConfig
             .getWriteDisposition()
             .equals(BigQueryIO.Write.WriteDisposition.WRITE_APPEND.toString())) {
       return null;
     }
     final Table destinationTable;

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

@@ -138,7 +138,8 @@ public static Builder builder() {
   @SchemaFieldDescription(
       "Specifies the action that occurs if the destination table already exists. "
           + "The following values are supported: "
-          + "WRITE_TRUNCATE (overwrites the table data), "
+          + "WRITE_TRUNCATE (overwrites the table data and schema), "
+          + "WRITE_TRUNCATE_DATA (overwrites the table data), "
           + "WRITE_APPEND (append the data to the table), "
           + "WRITE_EMPTY (job must fail if the table is not empty).")
   @Pure

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

@@ -590,6 +590,7 @@ class BigQueryDisposition(object):
 
   CREATE_NEVER = 'CREATE_NEVER'
   CREATE_IF_NEEDED = 'CREATE_IF_NEEDED'
+  WRITE_TRUNCATE_DATA = 'WRITE_TRUNCATE_DATA'
   WRITE_TRUNCATE = 'WRITE_TRUNCATE'
   WRITE_APPEND = 'WRITE_APPEND'
   WRITE_EMPTY = 'WRITE_EMPTY'
@@ -606,6 +607,7 @@ def validate_create(disposition):
   @staticmethod
   def validate_write(disposition):
     values = (
+        BigQueryDisposition.WRITE_TRUNCATE_DATA,
         BigQueryDisposition.WRITE_TRUNCATE,
         BigQueryDisposition.WRITE_APPEND,
         BigQueryDisposition.WRITE_EMPTY)
@@ -1517,6 +1519,7 @@ def __init__(
         - BigQueryDisposition.CREATE_NEVER: fail the write if does not exist.
       write_disposition: A string describing what happens if the table has
         already some data. Possible values are:
+        -  BigQueryDisposition.WRITE_TRUNCATE_DATA: delete rows, keep schema.
         -  BigQueryDisposition.WRITE_TRUNCATE: delete existing rows.
         -  BigQueryDisposition.WRITE_APPEND: add to existing rows.
         -  BigQueryDisposition.WRITE_EMPTY: fail the write if table not empty.
@@ -1558,6 +1561,7 @@ def __init__(
     self.create_disposition = create_disposition
     self.write_disposition = write_disposition
     if write_disposition in (BigQueryDisposition.WRITE_EMPTY,
+                             BigQueryDisposition.WRITE_TRUNCATE_DATA,
                              BigQueryDisposition.WRITE_TRUNCATE):
       raise ValueError(
           'Write disposition %s is not supported for'
@@ -2126,6 +2130,7 @@ def __init__(
       write_disposition (BigQueryDisposition): A string describing what happens
         if the table has already some data. Possible values are:
 
+        * :attr:`BigQueryDisposition.WRITE_TRUNCATE_DATA`: delete table data but preserve table metadata.
         * :attr:`BigQueryDisposition.WRITE_TRUNCATE`: delete existing rows.
         * :attr:`BigQueryDisposition.WRITE_APPEND`: add to existing rows.
         * :attr:`BigQueryDisposition.WRITE_EMPTY`: fail the write if table not

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

@@ -416,7 +416,11 @@ def process(self, element, schema_mod_job_name_prefix):
     # When writing to normal tables WRITE_TRUNCATE will overwrite the schema but
     # when writing to a partition, care needs to be taken to update the schema
     # even on WRITE_TRUNCATE.
-    if (self._write_disposition not in ('WRITE_TRUNCATE', 'WRITE_APPEND') or
+    if (self._write_disposition not in (
+        'WRITE_TRUNCATE',
+        'WRITE_TRUNCATE_DATA',
+        'WRITE_APPEND'
+      ) or
         not additional_parameters or
         not additional_parameters.get("schemaUpdateOptions")):
       # No need to modify schema of destination table
@@ -613,11 +617,12 @@ def _determine_write_disposition(self, copy_to_reference) -> tuple[bool, str]:
     Determines the write disposition for a BigQuery copy job,
      based on destination.
 
-    When the write_disposition for a job is WRITE_TRUNCATE, multiple copy jobs
-    to the same destination can interfere with each other, truncate data, and
-    write to the BigQuery table repeatedly. To prevent this, the first copy job
-    runs with the user's specified write_disposition, but subsequent jobs must
-    always use WRITE_APPEND. This ensures that subsequent copy jobs do not
+    When the write_disposition for a job is WRITE_TRUNCATE or 
+    WRITE_TRUNCATE_DATA, multiple copy jobs to the same destination can 
+    interfere with each other, truncate data, and write to the BigQuery 
+    table repeatedly. To prevent this, the first copy job runs with the 
+    user's specified write_disposition, but subsequent jobs must always
+    use WRITE_APPEND. This ensures that subsequent copy jobs do not
     clear out data appended by previous jobs.
 
     Args:
@@ -1231,7 +1236,9 @@ def _load_data(
                 load_job_project_id=self.load_job_project_id),
             schema_mod_job_name_pcv))
 
-    if self.write_disposition in ('WRITE_EMPTY', 'WRITE_TRUNCATE'):
+    if self.write_disposition in (
+        'WRITE_EMPTY', 'WRITE_TRUNCATE', 'WRITE_TRUNCATE_DATA'
+      ):
       # All loads going to the same table must be processed together so that
       # the truncation happens only once. See
       # https://github.com/apache/beam/issues/24535.

sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py

@@ -5440,12 +5440,15 @@ class JobConfigurationTableCopy(_messages.Message):
       destination table already exists. The following values are supported: *
       WRITE_TRUNCATE: If the table already exists, BigQuery overwrites the
       table data and uses the schema and table constraints from the source
-      table. * WRITE_APPEND: If the table already exists, BigQuery appends the
-      data to the table. * WRITE_EMPTY: If the table already exists and
-      contains data, a 'duplicate' error is returned in the job result. The
-      default value is WRITE_EMPTY. Each action is atomic and only occurs if
-      BigQuery is able to complete the job successfully. Creation, truncation
-      and append actions occur as one atomic update upon job completion.
+      table. * WRITE_TRUNCATE_DATA: If the table already exists, BigQuery
+      overwrites the table data, but does not use schema and table constraints
+      from the source table.* WRITE_APPEND: If the table already exists, 
+      BigQuery appends the data to the table. * WRITE_EMPTY: If the table 
+      already exists and contains data, a 'duplicate' error is returned in 
+      the job result. The default value is WRITE_EMPTY. Each action is atomic 
+      and only occurs if BigQuery is able to complete the job successfully. 
+      Creation, truncation and append actions occur as one atomic update 
+      upon job completion.
   """
 
   class OperationTypeValueValuesEnum(_messages.Enum):