如何基于PySpark DataFrame每行生成独立子DataFrame以解析JSON

原始DataFrame #

获取原始数据的代码及数据展示：

df_json = spark.sql("select id, type, json from process.base_json")
df_json.show(truncate=False)

输出结果：

+--------+--------+---------------------------------------------------------------------------------------------+
|id      |type    |json                                                                                         |
+--------+--------+---------------------------------------------------------------------------------------------+
|23016494|TAX     |{"Id":"253","RESULT":{"DATA":{"response":[{"message":"ID 253 invalid"}]}}}                   |
|23020867|WARRANTY|{"Id":"108","RESULT":{"DATA":{"result":[{"message":"Nomatches"}]}},"Type":"ID"}              |
|23021055|WARRANTY|{"Id":"332","RESULT":{"DATA":{"detail":{"cre":"BANK","nid":"332"}},"Type":"ID"}}             |
|23016497|TAX     |{"Id":"643","RESULT":{"DATA":{"registry":[{"dv":"5","st":"ACT","name":"MAY"}]}}}             |
+--------+--------+---------------------------------------------------------------------------------------------+

需求说明 #

将上述DataFrame的每一行数据单独拆分，生成对应数量的独立子DataFrame，每个子DataFrame仅包含原DataFrame中的一行数据，方便对每行的JSON字段单独解析处理。

实现方案 #

方式1：添加索引后过滤生成（推荐大数据量场景） #

# 为原DataFrame添加行索引
df_with_index = df_json.rdd.zipWithIndex().toDF(["row_data", "index"])

# 生成子DataFrame字典，键为子DataFrame名称，值为对应DataFrame
sub_dfs = {}
total_rows = df_json.count()
for i in range(total_rows):
    # 过滤指定索引的行，并展开原始字段
    sub_df = df_with_index.filter(df_with_index.index == i) \
                          .select("row_data.id", "row_data.type", "row_data.json")
    sub_dfs[f"df_json{i+1}"] = sub_df

# 验证df_json1
sub_dfs["df_json1"].show(truncate=False)

输出df_json1：

+--------+----+---------------------------------------------------------------------------------------------+
|id      |type|json                                                                                         |
+--------+----+---------------------------------------------------------------------------------------------+
|23016494|TAX |{"Id":"253","RESULT":{"DATA":{"response":[{"message":"ID 253 invalid"}]}}}                   |
+--------+----+---------------------------------------------------------------------------------------------+

方式2：收集行数据逐个创建（仅适用于小数据量场景） #

# 收集所有行数据到Driver节点
rows = df_json.collect()

# 逐个创建子DataFrame
sub_dfs = {}
for idx, row in enumerate(rows, start=1):
    sub_df = spark.createDataFrame([row], schema=df_json.schema)
    sub_dfs[f"df_json{idx}"] = sub_df

# 验证df_json2
sub_dfs["df_json2"].show(truncate=False)

输出df_json2：

+--------+--------+---------------------------------------------------------------------------------------------+
|id      |type    |json                                                                                         |
+--------+--------+---------------------------------------------------------------------------------------------+
|23020867|WARRANTY|{"Id":"108","RESULT":{"DATA":{"result":[{"message":"Nomatches"}]}},"Type":"ID"}              |
+--------+--------+---------------------------------------------------------------------------------------------+

注意事项 #

• 方式2的collect()会将全量数据拉取到Driver节点，大数据量场景下易引发内存溢出，优先使用方式1。
• 生成的子DataFrame可通过sub_dfs字典直接调用，后续可针对每个子DataFrame单独解析JSON字段。

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