数据处理神器可不止 Pandas 哦,还有 Polars,全方位解析 Polars

发布时间 2024-01-12 00:07:39作者: 古明地盆

楔子

Python 在数据处理领域有如今的地位,和 Pandas 的存在密不可分,然而除了 Pandas 之外,还有一个库也在为 Python 的数据处理添砖加瓦,它就是我们本次要介绍的 Polars。和 Pandas 相比,Polars 的速度更快,执行常见运算的速度是 Pandas 的 5 到 10 倍。 另外 Polars 运算的内存需求也明显小于 Pandas,Pandas 需要数据集大小的 5 到 10 倍左右的 RAM 来执行运算,而 Polars 需要 2 到 4 倍。

你可能会好奇,Polars 是怎么获得这种性能的?原因很简单,Polars 在设计上从一开始就以性能为宗旨,并通过多种方式实现。

1)用 Rust 编写

Rust 一种几乎与 C 和 C++ 一样快的低级语言,并且 Rust 天然允许安全并发,使并行性尽可能可预测。 这意味着 Polars 可以安全使用所有的 CPU 核心执行涉及多个列的复杂查询,甚至让 Ritchie Vink 将 Polar 的性能描述为过分并行(表示对并行的支持过于友好)。 所以 Polars 的性能远高于 Pandas,因为 Pandas 只使用一个核心执行运算。

2)基于 Arrow

Polars 具有惊人性能的一个因素是 Apache Arrow,一种独立于语言的内存格式。在 Arrow 上构建数据库的主要优点之一是互操作性,这种互操作性可以提高性能,因为它避免了将数据转换为不同格式以在数据管道的不同步骤之间传递的需要(换句话说它避免了对数据进行序列化和反序列化)。 此外 Arrow 还具有更高的内存效率,因为两个进程可以共享相同的数据,无需创建副本。 据估计,序列化/反序列化占数据工作流中 80-90% 的计算开销,Arrow 的通用数据格式为 Polars 带来了显著性能提升。

Arrow 还具有比 pandas 更广泛的数据类型内置支持,由于 Pandas 基于 NumPy,它在处理整数和浮点列方面非常出色,但难以应对其他数据类型。虽然 NumPy 的核心是以 C 编写,但它仍然受到 Python 某些类型的制约,导致处理这些类型时性能不佳,比如字符串、列表等等,因为 Numpy 本身就不是为 Pandas 而设计的。 相比之下,Arrow 对日期时间、布尔值、字符串、二进制甚至复杂的列类型(例如包含列表的列类型)提供了很好的支持。 另外,Arrow 能够原生处理缺失数据,这在 NumPy 中需要额外步骤。

最后,Arrow 使用列式数据存储,无论数据类型如何,所有列都存储在连续内存块中。 这不仅使并行更容易,也使数据检索更快。

3)查询优化

Polars 性能的另一个核心是评估代码的方式,Pandas 默认使用 Eager 执行,也就是按照代码编写的顺序执行运算。 相比之下,Polars 能够同时执行 Eager 和惰性执行,查询优化器将对所有必需运算求值并制定最有效的代码执行方式。,这可能包括重写运算的执行顺序或删除冗余计算。 例如,我们要基于列 Category 对列 Number 进行聚合求平均值,然后将 Category 中值 A 和 B 的记录筛选出来。

(
    df
    .groupby(by="Category").agg(pl.col("Number").mean())
    .filter(pl.col("Category").is_in(["A", "B"]))
)

如果表达式是 Eager 执行,则会多余地对整个 DataFrame 执行 groupby 运算,然后按 Category 筛选。 通过惰性执行,DataFrame 会先经过筛选,并仅对所需数据执行 groupby。

4)表达性 API

最后,Polars 拥有一个极具表达性的 API,基本上你想执行的任何运算都可以用 Polars 方法表达。 相比之下,Pandas 中更复杂的运算通常需要作为 lambda 表达式传递给 apply 方法。 apply 方法的问题是它循环遍历 DataFrame 的行,对每一行按顺序执行运算,这样效率很低,而 Polars 能够让你在列级别上通过 SIMD 实现并行。

以上就是 Polars 的优先,下面我们来安装它,直接 pip install polars 即可。当然啦, Polars 在安装时还提供了可选的依赖项。

可以根据自身情况选择安装,如果你需要所有的依赖项,那么直接 pip install 'polars[all]' 即可。

读取数据,创建 DataFrame

安装完 Polars 之后,我们来看如何读取数据并创建 DataFrame。Polars 支持读写所有的通用文件(如 CSV、JSON、Parquet、Excel),云存储(如 S3、Azure Blob、BigQuery)和数据库(如 Postgres、MySQL),我们分别介绍一下。

读取内置数据结构

最简单的方式,通过内置数据结构来创建。

from datetime import datetime
import polars as pl

df = pl.DataFrame(
    {
        "name": ["satori", "scarlet", "marisa"],
        "length": [155.3, 145.9, 152.1],
        "salary": [12000, 14000, 9000],
        "join_time": [
            datetime(1998, 12, 11, 12, 43, 18),
            datetime(1997, 8, 21),
            datetime(2005, 6, 18, 7, 22, 37),
        ]
    }
)
print(df)
"""
shape: (3, 4)
┌─────────┬────────┬────────┬─────────────────────┐
│ name    ┆ length ┆ salary ┆ join_time           │
│ ---     ┆ ---    ┆ ---    ┆ ---                 │
│ str     ┆ f64    ┆ i64    ┆ datetime[μs]        │
╞═════════╪════════╪════════╪═════════════════════╡
│ satori  ┆ 155.3  ┆ 12000  ┆ 1998-12-11 12:43:18 │
│ scarlet ┆ 145.9  ┆ 14000  ┆ 1997-08-21 00:00:00 │
│ marisa  ┆ 152.1  ┆ 9000   ┆ 2005-06-18 07:22:37 │
└─────────┴────────┴────────┴─────────────────────┘
"""

非常简单,并且每一列都有严格的类型,这里由于我们没有指定,因此 Polars 会自己推断。

那么我们如何手动指定类型呢?

import polars as pl

# 通过第二个参数 schema 指定列的类型
df = pl.DataFrame(
    {"col1": [0, 2], "col2": [3, 7]},
    schema={"col1": pl.Float32, "col2": pl.Int64}
)
print(df)
"""
shape: (2, 2)
┌──────┬──────┐
│ col1 ┆ col2 │
│ ---  ┆ ---  │
│ f32  ┆ i64  │
╞══════╪══════╡
│ 0.0  ┆ 3    │
│ 2.0  ┆ 7    │
└──────┴──────┘
"""
# schema 还可以接收一个列表,下面这种也是可以的
df = pl.DataFrame(
    {"col1": [0, 2], "col2": [3, 7]},
    schema=[("col1", pl.Float32), ("col2", pl.Int64)]
)

DataFrame 也可以基于 Series 创建,因为 DataFrame 本身就可以看作是多个 Series 的组合。

import polars as pl

# 通过第二个参数 schema 指定列的类型
df = pl.DataFrame(
    [
        pl.Series("col1", [0, 2], dtype=pl.Float32),
        pl.Series("col2", [3, 7], dtype=pl.Int64),
    ]
)
print(df)
"""
shape: (2, 2)
┌──────┬──────┐
│ col1 ┆ col2 │
│ ---  ┆ ---  │
│ f32  ┆ i64  │
╞══════╪══════╡
│ 0.0  ┆ 3    │
│ 2.0  ┆ 7    │
└──────┴──────┘
"""

由于 Series 里面已经包含了列名和类型,这时候就不用再指定 schema 参数了,如果指定了,还是以 schema 参数为准。

然后 DataFrame 里面还有一个比较重要的参数叫 orient,表示数据是按行解释,还是按列解释。该参数有三种选择:"row"、"col"、None,默认为 None,表示让 Polars 根据数据自己推断。这参数怎么理解呢?首先如果我们传递的数据是字典,那么一个键值对就是一列;如果传递的是包含多个 Series 的列表,那么一个 Series 就是一列,这很好理解,没有歧义。

但是问题来了,请看下面的例子。

import pandas as pd
import polars as pl

# Pandas 的 DataFrame
df1 = pd.DataFrame(
    [[0, 2], [3, 7]], columns=["col1", "col2"]
)
"""
   col1  col2
0     0     2
1     3     7
"""

# Polars 的 DataFrame
df2 = pl.DataFrame(
    [[0, 2], [3, 7]], schema=["col1", "col2"]
)
print(df1)
print(df2)
"""
shape: (2, 2)
┌──────┬──────┐
│ col1 ┆ col2 │
│ ---  ┆ ---  │
│ i64  ┆ i64  │
╞══════╪══════╡
│ 0    ┆ 3    │
│ 2    ┆ 7    │
└──────┴──────┘
"""

此时数据是一个二维列表,对于 Pandas 来说,内部的每个列表都是一行,而对于 Polars 来说,内部的每个列表都是一列。换句话说,此时 Polars 会按列来解释数据,如果想让它按行来解释,就需要 orient 参数了。

import polars as pl

# 将 orient 指定为 "row",那么内部每个列表都是一行
# 注意 schema,可以只指定列名,不指定类型(让 Polars 自己推断)
df = pl.DataFrame(
    [[0, 2], [3, 7]], schema=["col1", "col2"], orient="row"
)
print(df)
"""
shape: (2, 2)
┌──────┬──────┐
│ col1 ┆ col2 │
│ ---  ┆ ---  │
│ i64  ┆ i64  │
╞══════╪══════╡
│ 0    ┆ 2    │
│ 3    ┆ 7    │
└──────┴──────┘
"""
# 将 orient 指定为 "col",那么内部每个列表都是一列
df = pl.DataFrame(
    [[0, 2], [3, 7]], schema=["col1", "col2"], orient="col"
)
print(df)
"""
shape: (2, 2)
┌──────┬──────┐
│ col1 ┆ col2 │
│ ---  ┆ ---  │
│ i64  ┆ i64  │
╞══════╪══════╡
│ 0    ┆ 3    │
│ 2    ┆ 7    │
└──────┴──────┘
"""

默认情况下 orient 参数为 None,会由 Polars 自己推断,但还是建议手动指定此参数。当然,这只有在数据不明确的情况下,才需要这么做。如果传递的数据是字典,或者传递的列表里面的元素是 Series,那么 orient 参数就无需指定了,即使指定也没有意义。因为此时数据是明确的,对于字典来说,里面的一个键值对就是一列;对于包含 Series 的列表来说,里面的一个 Series 就是一列。

我们验证一下:

import polars as pl

df1 = pl.DataFrame(
    [
        pl.Series("col1", [0, 2], dtype=pl.Float32),
        pl.Series("col2", [3, 7], dtype=pl.Int64),
    ],
    orient="row"
)
df2 = pl.DataFrame(
    [
        pl.Series("col1", [0, 2], dtype=pl.Float32),
        pl.Series("col2", [3, 7], dtype=pl.Int64),
    ],
    orient="col"
)
print(df1)
"""
shape: (2, 2)
┌──────┬──────┐
│ col1 ┆ col2 │
│ ---  ┆ ---  │
│ f32  ┆ i64  │
╞══════╪══════╡
│ 0.0  ┆ 3    │
│ 2.0  ┆ 7    │
└──────┴──────┘
"""
print(df2)
"""
shape: (2, 2)
┌──────┬──────┐
│ col1 ┆ col2 │
│ ---  ┆ ---  │
│ f32  ┆ i64  │
╞══════╪══════╡
│ 0.0  ┆ 3    │
│ 2.0  ┆ 7    │
└──────┴──────┘
"""

可以看到,不管 orient 是啥,结果都没有变化。但如果传递的数据是列表,列表里面还是列表,此时 Polars 就不知道了,它会自己推断,因此这种情况下建议显式指定 orient 参数。当然啦,你也可以先将数据转成字典,然后再传进去。

最后,我们也可以不指定列名,让 Polars 自动为我们生成。

df = pl.DataFrame(
    [[0, 2], [3, 7]]
)
print(df)
"""
shape: (2, 2)
┌──────────┬──────────┐
│ column_0 ┆ column_1 │
│ ---      ┆ ---      │
│ i64      ┆ i64      │
╞══════════╪══════════╡
│ 0        ┆ 3        │
│ 2        ┆ 7        │
└──────────┴──────────┘
"""
# 由于我们没有指定列名,因此 Polars 会自动以 column_0、column_1、··· 的方式赋予列名
# 当然啦,我们肯定还是要手动指定列名的
df = pl.DataFrame(
    [[0, 2], [3, 7]],
    schema={"col1": pl.Float32, "col2": pl.Int64}
)
print(df)
"""
shape: (2, 2)
┌──────┬──────┐
│ col1 ┆ col2 │
│ ---  ┆ ---  │
│ f32  ┆ i64  │
╞══════╪══════╡
│ 0.0  ┆ 3    │
│ 2.0  ┆ 7    │
└──────┴──────┘
"""
# 会按照顺序依次赋值,col1 作为第一列的名称,col2 作为第二列的名称,因为 Python 从 3.6 开始,字典是有序的
# 不过我们不应该依赖字典有序这个特性,所以这种情况还是建议给 schema 传一个列表
# df = pl.DataFrame([[0, 2], [3, 7]], schema=[("col1", pl.Float32), ("col2", pl.Int64)])

以上就是通过内置数据结构创建 DataFrame,还是比较简单的,大部分情况下,我们传一个字典即可。

读取 CSV

读取 CSV 文件,Polars 有自己的快速实现,并且支持很多参数。

未完待续

本文参考自:

  • https://docs.pola.rs/user-guide/
  • https://blog.jetbrains.com/zh-hans/dataspell/2023/09/polars-vs-pandas-what-s-the-difference/