在 Spark 中修改 UDF 以创建附加键列答案

【问题标题】：Modifying UDF in Spark to Create additional key column在 Spark 中修改 UDF 以创建附加键列
【发布时间】：2022-01-23 12:14:08
【问题描述】：

我有一个由数据行和需要解析的 XML 列组成的数据框。我可以使用stack overflow solution 中的以下代码解析该 XML：

import xml.etree.ElementTree as ET
import pyspark.sql.functions as F

@F.udf('array<struct<id:string, age:string, sex:string>>')
def parse_xml(s):
    root = ET.fromstring(s)
    return list(map(lambda x: x.attrib, root.findall('visitor')))
    
df2 = df.select(
    F.explode(parse_xml('visitors')).alias('visitors')
).select('visitors.*')

df2.show()

这个函数为解析的 XML 数据创建一个新的数据框。

相反，我如何修改此函数以包含原始数据框中的一列，以便以后加入？

例如，如果原始数据框如下所示：

+----+---+----------------------+
|id  |a  |xml                   |
+----+---+----------------------+
|1234|.  |<row1, row2>          |
|2345|.  |<row3, row4>, <row5>  |
|3456|.  |<row6>                |
+----+---+----------------------+

如何在新创建的数据框的每一行中包含 ID？

【问题讨论】：

标签： python apache-spark pyspark user-defined-functions

【解决方案1】：

在构造df2 时，还需要select id 列。我认为您可以执行以下操作：

df2 = df.select('id',
    F.explode(parse_xml('visitors')).alias('visitors')
).select('id','visitors.*')

这里有一个独立的小例子来展示这个想法：

import pyspark.sql.functions as F
df = spark.createDataFrame([(1,["xml1", "xml2", "xml3"]), (2,["xml4", "xml5", "xml6"]),(3,["xml7", "xml8", "xml9"])], ["id", "xml"])
df.show()
df_exploded_with_id = df.select("id", F.explode(F.col("xml")))
df_exploded_with_id.show()

输出：

+---+------------------+
| id|               xml|
+---+------------------+
|  1|[xml1, xml2, xml3]|
|  2|[xml4, xml5, xml6]|
|  3|[xml7, xml8, xml9]|
+---+------------------+

+---+----+
| id| col|
+---+----+
|  1|xml1|
|  1|xml2|
|  1|xml3|
|  2|xml4|
|  2|xml5|
|  2|xml6|
|  3|xml7|
|  3|xml8|
|  3|xml9|
+---+----+

【讨论】：