如何在PySpark中从schema JSON文件创建DataFrame schema

你提供的是类BigQuery格式的JSON schema定义，可以通过递归解析该JSON文件生成PySpark对应的StructType schema对象，再用该schema直接加载JSON数据文件，避免Spark自动推断schema带来的类型误差和性能损耗。

步骤1：编写schema转换逻辑 #

首先导入依赖并定义类型映射、递归转换函数：

from pyspark.sql import SparkSession
from pyspark.sql.types import (
    StructType, StructField,
    IntegerType, StringType, DoubleType, BooleanType, TimestampType, ArrayType
)
import json

# 定义JSON中类型到Spark类型的映射，可根据实际需求扩展
TYPE_MAPPING = {
    "INTEGER": IntegerType(),
    "STRING": StringType(),
    "BOOLEAN": BooleanType(),
    "DOUBLE": DoubleType(),
    "TIMESTAMP": TimestampType()
}

def build_spark_schema(bq_schema_fields: list) -> StructType:
    """递归将BQ格式的JSON schema转换为PySpark StructType"""
    spark_fields = []
    for field in bq_schema_fields:
        field_name = field["name"]
        field_mode = field["mode"]
        field_type = field["type"]
        nullable = field_mode == "NULLABLE"
        
        if field_type == "RECORD":
            # 嵌套类型递归处理子字段
            nested_schema = build_spark_schema(field["fields"])
            current_type = nested_schema
        else:
            # 基础类型直接映射
            current_type = TYPE_MAPPING[field_type]
        
        # 处理数组类型（mode为REPEATED的场景）
        if field_mode == "REPEATED":
            current_type = ArrayType(current_type)
        
        spark_fields.append(StructField(field_name, current_type, nullable))
    return StructType(spark_fields)

步骤2：读取schema文件生成Spark schema #

# 初始化SparkSession
spark = SparkSession.builder.appName("load_with_custom_schema").getOrCreate()

# 读取本地的schema JSON文件，替换为你自己的文件路径
with open("/path/to/your/schema.json", "r", encoding="utf-8") as f:
    bq_schema = json.load(f)

# 生成Spark可用的schema对象
spark_schema = build_spark_schema(bq_schema)

步骤3：用生成的schema加载JSON数据文件 #

# 加载JSON数据，schema会严格匹配预先定义的规则，替换为你自己的数据文件路径
df = spark.read.schema(spark_schema).json("/path/to/your/data.json")

# 验证schema和数据
df.printSchema()
df.show()

注意事项 #

如果你的schema还包含其他数据类型（比如FLOAT、DATE等），直接在TYPE_MAPPING中补充对应Spark类型即可。

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