使用UDF解析PySpark DataFrame中的嵌套XML字段

嗨，我来帮你搞定这个PySpark解析XML列的需求！咱们直接上代码和步骤，保证你能跑通：

首先明确咱们的目标：把visitors列里的每个<visitor>节点解析成独立的行，同时提取节点属性作为新列，还要保留原有的sex和updated_at字段。

方法一：使用Python UDF解析XML #

这种方法适合需要自定义XML处理逻辑的场景，步骤如下：

1. 导入必要模块 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf, explode, col
from pyspark.sql.types import ArrayType, StructType, StructField, StringType, IntegerType
import xml.etree.ElementTree as ET

2. 初始化SparkSession并创建示例DataFrame #

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

# 模拟你的业务数据
data = [
    ("F", 1574264158, '<?xml version="1.0" encoding="utf-8"?> <visitors> <visitor id="9615" age="68" sex="F" /> <visitor id="1882" age="34" sex="M" /> <visitor id="5987" age="23" sex="M" /> </visitors>')
]
df = spark.createDataFrame(data, ["sex", "updated_at", "visitors"])

3. 定义UDF的返回Schema #

因为XML里有多个<visitor>节点，我们需要返回一个结构体数组来存储所有访客信息：

visitor_schema = ArrayType(
    StructType([
        StructField("visitor_id", StringType(), nullable=True),
        StructField("visitor_age", IntegerType(), nullable=True),
        StructField("visitor_sex", StringType(), nullable=True)
    ])
)

4. 编写XML解析函数并注册UDF #

def parse_visitor_xml(xml_str):
    root = ET.fromstring(xml_str)
    visitors = []
    # 遍历每个visitor节点，提取属性值
    for visitor in root.findall("visitor"):
        visitors.append({
            "visitor_id": visitor.attrib.get("id"),
            "visitor_age": int(visitor.attrib.get("age")) if visitor.attrib.get("age") else None,
            "visitor_sex": visitor.attrib.get("sex")
        })
    return visitors

parse_xml_udf = udf(parse_visitor_xml, visitor_schema)

5. 应用UDF并展开结果 #

# 生成包含访客数组的新列
df_with_visitors = df.withColumn("visitor_data", parse_xml_udf(col("visitors")))

# 展开数组，让每个访客单独占一行
df_exploded = df_with_visitors.select("sex", "updated_at", explode(col("visitor_data")).alias("visitor"))

# 提取结构体中的字段作为独立列
final_df = df_exploded.select(
    "sex", 
    "updated_at", 
    col("visitor.visitor_id").alias("visitor_id"),
    col("visitor.visitor_age").alias("visitor_age"),
    col("visitor.visitor_sex").alias("visitor_sex")
)

final_df.show()

运行后你会得到这样的结果：

+---+----------+----------+-----------+------------+
|sex|updated_at|visitor_id|visitor_age|visitor_sex|
+---+----------+----------+-----------+------------+
|  F|1574264158|      9615|         68|           F|
|  F|1574264158|      1882|         34|           M|
|  F|1574264158|      5987|         23|           M|
+---+----------+----------+-----------+------------+

方法二：使用Spark XML内置函数（推荐） #

如果处理大数据量，内置的from_xml函数性能比Python UDF好很多，步骤如下：

1. 添加依赖 #

运行Spark时需要带上spark-xml包，比如用spark-submit：

spark-submit --packages com.databricks:spark-xml_2.12:0.15.0 your_script.py

（版本号请匹配你的Spark版本，比如Spark 3.3.x对应0.15.0）

2. 代码实现 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import from_xml, explode, col
from pyspark.sql.types import ArrayType, StructType, StructField, StringType, IntegerType

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

# 定义XML对应的Schema
visitor_schema = ArrayType(
    StructType([
        StructField("id", StringType(), nullable=True),
        StructField("age", IntegerType(), nullable=True),
        StructField("sex", StringType(), nullable=True)
    ])
)

# 模拟数据
data = [
    ("F", 1574264158, '<?xml version="1.0" encoding="utf-8"?> <visitors> <visitor id="9615" age="68" sex="F" /> <visitor id="1882" age="34" sex="M" /> <visitor id="5987" age="23" sex="M" /> </visitors>')
]
df = spark.createDataFrame(data, ["sex", "updated_at", "visitors"])

# 解析XML列
df_with_visitors = df.withColumn("visitor_data", from_xml(col("visitors"), visitor_schema))

# 展开并提取字段
final_df = df_with_visitors.select(
    "sex", 
    "updated_at", 
    explode(col("visitor_data")).alias("visitor")
).select(
    "sex", 
    "updated_at", 
    col("visitor.id").alias("visitor_id"),
    col("visitor.age").alias("visitor_age"),
    col("visitor.sex").alias("visitor_sex")
)

final_df.show()

这个方法和UDF的结果完全一致，但性能更优，适合生产环境的大数据处理。

内容的提问来源于stack exchange，提问作者Mr.Tssou