首页 > 解决方案 > 查找互连行值的列名 - Spark

问题描述

我有一个遵循以下结构的 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,然后使用映射函数并将结果转换回 DataFrame。

欢迎任何想法。

标签: pythonapache-sparkpyspark

解决方案


您可能可以尝试这段代码:

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

推荐阅读