diff --git a/CHANGES.md b/CHANGES.md index 537c3a32447c..7ef1c252c060 100644 --- a/CHANGES.md +++ b/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 diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java index 0cefea4639fd..2f1d6b702691 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java @@ -3038,6 +3038,17 @@ public enum WriteDisposition { */ WRITE_TRUNCATE, + /** + * Specifies that write should replace a table, but not the table schema. + * + *

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. + * + *

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 @@ && getStorageApiTriggeringFrequency(bqOptions) != null) { 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 bigLakeConfiguration = getBigLakeConfiguration(); if (bigLakeConfiguration != null) { @@ -4229,6 +4241,9 @@ private WriteResult continueExpandTyped( 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()); diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpdateSchemaDestination.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpdateSchemaDestination.java index 536d318f0213..7f801ef38824 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpdateSchemaDestination.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpdateSchemaDestination.java @@ -257,6 +257,9 @@ public void finishBundle(FinishBundleContext context) throws Exception { 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())) { diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java index ec59db14501b..ec498c66aacc 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryWriteConfiguration.java +++ b/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 diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index 2a52b31edf7f..34051e98a992 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/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 diff --git a/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py b/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py index dabe80c5c27e..db51ecf22403 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py +++ b/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. diff --git a/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py b/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py index 3392bdaf9592..8b9fa0f883f4 100644 --- a/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/bigquery_v2_messages.py +++ b/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):