How to duplicate the row with value in one column and 0 in another column and vice versa in pyspark

SaurabhShelar 0
Viewed 45

I am trying to achieve below scenario in pyspark. Can someone please help me with it?

Source Dataframe:

ID      Col1    Flag1   Flag2
1234    Value1  1       1
5678    Value2  0       1

Desired results:

ID      Col1    Flag1   Flag2
1234    Value1  1       0
1234    Value1  0       1
5678    Value2  0       1
pyspark
2 Answers

If you have only two flags you can take the df, filter only the flag1 then the same for flag2 and do a union operation on the two dfs.

flag_one = df \
.withColumn("flag2", lit(0)) \
.filter(col("flag1") == 1)

flag_two = df \
.withColumn("flag1", lit(0)) \
.filter(col("flag2") == 1)

flag_one.union(flag_two).show()
walking0

It worked like charm, thanks. The logic is awesome! I have multiple flags though as you correctly pointed out, to tackle them I used a for loop.

def compute(df):

for_BoW = df.select("Flag3","Flag2")
Flag1 = df
for col_name in for_BoW.columns:
    Flag1 = Flag1.withColumn(col_name, F.lit(0)).filter(F.col("Flag1") == 1)

for_Flag2 = df.select("Flag3","Flag1")
Flag2 = df
for col_name in for_Flag2.columns:
    Flag2 = Flag2.withColumn(col_name, F.lit(0)).filter(F.col("Flag2") == 1)

for_Flag3 = df.select("Flag2","Flag1")
Flag3 = df
for col_name in for_Flag3.columns:
    Flag3 = Flag3.withColumn(col_name, F.lit(0)).filter(F.col("Flag3") == 1)

output = (Flag1.union(Flag2)).union(Flag3)
    
return output
SaurabhShelar0
Related