首页 > 解决方案 > Spark UDF:如何在每一行上编写一个 UDF 以提取嵌套结构中的特定值?

问题描述

我在 Java 中使用 Spark 来处理 XML 文件。来自 databricks 的包 spark-xml 包用于将 xml 文件读入数据框。

示例 xml 文件是:

<RowTag>
    <id>1</id>
    <name>john</name>
    <expenses>
        <travel>
            <details>
                <date>20191203</date>
                <amount>400</amount>
            </details>
        </travel>
    </expenses>
</RowTag>
<RowTag>
    <id>2</id>
    <name>joe</name>
    <expenses>
        <food>
            <details>
                <date>20191204</date>
                <amount>500</amount>
            </details>
        </food>
    </expenses>
</RowTag>

结果 sparkDataset<Row> df如下图,每一行代表一个 xml 文件。

+--+------+----------------+
|id| name |expenses        |
+---------+----------------+
|1 | john |[[20191203,400]]|
|2 | joe  |[[20191204,500]]|
+--+------+----------------+

df.printSchema();如下图所示:

root
|-- id: int(nullable = true)
|-- name: string(nullable = true)
|-- expenses: struct (nullable = true)
|    |-- travel: struct (nullable = true)
|    |    |-- details: struct (nullable = true)
|    |    |    |-- date: string (nullable = true)
|    |    |    |-- amount: int (nullable = true)
|    |-- food: struct (nullable = true)
|    |    |-- details: struct (nullable = true)
|    |    |    |-- date: string (nullable = true)
|    |    |    |-- amount: int (nullable = true)

所需的输出数据帧如下:

+--+------+-------------+
|id| name |expenses_date|
+---------+-------------+
|1 | john |20191203     |
|2 | joe  |20191204     |
+--+------+-------------+

基本上我想要一个通用的解决方案来从具有以下结构的 xml 中获取日期,其中只有标签<X>会有所不同。

<RowTag>
    <id>1</id>
    <name>john</name>
    <expenses>
        **<X>**
            <details>
                <date>20191203</date>
                <amount>400</amount>
            </details>
        **</X>**
    </expenses>
</RowTag>

我试过的:

spark.udf().register("getDate",(UDF1 <Row, String>) (Row row) -> {
            return row.getStruct(0).getStruct(0).getAs("date").toString();
        }, DataTypes.StringType);

df.select(callUDF("getDate",df.col("expenses")).as("expenses_date")).show();

但它不起作用,因为 row.getStruct(0) 路由到<travel>,但是对于 row joe,在 下没有<travel>标签<expenses>,所以它返回了一个java.lang.NullPointerException. 我想要的是一个通用的解决方案,对于每一行,它可以自动获取下一个标签名称,例如row.getStruct(0)路由到<travel>for row john 和 to <food>for row joe。

所以我的问题是:我应该如何重新制定我的 UDF 来实现这一目标?

提前致谢!!:)

标签: javaxmlapache-sparkapache-spark-sqluser-defined-functions

解决方案


spark-xml包允许您直接在选择表达式中访问嵌套字段。为什么要寻找 UDF?

df.selectExpr("id", "name", "COALESCE(`expenses`.`food`.`details`.`date`, `expenses`.`travel`.`details`.`date`) AS expenses_date" ).show()

输出:

+---+----+-------------+
| id|name|expenses_date|
+---+----+-------------+
|  1|john|     20191203|
|  2| joe|     20191204|
+---+----+-------------+

编辑

如果唯一改变的标签是expensesstruct 之后的标签,那么您可以搜索下的所有字段expenses,然后搜索coalesce列:expenses.X.details.date。在 Spark 中是这样的:

val expenses_fields = df.select(col("expenses.*")).columns
val date_cols = expenses_fields.map(f => col(s"`expenses`.`$f`.`details`.`date`"))

df.select(col("id"), col("name"), coalesce(date_cols: _*).alias("expenses_date")).show()

不过,您不需要使用 UDF!


推荐阅读