在Databricks中使用PySpark扁平化嵌套JSON结构的问题

核心问题分析 #

你遇到的问题根源是location.address字段存储的是JSON格式的字符串，而非原生嵌套JSON结构，因此spark.read.json只能将其识别为字符串类型，无法自动展开。需要先将该字符串解析为结构化的StructType，再进行全表扁平化处理。

步骤1：解析JSON字符串为结构化字段 #

方法A：自动推断address字段的Schema（推荐）

若不清楚address的具体结构，可从样本数据中提取并自动推断Schema：

from pyspark.sql.functions import col, from_json

# 提取非空的address字符串样本
sample_address = df.select("location.address").filter(col("location.address").isNotNull()).first()[0]
# 自动推断address的Schema
address_schema = spark.read.json(sc.parallelize([sample_address])).schema

# 将字符串类型的address转换为结构化字段
df = df.withColumn(
    "location_address_struct",
    from_json(col("location.address"), address_schema)
)
# 替换原字符串字段为结构化字段并清理临时列
df = df.withColumn("location", col("location").dropField("address")) \
       .withColumn("location", col("location").withField("address", col("location_address_struct"))) \
       .drop("location_address_struct")

方法B：手动定义address字段的Schema

若已知address的结构，可直接手动定义Schema以提升效率：

from pyspark.sql.types import StructType, StructField, StringType, IntegerType

# 手动定义address的Schema示例（根据实际结构调整）
address_schema = StructType([
    StructField("street", StringType(), nullable=True),
    StructField("city", StringType(), nullable=True),
    StructField("state", StringType(), nullable=True),
    StructField("zip_code", IntegerType(), nullable=True)
])

# 转换字符串为结构化字段
df = df.withColumn(
    "location.address",
    from_json(col("location.address"), address_schema)
)

步骤2：扁平化所有嵌套列 #

使用递归函数自动处理所有嵌套的StructType字段，展开为扁平列：

def flatten_dataframe(nested_df):
    # 区分扁平列与嵌套列
    flat_columns = [col_name for col_name, dtype in nested_df.dtypes if dtype[:6] != "struct"]
    nested_columns = [col_name for col_name, dtype in nested_df.dtypes if dtype[:6] == "struct"]

    # 展开嵌套列，列名格式为"父列_子列"
    expanded_df = nested_df.select(
        flat_columns +
        [col(f"{parent_col}.{child_col}").alias(f"{parent_col}_{child_col}")
         for parent_col in nested_columns
         for child_col in nested_df.select(f"{parent_col}.*").columns]
    )

    # 递归处理剩余嵌套列
    if len(nested_columns) > 0:
        return flatten_dataframe(expanded_df)
    else:
        return expanded_df

# 执行扁平化
flat_df = flatten_dataframe(df)
flat_df.show()

额外优化建议 #

• 如果API返回的原始数据中address就是转义后的JSON字符串，建议在API调用后、传入Spark前先用Python的json.loads解析为结构化数据，减少后续Spark端的处理步骤。
• 若存在多个包含address字符串字段的location类列，可通过循环遍历处理，避免重复代码。

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