Apache Beam Python：自动推断Schema并按Excel列序写入BigQuery

要保持Excel的列顺序写入BigQuery，不能依赖SCHEMA_AUTODETECT（它会按字母序或其他规则重排字段），需要手动生成与Excel列顺序完全一致的BigQuery Schema，再传入写入组件。具体实现步骤如下：

核心思路 #

1. 提前读取Excel文件的列名和数据类型，生成严格遵循Excel列顺序的BigQuery格式Schema
2. 在Beam写入BigQuery时，使用该自定义Schema替代自动推断选项

修改后的完整代码 #

import pandas as pd
from apache_beam.options.pipeline_options import PipelineOptions
import apache_beam as beam

def handle_nan(record):
    # 将NaN转换为BigQuery可识别的NULL（None）
    return {k: None if pd.isna(v) else v for k, v in record.items()}

def generate_bigquery_schema_from_excel(file_path, sheet_name):
    # 读取Excel前100行推断数据类型（可根据数据量调整行数）
    df = pd.read_excel(file_path, sheet_name=sheet_name, nrows=100)
    # 映射pandas数据类型到BigQuery标准类型
    dtype_mapping = {
        'int64': 'INTEGER',
        'float64': 'FLOAT',
        'datetime64[ns]': 'DATETIME',
        'bool': 'BOOLEAN',
        'string': 'STRING'
    }
    # 严格按Excel列顺序生成Schema
    schema = []
    for col in df.columns:
        # 优化类型推断，避免将文本识别为object类型
        col_dtype = df[col].convert_dtypes().dtype.name
        bq_type = dtype_mapping.get(col_dtype, 'STRING')  # 默认使用字符串类型兼容特殊格式
        schema.append({
            'name': col,
            'type': bq_type,
            'mode': 'NULLABLE'
        })
    return schema

def run_pipeline(input_file, tab_bottlenecks, output_table_bottlenecks, temp_location):
    # 关键步骤：提前生成带列顺序的BigQuery Schema
    bq_schema = generate_bigquery_schema_from_excel(input_file, tab_bottlenecks)

    # 配置管道选项
    pipeline_options = {
        'project': 'xxxx'
    }
    pipeline_options = PipelineOptions.from_dictionary(pipeline_options)
    
    with beam.Pipeline(options=pipeline_options) as pipeline:
        # 读取Excel数据并预处理
        input_data = (
            pipeline
            | 'Read From GCS' >> beam.Create([input_file])
            | 'Read Excel' >> beam.FlatMap(lambda file: pd.read_excel(file, tab_bottlenecks).to_dict('records'))
            | 'Handle NaN' >> beam.Map(handle_nan)
        )

        # 写入BigQuery时使用自定义Schema
        input_data | 'Write To BigQuery -> Bottlenecks' >> beam.io.WriteToBigQuery(
            output_table_bottlenecks,
            schema=bq_schema,  # 替换为手动生成的带顺序Schema
            create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
            write_disposition=beam.io.BigQueryDisposition.WRITE_TRUNCATE,
            custom_gcs_temp_location=temp_location
        )
        
        pipeline.run().wait_until_finish()

if __name__ == '__main__':
    input_file = 'gs://xxx/xxx.xlsx'
    tab_bottlenecks = 'Lista'
    output_table_bottlenecks = 'xx:xx.xx'
    temp_location = 'gs://xxx'
    # 需安装依赖库支持pandas读取GCS文件：pip install gcsfs
    run_pipeline(input_file, tab_bottlenecks, output_table_bottlenecks, temp_location)

关键注意事项 #

• 依赖安装：必须安装gcsfs库，让pandas可以直接读取GCS上的Excel文件
• 类型推断优化：使用df.convert_dtypes()可以更准确地识别字符串类型，避免将文本列识别为object类型
• Schema生成时机：在管道启动前生成Schema，避免分布式环境下重复读取Excel表头，提升运行效率
• 特殊列名处理：如果Excel列名包含空格、特殊字符，需额外处理（比如替换为下划线），否则BigQuery会抛出列名格式错误

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