使用Scala和Spark将数组列的行与另一个数据框的标题进行比较

问题描述

我正在使用Scala和Spark。 我有两个数据框。

一个如下:

i

在第二个中,数据帧头是

+------+------+-----------+
| num1 | num2 |    arr    |
+------+------+-----------+
|   25 |   10 | [a,c]     |
|   35 |   15 | [a,b,d]   |
+------+------+-----------+

我通过添加所有可能的标题列创建了一个案例类。

现在我想要的是,通过匹配num1和num2列,我必须检查是否 arr列中的数组包含第二个数据帧的标题。 如果是这样,则该值应为1,否则为0。

因此所需的输出是:

num1,num2,a,c,d

解决方法

如果我理解正确,则希望将数组列arr转换为每个可能值的一列,该值将包含数组是否包含该值。

如果是这样,您可以像这样使用array_contains函数:

val df = Seq((25,10,Seq("a","c")),(35,15,"b","d")))
             .toDF("num1","num2","arr")

val values = Seq("a","c","d")
df
    .select(Seq("num1","num2").map(col) ++
            values.map(x => array_contains('arr,x) as x) : _*)
    .show
+----+----+---+---+---+---+
|num1|num2|  a|  b|  c|  d|
+----+----+---+---+---+---+
|  25|  10|  1|  0|  1|  0|
|  35|  15|  1|  1|  0|  1|
+----+----+---+---+---+---+