Alternative of for loop on a column subsetting in pandas

Question

I am stuck at a time consumption related problem in pandas:

The code looks like this:

df = pd.DataFrame({"IDs": [1, 1, 1, 2, 2, 2, 3, 3, 3],
                   "Month": ["01", "02", "01", "01", "02", "01", "01", "02", "01"],
                   "column1": [0.9, 0.5, 0.3, 0.8, 0.5, 0.1, 0.6, 0.2, 0.8]})

df_list = []
for id in df.IDs.unique():
    temp = df[df.IDs == id]
    temp = temp.groupby("Month").mean()
    temp2 = temp['column1'].ewm(span=3, adjust=True).sum()
    df_list.append(temp2)

Note that unique IDs contains around 500k elements, the original dataframe df contains around 6 mil records.

Now I used tqdm to check the estimated time and it's taking 14-15 hrs to complete. If I have even loop only for temp = df[df.IDs == id] line and the estimated time is same (basically those are pandas function, so it should not make any performance issue). So the problem lies in that line.

Is there any alternative way to do this? Thanks for any suggestion.

Please provide data to make this reproducible on a small scale~ See: [How to make good reproducible pandas examples](https://stackoverflow.com/questions/20109391/how-to-make-good-reproducible-pandas-examples) — BeRT2me, Jul 23 '22 at 17:31
I take it from your question that your issue is related to the length of time it takes to perform a function. I am just not sure what it is you are trying to accomplish and "Is there any alternative way to do" doesn't help me understand. — itprorh66, Jul 23 '22 at 17:44
Hi @itprorh66 I am not sure what following operation actually called: `df[df.column = x]`... — Stat_prob_001, Jul 23 '22 at 17:47
also the main time consumption occurring when I am doing 'df[df.column = x]`. So I need an alternative of that in this situation. — Stat_prob_001, Jul 23 '22 at 17:49
@Stat_prob_001 how close does `df.groupby(['IDs', 'Month']).mean().ewm(span=3, adjust=True).sum()` get you? — Jon Clements, Jul 23 '22 at 17:50
@JonClements That's not giving the same numbers as OPs loop, and I can't figure out why. — BeRT2me, Jul 23 '22 at 17:53
thanks a lot @JonClements, I just have to reindex, thanks a lot! — Stat_prob_001, Jul 23 '22 at 18:06

score 1 · Accepted Answer · answered Jul 23 '22 at 18:06

out = (df.groupby(['IDs', 'Month'])
         .mean()['column1']
         .ewm(span=3, adjust=True)
         .sum())
print(out.reset_index())

Output:

   IDs Month  column1
0    1    01  0.60000
1    1    02  0.80000
2    2    01  0.85000
3    2    02  0.92500
4    3    01  1.16250
5    3    02  0.78125

out = (df.groupby('IDs')
         .apply(lambda x: x.groupby('Month')
                           .mean()['column1']
                           .ewm(span=3, adjust=True)
                           .sum())
         .stack())
print(out.reset_index(name='column1'))

Output:

   IDs Month  column1
0    1    01    0.600
1    1    02    0.800
2    2    01    0.450
3    2    02    0.725
4    3    01    0.700
5    3    02    0.550

I do not know why these give different results, but that probably comes from not understanding what `ewm` does. — BeRT2me, Jul 23 '22 at 18:44
"ewm" is exponential weighted moving average. The result I got is fine, it's same when I was using loop. — Stat_prob_001, Jul 23 '22 at 18:49

Alternative of for loop on a column subsetting in pandas

1 Answers1