前文介绍了4种提高python运算速度的思路,1是向量化函数(一文搞懂所有Python金融场景提速神器),2是多进程,3是内存映射惰性计算(续:Python加速之多进程(Swifter, Modin, Multiprocess, Pandarallel)和Vaex),4是JIT编译器续:python提速之JIT即时编译 (Numba)。本篇文章将会介绍一个融合了前述多种方法的库——Polars!
【关注公众号,回复 Polars 就可以获得本文全部Python代码!】
Polars的目标是提供一个快速的DataFrame库,从而替代Pandas。为什么Polars会比pandas快?以下几个原因:
1、可以多核心并行运算 (不需要用multiprocessing)
2、支持类似于Vaex的惰性计算和内存映射,通过Lazy API实现
3、使用了Apache Arrow数据格式,列式存储,减少了IO操作和内存占用,numpy是行式存储
4、底层Rust编写,接近C语言编译的速度 (不再需要JIT)
接下来,让我们具体看如何使用polars,以及比pandas快多少
直接pip install polars 就可以,本文使用的是最新的0.20.6
测试数据集:5000*4500的pandas Dataframe, index是日期,value是float,用pl.from_pandas把pandas.Dataframe转化成Polars.dataframe,但是注意polars不支持index索引,没有.index的attribute,所以需要reset_index后再转化就可以保留index列,不然直接转化,index会消失。
功能目录
1、行筛选
polars行筛选用filter,没有iloc,loc等索引相关的切片操作,但是可以通过slice(row,col)来切片,速度比pandas快,polars也支持is_in()筛选运算
2、列筛选
列筛选polars的表达式是select, filter和select可以串行使用,这和pandas性质一样。
注意,多个条件筛选最好是放在一个filter内部用与或非连接,而不是多个filter串行,因为多个filter串行实际是在内存上复制了多个原始数据,占用了很多内存,速度也会变慢。
3、基础运算
pandas的统计运算都支持,字符串操作也支持,还支持类似pandas case...when的语法,when... then..otherwise
4、拼接concat,join,hstack,vstack
各类join都支持
也支持pandas的merge_asof
还提供了merge_sorted函数,类似concat(axis=0)
5、特殊数据结构Struct
struct类似C语言的结构体,是apache arrow的一种数据类型,在rust编程中也会遇到
函数unnest()可以把struct拆开
6、添加列,删除列
用with_columns或者select都可以添加列,类似SQL语言;用pl.exclude()删除指定列,这些内部都是并行计算
7、Rolling
我们测试了常见的rolling窗口函数,求mean和max。Pandas耗时3s,但是polars只需要0.18s,速度是pandas的接近20倍!具体看Polars语法,
1、pldf.lazy()是启动惰性计算,如果没有最后的.collect()则并不会实际执行,只有.collect()后才会实际执行
2、rolling前需要确保index是排序后的,所以用了sort()。rolling()函数有两个重要参数,一是按照哪一列rolling,金融场景下多是时间,period则是window,polars支持1d,1s,1m等时间单位,也支持index count,1i。
3、agg()的用法和pandas相同。pl.exclude()是排除index日期列,对剩余的值来求mean和max
如果是复杂的函数计算,需要用到自定函数,可以直接python普通语法写函数,如下图所示,我们自定义了一个回归函数,rolling(window=4)滚动回归。按照这个方式调用self_func(x,y).rolling(),使用with_columns是在原始df中新增列,也可以使用select()来只获得计算结果。
也可以使用上述第一步的rolling().agg(result=func(x,y))实现
但是不支持在agg里面使用rolling(),agg里也不支持over()
polars的rolling()有一个参数by,意味着rolling里面可以实现groupby的功能,这个我们在下面Groupby+Rolling部分重点介绍
8、Groupby
当使用groupby.mean()等简单聚合函数操作时候,不需要apply的时候,polars是pandas的2倍速度
如果是自定义了复杂的函数运算,pandas就需要用Apply,polars并不需要apply,直接类SQL的写法就可以,是pandas的6倍速度
9、Groupby+Rolling+Agg
polars不支持groupby().rolling(),也不支持rolling().groupby(),最简便的写法是利用over()窗口函数实现groupby,再在前面叠加rolling_sum()或rolling_map(),但是不支持普通的rolling().over()
对于简单的聚合函数,polars提供了直接的rolling_sum()等函数,对于其他的算子,可以调用rolling_map(func,window_size)来实现,注意这里的func并不支持复杂的自定义函数,仅支持numpy的通用函数,元素层面的 (element-wise) 函数,具体的通用函数见链接http:/numpy.org/doc/stable/reference/ufuncs.html#methods
如果要实现自定义复杂函数怎么办呢?polars的rolling函数自带了参数by,0.19版本之前是group_by_rolling。所以可以直接rolling().agg()就可以计算了。
我们在测试数据集:5000*4500的pandas Dataframe, index是日期,value是float上做了比较测试。按照日期的年份groupby,然后rolling(window=20)计算每一列的mean()/max()
polars计算需要0.345s
pandas计算,需要9s,polars是pandas的20倍+速度。在本系列的第一篇文章中,同样的数据集用numpy stride实现rolling window=20计算mean()/max()用时1.8s,如果加上groupby耗时会更久,为了提速还要叠加numba,代码会更加复杂