查找互连行值的列名 - Spark

Question

我有一个遵循以下结构的 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 如何以及是否有效。

欢迎提出任何想法。

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