Spark：将多列分解为一列

Posted 2023-04-13

【中文标题】Spark：将多列分解为一列

【英文标题】：Spark: explode multiple columns into one

【发布时间】：2020-10-16 17:41:25

【问题描述】：

是否可以在 spark 中将多个列分解为一个新列？我有一个如下所示的数据框：

userId         varA     varB
   1         [0,2,5]   [1,2,9]

想要的输出：

userId     bothVars
   1         0     
   1         2     
   1         5  
   1         1  
   1         2  
   1         9  

到目前为止我所尝试的：

val explodedDf = df.withColumn("bothVars", explode($"varA")).drop("varA")
                   .withColumn("bothVars", explode($"varB")).drop("varB")

这是行不通的。非常感谢任何建议。

【参考方案1】：

你可以将两个数组合二为一，flatten是嵌套数组，然后再分解，如下所示：

val df = Seq(
  (1, Seq(0, 2, 5), Seq(1, 2, 9)),
  (2, Seq(1, 3, 4), Seq(2, 3, 8))
).toDF("userId", "varA", "varB")

df.
  select($"userId", explode(flatten(array($"varA", $"varB"))).as("bothVars")).
  show
// +------+--------+
// |userId|bothVars|
// +------+--------+
// |     1|       0|
// |     1|       2|
// |     1|       5|
// |     1|       1|
// |     1|       2|
// |     1|       9|
// |     2|       1|
// |     2|       3|
// |     2|       4|
// |     2|       2|
// |     2|       3|
// |     2|       8|
// +------+--------+

请注意，flatten 在 Spark 2.4+ 上可用。

【参考方案2】：

使用array_union，然后使用explode函数。

scala> df.show(false)
+------+---------+---------+
|userId|varA     |varB     |
+------+---------+---------+
|1     |[0, 2, 5]|[1, 2, 9]|
|2     |[1, 3, 4]|[2, 3, 8]|
+------+---------+---------+

scala> df
.select($"userId",explode(array_union($"varA",$"varB")).as("bothVars"))
.show(false)

+------+--------+
|userId|bothVars|
+------+--------+
|1     |0       |
|1     |2       |
|1     |5       |
|1     |1       |
|1     |9       |
|2     |1       |
|2     |3       |
|2     |4       |
|2     |2       |
|2     |8       |
+------+--------+

array_union 在 Spark 2.4+ 中可用

【讨论】：

请注意array_union 将只保留数组之间任何重叠元素的一份副本。