BigQueryOperator changes the table schema and column modes when write_disposition='WRITE_TRUNCATE'

user_12317609 0
Viewed 2531

I am using Airflow's BigQueryOperator to populate the BQ table with write_disposition='WRITE_TRUNCATE'. The problem is that every time the task runs, it alters the table schema and also the column mode from Required to Nullable. The create_disposition I am using is 'CREATE_NEVER'. Since my tables are pre-created, I don't want the schemas or column modes to be altered. Using write_disposition='WRITE_APPEND' fixes the issue but my requirement is to use WRITE_TRUNCATE. Any idea why BigQueryOperator alters the schema and mode?

google-bigquery python
1 Answers

I had a similar issue, not with the required/nullable shcema value, but on policy tags, and the behavior is the same: the policy tags are overriden (and lost). Here the answer of the Google support team:

If you overwrite to a destination table, any existing policy tags are removed from the table, unless you use the --destination_schema flag to specify a schema with policy tags. For WRITE_TRUNCATE, the disposition overwrites the existing table and the schema. If you want to keep the policy tags, you can use "--destination_schema" to specify a schema with policy tags.

However, with my test in python, I observed 2 different behaviors between the QueryJob (job based on sql query and that sinks the result in a table) and the LoadJob (job that loads data from a file and that sinks the data in a table).

  • If you perform a LoadJob,
    • Remove the schema autodetec
    • Get the schema of the original table
    • Perform the load job

Like this in Python

    job_config = bigquery.job.LoadJobConfig()
    job_config.create_disposition = bigquery.job.CreateDisposition.CREATE_IF_NEEDED
    job_config.write_disposition = bigquery.job.WriteDisposition.WRITE_TRUNCATE
    job_config.skip_leading_rows = 1
    # job_config.autodetect = True
    job_config.schema = client.get_table(table).schema

    query_job = client.load_table_from_uri(uri, table, job_config=job_config)
    res = query_job.result()

This solution, to copy the schema, doesn't work with a QueryJob

The workaround is the following (works for LoadJob and QueryJob)

  • Truncate the table
  • Perform a job in WRITE_EMPTY mode

Trade-off:

  • the WRITE_TRUNCATE is atomic: if the write write failed, the data aren't truncated
  • the workaround is in 2 steps: if the write write failed, the data are already deleted
    config = bigquery.job.QueryJobConfig()
    config.create_disposition = bigquery.job.CreateDisposition.CREATE_IF_NEEDED
    config.write_disposition = bigquery.job.WriteDisposition.WRITE_EMPTY
    # Don't work
    # config.schema = client.get_table(table).schema
    config.destination = table

    # Step 1 truncate the table
    query_job = client.query(f'TRUNCATE TABLE `{table}`')
    res = query_job.result()
    # Step 2: Load the new data
    query_job = client.query(request, job_config=job_config)
    res = query_job.result()

All of this to tell you that the BigQuery operator on Airflow isn't the problem. It's a BigQuery issue. You have a workaround to achieve what you want.

user_113725930
Related