首页 > 解决方案 > 如果第一列的顺序无关紧要,Pandas 进行多列排序的最有效方法

问题描述

我有一个大约 3 亿行和两列的 DataFrame。我需要对第二列进行排序,同时将第一列中具有相同值的所有行放在一起,但第一列的顺序无关紧要。

如果我开始:

 |ID  |Date
0|DEF |2000-01-01
1|ABC |2000-01-01
2|DEF |2000-02-01
3|ABC |2000-02-01

我需要结束:

 |ID  |Date
0|DEF |2000-01-01
2|DEF |2000-02-01
1|ABC |2000-01-01
3|ABC |2000-02-01

或者

 |ID  |Date
1|ABC |2000-01-01
3|ABC |2000-02-01
0|DEF |2000-01-01
2|DEF |2000-02-01

但我得到第一个还是第二个都没关系。跑步

df.sort_values(by=["ID", "Date"])

完整的数据集大约需要 8 分钟,但其中一部分时间用于确保 ID 列的顺序正确,这对我来说无关紧要。我认为在 ID 组中对日期进行排序而不对 ID 进行排序的唯一方法是

df.groupby("ID", sort=False).apply(lambda x: x.sort_values(by=["Date"])

但是 groupby 代码的运行时间比仅对 ID 列进行排序要长得多。在 10,000,000 行的较小样本中,groupby 排序需要 10 多分钟,而两列排序需要 17.6 秒。

有没有一种聪明的方法可以在第一列的组中对第二列进行排序,这比仅对两列进行排序更有效,还是双列排序是我最快的?

标签: pythonpandasperformancesorting

解决方案


免责声明:我为bodo.ai工作

我用 Bodo 尝试了你的代码,这就是它对我的影响。

首先生成一些数据并将其存储在parquet文件中(我测试了两种数据帧大小,1m 和 100m 行):


import pandas as pd
import numpy as np



def gen_data(n):
    df = pd.DataFrame(
        {
        "Date" : np.repeat(pd.date_range("2013-01-03", periods=1000), n),
        "id" : np.random.choice(['ABC','DEF','GHI'], n * 1000)
        }
        )
    df.to_parquet("pd_example"+str(n)+".pq", row_group_size=100000)

gen_data(100000)

然后我在你的代码中添加了一个 Bodo JIT 装饰器,打印出头部和尾部以查看它的外观,并添加了一些计时器来查看 Bodo 花费了多少时间。


import bodo
import time

@bodo.jit
def get_output(input_file):

    # add timer here to start timing after compilation
    t1 = time.time()

    df = pd.read_parquet(input_file)
    # your original code
    df1 = df.sort_values(by=["id", "Date"]) 

    print("time taken:",time.time()-t1)
    print(df1.head(10), df1.tail(10)) 

    return df1

Bodo 使用MPI来并行化 Pandas 代码,因此我也能够在 4 个内核上运行此代码。以下是我在 2020 Macbook Pro 上的结果:

核心 用的时间
1m 1 3.22s
1m 4 0.97s
10m 1 41s
10m 4 12.5s

Bodo 可作为conda 安装使用,最多可在 4 个内核上免费使用。


推荐阅读