在数据分析相关的工作中,Pandas无疑是一个强大的工具,它的易用性和灵活性广受青睐。
然而,随着数据量的不断增长和计算需求的日益复杂,Pandas代码的性能问题也逐渐浮出水面。
如何让Pandas代码运行得更快、更高效,成为了每一个人使用者都需要面对的挑战。

今天,本文就一个简化版的实际分析案例,来一起探讨Pandas代码如何写才能性能更好,让你的数据处理流程更加顺滑,不再为漫长的等待运行而烦恼。

1. 案例介绍

假设有个股票分析的场景,我们需要给所有的股票每天的交易情况做一个评估,评估的结果有3个级别:
其中,“优”的条件是当天成交额大于10亿收盘价大于开盘价,也就是股票上涨;
“差”的条件是当天成交额小于1亿收盘价小于开盘价,也就是股票下跌;
除此之外的情况就是“中”

再次强调一下,这是一个简化的评估方法,主要是为了下面演示Pandas的代码性能,真实的评估不会如此粗糙。

股票交易的数据来自A股2024年1,2月份的日交易数据,大约20多万条。
数据可从地址 https://databook.top/stock/2024 下载。

导入数据:

import pandas as pd

# 这个路径根据实际情况修改
fp = r'D:\data\2024\历史行情数据-东财-不复权-2024.csv'

df = pd.read_csv(fp)
df = df.loc[:, ["股票代码", "日期", "开盘", "收盘", "最高", "最低", "成交量"]]
df

2. 不同写法的性能比较

下面是3种代码的写法都是基于pandas的,完成的功能也是一样的。

2.1. 循环遍历

给每条数据加一个评估的指标,最直接想到的方法就是遍历所有的数据,然后根据每条数据的情况,
给予一个评估指标(优,中,差)。

首先,封装一个评估一条数据的函数:

def eval_stock(row):
    """
    评估一条的数据
    """
    # 成交额
    volumn = row["收盘"] * row["成交量"]
    
    if volumn > 1_000_000_000 and row["收盘"] > row["开盘"]:
        return "优"

    if volumn < 100_000_000 and row["收盘"] < row["开盘"]:
        return "差"

    return "中"

然后用遍历的方式评估我们准备的数据(A股2024年1,2月份的日交易数据)。

for idx, row in df.iterrows():
    df.loc[idx, "评估"] = eval_stock(row)

df

虽然只有20多万条数据,但是执行时间还挺长的,在jupyter notebook 中用 %%timeit 魔法函数测试性能如下。

%%timeit
for idx, row in df.iterrows():
    df.loc[idx, "评估"] = eval_stock(row)

运行结果:

36.4 s ± 367 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

执行一次要36秒多,效率很低。

2.2. apply方法

apply方法是pandas提供的一种灵活处理数据的接口,它允许我们传入一个自定义函数来处理数据。
下面我们看看这种方式的性能如何。

%%timeit
df["评估"] = df.apply(eval_stock, axis=1)

运行结果:

4.9 s ± 86.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

性能提升非常明显(36.4秒->4.9秒),代码也变得更加简洁。

2.3. 向量化方法

最后,我们看看终极的写法,这种写法把pandas的数据结构DataFrame看成是一个一维的向量数组(每列一个向量),而不是一个二维的数值数组。
这样,我们操作数据是以为单位来操作,看看这样写的性能如何:

%%timeit
df["评估"] = "中"
df.loc[
    (df["收盘"] * df["成交量"] > 1_000_000_000) & (df["收盘"] > df["开盘"]),
    "评估",
] = "优"
df.loc[
    (df["收盘"] * df["成交量"] < 100_000_000) & (df["收盘"] < df["开盘"]), "评估"
] = "差"

运行结果:

8.22 ms ± 434 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

运行时间几乎可以忽略不计

3. 总结

同样使用pandas写数据分析的代码,性能差距居然会如此天差地别。

可见,学习pandas,不仅仅是学习它的各种接口和函数,
更重要的是了解从pandas执行的角度应该如何看待数据,是把数据看成一个一个独立的值,还是一行一行或一列一列的向量。
只有这样,才能用pandas高效的处理数据,这在数据量膨胀之后,会大大提高我们分析的效率。