如何用Pandas向量化函数按首尾日期聚合DataFrame数据

给定如下百万级规模的示例数据：

import pandas as pd
import numpy as np
import datetime as dt

df_data = pd.DataFrame([
    [dt.date(2023, 5, 8),  'Firm A', 'AS', 250.0, -1069.1],
    [dt.date(2023, 5, 8),  'Firm A', 'JM', 255.0, -1045.5],
    [dt.date(2023, 5, 8),  'Firm A', 'WC', 250.0, -1068.8],
    [dt.date(2023, 5, 11), 'Firm A', 'WC', 250.0, -1068.8],
    [dt.date(2023, 5, 8),  'Firm B', 'AS',  31.9,  -317.9],
    [dt.date(2023, 5, 8),  'Firm B', 'JM',  33.5,  -310.7],
    [dt.date(2023, 5, 8),  'Firm B', 'WC',  34.5,  -305.9],
    [dt.date(2023, 5, 11), 'Firm B', 'AS',  33.0,  -313.1],
    [dt.date(2023, 5, 11), 'Firm B', 'JM',  33.5,  -310.7],
    [dt.date(2023, 5, 11), 'Firm B', 'WC',  35.0,  -303.5],
    [dt.date(2023, 5, 10), 'Firm C', 'BC', 167.0,   301.0],
    [dt.date(2023, 5, 9),  'Firm D', 'BA', 791.9,  1025.0],
    [dt.date(2023, 5, 9),  'Firm D', 'CT', 783.8,  1000.0],
    [dt.date(2023, 5, 11), 'Firm D', 'BA', 783.8,  1000.0],
    [dt.date(2023, 5, 11), 'Firm D', 'CT', 767.9,   950.0]],
    columns=['Date', 'Name', 'Source', 'Value1', 'Value2'])

需求要求按Name分组，完成以下计算：

• 找到每个分组的最早（Min Date）和最晚（Max Date）日期
• 计算这两个日期下Value1、Value2的均值
• 统计对应日期的Source数量
• 计算首尾日期间Value1、Value2的变化量
• 需兼容仅存在单日期数据的分组

当前使用分组循环的方法可以得到目标格式，但百万级数据下速度极慢：

def compute_entry(df: pd.DataFrame) -> dict:
    dt_min  = df.Date.min()
    dt_max  = df.Date.max()
    idx_min = df.Date == dt_min
    idx_max = df.Date == dt_max
    data = {
        'Min Date':        dt_min,
        'AvgValue1 (Min)': df[idx_min].Value1.mean(),
        'AvgValue2 (Min)': df[idx_min].Value2.mean(),
        '#Sources (Min)':  df[idx_min].Value2.count(),
        'Max Date':        dt_max,
        'AvgValue1 (Max)': df[idx_max].Value1.mean(),
        'AvgValue2 (Max)': df[idx_max].Value2.mean(),
        '#Sources (Max)':  df[idx_max].Value2.count(),
        'Value1 Change':   df[idx_max].Value1.mean() - df[idx_min].Value1.mean(),
        'Value2 Change':   df[idx_max].Value2.mean() - df[idx_min].Value2.mean()
    }
    return data

df_pivot = pd.DataFrame.from_dict({sn_id: compute_entry(df_sub)
                                   for sn_id, df_sub in df_data.groupby('Name')}, orient='index')

尝试用pd.pivot_table提升效率，但输出格式不符合需求，难以转换：

pd.pivot_table(df_data,
               index=['Name', 'Date'],
               aggfunc={'Value1': np.mean, 'Value2': np.mean, 'Source': len})

需要用Pandas内置的向量化函数实现符合要求的输出格式，同时保证处理百万级数据的效率。

可以通过分组标记首尾日期 + 聚合 + 重塑格式的方式实现全向量化操作，避免循环，大幅提升效率：

步骤1：为每个分组标记最小/最大日期 #

先给每条数据标记属于该分组的Min Date或Max Date，同时处理单日期的情况（此时最小和最大日期相同）：

# 按Name分组，计算每个分组的最小和最大日期
grouped = df_data.groupby('Name')['Date']
df_data['is_min_date'] = df_data['Date'] == grouped.transform('min')
df_data['is_max_date'] = df_data['Date'] == grouped.transform('max')

步骤2：分别聚合最小/最大日期的统计量 #

分别筛选最小、最大日期的数据，按Name聚合得到所需的均值和计数：

# 聚合最小日期的统计数据
agg_min = df_data[df_data['is_min_date']].groupby('Name').agg(
    Min_Date=('Date', 'first'),
    AvgValue1_Min=('Value1', 'mean'),
    AvgValue2_Min=('Value2', 'mean'),
    Sources_Min=('Source', 'count')
)

# 聚合最大日期的统计数据
agg_max = df_data[df_data['is_max_date']].groupby('Name').agg(
    Max_Date=('Date', 'first'),
    AvgValue1_Max=('Value1', 'mean'),
    AvgValue2_Max=('Value2', 'mean'),
    Sources_Max=('Source', 'count')
)

步骤3：合并数据并计算变化量 #

将两个聚合结果合并，然后计算Value1和Value2的变化量：

# 合并最小和最大日期的统计结果
result = agg_min.join(agg_max)

# 计算变化量
result['Value1 Change'] = result['AvgValue1_Max'] - result['AvgValue1_Min']
result['Value2 Change'] = result['AvgValue2_Max'] - result['AvgValue2_Min']

# 重命名列名以匹配需求格式
result = result.rename(columns={
    'Min_Date': 'Min Date',
    'AvgValue1_Min': 'AvgValue1 (Min)',
    'AvgValue2_Min': 'AvgValue2 (Min)',
    'Sources_Min': '#Sources (Min)',
    'Max_Date': 'Max Date',
    'AvgValue1_Max': 'AvgValue1 (Max)',
    'AvgValue2_Max': 'AvgValue2 (Max)',
    'Sources_Max': '#Sources (Max)'
})

一步到位的简化写法 #

也可以用pivot结合聚合来实现更紧凑的代码：

# 给每条数据标记类型：min或max，单日期的同时标记为两者
df_data['type'] = np.where(df_data['is_min_date'], 'min', '')
df_data['type'] = np.where(df_data['is_max_date'], df_data['type'] + '_max', df_data['type'])
df_data['type'] = df_data['type'].str.strip('_')

# 按Name和type聚合，然后重塑为宽表
agg_df = df_data.groupby(['Name', 'type']).agg(
    Date=('Date', 'first'),
    AvgValue1=('Value1', 'mean'),
    AvgValue2=('Value2', 'mean'),
    Sources=('Source', 'count')
).unstack()

# 整理列名并计算变化量
agg_df.columns = [f'{col[0]} ({col[1].capitalize()})' for col in agg_df.columns]
result = agg_df.rename(columns={
    'Date (Min)': 'Min Date',
    'Date (Max)': 'Max Date',
    'Sources (Min)': '#Sources (Min)',
    'Sources (Max)': '#Sources (Max)'
})
result['Value1 Change'] = result['AvgValue1 (Max)'] - result['AvgValue1 (Min)']
result['Value2 Change'] = result['AvgValue2 (Max)'] - result['AvgValue2 (Min)']

效率说明 #

这种方法完全使用Pandas的向量化分组和聚合操作，避免了Python层面的循环遍历，处理百万级数据时的效率会比原循环方法提升数倍甚至数十倍，同时保证输出格式完全符合需求。

内容的提问来源于stack exchange，提问作者Phil-ZXX