Partitioning dataframe based on certain range in PySpark - dataframe

Question

Suppose, my dataframe looks like the one below. It has three different accounts (column "accnt") with the month (column "mn") ranging from 1 to 4. I want to precisely extract the "val2" for each account which has Null values in months 1 to 3.
---------------------------
accnt | mn | value | val2
---------------------------
1 | 1 | Null | 5
1 | 2 | Null | 5
1 | 3 | Null | 5
1 | 4 | 5 | 5
2 | 1 | 3 | 4.5
2 | 2 | 2 | 4.5
2 | 3 | Null | 4.5
2 | 4 | 4 | 4.5
3 | 1 | Null | Null
3 | 2 | Null | Null
3 | 3 | Null | Null
3 | 4 | Null | Null
So, my output would be like this:
---------------------------
accnt | mn | value | val2
---------------------------
1 | 1 | Null | 5
1 | 2 | Null | 5
1 | 3 | Null | 5
3 | 1 | Null | Null
3 | 2 | Null | Null
3 | 3 | Null | Null
My code:
df_data = {'accnt': [1,1,1,1,2,2,2,2,3,3,3,3,],
'mn': [1,2,3,4,1,2,3,4,1,2,3,4,],
'value': [None,None,None,5,3,2,None,4, None, None, None, None],
'val2': [5,5,5,5,4.5,4.5, 4.5, 4.5, None, None, None, None],
}
import pandas as pd
from pyspark.sql.function import *
df_pandas = pd.DataFrame.from_dict(df_data)
df = spark_session.createDataFrame(df_pandas)
df_final = df.where(isnan(col("value")) & \
(col("mn").isin(1,2,3))).select('accnt','mn','value', 'val2')
But the output also contains "accnt" value 2 which should not be there. I guess window functions would be handy here. Can anyone help me, please?

Answer 1

You're correct - you will need window functions. Your current filter does not take into account other rows for the same "accnt".
from pyspark.sql import functions as F, Window as W
cond = (F.isnan("value") & F.col("mn").isin(1,2,3)).cast('long')
df2 = df.withColumn('_flag', F.sum(cond).over(W.partitionBy('accnt')) >= 3)
df_final = df2.filter('_flag').select(df['*'])
df_final.show()
# +-----+---+-----+----+
# |accnt| mn|value|val2|
# +-----+---+-----+----+
# | 1| 1| NaN| 5.0|
# | 1| 2| NaN| 5.0|
# | 1| 3| NaN| 5.0|
# | 1| 4| 5.0| 5.0|
# | 3| 1| NaN| NaN|
# | 3| 2| NaN| NaN|
# | 3| 3| NaN| NaN|
# | 3| 4| NaN| NaN|
# +-----+---+-----+----+