Faster calculation of moving averages in Koalas.

[sushmit86]

I am trying to calculate the exponential moving average to a koalas dataframe. I am able to achieve this as below

import pandas as pd
import databricks.koalas as ks
from databricks.koalas import pandas_wraps

df = ks.DataFrame({'cust_id':['a', 'a', 'a', 'b', 'b'],
                   'sales': [100, 200, 300, 400, 500]})
def fun(col1) -> ks.Series[np.float64]:
    return col1.apply(lambda x: x.ewm(alpha=0.5, adjust=False).mean())  # Arbitrary pandas code.
df['moving_average'] = fun(df.groupby('cust_id').sales)
df.head()

However, when I try to implement the above in a dataset that has 30M records it takes 4 hrs to complete. Is there any way to speed this up

[HyukjinKwon]

@sushmit86, can you check here https://koalas.readthedocs.io/en/latest/user_guide/best_practices.html? In particular, Use distributed or distributed-sequence default index

[HyukjinKwon]

Also, it would be help analyze further if you are able to provide full codes probably with data if possible.

[sushmit86]

@HyukjinKwon I will check and let you know. The data is proprietary so I will cook up some data and test it. Will share once I have it ready

[sushmit86]

@HyukjinKwon While investigating I encountered another problem Here is the code

import pandas as pd
import databricks.koalas as ks
from databricks.koalas import pandas_wraps

kdf = ks.DataFrame({'cust_id':['a', 'a', 'a', 'b', 'b'],
                   'sales': [100, 200, 300, 400, 500],
                  'days':[12.1,13.1,14.1,78.1,87.2]})
@pandas_wraps
def fun(col1) -> ks.Series[np.float64]:
    return col1.apply(lambda x: x.ewm(alpha=0.5, adjust=False).mean())  # Arbitrary pandas code.
kdf['koalas_moving_average'] = fun(kdf.groupby('cust_id').days)
display(kdf)

The output to this is

cust_id sales days koalas_moving_average
3 b 400 78.1 400.0
2 a 300 14.1 300.0
4 b 500 87.2 450.0
0 a 100 12.1 175.0
1 a 200 13.1 250.0

while if we use pandas as below

pdf = kdf.to_pandas() pdf['pandas_moving_average'] = pdf.groupby('cust_id')['days'].apply(lambda x: x.ewm(alpha=0.5, adjust=False).mean()) display(pdf)

The output in this case is

cust_id sales days koalas_moving_average pandas_moving_average
0 a 100 12.1 175.0 12.10
1 a 200 13.1 200.0 12.60
3 b 400 78.1 450.0 78.10
2 a 300 14.1 250.0 13.35
4 b 500 87.2 500.0 82.65

As we can see when using Koalas the moving average seems to be getting calculated on sales column instead of the days column as desired

[HyukjinKwon]

^ this seems like a bug. Let me fix it.

[HyukjinKwon]

The issue above will be fixed by #1224

[HyukjinKwon]

I am leaving this resolved. The blocker bug was fixed. Let's track it in a different ticket if you still face the performance issue.