从 Spark Dataframe 中的数组中提取单个元素

Posted 2023-04-13

【中文标题】从 Spark Dataframe 中的数组中提取单个元素

【英文标题】：Extracting single element from array in Spark Dataframe

【发布时间】：2020-08-11 15:50:27

【问题描述】：

我有一个Dataset<Row>，其结构如下：

"name": "Ben",
"lastHolidayDestination": "Florida",
"holidays": [
    "destination": "Florida",
     "year": 2020,
    "destination": "Lille",
     "year": 2019
]

我想使用 Spark SQL 将新列 lastHolidayYear 添加到数据集的根，通过查找连接到 lastHolidayDestination 的 holidays 元素来填充（假设永远只有一个）。所以输出数据集将是：

"name": "Ben",
"lastHolidayDestination": "Florida",
"lastHolidayYear": 2020,
"holidays": [
    "destination": "Florida",
     "year": 2020,
    "destination": "Lille",
     "year": 2019
]

我一直在玩 dataset.withColumn() 和 when()（使用 Java，但 Scala/Python 的答案很好），但到目前为止我一无所获。除非必须，否则我真的不想使用 UDF。有什么建议吗？

【问题讨论】：

您可以使用 SQL 来查询数组和对象...这些不是“嵌套数据帧”。例如，选择exploded holidays中的年份列的最大值

谢谢，我去看看。需要明确的是，我特别希望根据值而不是最大/最小值来选择元素。我的实际用例更复杂（你说得对，在这个简化的示例中它没有嵌套，所以我将重命名它）但我只是在这里将它简化为更简单的东西。

对我来说，您似乎想设置lastHolidayYear = max(year) in holidays...这不正确吗？

“通过查找连接到 lastHolidayDestination 的假期元素填充”所以不，不是最大值

啊。我假设“最后一个”您正在寻找数组中年份最大的元素。在这种情况下，您不需要单独存储 Florida，因为这将是来自原始数据集的派生视图

【参考方案1】：

要模拟与数组的连接，您可以使用 flatten 和 filter 组合：

val result = ds.withColumn("expl", explode(col("holidays")))
               .filter("lastHolidayDestination = expl.destination")
               .withColumn("lastHolidayYear", col("expl.year"))
               .drop("expl")

【参考方案2】：

从Spark 3.0开始，可以先过滤数组，然后用下面的表达式获取数组的第一个元素：

import org.apache.spark.sql.functions.element_at, filter, col

val extractElementExpr = element_at(filter(col("myArrayColumnName"), myCondition), 1)

其中"myArrayColumnName" 是包含数组的列的名称，myCondition 是条件，它是一个Column => Column 表达式。

对于您的具体示例，代码为：

import org.apache.spark.sql.functions.col, element_at, filter
import org.apache.spark.sql.Column

val isLastHoliday = (c: Column) => c.getField("destination") === col("lastHolidayDestination")
val getLastHoliday = element_at(filter(col("holidays"), isLastHoliday), 1)

val result = df.withColumn("lastHolidayYear", getLastHoliday.getField("year"))

使用此代码，如果您的输入数据框包含以下值：

+------+----------------------+--------------------------------+
|name  |lastHolidayDestination|holidays                        |
+------+----------------------+--------------------------------+
|Ben   |Florida               |[[Florida, 2020], [Lille, 2019]]|
|Alice |Peru                  |[[Florida, 2020], [Lille, 2019]]|
|Robert|Lille                 |[[Florida, 2020], [Lille, 2019]]|
+------+----------------------+--------------------------------+

输出将是：

+------+----------------------+--------------------------------+---------------+
|name  |lastHolidayDestination|holidays                        |lastHolidayYear|
+------+----------------------+--------------------------------+---------------+
|Ben   |Florida               |[[Florida, 2020], [Lille, 2019]]|2020           |
|Alice |Peru                  |[[Florida, 2020], [Lille, 2019]]|null           |
|Robert|Lille                 |[[Florida, 2020], [Lille, 2019]]|2019           |
+------+----------------------+--------------------------------+---------------+