PySpark中从嵌套DataFrame提取指定字段并匹配Type值

给定如下嵌套结构的JSON样本数据：

{
    "GeneralInformation":{
        "ID":"00001",
        "WebLinksInfo":{
            "LastUpdated":"2019-10-27",
            "WebSite":{
                "Type":"Home Page",
                "text":"https://www.aaaa.com/"
            }
        },
        "TextInfo":{
            "Text":[
                {
                    "Type":"Business",
                    "updated_at":"2018-09-14",
                    "unused_field":"en-US",
                    "Description":"Lorem ipsum dolor sit amet, consectetur adipiscing elit, laborum."
                },
                {
                    "Type":"Financial",
                    "updated_at":"2022-08-26",
                    "unused_field":"en-US",
                    "Description":"Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat."
                }
            ]
        },
        "Advisors":{
            "Auditor":{
                "Code":"AAA",
                "Name":"Aristotle"
            }
        }
    }
}

对应的PySpark DataFrame Schema为：

root
 |-- GeneralInformation: struct (nullable = true)
 |    |-- Advisors: struct (nullable = true)
 |    |    |-- Auditor: struct (nullable = true)
 |    |    |    |-- Code: string (nullable = true)
 |    |    |    |-- Name: string (nullable = true)
 |    |-- ID: string (nullable = true)
 |    |-- TextInfo: struct (nullable = true)
 |    |    |-- Text: array (nullable = true)
 |    |    |    |-- element: struct (containsNull = true)
 |    |    |    |    |-- Description: string (nullable = true)
 |    |    |    |    |-- updated_at: string (nullable = true)
 |    |    |    |    |-- Type: string (nullable = true)
 |    |    |    |    |-- unused_field: string (nullable = true)
 |    |-- WebLinksInfo: struct (nullable = true)
 |    |    |-- LastUpdated: string (nullable = true)
 |    |    |-- WebSite: struct (nullable = true)
 |    |    |    |-- text: string (nullable = true)
 |    |    |    |-- Type: string (nullable = true)

需求说明 #

提取GeneralInformation.ID，以及TextInfo.Text数组中Type为Business的Description和updated_at字段，将这些字段移至DataFrame根层级，且每个ID仅保留一条符合条件的记录。

用户尝试过的代码（存在局限性）：

df = spark.read.json('my_file_path')

new_df = df.select(col('GeneralInformation.ID'))
new_df = new_df.join(df.select(col('GeneralInformation.TextInfo.Text')))

以及：

df = df.withColumn('description', F.when(
    df['Text'][0]['Type'] == 'Business', lit(df['Text'][0]['Description'])))

上述方法无法遍历数组所有元素，无法确保获取到Type为Business的元素，期望输出如下：

{
    "ID": "00001",
    "Description": "Lorem ipsum dolor sit amet, consectetur adipiscing elit, laborum.",
    "Updated_at": "2018-09-14"
}

正确实现方案 #

方案1：展开数组过滤后聚合 #

先将数组展开，过滤出符合条件的元素，再按ID聚合确保单条记录：

from pyspark.sql import functions as F
from pyspark.sql.functions import col

# 读取数据
df = spark.read.json('my_file_path')

# 提取ID并展开Text数组
exploded_df = df.select(
    col("GeneralInformation.ID").alias("ID"),
    F.explode(col("GeneralInformation.TextInfo.Text")).alias("text_element")
)

# 过滤Type为Business的记录，提取目标字段
filtered_df = exploded_df.filter(col("text_element.Type") == "Business") \
    .select(
        "ID",
        col("text_element.Description").alias("Description"),
        col("text_element.updated_at").alias("Updated_at")
    )

# 按ID聚合，确保每个ID仅保留一条记录（若存在多条Business类型，取第一条）
final_df = filtered_df.groupBy("ID") \
    .agg(
        F.first("Description").alias("Description"),
        F.first("Updated_at").alias("Updated_at")
    )

# 查看结果
final_df.show()
final_df.printSchema()

方案2：直接过滤数组元素（性能更优） #

无需展开数组，用array_filter直接筛选符合条件的元素，再提取字段：

from pyspark.sql import functions as F
from pyspark.sql.functions import col

df = spark.read.json('my_file_path')

# 过滤数组并提取第一个符合条件的元素，再提取目标字段
final_df = df.select(
    col("GeneralInformation.ID").alias("ID"),
    F.element_at(
        F.array_filter(
            col("GeneralInformation.TextInfo.Text"),
            lambda x: x["Type"] == "Business"
        ),
        1  # 取数组中第一个符合条件的元素
    ).alias("business_element")
) \
.select(
    "ID",
    col("business_element.Description").alias("Description"),
    col("business_element.updated_at").alias("Updated_at")
)

# 查看结果
final_df.show()

两种方案均可实现需求：方案2无需展开数组，性能更适合大数据场景；方案1适合需要处理多条符合条件记录的场景。

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