Pandas分组填充DataFrame至指定长度倍数 排查reindex生效异常

初始数据 #

测试用DataFrame构造代码：

import pandas as pd
df = pd.DataFrame({"my_label":["A","A","B","B","B","C"], "value":[1,85,65,41,21,3]})

初始数据结构：

my_label   value
0        A   1
1        A   85
2        B   65
3        B   41
4        B   21
5        C   3

需求说明 #

按my_label字段分组，将每个分组的长度填充为指定长度的整数倍，新增填充行取对应分组的最后一个值前向填充。指定目标长度倍数为4时，预期输出如下：

my_label   value
0        A   1
1        A   85
2        A   85
3        A   85
4        B   65
5        B   41
6        B   21
7        B   21
8        C   3
9        C   3
10       C   3
11       C   3

原有失效代码 #

def _pad(group, seq_len):
    pad_number = seq_len - (len(group) % seq_len)
    if pad_number != seq_len:
        group = group.reindex(range(len(group)+pad_number)).ffill()
    return group
df = (df.groupby('my_label')
        .apply(_pad, (4))
        .reset_index(drop = True))

运行上述代码无法得到预期结果，reindex操作未在分组末尾正确生效。

分组对象会保留原DataFrame的全局索引，比如A组索引为[0,1]、B组索引为[2,3,4]、C组索引为[5]，并非分组内从0开始的连续序号。直接用从0开始的range序列做reindex时，新索引和分组原有索引不匹配，不仅无法正确在末尾补行，原有数据也会因为索引对齐失败生成空值。
另外groupby加apply时默认会把分组键作为额外索引层级，也会干扰最终结果输出。

只需要在reindex前先重置分组内索引为从0开始的连续值，同时给groupby加上group_keys=False参数避免额外索引层级即可：

def _pad(group, seq_len):
    pad_number = seq_len - (len(group) % seq_len)
    if pad_number != seq_len:
        # 先重置分组索引为0开始的连续值，再扩展索引做前向填充
        group = group.reset_index(drop=True).reindex(range(len(group) + pad_number)).ffill()
    return group

df = (df.groupby('my_label', group_keys=False)
        .apply(_pad, 4)
        .reset_index(drop=True))

运行后输出和预期完全一致。

内容的提问来源于stack exchange，提问作者Mrofsnart