ffill pandas dataframe with a strict fill limit

Question

This questions builds upon the old question: pandas ffill/bfill for specific amount of observation

Where the following answer is given.

df['filled'] = df.groupby("id")["indicator"].ffill(limit=2)
print (df)
    id  indicator  filled
0    1        NaN     NaN
1    1        NaN     NaN
2    1        1.0     1.0
3    1        NaN     1.0
4    1        NaN     1.0
5    1        NaN     NaN
6    1        NaN     NaN
7    1        NaN     NaN
8    1        4.0     4.0
9    1        NaN     4.0
10   1        NaN     4.0
11   1        NaN     NaN
12   1        NaN     NaN
13   2        NaN     NaN
14   2        NaN     NaN
15   2        1.0     1.0
16   2        NaN     1.0
17   2        NaN     1.0
18   2        NaN     NaN
19   2        5.0     5.0
20   2        NaN     5.0
21   3        3.0     3.0
22   3        NaN     3.0
23   3        NaN     3.0
24   3        NaN     NaN
25   3        NaN     NaN

However in my case only when the number of consec nan < limit it should forward fill. To illustrate the result for my application should look like:

df['filled'] = df.groupby("id")["indicator"].ffill(limit=2)
print (df)
    id  indicator  filled
0    1        NaN     NaN
1    1        NaN     NaN
2    1        1.0     1.0
3    1        NaN     NaN
4    1        NaN     NaN
5    1        NaN     NaN
6    1        NaN     NaN
7    1        NaN     NaN
8    1        4.0     4.0
9    1        NaN     NaN
10   1        NaN     NaN
11   1        NaN     NaN
12   1        NaN     NaN
13   2        NaN     NaN
14   2        NaN     NaN
15   2        1.0     1.0
16   2        NaN     NaN
17   2        NaN     NaN
18   2        NaN     NaN
19   2        5.0     5.0
20   2        NaN     5.0
21   3        3.0     3.0
22   3        NaN     NaN
23   3        NaN     NaN
24   3        NaN     NaN
25   3        NaN     NaN

Answer 1

In your case, you can check the consecutive NaN blocks and mask the filled column:

forward=2

# we groupby on `.notna().cumsum()` to find block sizes
# then compare to number of forward limit
valid_blocks = (df.groupby([df['indicator'].notna().cumsum(), 'id'])
                  ['id'].transform('size') <= forward
               )

# ffill as usual, then mask those invalid with `NaN`
df['filled'] = (df.groupby("id")["indicator"].ffill(limit=forward)
                  .where(valid_blocks | df['indicator'].notna())
               )  

Output:

    id  indicator  filled
0    1        NaN     NaN
1    1        NaN     NaN
2    1        1.0     1.0
3    1        NaN     NaN
4    1        NaN     NaN
5    1        NaN     NaN
6    1        NaN     NaN
7    1        NaN     NaN
8    1        4.0     4.0
9    1        NaN     NaN
10   1        NaN     NaN
11   1        NaN     NaN
12   1        NaN     NaN
13   2        NaN     NaN
14   2        NaN     NaN
15   2        1.0     1.0
16   2        NaN     NaN
17   2        NaN     NaN
18   2        NaN     NaN
19   2        5.0     5.0
20   2        NaN     5.0
21   3        3.0     3.0
22   3        NaN     NaN
23   3        NaN     NaN
24   3        NaN     NaN
25   3        NaN     NaN