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

时间:2018-04-06 04:16:33

标签: scala apache-spark

假设我有以下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文档,但找不到允许我这样做的任何内容。我可以使用splitconcat_ws来一起破解某些东西,但如果有可能的话,我宁愿选择更优雅的解决方案。感谢。

3 个答案:

答案 0 :(得分:3)

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()是在您问问题7个月后的2018年1月2日spark 2.4.0版本发布之后添加的:)请参见https://spark.apache.org/news/index.html

答案 1 :(得分: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]|
+---+------+

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

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

我希望答案很有帮助

答案 2 :(得分:0)

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

%sql

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

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

您可以输入数组(“字符串”)或数组(您的列)之类的值。