CODEWITHSUNDEEP
pythonapache-sparkpyspark
iJup
iJup

Reputation: 281

Find column names of interconnected row values - Spark

I have a Spark dataframe that adheres to the following structure:

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

For every row I want to find in which column the value of Y is denoted as the first element. So, ideally I want to retrieve a list like: [Name2,Name1,Name2].

I am not sure how and whether it works to convert first to a RDD, then use a map function and convert the result back to DataFrame.

Any ideas are welcome.

Upvotes: 0

Views: 129

Answers (1)

Steven
Steven

Reputation: 15273

You can probably try this piece of code : 

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

Upvotes: 1

Related Questions