How can we set a flag for last occurence of a value in a column of Pyspark Dataframe

Viewed 270

Requirement : Here when last occurence of loyal with value is 1 then set flag as 1 else 0

Input:

+-----------+----------+----------+-------+-----+---------+-------+---+
|consumer_id|product_id|    TRX_ID|pattern|loyal| trx_date|row_num| mx|
+-----------+----------+----------+-------+-----+---------+-------+---+
|         11|         1|1152397078|  VVVVM|    1| 3/5/2020|      1|  5|
|         11|         1|1152944770|  VVVVV|    1| 3/6/2020|      2|  5|
|         11|         1|1153856408|  VVVVV|    1|3/15/2020|      3|  5|
|         11|         2|1155884040|  MVVVV|    1| 4/2/2020|      4|  5|
|         11|         2|1156854301|  MMVVV|    0|4/17/2020|      5|  5|
|         12|         1|1156854302|  VVVVM|    1| 3/6/2020|      1|  3|
|         12|         1|1156854303|  VVVVV|    1| 3/7/2020|      2|  3|
|         12|         2|1156854304|  MVVVV|    1|3/16/2020|      3|  3|
+-----------+----------+----------+-------+-----+---------+-------+---+

df = spark.createDataFrame(
[('11','1','1152397078','VVVVM',1,'3/5/2020',1,5),
('11','1','1152944770','VVVVV',1,'3/6/2020',2,5),
('11','1','1153856408','VVVVV',1,'3/15/2020',3,5),
('11','2','1155884040','MVVVV',1,'4/2/2020',4,5),
('11','2','1156854301','MMVVV',0,'4/17/2020',5,5),
('12','1','1156854302','VVVVM',1,'3/6/2020',1,3),
('12','1','1156854303','VVVVV',1,'3/7/2020',2,3),
('12','2','1156854304','MVVVV',1,'3/16/2020',3,3)
]
,["consumer_id","product_id","TRX_ID","pattern","loyal","trx_date","row_num","mx"])
df.show()

Output Required:

Note : Here Flag only checks whether last loyal value containing 1 and set the flag.

+-----------+----------+----------+-------+-----+---------+-------+---+----+
|consumer_id|product_id|    TRX_ID|pattern|loyal| trx_date|row_num| mx|Flag|
+-----------+----------+----------+-------+-----+---------+-------+---+----+
|         11|         1|1152397078|  VVVVM|    1| 3/5/2020|      1|  5|   0|
|         11|         1|1152944770|  VVVVV|    1| 3/6/2020|      2|  5|   0|
|         11|         1|1153856408|  VVVVV|    1|3/15/2020|      3|  5|   0|
|         11|         2|1155884040|  MVVVV|    1| 4/2/2020|      4|  5|   1|
|         11|         2|1156854301|  MMVVV|    0|4/17/2020|      5|  5|   0|
|         12|         1|1156854302|  VVVVM|    1| 3/6/2020|      1|  3|   0|
|         12|         1|1156854303|  VVVVV|    1| 3/7/2020|      2|  3|   0|
|         12|         2|1156854304|  MVVVV|    1|3/16/2020|      3|  3|   1|
+-----------+----------+----------+-------+-----+---------+-------+---+----+

What i tried :

from pyspark.sql import functions as F
from pyspark.sql.window import Window
w2 = Window().partitionBy("consumer_id").orderBy('row_num')
df = spark.sql("""select * from inter_table""")
df = df.withColumn("Flag",F.when(F.last(F.col('loyal') == 1).over(w),1).otherwise(0))

Here there are two scenarios :

1. Value 1 with preceding 0 (for your reference row_num 4 for consumer_id 11)

2. Value 1 with no preceding (for your reference row_num 3 for consumer_id 12)

2 Answers

Try this.

from pyspark.sql import functions as F
from pyspark.sql.window import Window

w = Window().partitionBy("product_id").orderBy('row_num')
df.withColumn("flag", F.when((F.col("loyal")==1)\
                             &(F.lead("loyal").over(w)==0),F.lit(1))\
                       .otherwise(F.lit(0))).show()

#+-----------+----------+----------+-------+-----+---------+-------+---+----+
#|consumer_id|product_id|    TRX_ID|pattern|loyal| trx_date|row_num| mx|flag|
#+-----------+----------+----------+-------+-----+---------+-------+---+----+
#|         11|         1|1152397078|  VVVVM|    1| 3/5/2020|      1|  5|   0|
#|         11|         1|1152944770|  VVVVV|    1| 3/6/2020|      2|  5|   0|
#|         11|         1|1153856408|  VVVVV|    1|3/15/2020|      3|  5|   0|
#|         11|         2|1155884040|  MVVVV|    1| 4/2/2020|      4|  5|   1|
#|         11|         2|1156854300|  MMVVV|    0|4/17/2020|      5|  5|   0|
#+-----------+----------+----------+-------+-----+---------+-------+---+----+

UPDATE:

from pypsark.sql import functions as F
from pyspark.sql.window import Window

w = Window().partitionBy("consumer_id").orderBy('row_num')
lead=F.lead("loyal").over(w)
df.withColumn("Flag", F.when(((F.col("loyal")==1)\
                             &((lead==0)|(lead.isNull()))),F.lit(1))\
                       .otherwise(F.lit(0))).show()

To add in the Murtaza's answer

We can add a new column which will check your second scenario for preceding null

window = Window.partitionBy('Consumer_id').orderBy('row_num')
df.withColumn('Flag',f.when((f.col('loyal')==1) 
                            & ((f.lead(f.col('loyal')).over(window)==0)
                              | (f.lead(f.col('loyal')).over(window).isNull())), f.lit('1')).otherwise(f.lit('0'))).show()

+-----------+----------+----------+-------+-----+---------+-------+---+----+
|consumer_id|product_id|    TRX_ID|pattern|loyal| trx_date|row_num| mx|Flag|
+-----------+----------+----------+-------+-----+---------+-------+---+----+
|         11|         1|1152397078|  VVVVM|    1| 3/5/2020|      1|  5|   0|
|         11|         1|1152944770|  VVVVV|    1| 3/6/2020|      2|  5|   0|
|         11|         1|1153856408|  VVVVV|    1|3/15/2020|      3|  5|   0|
|         11|         2|1155884040|  MVVVV|    1| 4/2/2020|      4|  5|   1|
|         11|         2|1156854300|  MMVVV|    0|4/17/2020|      5|  5|   0|
|         12|         1|1156854300|  VVVVM|    1| 3/6/2020|      1|  4|   0|
|         12|         1|1156854300|  VVVVV|    1| 3/7/2020|      2|  4|   0|
|         12|         2|1156854300|  MVVVV|    1|3/16/2020|      3|  4|   0|
|         12|         1|1156854300|  MVVVV|    1| 4/3/2020|      4|  4|   1|
+-----------+----------+----------+-------+-----+---------+-------+---+----+
Related