Groupby and moving average function in pandas works but is slow

Question

I have some time series data collected for a lot of people (over 50,000) over a two year period on 1 day intervals. I want to applying a exponential weighted moving average function for each person and each metric in the dataset. After calculating the moving average, I want to join the new values up with the existing values in the dataframe. I have figured out how to do this on a small sample dataset, but I'm afraid that it's not optimized and therefore won't scale to my actual dataset. I have plenty of RAM available, so I'm not afraid of running out of space, but I am concerned with how long this could take to calculate for my entire dataset. Especially as I experiment with different spans and aggregation levels.

Thanks to issue #13966 for the starting place!

How can I optimize this code to better scale to a larger dataset?

import pandas as pd
import random
from datetime import datetime

# create some sample data
vals = ['A,' * 8, 'B,' * 8, 'C,' *8]
person = []
for x in vals:
    for y in x.split(','):
        if y != '':
            person.append(y)

metric1 = [random.random() for _ in range(0, 24)]
metric2 = [random.random() for _ in range(0, 24)]
dates = [datetime(2017,1,1), datetime(2017, 2, 1), datetime(2017, 3, 1), datetime(2017, 4, 1), 
         datetime(2017,5,1), datetime(2017, 6, 1), datetime(2017, 7, 1), datetime(2017, 8, 1)] * 3

# load into a dataframe
df = pd.DataFrame({'person': person, 
                   'metric1': metric1, 
                   'metric2': metric2,
                   'timestamp': dates})

def run_ewm(df):

    # sort the dataframe by the timestamp
    df.sort_values('timestamp', inplace=True)

    # group the df by person
    grouped = df.groupby('person')

    # create a temporary list to hold frames
    frames = []

    # iterate over the groups and apply exp. weighted moving average
    for group in grouped.groups:
        frame = grouped.get_group(group)
        frame['metric1_emw'] = frame['metric1'].ewm(span=60).mean()
        frames.append(frame)

    # concat the frames for a new dataframe
    df_new = pd.concat(frames)

    return df_new

%timeit df_new = run_ewm(df)

/home/curtis/Program_Files/miniconda2/envs/py35/lib/python3.5/site-packages/ipykernel_launcher.py:15: SettingWithCopyWarning: 
A value is trying to be set on a copy of a slice from a DataFrame.
Try using .loc[row_indexer,col_indexer] = value instead

See the caveats in the documentation: http://pandas.pydata.org/pandas-docs/stable/indexing.html#indexing-view-versus-copy
  from ipykernel import kernelapp as app

10 loops, best of 3: 101 ms per loop

Answer 1

Iterating in Python is slow, iterating in C is fast. So, it's best to keep as much as possible within Pandas to take advantage of its C implementation and avoid Python.

You can use apply on groupby objects to apply a function over every group in Pandas instead of iterating over them individually in Python.

df["metric1_ewm"] = df.groupby("person").apply(lambda x: x["metric1"].ewm(span=60).mean())

This gives me a run time of

3.21 ms ± 181 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

You can then also apply this over multiple columns:

df[["metric1_ewm", "metric2_ewm"]] = df.groupby("person").apply(lambda x: x[["metric1", "metric2"]].ewm(span=60).mean())

With that being said, you could generalize and reduce your run_ewm function by passing it a dataframe and a list of columns to apply ewm to:

def set_ewm(df, cols):
   df.sort_values("timestamp",inplace=True)
   df[[c+"_ewm" for c in cols]] = df.groupby("person").apply(lambda x: x[cols].ewm(span=60).mean())

You could also just return the columns or copy the dataframe before setting ewm, then returning the new dataframe to avoid modifying the dataframe inplace.