Spark Window - Как сравнить первую строку с n-й строкой фрейма данных?

Question

У меня есть датафрейм, как показано ниже. Я должен взять ранг строки, в которой для request_typ указано «PD», и сделать разницу с рангом текущей строки и добавить его в качестве нового столбца.

Исходный кадр данных:

Id svc_dt clm_typ rank 48115882 20180209 RV 1 48115882 20180209 RJ 2 48115882 20180216 RJ 3 48115882 20180302 RJ 4 48115882 20180402 PD 5 48115882 20180502 RJ 6

Ожидаемый результирующий кадр данных:

Id svc_dt clm_typ rank diff_PD_Rank 48115882 20180209 RV 1 4 (Current rank - rank of column with 'PD') 48115882 20180209 RJ 2 3 48115882 20180216 RJ 3 2 48115882 20180302 RJ 4 1 48115882 20180402 PD 5 null 48115882 20180502 RJ 6 null

Answer 1

Решение PySpark.

Предполагая, что для каждого идентификатора clm_type 'PD' имеется одна строка, вы можете использовать условное агрегирование с max(when...)) для получения необходимой разницы.

# necessary imports 
w1 = Window.partitionBy(df.id).orderBy(df.svc_dt)
df = df.withColumn('rnum',row_number().over(w1))
w2 = Window.partitionBy(df.id)
res = df.withColumn('diff_pd_rank',max(when(df.clm_typ == 'PD',df.rnum)).over(w2) - rnum)
res.show()