使用PyArrow/Pandas在分组后补全缺失的[日期,麻雀类型]分组

问题描述 #

原始数据集：

需要补全所有缺失的[日期,麻雀类型]组合，遵循以下规则：

• 若该麻雀类型在缺失日期之前已出现，用最近的有效值填充；
• 若仅在未来出现，则填充Null。

期望得到的结果：

使用Pandas实现 #

步骤说明 #

1. 将date列转换为日期类型并按时间排序；
2. 生成所有日期与麻雀类型的笛卡尔积，构建完整的索引组合；
3. 将原始数据重新映射到完整索引上；
4. 按麻雀类型分组后，用向前填充（ffill）补全最近的有效值，未提前出现的类型保留Null。

代码示例 #

import pandas as pd

# 加载原始数据
df = pd.DataFrame({
    'date': ['2022-05-01', '2022-05-01', '2022-05-02', '2022-05-03', '2022-05-03'],
    'sparrow type': ['red', 'blue', 'red', 'blue', 'green'],
    'var': [100, 100, 20, 30, 40]
})

# 转换日期类型并按类型+日期排序
df['date'] = pd.to_datetime(df['date'])
df = df.sort_values(['sparrow type', 'date'])

# 生成所有可能的日期-类型组合
all_dates = df['date'].unique()
all_types = df['sparrow type'].unique()
full_index = pd.MultiIndex.from_product([all_dates, all_types], names=['date', 'sparrow type'])

# 重新索引并按类型向前填充
result = df.set_index(['date', 'sparrow type']).reindex(full_index).groupby('sparrow type').ffill().reset_index()

# 按日期排序对齐期望结果
result = result.sort_values('date').reset_index(drop=True)
print(result)

使用PyArrow实现 #

步骤说明 #

1. 将原始数据转换为PyArrow Table；
2. 生成日期与麻雀类型的笛卡尔积，构建完整的基础表；
3. 通过左连接将原始数据与基础表关联，得到带Null值的完整结构；
4. 使用窗口函数按类型分组、日期排序，取当前行及之前的最后有效值填充Null，未来未出现的类型保留Null。

代码示例 #

import pyarrow as pa
import pyarrow.compute as pc
from pyarrow.dataset import field

# 加载原始数据为PyArrow Table
data = pa.table({
    'date': pa.array(['2022-05-01', '2022-05-01', '2022-05-02', '2022-05-03', '2022-05-03'], type=pa.date32()),
    'sparrow type': pa.array(['red', 'blue', 'red', 'blue', 'green']),
    'var': pa.array([100, 100, 20, 30, 40], type=pa.int64())
})

# 生成所有日期-类型的笛卡尔积
all_dates = pc.unique(data['date'])
all_types = pc.unique(data['sparrow type'])
full_dates = pa.repeat(all_dates, len(all_types))
full_types = pa.concat_arrays([pa.repeat(t, len(all_dates)) for t in all_types])
full_table = pa.table({'date': full_dates, 'sparrow type': full_types})

# 左连接原始数据与完整表
joined = full_table.join(data, keys=['date', 'sparrow type'], join_type='left')

# 定义窗口：按类型分组，按日期排序，取当前行及之前的所有数据
window = pa.window(
    partition_by=field('sparrow type'),
    order_by=field('date'),
    preceding=pa.window.unbounded_preceding(),
    following=0
)

# 用窗口函数取最后有效值填充Null，未来未出现的类型保留Null
filled_var = pc.fill_null(pc.last(joined['var'], window=window), None)

# 构建结果表
result_table = pa.table({
    'date': joined['date'],
    'sparrow type': joined['sparrow type'],
    'var': filled_var
})

# 转换为Pandas查看结果（也可直接处理PyArrow表）
print(result_table.to_pandas().sort_values('date').reset_index(drop=True))

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