基于范围聚合DataFrame数值列：高效计算min/max/mean/sum

数据定义 #

我有两个DataFrame，分别为df和ranges，具体数据如下：

import pandas as pd

# 构建df
data_df = {
    'group': ['A', 'B', 'A', 'C', 'B'],
    'start': [10, 20, 15, 30, 25],
    'end': [50, 40, 60, 70, 45],
    'val1': [5, 10, 11, 12, 6],
    'val2': [5, 2, 1, 1, 0],
}
df = pd.DataFrame(data_df)

# 构建ranges
data_ranges = {
    'group': ['A', 'B', 'C'],
    'start': [0, 5, 25],
    'end': [50, 7, 35],
}
ranges = pd.DataFrame(data_ranges)

需求说明 #

需要根据ranges中定义的分组规则（同group下，df的行区间需落在对应ranges的区间范围内），对df的行进行分组聚合。针对val1、val2两列，分别计算每组的min、max、mean、sum统计量。

实际场景中，ranges约有5000个区间，df约有50万行，因此需要高效且内存友好的实现方案，可使用vaex等大数据处理框架。

预期输出 #

range_id val1              val2             
            min max mean sum  min max mean sum
0        0    5   5  5.0   5    5   5  5.0   5

解决方案 #

方案1：Pandas高效实现 #

避免笛卡尔积式的全量匹配，通过分组匹配减少计算量：

import pandas as pd

# 给ranges添加唯一标识range_id
ranges['range_id'] = ranges.index

# 定义分组匹配函数：同group下，匹配df行是否落在ranges的区间内
def match_interval(df_group, ranges):
    # 筛选当前group对应的ranges规则
    range_subset = ranges[ranges['group'] == df_group.name]
    if range_subset.empty:
        return pd.DataFrame()
    # 生成匹配掩码：df行的start >= range.start 且 df行的end <= range.end
    mask = (df_group['start'].values >= range_subset['start'].values[:, None]) & \
           (df_group['end'].values <= range_subset['end'].values[:, None])
    # 找到每个df行对应的range_id（取第一个匹配项，可根据需求调整多匹配逻辑）
    match_indices = mask.argmax(axis=0)
    df_group['range_id'] = range_subset['range_id'].iloc[match_indices].values
    # 过滤无匹配的行
    return df_group[mask.any(axis=0)]

# 分组匹配
df_matched = df.groupby('group', group_keys=False).apply(match_interval, ranges=ranges)

# 按range_id聚合计算统计量
agg_result = df_matched.groupby('range_id')[['val1', 'val2']].agg(['min', 'max', 'mean', 'sum'])

print(agg_result)

方案2：Vaex内存友好实现 #

Vaex采用懒加载机制，适合处理超大数据集，无需将全量数据载入内存：

import vaex

# 将Pandas DataFrame转为Vaex DataFrame
df_vaex = vaex.from_pandas(df)
ranges_vaex = vaex.from_pandas(ranges)
# 给ranges添加range_id
ranges_vaex['range_id'] = ranges_vaex.index

# 条件关联：同group + df区间落在range区间内
joined = df_vaex.join(ranges_vaex, on='group', how='inner', left_prefix='df_', right_prefix='range_')
joined = joined[(joined.df_start >= joined.range_start) & (joined.df_end <= joined.range_end)]

# 分组聚合计算统计量
agg_result = joined.groupby('range_range_id').agg({
    'df_val1': ['min', 'max', 'mean', 'sum'],
    'df_val2': ['min', 'max', 'mean', 'sum']
})

# 调整列名格式以匹配预期输出
agg_result.columns = pd.MultiIndex.from_tuples([
    ('val1', 'min'), ('val1', 'max'), ('val1', 'mean'), ('val1', 'sum'),
    ('val2', 'min'), ('val2', 'max'), ('val2', 'mean'), ('val2', 'sum')
])
agg_result.index.name = 'range_id'

print(agg_result)

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