基于Source列补全Pandas DataFrame缺失日期的实现问题

需求是补全DataFrame中缺失的日期，但要按source列分组，每个source都要覆盖数据中的最小到最大日期区间，缺失的数值列填充为0。

原始数据集 #

date_found  source        count_unique_uuids    count_unique_uuids_raw
2021-05-13  source_1      20                    20
2021-05-14  source_2      1829                  1829
2021-05-14  source_3      2245                  2245
2021-05-14  source_1      40                    40
2021-05-15  source_1      903                   903
2021-05-16  source_2      20                    20
2021-05-18  source_3      89                    89

目标数据集 #

date_found  source        count_unique_uuids    count_unique_uuids_raw
    2021-05-13  source_1      20                    20
    2021-05-13  source_2      0                     0
    2021-05-13  source_3      0                     0
    2021-05-14  source_1      40                    40
    2021-05-14  source_2      1829                  1829
    2021-05-14  source_3      2245                  2245 
    2021-05-15  source_1      903                   903
    2021-05-15  source_2      0                     0
    2021-05-15  source_3      0                     0
    2021-05-16  source_1      0                     0
    2021-05-16  source_2      20                    20
    2021-05-16  source_3      0                     0
    2021-05-17  source_1      0                     0
    2021-05-17  source_2      0                     0
    2021-05-17  source_3      0                     0
    2021-05-18  source_1      0                     0
    2021-05-18  source_2      0                     0
    2021-05-18  source_3      89                    89

错误代码及问题 #

尝试了以下代码，报错ValueError: StringArray requires a sequence of strings or pandas.NA：

def add_missing_dates(df: pd.DataFrame) -> pd.DataFrame:
    df['date_found'] = pd.to_datetime(df['date_found'], format='%Y-%m-%d')
    min_date = df['date_found'].min()
    max_date = df['date_found'].max()

    (df.set_index(['date_found'])
     .groupby(['source'], as_index=False, group_keys=False)
     .apply(lambda x: x.reindex(pd.date_range(min_date, max_date)))
     .reset_index().rename(columns={'index': 'date_found'})
     .fillna(0)
     )
    return df


def add_dates_to_source(df: pd.DataFrame, source: str = 'source') -> pd.DataFrame:
    sources = df[source].tolist()
    dfs_to_concat = []

    for source_value in sources:
        filtered_df = df.loc[df[source] == source_value]
        df_ = add_missing_dates(filtered_df)
        df_[source] = source_value
        dfs_to_concat.append(df_)

    return pd.concat(dfs_to_concat)

运行方式：

df = add_dates_to_source(df)

解决方案 #

错误核心原因：

1. add_missing_dates函数未将处理后的结果赋值，直接返回原始DataFrame，等于未执行补全逻辑
2. 循环处理时使用了包含重复值的source列表，且重索引过程丢失了source列的字符串信息，导致后续填充时类型不匹配

推荐两种简洁的实现方式：

方法一：透视表+重索引

import pandas as pd

def fill_missing_dates(df):
    # 转换日期格式
    df['date_found'] = pd.to_datetime(df['date_found'])
    # 获取全量日期区间和所有唯一source值
    all_dates = pd.date_range(df['date_found'].min(), df['date_found'].max())
    all_sources = df['source'].unique()
    
    # 透视成宽表，补全日期后转回长表
    pivot_df = df.pivot(index='date_found', columns='source', values=['count_unique_uuids', 'count_unique_uuids_raw'])
    pivot_df = pivot_df.reindex(all_dates).fillna(0)
    # 转回长表并整理列结构
    result = pivot_df.stack('source').reset_index()
    # 恢复原始列顺序
    result = result[['date_found', 'source', 'count_unique_uuids', 'count_unique_uuids_raw']]
    # 确保数值列为整数类型
    result[['count_unique_uuids', 'count_unique_uuids_raw']] = result[['count_unique_uuids', 'count_unique_uuids_raw']].astype(int)
    return result

方法二：分组后重索引

import pandas as pd

def fill_missing_dates(df):
    df['date_found'] = pd.to_datetime(df['date_found'])
    min_date = df['date_found'].min()
    max_date = df['date_found'].max()
    all_dates = pd.date_range(min_date, max_date)
    
    # 定义分组补全逻辑
    def fill_group(group):
        # 重索引到全量日期，保留source列
        group = group.set_index('date_found').reindex(all_dates)
        group['source'] = group['source'].ffill().bfill()  # 填充source列的缺失值
        group[['count_unique_uuids', 'count_unique_uuids_raw']] = group[['count_unique_uuids', 'count_unique_uuids_raw']].fillna(0)
        return group.reset_index().rename(columns={'index': 'date_found'})
    
    # 分组执行补全并合并结果
    result = df.groupby('source', group_keys=False).apply(fill_group)
    # 按日期和source排序
    result = result.sort_values(['date_found', 'source']).reset_index(drop=True)
    # 转换数值列为整数
    result[['count_unique_uuids', 'count_unique_uuids_raw']] = result[['count_unique_uuids', 'count_unique_uuids_raw']].astype(int)
    return result

使用示例 #

调用函数即可得到目标数据集：

df = pd.DataFrame([
    ['2021-05-13', 'source_1', 20, 20],
    ['2021-05-14', 'source_2', 1829, 1829],
    ['2021-05-14', 'source_3', 2245, 2245],
    ['2021-05-14', 'source_1', 40, 40],
    ['2021-05-15', 'source_1', 903, 903],
    ['2021-05-16', 'source_2', 20, 20],
    ['2021-05-18', 'source_3', 89, 89]
], columns=['date_found', 'source', 'count_unique_uuids', 'count_unique_uuids_raw'])

result_df = fill_missing_dates(df)
print(result_df)

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