如何使用PySpark DataFrame的toJSON()保留含Null值的键

问题说明 #

将Spark DataFrame转为JSON文件时，含Null值的列会被自动丢弃，需要保留对应键并将值设为null。

解决方案 #

方法一：使用DataFrameWriter直接写入（推荐） #

Spark的DataFrame.write.json()方法支持通过ignoreNulls参数控制Null值键的保留逻辑，默认该参数为True（丢弃Null键），设置为False即可保留：

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("DataFrameToJson").getOrCreate()

# 示例数据
data = [
    {"first": "1", "second": "2", "third": "3", "fourth": "4", "fifth": 5, "six": 6},
    {"first": "11", "second": "22", "third": "33", "fourth": "44", "fifth": 55, "six": None}
]

final_df = spark.createDataFrame(data)

# 写入JSON文件，保留Null值对应的键
final_df.write.option("ignoreNulls", "false").json("output_dir")

注：该方法会将数据写入目录下的多个分区文件，若需要单个文件，可先调用coalesce(1)再写入，但大数据量场景不建议使用此操作。

方法二：处理toJSON()返回的结果 #

如果必须使用toJSON().collect()的方式，可以手动补全每个JSON对象中缺失的列键：

from pyspark.sql import SparkSession
import json

spark = SparkSession.builder.appName("DataFrameToJson").getOrCreate()

data = [
    {"first": "1", "second": "2", "third": "3", "fourth": "4", "fifth": 5, "six": 6},
    {"first": "11", "second": "22", "third": "33", "fourth": "44", "fifth": 55, "six": None}
]

final_df = spark.createDataFrame(data)
# 获取DataFrame所有列名
all_columns = final_df.columns

df_json = final_df.toJSON().collect()
df_list_of_dicts = []
for json_str in df_json:
    row_dict = json.loads(json_str)
    # 补全缺失的列，设置值为None
    for col_name in all_columns:
        if col_name not in row_dict:
            row_dict[col_name] = None
    df_list_of_dicts.append(row_dict)

# 写入单个JSON文件
json_object = json.dumps(df_list_of_dicts, indent=4)
with open("output.json", "w") as f:
    f.write(json_object)

运行后生成的output.json会保留所有列键，Null值列将显示为null，符合预期输出。

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