Calculate the autocorrelation in pyspark - pyspark

Question

I am currently migrate my scripts from pandas to pyspark. I want to calculate the autocorrelation of return for each stock in each day. My data looks like:
+-----+--------+-------+----------+----------+
|stock| date | hour | minute | return |
+-----+--------+-------+----------+----------+
VOD | 01-02 | 10 | 13 | 0.05 |
VOD | 01-02 | 10 | 14 | 0.02 |
VOD | 01-02 | 10 | 16 | -0.02 |
VOD | 01-02 | 11 | 13 | 0.05 |
VOD | 01-02 | 12 | 03 | 0.02 |
VOD | 01-02 | 13 | 45 | -0.02 |
... ... ... .... ...
ABC | 01-02 | 11 | 13 | 0.01 |
ABC | 01-02 | 11 | 14 | 0.02 |
ABC | 01-02 | 11 | 15 | 0.03 |
The desired output should be like
+-----+--------+-------+
|stock| date | auto |
+-----+--------+-------+
VOD | 01-02 | 0.04 |
VOD | 01-03 | 0.07 |
VOD | 01-04 | 0.01 |
VOD | 01-05 | 0.05 |
Its very simple to do this in pandas
df_auto=df.groupby[('stock','date')]['return'].apply(pd.Series.autocorr,lag=1).reset_index(name='auto')
However, can someone let me know how to get the autocorrelation of a factor in pyspark? Thanks.

Answer 1

sample data:
df.show()
+-----+-----+----+------+------+
|stock| date|hour|minute|return|
+-----+-----+----+------+------+
| VOD|01-02| 10| 13| 0.05|
| VOD|01-02| 10| 14| 0.02|
| VOD|01-02| 10| 16| -0.02|
| VOD|01-02| 11| 13| 0.05|
| VOD|01-02| 12| 3| 0.02|
| VOD|01-02| 13| 45| -0.02|
+-----+-----+----+------+------+
Use groupby with collectist and apply udf on collected list. There is no autocorr function in spark so we have to use pandas/series:
from pyspark.sql import functions as F
def autocorr(ret):
import pandas as pd
s = pd.Series(ret)
return float(s.autocorr(lag=1))
auto=F.udf(autocorr, FloatType())
df.groupBy("stock","date").agg(F.collect_list(F.col("return")).alias("return")).withColumn("auto", auto("return")).select("stock","date","auto").show(truncate=False)
+-----+-----+-----------+
|stock|date |auto |
+-----+-----+-----------+
|VOD |01-02|-0.28925422|
+-----+-----+-----------+