查找互连的行值的列名-Spark

时间:2019-05-07 11:25:47

标签: python apache-spark pyspark

我有一个遵循以下结构的Spark数据框:

    +------+-----------+-----------+-----------+------+
    |ID    |   Name1   |   Name2   |   Name3   |   Y  |
    +------+-----------+-----------+-----------+------+
    |   1  |       A,1 |       B,1 |       C,4 |   B  |
    |   2  |       D,2 |       E,2 |       F,8 |   D  |
    |   3  |       G,5 |       H,2 |       I,3 |   H  |
    +------+-----------+-----------+-----------+------+

对于我想查找的每一行,Y的值都表示为第一个元素。因此,理想情况下,我想检索类似[Name2,Name1,Name2]的列表。

我不确定先转换为RDD,然后使用map函数并将结果转换回DataFrame的方式和方式。

欢迎提出任何想法。

1 个答案:

答案 0 :(得分:1)

您可能可以尝试以下代码:

df.show()                                                                                                         
+---+-----+-----+-----+---+                                                     
| ID|Name1|Name2|Name3|  Y|
+---+-----+-----+-----+---+
|  1|  A,1|  B,1|  C,4|  B|
|  2|  D,2|  E,2|  F,8|  D|
|  3|  G,5|  H,2|  I,3|  H|
+---+-----+-----+-----+---+

from pyspark.sql import functions as F

name_cols = ["Name1", "Name2", "Name3"]
cond = F

for col in name_cols: 
     cond = cond.when(F.split(F.col(col),',').getItem(0) == F.col("Y"), col)

df.withColumn("whichName", cond).show()

+---+-----+-----+-----+---+---------+
| ID|Name1|Name2|Name3|  Y|whichName|
+---+-----+-----+-----+---+---------+
|  1|  A,1|  B,1|  C,4|  B|    Name2|
|  2|  D,2|  E,2|  F,8|  D|    Name1|
|  3|  G,5|  H,2|  I,3|  H|    Name2|
+---+-----+-----+-----+---+---------+