GCP Dataflow中TensorFlow Transform结果写入BigQuery报错咨询

问题背景 #

在GCP Dataflow流水线中，尝试将TensorFlow Transform（TFT）组件输出的转换数据写入BigQuery时，先后遇到两次序列化错误，需要适配BigQuery的写入要求进行代码修正。

第一次报错：RecordBatch无法序列化 #

File "/usr/local/lib/python3.7/site-packages/apache_beam/io/gcp/bigquery_file_loads.py", line 261, in process
    writer.write(row)
  File "/usr/local/lib/python3.7/site-packages/apache_beam/io/gcp/bigquery_tools.py", line 1400, in write
    return self._file_handle.write(self._coder.encode(row) + b'\n')
  File "/usr/local/lib/python3.7/site-packages/apache_beam/io/gcp/bigquery_tools.py", line 1351, in encode
    default=default_encoder).encode('utf-8')
  File "/usr/local/lib/python3.7/json/__init__.py", line 238, in dumps
    **kw).encode(obj)
  File "/usr/local/lib/python3.7/json/encoder.py", line 199, in encode
    chunks = self.iterencode(o, _one_shot=True)
  File "/usr/local/lib/python3.7/json/encoder.py", line 257, in iterencode
    return _iterencode(o, 0)
  File "/usr/local/lib/python3.7/site-packages/apache_beam/io/gcp/bigquery_tools.py", line 134, in default_encoder
    "Object of type '%s' is not JSON serializable" % type(obj).__name__)
TypeError: Object of type 'RecordBatch' is not JSON serializable [while running 'train - Write to BigQuery/BigQueryBatchFileLoads/ParDo(WriteRecordsToFile)/ParDo(WriteRecordsToFile)-ptransform-128']

原因：代码中设置了output_record_batches=True，TFT输出的是pyarrow.RecordBatch对象，而BigQuery的WriteToBigQuery组件仅支持可JSON序列化的行字典。

第二次报错：float32无法序列化 #

将output_record_batches改为False后，出现新错误：

File "/usr/local/lib/python3.7/site-packages/apache_beam/io/gcp/bigquery_file_loads.py", line 261, in process
    writer.write(row)
  File "/usr/local/lib/python3.7/site-packages/apache_beam/io/gcp/bigquery_tools.py", line 1400, in write
    return self._file_handle.write(self._coder.encode(row) + b'\n')
  File "/usr/local/lib/python3.7/site-packages/apache_beam/io/gcp/bigquery_tools.py", line 1351, in encode
    default=default_encoder).encode('utf-8')
  File "/usr/local/lib/python3.7/json/__init__.py", line 238, in dumps
    **kw).encode(obj)
  File "/usr/local/lib/python3.7/json/encoder.py", line 199, in encode
    chunks = self.iterencode(o, _one_shot=True)
  File "/usr/local/lib/python3.7/json/encoder.py", line 257, in iterencode
    return _iterencode(o, 0)
  File "/usr/local/lib/python3.7/site-packages/apache_beam/io/gcp/bigquery_tools.py", line 134, in default_encoder
    "Object of type '%s' is not JSON serializable" % type(obj).__name__)
TypeError: Object of type 'float32' is not JSON serializable [while running 'train - Write to BigQuery/BigQueryBatchFileLoads/ParDo(WriteRecordsToFile)/ParDo(WriteRecordsToFile)-ptransform-110']

原因：TFT输出的数值类型是TensorFlow的float32等张量类型，无法直接被JSON序列化。

解决方案 #

方案1：处理RecordBatch输出（保留output_record_batches=True） #

将RecordBatch转换为BigQuery可接受的行字典：

import pyarrow as pa

def record_batch_to_dicts(record_batch):
    # 将RecordBatch转为pandas DataFrame，再生成字典列表
    df = record_batch.to_pandas()
    return df.to_dict('records')

def write_to_bigquery(transformed_dataset, step):
    transformed_data, transformed_metadata = transformed_dataset
    schema_string = convert_schema_to_string(transformed_metadata.schema)
    
    # 转换RecordBatch为可序列化的行字典
    formatted_data = (
        transformed_data
        | f'{step} - Convert RecordBatch to Dicts' >> beam.FlatMap(record_batch_to_dicts)
    )
    
    formatted_data | f'{step} - Write to BigQuery' >> beam.io.WriteToBigQuery(
        table=f'{PROJECT}.{OUT_DATASET_ID}.{OUT_TABLE_NAME}',
        schema=schema_string,
        create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
        write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND
    )

方案2：处理张量类型（使用output_record_batches=False） #

将TFT输出的张量类型转换为Python原生类型：

def convert_tensor_types(row):
    converted_row = {}
    for key, value in row.items():
        # 将TensorFlow张量转为Python原生数值
        if hasattr(value, 'numpy'):
            # 单个值取item，数组转列表
            converted_row[key] = value.numpy().item() if value.shape == () else value.numpy().tolist()
        else:
            converted_row[key] = value
    return converted_row

def write_to_bigquery(transformed_dataset, step):
    transformed_data, transformed_metadata = transformed_dataset
    schema_string = convert_schema_to_string(transformed_metadata.schema)
    
    # 转换张量类型为Python原生类型
    formatted_data = (
        transformed_data
        | f'{step} - Convert Tensor Types' >> beam.Map(convert_tensor_types)
    )
    
    formatted_data | f'{step} - Write to BigQuery' >> beam.io.WriteToBigQuery(
        table=f'{PROJECT}.{OUT_DATASET_ID}.{OUT_TABLE_NAME}',
        schema=schema_string,
        create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
        write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND
    )

额外优化：Schema转换函数改进 #

原Schema转换函数可优化为更健壮的字段列表格式，避免手动拼接字符串的错误：

from tensorflow_metadata.proto.v0 import schema_pb2

def convert_schema_to_bigquery_fields(schema):
    fields = []
    for feature in schema.feature:
        field = {"name": feature.name}
        if feature.type == schema_pb2.FLOAT:
            field["type"] = "FLOAT"
        elif feature.type == schema_pb2.INT:
            field["type"] = "STRING" if feature.int_domain.is_categorical else "INTEGER"
        elif feature.type == schema_pb2.STRING:
            field["type"] = "STRING"
        elif feature.type == schema_pb2.BOOL:
            field["type"] = "BOOLEAN"
        fields.append(field)
    return fields

# 在write_to_bigquery中使用：
schema_fields = convert_schema_to_bigquery_fields(transformed_metadata.schema)
beam.io.WriteToBigQuery(
    ...,
    schema=schema_fields,  # 直接传入字段列表
    ...
)

标准流程说明 #

TFT与BigQuery集成的核心是确保数据格式适配BigQuery的序列化要求：

1. 若使用RecordBatch输出，需转换为行字典；Beam 2.30+版本支持use_arrow=True参数，可直接写入Arrow格式数据到BigQuery
2. 若使用字典输出，必须将TensorFlow张量类型转换为Python原生类型
3. Schema需正确映射TFT输出类型到BigQuery支持的类型

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