如何将元素附加到 Spark Dataframe 的数组列？答案

【问题标题】：How to append an element to an array column of a Spark Dataframe?如何将元素附加到 Spark Dataframe 的数组列？
【发布时间】：2021-04-13 22:43:03
【问题描述】：

假设我有以下 DataFrame：

scala> val df1 = Seq("a", "b").toDF("id").withColumn("nums", array(lit(1)))
df1: org.apache.spark.sql.DataFrame = [id: string, nums: array<int>]

scala> df1.show()
+---+----+
| id|nums|
+---+----+
|  a| [1]|
|  b| [1]|
+---+----+

我想在nums 列中向数组中添加元素，以便得到如下内容：

+---+-------+
| id|nums   |
+---+-------+
|  a| [1,5] |
|  b| [1,5] |
+---+-------+

有没有办法使用 DataFrame 的 .withColumn() 方法来做到这一点？例如。

val df2 = df1.withColumn("nums", append(col("nums"), lit(5)))

我查看了 Spark 的 API 文档，但找不到任何可以让我执行此操作的内容。我可能可以使用split 和concat_ws 一起破解一些东西，但如果可能的话，我更喜欢更优雅的解决方案。谢谢。

【问题讨论】：

标签： scala apache-spark

【解决方案1】：

import org.apache.spark.sql.functions.{lit, array, array_union}

val df1 = Seq("a", "b").toDF("id").withColumn("nums", array(lit(1)))
val df2 = df1.withColumn("nums", array_union($"nums", lit(Array(5))))
df2.show

+---+------+
| id|  nums|
+---+------+
|  a|[1, 5]|
|  b|[1, 5]|
+---+------+

array_union() 是自 2018 年 11 月 2 日 spark 2.4.0 发布以来添加的，在您提出问题 7 个月后，:) 请参阅 https://spark.apache.org/news/index.html

【讨论】：

【解决方案2】：

您可以使用udf 函数作为

def addValue = udf((array: Seq[Int])=> array ++ Array(5))

df1.withColumn("nums", addValue(col("nums")))
  .show(false)

你应该得到

+---+------+
|id |nums  |
+---+------+
|a  |[1, 5]|
|b  |[1, 5]|
+---+------+

更新另一种方法是使用数据集方式并使用地图作为

df1.map(row => add(row.getAs[String]("id"), row.getAs[Seq[Int]]("nums")++Seq(5)))
  .show(false)

其中 add 是一个案例类

case class add(id: String, nums: Seq[Int])

希望回答对你有帮助

【讨论】：

对于大数组有没有更有效的方法？因为默认情况下，Seq 是 List 并且附加 List 很慢
Majarjan 谢谢，但就性能而言，它是一样的。所以这与我所说的无关
感谢您的支持，让我知道性能是一样的。您可以尝试 rdd 方式来检查性能。 @MinhThai

【解决方案3】：

如果您和我一样，正在寻找如何在 Spark SQL 语句中执行此操作；方法如下：

%sql

select array_union(array("value 1"), array("value 2"))

您可以使用 array_union 连接两个数组。为了能够使用它，您必须将附加值转换为数组。使用 array() 函数来做到这一点。

您可以输入一个值，如 array("a string") 或 array(yourColumn)。

【讨论】：