Apache Beam写入BigQuery：Storage Write API不遵循主键问题

我使用以下DDL在BigQuery中创建了表：

CREATE TABLE mytable AS
(
  id STRING,
  source STRING,
  PRIMARY KEY (id) NOT ENFORCED
);

可以看到id被设为表的主键。我的Beam流水线定义如下：

def process_message(message):  
    import apache_beam as beam
    import struct
    data = json.loads(message.decode("utf-8"))
    if data == {}:
        print(f"Running DELETE operation on row {this_message['Key']}")
        data['_CHANGE_TYPE'] = 'DELETE'
    else:   
        print(f"Running UPSERT operation on row {this_message['Key']}")
        data['_CHANGE_TYPE'] = 'UPSERT'
    data['_CHANGE_SEQUENCE_NUMBER'] = str(struct.pack('d', int(round(float(this_message['Value']['updated'])))).hex())
    return [data]


with beam.Pipeline(options=PipelineOptions([
    f"--project={project_id}",
    "--region=europe-west2",
    "--runner=DataflowRunner",
    "--streaming",
    "--temp_location=gs://tmp/cdc",
    "--staging_location=gs://tmp/cdc",
])) as pipeline:
  data = pipeline | 'ReadFromPubSub' >> ReadFromPubSub(subscription=f'projects/{project_id}/subscriptions/{bq_table_name}')

  data | 'ProcessMessages' >> beam.ParDo(process_message) | 'WriteToBigQuery' >> WriteToBigQuery(
    f'{project_id}:{bq_dataset}.{bq_table_name}',
    schema=schema,
    method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API,
    triggering_frequency=5
  )

消费若干条记录后查询BigQuery表，发现id字段存在大量重复值。请问如何让流水线遵循主键约束，正确执行UPSERT操作？

要让Beam流水线正确执行UPSERT/DELETE并遵循BigQuery主键约束，需要做以下关键调整：

1. 启用Storage Write API的CDC模式 #

你已使用STORAGE_WRITE_API方法，但需显式开启CDC（变更数据捕获）模式并指定主键字段，让BigQuery根据主键处理变更操作。修改WriteToBigQuery参数：

WriteToBigQuery(
    f'{project_id}:{bq_dataset}.{bq_table_name}',
    schema=schema,
    method=beam.io.WriteToBigQuery.Method.STORAGE_WRITE_API,
    triggering_frequency=5,
    write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND,
    additional_bq_parameters={
        'enableStreamingInserts': True,
        'ignoreUnknownValues': True,
        # 指定表的主键字段，匹配DDL定义
        'primaryKeyFields': ['id'],
        # 开启CDC模式，识别_CHANGE_TYPE字段
        'useCdc': True
    }
)

2. 修正数据处理逻辑的变量错误 #

原函数存在未定义变量this_message的问题，同时DELETE操作必须包含主键字段才能定位目标行，修正后的处理逻辑：

def process_message(message):  
    import apache_beam as beam
    import struct
    import json
    # 正确解析PubSub消息内容
    this_message = json.loads(message.decode("utf-8"))
    data = this_message.get('Value', {})
    
    if data == {}:
        print(f"Running DELETE operation on row {this_message['Key']}")
        # DELETE操作必须携带主键id
        data['id'] = this_message['Key']
        data['_CHANGE_TYPE'] = 'DELETE'
    else:   
        print(f"Running UPSERT operation on row {this_message['Key']}")
        data['id'] = this_message['Key']
        data['_CHANGE_TYPE'] = 'UPSERT'
    
    data['_CHANGE_SEQUENCE_NUMBER'] = str(struct.pack('d', int(round(float(this_message['Value']['updated'])))).hex())
    return [data]

3. 优化主键约束配置 #

若需要严格强制主键唯一性，可将表DDL中的NOT ENFORCED改为ENFORCED（注意：强制约束会增加写入开销，且要求表使用BigQuery标准SQL的分区/集群特性）。即使保留NOT ENFORCED，结合CDC模式的primaryKeyFields配置，BigQuery仍会按主键自动处理UPSERT/DELETE，避免重复值。

4. 确保变更序列的正确性 #

_CHANGE_SEQUENCE_NUMBER用于保证变更的顺序性，BigQuery会根据该字段的顺序应用变更，确保最新的操作覆盖旧数据。需确保该字段的生成逻辑能准确反映事件的时间先后顺序。

内容的提问来源于stack exchange，提问作者SeeBeeOss