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