Spark动态解析DataFrame JSON列性能优化求助

在Databricks中通过Spark JDBC关联查询PostgreSQL，得到包含17万条记录的DataFrame，其中data列存储结构不固定的JSON字符串。由于Schema非静态，未预先定义结构，使用以下代码动态推导并解析JSON：

json_schema = spark.read.json(df_code.select("data").rdd.map(lambda row: row.data)).schema
df_json_data = df_code.withColumn('json_data', F.from_json(F.col('data'), json_schema)).drop('data')

处理17万条数据时耗时极长，但小批量（数千条）数据无性能问题。当前集群配置为1个驱动节点+4个工作节点，每个工作节点拥有56GB内存与8核CPU，且为独享状态，无其他任务运行。

示例数据 #

注：部分记录包含大体积Base64编码内容（如第一条记录中的867506个随机字符，推测是网页捕获的JPG图片），且各记录Schema无统一结构：

Record1 #

{"recomend_portal":2,
"upload":[{"storage":"base64","name":"dbeaver.exe"
,"url":"data:application/x-msdownload;base64,867506 Random Characters","size":532488,"type":"application/x-msdownload","originalName":"dbeaver.exe","hash":"0fbee5a6f48b20225eb23bb59d870147","fileType":"JPG"}]
,"secondUploadPdf":[{"storage":"base64","name":"My HR Tool-94a42a23-6239-42ca-b5d6-ea13c5969b24.url","url":"data:application/octet-stream;base64,116 Random Characters",
"size":86,"type":null,"originalName":"My HR Tool.url","hash":"0fbee5a6f48b20225eb23bb59d870147","fileType":"PDF"}],
"aFieldWithInlineValidation":"df","thePortalsCapabilitiesMetMyNeeds":4}

Record2 #

{"recomend_portal":0,"thePortalsCapabilitiesMetMyNeeds":8,"upload":[],"secondUploadPdf":[]}

Record3 #

{"start_time":"2023-06-28 16:00:00","end_time":"2023-06-28 17:00:00"}

1. 优化Schema推导逻辑 #

• 避免全量数据推导Schema：从17万条数据中抽样（比如取10%或固定1-2万条）来推导Schema，全量推导会遍历所有JSON字符串，尤其包含大Base64数据时，IO和计算成本极高。示例代码：

  # 抽样推导Schema
  sampled_df = df_code.sample(fraction=0.1, seed=42)
  json_schema = spark.read.json(sampled_df.select("data").selectExpr("data as value")).schema
  

• 替换RDD转换：用selectExpr直接传递列给spark.read.json，避免RDD的序列化开销，减少性能损耗。

2. 分离大体积Base64数据处理 #

大Base64字符串是核心性能瓶颈，可拆分数据集分别处理：

from pyspark.sql import functions as F
import json
import base64
import uuid

# 标记包含大Base64的记录
df_with_flag = df_code.withColumn(
    "has_large_base64",
    F.when(F.length(F.regexp_extract(F.col("data"), r'"url":"data:[^;]+;base64,([^"]+)"', 1)) > 100000, 1).otherwise(0)
)

# 小数据集正常解析
small_df = df_with_flag.filter(F.col("has_large_base64") == 0)
small_parsed = small_df.withColumn('json_data', F.from_json(F.col('data'), json_schema)).drop('data')

# 大数据集单独存储Base64文件，保留路径
def save_large_base64(row):
    data_json = json.loads(row.data)
    # 处理upload中的大文件
    for item in data_json.get("upload", []):
        url = item.get("url", "")
        if "," in url and len(url.split(",")[-1]) > 100000:
            file_id = str(uuid.uuid4())
            file_path = f"/dbfs/large_files/{file_id}.bin"
            with open(file_path, "wb") as f:
                f.write(base64.b64decode(url.split(",")[-1]))
            item["url"] = file_path
    # 同理处理secondUploadPdf
    for item in data_json.get("secondUploadPdf", []):
        url = item.get("url", "")
        if "," in url and len(url.split(",")[-1]) > 100000:
            file_id = str(uuid.uuid4())
            file_path = f"/dbfs/large_files/{file_id}.bin"
            with open(file_path, "wb") as f:
                f.write(base64.b64decode(url.split(",")[-1]))
            item["url"] = file_path
    return row._replace(data=json.dumps(data_json))

large_df_processed = df_with_flag.filter(F.col("has_large_base64") == 1).rdd.map(save_large_base64).toDF()
large_parsed = large_df_processed.withColumn('json_data', F.from_json(F.col('data'), json_schema)).drop('data')

# 合并结果
final_df = small_parsed.unionByName(large_parsed, allowMissingColumns=True)

3. 调优Spark配置 #

• 调整并行度：当前集群有32可用核（4节点×8核），设置spark.sql.shuffle.partitions为64（核数的2倍），避免分区过少导致资源闲置：

  spark.conf.set("spark.sql.shuffle.partitions", "64")
  

• 优化内存配置：开启堆外列存储，调整堆外内存预留，应对大Base64数据的内存压力：

  spark.conf.set("spark.sql.columnVector.offheap.enabled", "true")
  spark.conf.set("spark.executor.memoryOverhead", "10g")
  

• 使用Kryo序列化：替换默认Java序列化，提升复杂数据的序列化效率：

  spark.conf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
  

4. 分阶段处理数据 #

• 缓存原始DataFrame：在推导Schema前缓存df_code，避免重复从JDBC读取数据：

  df_code.cache()
  df_code.count()  # 触发缓存
  

• 按需解析字段：先将JSON转为Map类型，只提取已知字段，未知字段保留为Map，减少不必要的解析开销：

  df_map = df_code.withColumn("json_data", F.from_json(F.col("data"), F.map_type(F.stringType(), F.stringType())))
  df_final = df_map.select(
      F.col("json_data.recomend_portal").cast("int"),
      F.col("json_data.thePortalsCapabilitiesMetMyNeeds").cast("int"),
      F.col("json_data.start_time"),
      F.col("json_data.end_time"),
      F.expr("json_data - array('recomend_portal', 'thePortalsCapabilitiesMetMyNeeds', 'start_time', 'end_time')").alias("other_fields")
  )
  

5. 推到数据库侧预处理 #

利用PostgreSQL的JSON函数，提前解析常用字段并剥离大Base64内容，减少Spark侧处理的数据量：

SELECT 
    id,
    (data::json->>'recomend_portal')::int as recomend_portal,
    (data::json->>'thePortalsCapabilitiesMetMyNeeds')::int as thePortalsCapabilitiesMetMyNeeds,
    -- 只返回文件元数据，不返回Base64内容
    json_build_object(
        'storage', data::json->'upload'->0->>'storage',
        'name', data::json->'upload'->0->>'name',
        'size', (data::json->'upload'->0->>'size')::bigint
    ) as upload_meta,
    data  -- 保留原始JSON用于其他字段解析
FROM your_table

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