Pyspark iterate over rows and compute counter with logic on result column - pyspark

Question

I have data in pyspark dataframe (it is a very big table with 900M rows)
This is the data that I have
+-------+---------+----------+
| key| time| cond|
+-------+---------+----------+
| 6| 3704| null|
| 6| 74967| 1062|
| 6|151565068| null|
| 6|154999554| null|
| 6|160595800| null|
| 6|166192324| null|
| 6|166549533| null|
| 6|171318946| null|
| 6|754759092| null|
| 6|754999359| 18882624|
| 6|755171746| 11381128|
| 6|761097038| null|
| 6|774496554| null|
| 6|930609982| null|
| 6|930809622| null|
| 1| 192427| null|
| 1| 192427| 2779|
| 1| 717931| null|
| 1| 1110573| null|
| 1| 1155854| null|
| 1| 70049289| null|
| 1| 70687548| null|
| 1| 71222733| null|
| 1| 85006084| null|
| 1| 85029676| null|
| 1| 85032605| 1424537|
| 1| 85240114| null|
| 1| 85573757| null|
| 1| 85710915| null|
| 1| 85870370| null|
+-------+---------+----------+
This is what I need to be doing with the dataframe (intermediate step):
+-------+---------+----------+--------+
| key| time| cond| result|
+-------+---------+----------+--------+
| 6| 3704| null| 0|
| 6| 74967| 1062| 1|
| 6|151565068| null| 0|
| 6|154999554| null| 1|
| 6|160595800| null| 2|
| 6|166192324| null| 3|
| 6|166549533| null| 4|
| 6|171318946| null| 5|
| 6|754759092| null| 6|
| 6|754999359| 18882624| 7|
| 6|755171746| 11381128| 0|
| 6|761097038| null| 0|
| 6|774496554| null| 1|
| 6|930609982| null| 2|
| 6|930809622| null| 3|
| 1| 192427| null| 0|
| 1| 192427| 2779| 1|
| 1| 717931| null| 0|
| 1| 1110573| null| 1|
| 1| 1155854| null| 2|
| 1| 70049289| null| 3|
| 1| 70687548| null| 4|
| 1| 71222733| null| 5|
| 1| 85006084| null| 6|
| 1| 85029676| null| 7|
| 1| 85032605| 1424537| 8|
| 1| 85240114| null| 0|
| 1| 85573757| null| 1|
| 1| 85710915| null| 2|
| 1| 85870370| null| 3|
+-------+---------+----------+--------+
The logic for 'result' column is as follows: have a running counter per key, zero the counter if 'cond' column is not null.
We can assume that table is orderBy("key",asc("time"))
My end results is actually average the result (per key) on rows were condition is not null.
It should look like this for above data (final result):
+--------+--------------+
| key | avg_per_key |
+--------+--------------+
| 6| 2.66666665| ==> (1+7+0)/3
| 1| 4.5| ==> (1+8)/2
+--------+--------------+
I plan to do it like this:
df_results = df3[df3.cond.isNotNull()].groupby(['key']).agg(
F.expr("avg(result)").alias("avg_per_key")
)
I assume it should work, but maybe there is a better way of doing it without the intermediate step in the middle.
How can that be done efficiently in pyspark? (remember that the dataset is huge)

Answer 1

Try this. Result is calculated by using an incremental sum over conditions, and then using those groupings as partitionBy in another window for row_number() - 1 to get desired result. Filter before groupBy should be performant by reducing shuffle.
from pyspark.sql import functions as F
from pyspark.sql.window import Window
w=Window().partitionBy("key").orderBy("time")
w1=Window().partitionBy("key","result").orderBy("time")
conditions=F.when((F.col("cond").isNotNull())&(F.col("lag").isNotNull()&\
F.col("lead").isNull()),F.lit(1))\
.when((F.col("cond").isNull())&(F.col("lag").isNotNull()),F.lit(1))\
.otherwise(F.lit(0))
df.withColumn("lag", F.lag("cond").over(w))\
.withColumn("lead", F.lead("cond").over(w))\
.withColumn("result",F.sum(conditions).over(w))\
.withColumn("result", F.row_number().over(w1)-1).filter("cond is not null")\
.groupBy("key").agg(F.mean(F.col("result")).alias("avg_per_key")).show()
#+---+------------------+
#|key| avg_per_key|
#+---+------------------+
#| 6|2.6666666666666665|
#| 1| 4.5|
#+---+------------------+

Answer 2

This was my solution for it, I am not saying that it is optimal, but worked for my case, when other attempts crashed the cluster.
I am a beginner in spark, so I understand that this approach might cause issues since it takes the dataset into memory.
If I had more time to play with it, I would try something with sortWithinPartitions
def handleRow(row):
temp = list(row[1])
temp = np.array([temp[x:x+2] for x in range(0, len(temp),2)])
temp[:,0] = temp[:,0].astype(float)
temp = temp[temp[:,0].argsort()]
avg_per_key= []
counter=0
for time,cond in temp:
if cond!=None:
avg_per_key.append(counter)
counter=0
else:
counter=counter+1
return [(row[0],-1 if len(avg_per_key)==0 else np.mean(avg_per_key))]
count = df3.rdd.map(lambda x: (x.key, (x.time, x.cond)))\
.reduceByKey(lambda a, b: a + b)\
.flatMap(handleRow)\
.collect()