Pandas按多列分组实现高效聚合的方法探究

原始数据与目标需求 #

原始DataFrame #

import pandas as pd

df = pd.DataFrame(
    {
        "group0": [1, 1, 2, 2, 3, 3],
        "group1": ["1", "1", "1", "2", "2", "2"],
        "relevant": [True, False, False, True, True, True],
        "value": [0, 1, 2, 3, 4, 5],
    }
)

目标DataFrame #

target = pd.DataFrame(
    {
        "group0": [1, 2, 2, 3],
        "group1": ["1","1", "2", "2"],
        "value": [0, 2, 3, 5],
    }
)

聚合规则 #

• 优先取分组内relevant为True的所有value的最大值
• 若分组内无relevant为True的行，则取分组内value的最大值

现有实现方式 #

通过自定义函数结合groupby.apply实现：

def fun(x):
    tmp = x["value"][x["relevant"]]
    if len(tmp):
        return tmp.max()
    return x["value"].max()

# 调用示例
result = df.groupby(["group0", "group1"]).apply(fun).reset_index(name="value")

更高效的实现方案及性能对比 #

替代方案：排序后取分组最后一行 #

利用sort_values调整行顺序，让符合优先条件的行排在分组末尾，再通过groupby.last()获取目标值：

out = (df
   .sort_values(by=['relevant', 'value'], ascending=[False, True])
   .groupby(['group0', 'group1'], as_index=False)
   ['value'].last()
)

性能测试对比 #

场景1：分组数量较少

测试数据：

import numpy as np
from time import perf_counter

df = pd.DataFrame(
    {
        "group0": np.random.randint(0, 30, size=10_000_000),
        "group1": np.random.randint(0, 30, size=10_000_000),
        "relevant": np.random.randint(0, 1, size=10_000_000).astype(bool),
        "value": np.random.random_sample(size=10_000_000) * 1000,
    }
)

耗时结果：

• 排序后取last方案：14.82秒
• apply自定义函数方案：1.50秒

场景2：分组数量大幅增加

调整分组基数后的测试数据：

df = pd.DataFrame(
    {
        "group0": np.random.randint(0, 500_000, size=10_000_000),
        "group1": np.random.randint(0, 100_000, size=10_000_000),
        "relevant": np.random.randint(0, 1, size=10_000_000).astype(bool),
        "value": np.random.random_sample(size=10_000_000) * 1000,
    }
)

耗时结果：

• 排序后取last方案：15.29秒
• apply自定义函数方案：1423.84秒

结论 #

排序后取分组最后一行的方案（由@mozway提出）在分组数量较多时性能优势极为明显；仅当明确数据分组数量很少时，apply自定义函数方案才会有性能表现上的优势。

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