Faster way to forward-fill and back-fill a groupby

user_2205969 0
Viewed 796

I want to ffill and bfill a specific column after a groupby.

My solution works:

import numpy as np
import pandas as pd

df = pd.DataFrame({
    "A": [1, 1, 1, 1, 2, 2, 2, 2],
    "B": [np.nan, 'f1', 'b1', np.nan, np.nan, 'f2', 'b2', np.nan]
})
df['B'] = df.groupby('A')['B'].apply(lambda _: _.ffill().bfill())

So this:

    A   B
0   1   NaN
1   1   f1
2   1   b1
3   1   NaN
4   2   NaN
5   2   f2
6   2   b2
7   2   NaN

Becomes this:

    A   B
0   1   f1
1   1   f1
2   1   b1
3   1   b1
4   2   f2
5   2   f2
6   2   b2
7   2   b2

Note that the sequences I want to ffill and bfill will always be in this format (Nan, x, y, Nan)

While this works, it's extremely slow on large dataframes.

I'm looking for some optimization to make this faster (ideally without resorting to using Dask or multiprocessing), maybe there's a Numpy optimization I can make?

I havn't had a lot of luck looking at other answers, like this one.

pandas optimization python
3 Answers

If you want speed, avoiding groupby and using numpy instead of pandas are good rules to follow. That's often not possible, but here you have special case with extremely regular data and all you need is subscript triplet of form [start:end:stride]:

df.iloc[0::4,1] = df.iloc[1::4,1].values
df.iloc[3::4,1] = df.iloc[2::4,1].values

Explanation: Most people know that you can use subscripts of the form [start:stop] but you can also add an optional stride argument. So the first line says to replace elements 0,4,8,... with elements 1,5,9,... The "values" is necessary to remove pandas indexing which is actually detrimental here.

This ought to be a bit faster just by avoiding groupby. For a little more speed, you could output column B to numpy, do work in numpy (basically the same code), and then reimport to pandas:

arr = df.B.values
arr[0::4] = arr[1::4]  
arr[3::4] = arr[2::4]
df.B = arr

Another thing you could do if you wanted to stay in pandas would be to unstack, copy entire columns, then re-stack. That's essentially what the above code is doing anyway. Honestly with such a rectangular sort of problem any array-style approach is going to be fairly fast.

user_38773380

If your data really well structured with continuous groups, then you can avoid groupby by using the limit parameter in ffill and bfill like:

print (df['B'].ffill(limit=1).bfill(limit=1))
0    f1
1    f1
2    b1
3    b1
4    f2
5    f2
6    b2
7    b2
Name: B, dtype: object
user_92747320

If you format is pre-fix as (Nan, x, y, Nan), when can do 

df.B=df.groupby([df.A,df.index//2]).B.transform('first')
Out[169]: 
    B
0  f1
1  f1
2  b1
3  b1
4  f2
5  f2
6  b2
7  b2
user_79645270
Related