Pandas：如何让你的代码性能飙升

在数据分析相关的工作中，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高效的处理数据，这在数据量膨胀之后，会大大提高我们分析的效率。