最快的编写大型pandas DataFrame的方法

4

我进行了一项测试,测试了10种编写DataFrame和10种读取DataFrame的方法。我在这里发现了测试(我进行了一些调整并将Parquet添加到了列表中)。最好的方法包括:

df.to_feather('test.feather') :
39.34544535900204s

table=pyarrow.Table.from_pandas(df)
pq.write_table(table, "test_parquet_write_snappy_dict.parquet",
               use_dictionary=True, version='2.0', compression='snappy') :
40.6873751259991s

table=pyarrow.Table.from_pandas(df, nthreads=4)
pq.write_table(table, "test_parquet_write_snappy_dict.parquet",
               use_dictionary=True, version='2.0', compression='snappy') :
41.051620177000586s

为了写作和其他相关事项

pd.read_hdf('test_fixed.hdf', 'test') :
1.5275615360005759

pd.read_feather('test.feather') :
20.635139821002667

pd.read_pickle('test.pkl') :
37.21131302599679

为了阅读。

这是数据框:

sz = 50000000
df = pd.DataFrame({'A': randn(sz), 'B': randn(sz), 'C': randn(sz), 'D': randn(sz)})

我有两个问题。为什么read_hdfread_feather快20倍,而to_hdf在写入测试中甚至没有进入前三名?
第二个问题是,40秒对我来说仍然太慢了。有没有办法提高速度?可以通过使用不同的参数与to_featherwrite_table或使用我不知道的函数/模块来实现吗?
我不是要求别人帮我查找,我自己可以做这件事,也不想浪费任何人的时间。我正在寻找已经知道这个问题并能引导我走最快捷方式的人。

1
对于HDF读取测试的结果,由于速度差异较大,我建议你要保持一定的怀疑态度。或许你还可以计算一下读取数据的“校验和”(比如某一列的总和),并将其与已知值进行比较。这样可以确保你的IO操作是可靠的,同时也能确保数据没有被懒加载。 - AKX
1
好的。HDF读取测试是我的错。它不快速。 - TheHaricover
1个回答

7

这不是一个确切的答案,但下面是各种针对您这类数据的方法的更详细的基准测试...

import timeit

import numpy as np
import pandas as pd
from numpy.random import randn, randint


def generate_data(n):
    df = pd.DataFrame(
        {
            "dt": randint(1_600_000_000, 1_700_000_000, size=n) * 1000,
            "a": randn(n),
            "b": randn(n),
            "c": randn(n),
        }
    )
    df.dt = pd.to_datetime(df.dt, unit="s")
    df.set_index("dt", inplace=True)
    return df


def benchmark(df, name, saver, loader):
    verify(df, loader, saver)
    save_timer = timeit.Timer(lambda: saver(df))
    load_timer = timeit.Timer(lambda: loader().a.sum())
    save_n, save_time = save_timer.autorange()
    load_n, load_time = load_timer.autorange()
    total_time = (load_time / load_n) + (save_time / save_n)
    print(
        f"{name:<15s} : "
        f"{save_n / save_time:>20.3f} save/s : "
        f"{load_n / load_time:>20.3f} load+sum/s : "
        f"{1 / total_time: >20.3f} total speed"
    )


def verify(df, loader, saver):
    saver(df)
    loaded = loader()
    assert np.allclose(loaded.a.sum(), df.a.sum())
    assert np.allclose(loaded.b.sum(), df.b.sum())
    assert list(loaded.columns) == list(df.columns), loaded.columns


def save_feather(df):
    df = df.reset_index()
    df.to_feather("dummy.feather")


def load_feather():
    df = pd.read_feather("dummy.feather")
    df.set_index("dt", inplace=True)
    return df


def main():
    df = generate_data(5_000_000)
    benchmark(df, "dummy", lambda df: None, lambda: df)
    benchmark(df, "csv", lambda df: df.to_csv("dummy.csv"), lambda: pd.read_csv("dummy.csv", index_col="dt"))
    benchmark(df, "hdf", lambda df: df.to_hdf("dummy.h5", "dummy"), lambda: pd.read_hdf("dummy.h5", "dummy"))
    benchmark(df, "pickle", lambda df: df.to_pickle("dummy.pickle"), lambda: pd.read_pickle("dummy.pickle"))
    benchmark(df, "feather", save_feather, load_feather)
    benchmark(
        df,
        "parquet",
        lambda df: df.to_parquet("dummy.parquet", allow_truncated_timestamps=True),
        lambda: pd.read_parquet("dummy.parquet"),
    )


if __name__ == "__main__":
    main()

在我的电脑上(Ryzen 7 3700X, SSD硬盘, Windows 10, Python 3.8, 最新的Pandas等)和一个百万行的数据框(没耐心等待CSV结果),我得到了:

dummy           :         10475677.467 save/s :              186.737 load+sum/s :              186.734 total speed
csv             :                0.185 save/s :                0.970 load+sum/s :                0.156 total speed
hdf             :               18.289 save/s :               28.514 load+sum/s :               11.142 total speed
pickle          :               14.058 save/s :               31.962 load+sum/s :                9.764 total speed
feather         :               34.766 save/s :               41.436 load+sum/s :               18.904 total speed
parquet         :                7.707 save/s :               19.603 load+sum/s :                5.532 total speed

因此,HDF5绝对不会快上几个数量级。


1
你怎么能这么快地进行基准测试啊 xD 而且,你的结果告诉我,feather确实是最快的,没有任何可以改进其他函数的参数。 我可能会坚持使用feather,并尝试加速它! - TheHaricover

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接