java - Spark UDF : How to write a UDF on each row to extract a specific value in a nested struct?

Question

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

示例 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>

Spark 结果Dataset<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>对于排约翰和<food>对于排乔。

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

提前致谢!! :)

Answer 1

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|
+---+----+-------------+

编辑

如果唯一发生变化的标签是 expenses 结构之后的标签，那么您可以搜索 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!