如何将元素追加到 Spark Dataframe 的数组列？

Question

假设我有以下数据框：

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 文档，但找不到任何可以让我执行此操作的内容。

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

自 Spark 2.4.0 于 2018 年 11 月 2 日发布以来添加了

array_union()

，即您提出问题 7 个月后，:) 请参阅 https://spark.apache.org/news/index.html

Answer 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])

希望我的回答对你有帮助

Answer 3

如果您像我一样正在搜索如何在 Spark SQL 语句中执行此操作；方法如下：

%sql

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

您可以使用 array_union 连接两个数组。为了能够使用它，您必须将要附加的值转换为数组。通过使用 array() 函数来完成此操作。

您可以输入一个值，例如数组（“字符串”）或数组（您的列）。

Answer 4

使用 Spark array_join 时要小心。它正在删除重复项。因此，如果数组中有重复的条目，您将不会获得预期的结果。而且至少要花费O(N)。因此，当我将它与数组聚合一起使用时，它变成了 O(N^2) 操作，并且对于某些大型数组来说需要很长时间。

Answer 5

Spark 3.4+ 有一个将元素追加到数组末尾的方法：

array_append($"nums", 5)

Spark 3.5+ 有一种将元素追加到数组开头的方法：

array_prepend($"nums", 5)

---

示例：

val df1 = Seq("a", "b").toDF("id").withColumn("nums", array(lit(1)))
df1.show()
// +---+----+
// | id|nums|
// +---+----+
// |  a| [1]|
// |  b| [1]|
// +---+----+

val df2 = df1.withColumn("nums", array_append($"nums", 5))
df2.show()
// +---+------+
// | id|  nums|
// +---+------+
// |  a|[1, 5]|
// |  b|[1, 5]|
// +---+------+

val df3 = df1.withColumn("nums", array_prepend($"nums", 5))
df3.show()
// +---+------+
// | id|  nums|
// +---+------+
// |  a|[5, 1]|
// |  b|[5, 1]|
// +---+------+