pandas加速利器—polars

17 天前

风暴之灵

百格血精风暴猫

2023年4月更新,由于polars需要增加新的学习成本，且与很多库的兼容性不及pandas因此也没有进一步的深入学习。目前仅仅把polars库作为一个pandas的读写增强引擎进行了封装，封装为pandasrw库，已经上次pypi可以pip 安装，方便使用。

pandas和excel、csv高效读写的增强库—pandasrw - 知乎 (zhihu.com)

polars可以有效加速pandas运算，它使用Apache Arrow数据格式，用Rust语言开发。但是作为一个2020的新项目，目前语法还不完备，且网上的教程较少。我计划试用一下并在本文逐步更新探索的用法。

使用polars替代pandas受限于 API和pandas不同以及资料较少，需要较高的学习成本，因此主要学习I/O，apply、groupby、lazy API以替代pandas的效率瓶颈。

2022年9月17日更新：研究发现polars的计算效率是建立在他的表达式编程上的，如果一旦使用python的原生表达则效率下降很快，和优化过的pandas相差无几。例如使用polars的apply应用python自定义函数比pandas快2倍，但是在pandas中使用迭代器itertruple也能达到类似效果。

对于一个新项目，最好的参考资料时它的官网。

一、安装Polars

使用百度pip源。

# 安装polars
pip install polars -i https://mirror.baidu.com/pypi/simple/

二、I/O加速

通过polars读写速度远快于pandas，读写较大的文件时，可以使用polars加速。读取速度约为3倍左右，且可避免过大的数据无法加载的问题。

df = pl.read_csv("path.csv")
#lazy模式
df = pl.scan_csv("path.csv")

读取excel 需要先安装库

pip install xlsx2csv -i https://pypi.tuna.tsinghua.edu.cn/simple

通过以下方式可以读取excl/csv，并转化为pandas

df=pl.read_excel(r"D:\data\xx.xlsx").to_pandas()
print(type(plf))
<class 'pandas.core.frame.DataFrame'>

三、与pandas互相转化

为了解决polars语法不完备和不熟悉的问题，计划使用polars加速部分代码，对于无法使用polars的部分，使用pandas解决。

polars 转为pandas

import pandas as pd
import polars as pl
pl=pl.read_excel(r"D:\data\yy.xlsx")
#需要加.T进行转置，否者会把index变为列名
df=pd.DataFrame(pl).T

但是用不能官网里的以下代码，会报错 'pyarrow' is required when using to_pandas()。估计包版本的问题。

df=pl.read_excel(r"D:\data\xx.xlsx")
df.to_pandas()
pf=pl.DataFrame(df)

注意：df=pd.DataFrame(pl).T及其耗时,虽然polars本身读取快，但是使用polas读取再转换的pandas的想法行不通，耗时是pandas直接读取的10多倍。

pandas 转化为 polars

import polars as pl
import pandas as pd
df=pd.read_excel(r"D:\data\xx.xlsx")
plf=pl.from_pandas(df)

*************来自五个月后的更新2023/3/22*************

之前估计到电脑里pyarrow包存在冲突，不能使用to_pandas。最近电脑重装后，重新安装包测试，正如评论区所说转换很快，且不需要向 df=pd.DataFrame(pl).T 需要对转换后的DataFrame转置。对于一个20多万行100多列17M的excel表读取后5ms即可转换完成。

不过近期pandas2.0将要发布，pandas 读取excel过慢以及一系列其他性能瓶颈问题将会得到解决，将给polars带来较大的压力。

我近期计划使用polars包装一个常见文件格式的表格数据读写库，希望能尽快完成。

四、数据选择

参考资料

分为

Selecting with expression 使用表达式选择数据和 Selecting with indexing 使用索引选择两种，其中polars并不鼓励使用索引。

要使用表达式选择数据，使用：

filter 选择行的方法
select 选择列的方法

1、选择行

方法一:

在 filter 方法中，将用于选择行的条件作为表达式传递。

官方文档示例：注意filter函数内是pl.col不是df.col

multi_filter_df = df.filter((pl.col("id") <= 2) & (pl.col("size") == "small"))

方法二：使用slice方法切片。通过起始行和偏移量来切片。

DataFrame.slice(offset:int,length:int|None=None)
plf.slice(1,2) #表示从第一行起，偏移2行 数据为闭区间选择，选择数据包括第一行以及偏移的第n行
plf.slice(1,) #表示从第一行起，选择之后所有数据

方法三：使用.is_in方法选择值在列表中的列。

注意:是 .is_in() ,他是polars的表达式的属性或者方法，不是单独的 is_in() 函数。

df.filter(pl.col("CGI").is_in (["ID7ID11481","ID7ID11381"]))

更新：该问题已解决。

如何选择值在列表中的列？polars 不支持 in 和isin()函数。

不支持python in 表达式

plf.filter(pl.col("CGI") in ["ID7ID11481"])
Since Expr are lazy, the truthiness of an Expr is ambiguous. 
Hint: use '&' or '|' to chain Expr together, not and/or.

不支持pandas的isin（）函数

df.filter(pl.col("CGI") isin (["ID7ID11481"])
SyntaxError: invalid syntax

2、选择列

select 我们使用该方法选择列。在该 select 方法中，我们可以指定列：

一个（字符串）列名
（字符串）列名列表
与列数长度相同的布尔列表
表达式，例如列名上的条件
一个 Series

官方文档示例：

选择单个列

single_select_df = df.select("id")

选择列列表

list_select_df = df.select(["id", "color"])

选择带有表达式的列，根据列名的条件进行选择：

condition_select_df = df.select(pl.col("^col.*$"))

要根据列的 dtype 进行选择

dtype_select_df = df.select(pl.col(pl.Int64))

3、选择行和列

结合 filter and select 使用链式表达式来进行

官方文档示例：

expression_df = df.filter(pl.col("id") <= 2).select(["id", "color"])

4、查询优化

使用懒加载来进行查询优化。

五、DataFrame与Series互相转化

参考资料

DataFrame转化为Series

polars.DataFrame.to_series将一列转化为Series

DataFrame.to_series(index:int=0)→ polars.internals.series.series.Series

polars.DataFrame.to_struct将全部数据转化为Series

DataFrame.to_struct( name:str )→ polars.internals.series.series.Series

六、apply

apply应用场景

select context-> 单个元素
groupby context-> 单组

在 select 上下文中， apply 表达式将列的元素传递给 python 函数。

请注意，您现在正在运行 python，这会很慢。

结论：不使用polars的表达式编程时；效率提升有限，约为2倍左右，需要经一步考虑和多进程结合。该效率与pandas中的itertruple大致相同，优点是写法简单，缺点是不如itertruple灵活以及可能存在的语法问题。

资料说，polars的高效来源于Apache Arrow 数内存据格式和它的并发。如果使用Python的原生代码则该部分代码会收到GIL影响。

我自己验证，使用了Python的原生代码后，polars的apply是pandas的两倍，但是 并不自动启动多进程 。

只用polars语法则可以加速15倍左右。

与pandas的差异：

1、polars 使用apply后行作为元组传递，pandas 是Series。

2、

官网说明：

DataFrame.apply(
    f:Callable[[tuple[Any,...]],Any],
    return_dtype:type[DataType]|None=None,
    inference_size:int=256)

在 DataFrame 的行上应用自定义函数。 行作为元组传递。

使用apply方法通常比使用表达式 API 实现相同的逻辑更慢且更占用内存，因为：

with .applyapply逻辑是用 Python 实现的，但使用表达式逻辑是用 Rust 实现的
with .apply应用DataFrame 在内存中实现
表达式可以并行化
表达式可以优化

如果可能，请使用表达式 API 以获得最佳性能。

参数F

自定义函数/ lambda 函数。

return_dtype

操作的输出类型。如果没有给出，Polars 会自动推断类型。

inference_size

仅在自定义函数返回行的情况下使用。这使用前n行来确定输出模式

具体用法：

参考

使用apply函数返回单列：

我们可以使用polars.struct将多个值传递给apply-by-stamp-coupling中的fun函数。在lambda函数中，值作为Python dict传递，列名称作为键。

df = pl.DataFrame({"a":[1,2,3,4,5],"b":[2,3,4,5,6],"c":[3,4,5,6,7]})
df.with_column(
    pl.struct(["a", "b"])
    .apply(lambda cols: fun(cols["a"], cols["b"], 3))
    .alias("result")
)

如果不选择数据，lambda函数会把整行作为数组进行传递,注意下面例子中sr_1 = sr[10:150]中，sr是数组。

import pymannkendall as mk
def mk_(data):
    result=mk.original_test(data)
    return (result.trend,result.z,result.slope)
def srmk(sr):
    sr_1 = sr[10:150]
    ls=1
    try:
        n=len(sr_1)
        if n>10:
            srnp=np.array(sr_1)
            trd=mk_(srnp)
            ls=[trd[0],trd[1],trd[2]]
    except:
    return ls
pf=pl.read_excel(r"xx")
pf=pf.apply(lambda x:srmk(x))

使用apply函数返回多列：

关键点1：Python的zip函数和具有所需名称的元组来实现键值对字典，其中键是所需的列名（在示例中为d和e）。示列中 dict(zip(("d", "e"), fun(cols["a"], cols["b"], 3))，fun函数的值通过zip函数打包，然后dict转化为字典。

关键点2：unnest()函数将列分开。该函数可以将struct组成的列分开。

df = pl.DataFrame({"a":[1,2,3,4,5],"b":[2,3,4,5,6],"c":[3,4,5,6,7]})
def fun(a,b,shift_len): 
    return a+b*shift_len,b-shift_len